从理论到实践，ViT-1.58b 在复杂 Transformer 架构中的应用 ！

视觉 Transformer （ViTs）通过利用注意力机制将图像块作为标记进行处理，在各种图像分类任务中取得了显著性能。然而，ViTs的高计算和内存需求对于在资源受限环境中部署提出了重大挑战。

本文介绍了ViT-1.58b，这是一种新颖的1.58位量化ViT模型，旨在大幅减少内存和计算开销，同时保持有竞争力的性能。

ViT-1.58b采用三元量化，通过将权重约束为{-1, 0, 1}并将激活量化为8位精度，在效率和准确度之间进行了精细平衡。作者的方法确保在内存和计算方面都能高效扩展。

在CIFAR-10和ImageNet-1k上的实验表明，ViT-1.58b在保持与全精度Vit相当的准确度的同时，显著减少了内存使用和计算成本。

本文强调了极端量化技术在开发可持续AI解决方案中的潜力，并为在实践应用中高效模型部署的更广泛讨论做出了贡献。

作者的代码和权重可在https://github.com/DLYuanGod/ViT-1.58b获取。

1 Introduction

Transformer架构的快速发展，特别是Vaswani等人（2017年）提出的Transformer Vaswani et al. (2017)，对计算机视觉领域产生了重大影响，尤其是随着Vision Transformers (ViTs) Dosovitskiy et al. (2020)的引入。通过将图像块视为标记，并利用注意力机制处理图像块，ViTs有效地捕捉整个图像中的复杂依赖关系，在各类图像分类任务中取得了显著性能，如Graham et al. (2021); Liu et al. (2021); Yuan et al. (2021)。然而，ViTs的计算和内存需求很大，这主要源于它们的注意力机制，其复杂度与标记数量成平方关系，如Dosovitskiy et al. (2020)所述。这种固有的复杂性为在资源受限的环境中部署ViTs带来了重大挑战，例如移动设备和嵌入式系统，在这些环境中能效和低延迟至关重要。

最近，为减轻这些需求，神经网络效率方面的研究探索了多种策略，包括模型剪枝Pan et al. (2023); Yu and Wu (2023); Song et al. (2022)，知识蒸馏，以及量化Li and Gu (2023); Sun et al. (2022); Lin et al. (2021)。在这些策略中，量化技术特别有效，因为它们直接降低了权重和激活的精度，从而显著降低了深度学习模型的内存和计算要求。传统上，由于不需要改变训练流程或重新训练模型，后训练量化因其简单性而受到青睐。然而，这种方法在较低精度水平下通常会导致显著的准确度损失，因为模型在训练期间没有针对量化表示进行优化，如Frantar et al. (2022); Chee et al. (2024)所述，这限制了其在高性能应用中的实用性。相比之下，量化感知训练（QAT）Park et al. (2018)将量化效应整合到训练过程本身中，模拟量化效果，通常比后训练方法获得更高的准确度。例如，极端量化，如同Binarized neural networks (BNNs) Courbariaux et al. (2016)中使用的1位模型，利用二进制权重和激活，显著降低了计算和内存需求。最近，将1位量化技术适应于Transformer模型，例如BitNet (Wang et al., 2023)和BiVit (He et al., 2023)，表明即使严重的量化也能在大幅降低资源消耗的同时保持性能。

然而，使用二进制量化的1位模型通常面临在极端降低权重精度的情况下保持模型准确度的挑战。为了解决这一限制，研究转向了三元量化，这种方法提供了更加平衡的方案。一个值得注意的例子是“1位LLM时代：所有大型语言模型都在1.58位中”(Ma et al., 2024)，该研究探讨了1.58位量化，其中权重可以取-1、0或1的值。这种方法通过引入零值，在性能和计算需求之间达到了更好的平衡，显著减少了计算开销。非极性权重（零）的包含不仅允许稀疏性，从而减少活跃计算的数量，而且比二进制权重保持了更丰富的表示，可能导致信息损失更少。这种量化策略在大规模语言模型中的成功表明，它有望应用于Vision Transformers，使得在保持模型准确度的同时更高效地处理视觉数据。

为了应对大规模ViT模型独特的需求，作者引入了ViT-1.58b，这是一个1.58位量化的ViT模型。ViT-1.58b利用三元量化来优化内存和计算效率，同时保持有竞争力的性能。这种方法利用了在语言模型中展示的极端量化的好处，将其适应到计算机视觉的背景中。作者的贡献总结如下：

