TurboESM：针对蛋白质语言模型的超高效3位KV缓存量化
定量方法
2026-03-30 v1
摘要
蛋白质语言模型(PLMs)的快速规模化已在蛋白质结构预测和设计中实现了前所未有的准确率，但在推理过程中键值(KV)缓存的二次内存增长仍是单GPU部署和高吞吐量生成的禁锢瓶颈。虽然8位量化已成为标准，但3位量化仍面临因激活值中严重数值异常导致的挑战。本文提出TurboESM，将Google的TurboQuant方法应用于PLM领域。我们推导了解决旋转位置嵌入(RoPE)与正交变换根本不兼容问题的方法，提出了面向氨基酸激活流形的分头SVD校准方法，采用双查找表(LUT)策略处理K/V分布的非对称性，以及引入1位量化Johnson-Lindenstrauss(QJL)残差校正。所有实验均在ESM-2 650M模型上进行，本实现实现了从330MB降至47MB的7.1倍内存减少，同时在多样化蛋白族（包括短肽、跨膜螺旋、酶活性位点片段及本征无序区域）的自回归解码中保持余弦相似度>0.96。我们进一步实现了一个基于Triton的融合解码注意力内核，消除了中间去量化的内存分配，仅针对KV获取操作即可实现1.96倍的PyTorch双步骤路径加速；然而，TurboESM因KV量化和打包在预填充阶段产生21-27ms的开销，因而最适用于内存受限场景而非延迟敏感的短序列工作负载。分析表明，PLMs的异常轮廓比大型语言模型(LLM)更尖锐，这源于氨基酸词汇稀疏性，我们的方法有效解决了这些分布问题。
引用
@article{arxiv.2603.26110,
title = {TurboESM: Ultra-Efficient 3-Bit KV Cache Quantization for Protein Language Models with Orthogonal Rotation and QJL Correction},
author = {Yue Hu and Junqing Wang and Yingchao Liu},
journal= {arXiv preprint arXiv:2603.26110},
year = {2026}
}
备注
16 pages, 7 tables