摘要

本文介绍了一种名为BitNet b1.58的1.58位量化感知训练方法，该方法在小型语言模型和视觉模型上实现了与大型语言模型相媲美的性能。通过引入一种依赖于中位数而非均值的量化过程，BitNet b1.58能够在模型参数从100K到48M的范围内保持高性能。研究结果表明，1.58位量化感知训练不仅在小型语言模型上接近最先进水平，而且在视觉模型上甚至超越了现有技术，为低资源应用中的模型部署提供了新的可能性。

原理

BitNet b1.58的核心在于其量化感知训练架构，该架构通过将16位权重层替换为仅假设值-1、0和1的层来实现。具体来说，BitLinear层作为PyTorch的torch.nn.Linear层的直接替代，通过五个步骤进行计算：激活归一化、量化到k位精度、16位影子权重量化到1.58位权重、量化激活与1.58位权重相乘、结果通过重缩放进行反量化。这种方法通过使用中位数而非均值进行权重量化，显著提高了模型的性能和稳定性。

流程

BitNet b1.58的工作流程包括：首先，对输入进行层归一化；其次，将归一化的激活量化到k位精度；然后，将16位影子权重量化为1.58位权重；接着，将量化后的激活与1.58位权重相乘；最后，通过重缩放对结果进行反量化。实验中，使用Adam优化器和批量大小为128进行训练，模型参数略有增加，但实际可训练参数数量不变。通过在小型Mistral-like模型和视觉模型上的广泛实验，验证了该方法的有效性和先进性。

应用

BitNet b1.58的1.58位量化感知训练方法在小型语言模型和视觉模型上的成功应用，预示着其在低资源环境中的广泛应用前景。这种方法不仅有助于减少模型的内存和计算需求，还能在保持高性能的同时，推动深度学习模型在资源受限环境中的部署。未来，该技术可能进一步扩展到其他类型的网络，如对象检测网络和编码器语言模型，为更广泛的AI应用提供支持。