• 作者提出了ViT-1.58b，这是第一个针对内存和计算效率优化的1.58位量化ViT模型。
• 作者在CIFAR-10和ImageNet-1k数据集上证明了ViT-1.58b的有效性，展示了与全精度ViT模型相比具有竞争力的准确度，同时显著降低了内存和计算成本。
• 作者提供了全面评估，将ViT-1.58b与最先进的量化方法和全精度ViT模型进行了比较，突出了在资源受限环境中的优势。

2 Methods

如图1所示，作者的ViT-1.58b架构主要基于用于图像分类的ViT。该过程首先将输入图像划分为多个块，然后对展平的块应用线性投影。这些块接着进行块和位置嵌入以及类别嵌入，之后被送入 Transformer 编码器。 Transformer 编码器的输出通过MLP进行处理，以产生最终的分类预测。作者采用来自BitNet b1.58（Ma等人，2024年）的BitLinear层来替换 Transformer 编码器中的传统_nn.Linear_层。采用不同的量化函数将权重量化到1.58位精度，将激活量化到8位精度，确保效率和性能。接下来，作者将介绍如何实现BitLinear层以实现1.58位权重和8位激活。

权重量化 在ViT-1.58b中，作者使用_absmean_量化函数来实现1.58位权重，该函数将权重限制为{-1, 0, 1}。具体来说，作者首先按权重矩阵 的平均绝对值进行缩放，从而得到 ，其中 是一个很小的浮点数，用来避免除以零并确保数值稳定性， 是权重矩阵的平均绝对值，计算如下：

然后，作者将缩放矩阵中的每个值四舍五入到最近的-1、0、+1之间的整数：

这个_RoundClip_函数将输入值 四舍五入到最近的整数，并将其限制在指定的范围 内。

使用这种方法，作者可以将权重矩阵 转换为三元矩阵 ，其中每个权重值被限制为{-1, 0, 1}。

激活量化。遵循（Dettmers等人，2022年）的方法，通过_absmax_量化函数以b位精度对激活进行量化，该函数将激活缩放到范围 ，其中 。在所提出的ViT-1.58b中，作者设置 ，这意味着激活被量化到8位精度，在计算效率和保持有效性能所需的足够精度之间取得了平衡，如（Ma等人，2024年）所述。激活量化过程如下所示：

其中 ，表示激活的最大绝对值，_Clip_函数确保缩放后的值被限定在指定的范围 内，定义为：

ViT-1.58b的核心组件是BitLinear层，它取代了Transformer中的传统_nn.Linear_层。最初，激活通过一个归一化层（Ba等人，2016年）确保激活值具有方差为1、均值为0，从数学上表示为 。在此之后，归一化的激活使用_absmax_量化函数进行量化，得到 。权重被量化到1.58位精度，矩阵乘法得到输出 ，然后输出随后去量化并通过 和 重新缩放至原始精度，表示为 ，其中

这些步骤共同使得_BitLinear_层在保持Transformer模型性能的同时，显著降低了计算成本和存储需求。

训练策略 在训练过程中，作者采用直通估计器（Straight-Through Estimator，STE）（Bengio等人，2013年）来处理反向传播中的不可微函数，如符号和裁剪函数。STE使得梯度能够通过这些不可微函数，从而有效训练量化的模型。

此外，作者还使用混合精度训练，在这种训练中，权重和激活被量化为低精度，但梯度和优化器状态保持高精度，以确保训练的稳定性和准确性。

3 Experiments and Results

实验设置。 作者在两个数据集CIFAR-10（Krizhevsky等人，2009年）和ImageNet-1k（Deng等人，2009年）上评估了作者的ViT-1.58b模型，并将其与几种版本的Vision Transformer Large（ViT-L）进行了比较：全精度ViT-L（即32位精度），以及16位、8位和4位推理版本，从内存成本、训练损失、CIFAR-10的测试准确度以及ImageNet-1k的Top-1和Top-3准确度进行对比。

本研究采用的计算框架包括四个NVIDIA TESLA A100 GPU，每个GPU有80 GB的VRAM。系统的处理核心使用AMD EPYC 7552 48核心处理器，配备80 GB的系统RAM，以有效管理大量数据集。作者使用了PyTorch 2.0.0版本，集成了CUDA 11.8，以优化跨GPU的张量操作。

