蛋白质语言模型太大?用3-bit量化把它缩小7倍
蛋白质语言模型太大?用3-bit量化把它缩小7倍
蛋白质语言模型在最近几年变得越来越大。ESM-2有650M参数,如果是更大的版本,参数量可以达到几十亿。这些大模型在蛋白质结构预测、功能注释、药物发现等任务上效果确实好,但有一个实际问题:它们太大了,大到几乎不可能在单个GPU上部署。
更大的模型意味着更大的KV cache——在推理过程中,需要存储attention机制里Key和Value矩阵的缓存。这个缓存的大小和序列长度的平方成正比。处理一个几百个氨基酸的蛋白质序列,KV cache可能就已经是几百MB甚至几个GB。对于想要在笔记本电脑上跑模型的人来说,这显然不现实。
解决这个问题的一个常见方法是量化——把原本用32位浮点数存储的参数压缩到更少的位数。但这个方向有一个难点:3-bit量化(在每个参数上用3个比特表示)一直很难做成功。
这篇论文解决了这个问题。
论文标题是 “TurboESM: Ultra-Efficient 3-Bit KV Cache Quantization for Protein Language Models with Orthogonal Rotation and QJL Correction”,2026年3月27日提交到arXiv,作者是Yue Hu、Junqing Wang、Yingchao Liu。
方法的核心是把Google的一个叫TurboQuant的量化技术迁移到蛋白质语言模型领域,同时针对蛋白质数据的特点做了一些定制。
先解释一下技术细节。
研究者面临的最大困难是:Rotary Position Embeddings(RoPE)和正交变换不兼容。RoPE是现在很多大语言模型里用的位置编码方式,它通过旋转来编码位置信息。但对KV cache做量化需要做正交旋转——这两个操作冲突。
TurboESM的解决方案是设计了一个”先RoPE再旋转”的流水线,先做位置编码,再做量化,解决了这个兼容性问题。
另外,蛋白质模型里的activation有一个特点:它们的outlier比大语言模型更尖锐。这是因为氨基酸的词表只有20个,比自然语言的词表小得多,所以某些氨基酸的激活值会特别大。这些outlier会让量化效果变差。
TurboESM用了一个叫head-wise SVD calibration的方法,针对氨基酸激活值的分布特点来调整量化参数。还用了一个双查找表(dual LUT)策略来处理K和V分布不对称的问题。
效果如何。
研究者在ESM-2 650M上测试,结果显示:KV cache的内存从330MB降低到47MB,减少了7.1倍。而且在各种不同类型的蛋白质序列上——短肽、跨膜螺旋、酶活性位点、无序区域——做自回归解码时,余弦相似度都能保持在0.96以上。这意味着量化没有明显影响模型的输出质量。
研究者还实现了一个基于Triton的融合解码attention kernel,消除了中间的反量化内存分配,在KV获取操作上实现了1.96倍的加速。
当然也有代价:因为KV量化和packing的开销,预填充(prefill)阶段多了21-27毫秒的延迟。所以这个方法最适合的场景是内存受限制的情况,而不是对延迟极其敏感的超短序列任务。
这对做蛋白质研究的人意味着什么。
如果你想在笔记本电脑或者消费级GPU上跑蛋白质语言模型,TurboESM提供了一个可行的方案。7倍的内存压缩意味着原来需要几GB显存的任务,现在可能只需要几百MB。
对于那些需要处理大量蛋白质序列的应用场景——比如虚拟筛选、蛋白质设计的大规模采样——这种方法可以显著降低硬件门槛。
参考文献
- TurboESM: Ultra-Efficient 3-Bit KV Cache Quantization for Protein Language Models with Orthogonal Rotation and QJL Correction. Yue Hu, Junqing Wang, Yingchao Liu. arXiv:2603.26110. https://arxiv.org/abs/2603.26110