探索蛋白质语言模型的后训练量化
机器学习
2023-10-31 v1 人工智能
定量方法
摘要
无监督蛋白质语言模型（ProteinLM，如 ESM-1b 和 ESM-2）近期的进展在不同蛋白质预测任务中展现出前景。然而，这些模型因其高计算需求、显著内存需求和延迟而面临挑战，限制了它们在资源受限设备上的使用。为此，我们探索了蛋白质语言模型的后训练量化（PTQ），重点关注基于 ESM-2 ProteinLM 的 AlphaFold 简化版 ESMFold。我们的研究是首次尝试量化蛋白质语言模型的所有权重和激活。我们观察到，典型的均匀量化方法在 ESMFold 上表现不佳，在使用 8-bit 量化时导致 TM-Score 显著下降。我们进行了广泛的量化实验，揭示了与 ESMFold 相关的独特挑战，特别是 Layer Normalization 之前高度非对称的激活范围，使得用低比特定点格式难以表示。为应对这些挑战，我们提出了一种针对蛋白质语言模型的新 PTQ 方法，利用分段线性量化处理非对称激活值以确保精确近似。我们在蛋白质结构预测任务中证明了方法的有效性，表明 ESMFold 可被准确量化为低比特宽度而不损害精度。此外，我们将方法应用于接触预测任务，展示了其通用性。总之，我们的研究引入了针对蛋白质语言模型的创新 PTQ 方法，解决了特定量化挑战，并可能催生更高效的蛋白质语言模型，对各类蛋白质相关应用具有重要意义。
引用
@article{arxiv.2310.19624,
title = {Exploring Post-Training Quantization of Protein Language Models},
author = {Shuang Peng and Fei Yang and Ning Sun and Sheng Chen and Yanfeng Jiang and Aimin Pan},
journal= {arXiv preprint arXiv:2310.19624},
year = {2023}
}
备注
8 pages, 4 figures