结果。 表1展示了在不同位精度下评估的两个广泛认可的数据集CIFAR-10和ImageNet-1k上ViT-L模型的性能。结果表明，从全精度（ViT-L）到较低位配置（16位、8位、4位和1.58位）的位精度降低如何影响模型在训练损失、测试准确度以及top-k准确度指标方面的有效性。

对于CIFAR-10，全精度ViT-L模型的测试准确度为76.28%，作为比较的 Baseline 。当位精度降低到16位时，准确度适度下降到74.61%。进一步降低到8位和4位，准确度分别显著下降到72.20%和70.69%。这一趋势表明，较低位精度通常会降低模型的性能，可能是由于量化误差增加以及捕捉数据中变异性的能力降低。所提出的1.58-ViT-1.58b-L模型，在比其他量化版本更低的位精度下运行，记录的测试准确度为72.27%。有趣的是，这个模型的性能更接近于8位版本，而不是4位。

在更复杂和多样的ImageNet-1k数据集上，也观察到了类似的模式。全精度ViT-L模型取得了76.54%的Top-1准确度和90.23%的Top-3准确度。精度降低到16位时，Top-1和Top-3准确度分别下降到75.29%和87.44%。在8位（Top-1：74.11%，Top-3：85.26%）和4位（Top-1：70.88%，Top-3：82.50%）时，这些下降变得更加显著。1.58位模型取得了74.25%的Top-1准确度和85.78%的Top-3准确度，展示了一个在Top-1准确度上超过8位版本，在Top-3准确度上几乎与之持平的性能。

如图2所示，与全精度ViT-L及其量化版本相比，ViT-1.58b-L在训练损失和内存消耗方面显示出有希望的结果。左图显示了训练损失曲线，其中ViT-1.58b-L的损失紧随全精度ViT-L之后。这表明1.58位量化有效地保留了模型的学习能力。右图突出了内存消耗。全精度ViT-L需要1.14 GB的内存，而ViT-1.58b-L将其大幅降低到仅57 MB。这种内存使用的显著减少使作者的模型非常适合部署在资源受限的环境中。总的来说，作者的ViT-1.58b-L模型在保持竞争性能的同时，实现了大量的内存节省，展示了其效率和在实际应用中的实用性。

在图3中，随着位精度的降低，计算性能（以TFLOPs测量）显著提高。全精度达到了0.692 TFLOPs，而作者的模型达到了11.069 TFLOPs。这种在较低精度下性能的显著提升，突显了通过极端量化可实现的效率增益。作者的模型在计算吞吐量上提供了巨大的提升，使其非常适合在既需要速度又需要资源效率的高性能计算环境中。

4 Conclusion

在本文中，作者介绍了ViT-1.58b，这是一个通过三元量化有效解决视觉模型部署中的计算和内存挑战的1.58位量化视觉 Transformer 。作者的结果显示，ViT-1.58b在CIFAR-10和ImageNet-1k等基准测试中取得了具有竞争力的准确度，同时显著降低了资源需求。这个模型展示了高级量化技术在复杂 Transformer 架构中的潜力，突显了它们在开发可持续AI解决方案中的作用。未来的工作将探索ViT-1.58b的可扩展性及其在更大系统中的集成，提高其实用性和环境效益。

5 Limitations

尽管ViT-1.58b模型展现出有前景的性能和计算效率，但它也存在一些需要解决的局限性。本节概述了与作者的方法相关的主要挑战。

对预训练的需求。 ViT-1.58b架构的一个显著局限性是在应用作者的三元量化技术时，需要从头开始预训练ViT模型。这一要求可能会显著提高使用作者提出的模型的门槛，因为预训练需要大量的计算资源和时间。对于资源有限的组织或个人来说，如果没有现成的预训练模型，他们可能会发现采用这项技术具有挑战性。此外，预训练在确保模型在有效部署之前的鲁棒性和泛化能力方面引入了额外的复杂性。

训练难度。 另一个关键挑战是与全精度版本相比，1.58位ViT模型的训练难度增加了。在作者的实验中，如图2所示，作者观察到在CIFAR-10上训练1.58位ViT大约需要250个周期才能使训练损失达到约0.026，而标准ViT在仅200个周期内就能达到类似的0.024的损失。对于1.58位模型，这种增加的训练周期表明了较低的学习效率，这可能是由于精度降低以及模型捕捉详细特征表示能力的相应限制所致。

参考

[1].ViT-1.58b: Mobile Vision Transformers in the 1-bit Era.