TWLA: Achieving Ternary Weights and Low-Bit Activations for LLMs via Post-Training Quantization

阅读笔记

以下为个人阅读笔记，原文见 arXiv:2606.13054v2。

论文背景与动机

大语言模型（LLMs）虽然在语言处理任务上表现卓越，但其巨大的参数量导致了高昂的内存和计算成本，阻碍了在边缘设备等资源受限平台上的部署。论文指出，三元化（Ternarization）是一种极具潜力的压缩技术，通过将权重约束为 $\{-1, 0, +1\}$，可以显著减少模型大小并降低推理复杂度。然而，现有的三元化方法主要面临以下挑战：

1. 激活值分布问题：现有的方法难以处理长尾激活分布（heavy-tailed activation distributions），因此通常保持激活值的高精度，这从根本上限制了端到端的推理加速。
2. 权重分布不匹配：经验表明，LLM 的通道权重分布通常接近单峰高斯分布，这与三元码本 $\{-1, 0, +1\}$ 不匹配，导致投影到三元空间时产生较大的量化误差。
3. 训练成本高昂：虽然 BitNet v2 展示了三元权重与 4 位激活的联合部署，但其依赖于昂贵的量化感知训练（QAT），训练成本超过 $10^4$ GPU 小时。

为了解决上述问题，作者提出了一种名为 TWLA（Ternarized Weights and Low-bit Activations）的后训练量化（PTQ）框架。该框架旨在实现 1.58 位权重压缩和 4 位激活量化（W1.58A4），同时保持高精度，且无需重新训练。

核心方法：TWLA 框架

TWLA 框架由三个紧密耦合的模块组成，旨在协同优化权重三元化和激活低比特量化。论文通过图 3 展示了整体工作流程，该流程包括初始化、权重三元化优化、正交旋转重塑以及层间感知的混合精度分配。

如图 3 所示，TWLA 的整体设计遵循后训练量化的范式，但引入了特定的优化策略来适应三元权重和低比特激活的联合需求。作者指出，TWLA 的核心在于解决两个关键问题：一是如何将权重分布重塑为适合三元化的三模态形式，二是如何抑制激活值的异常值以支持低比特量化。

关键模块详解

为了实现上述目标，论文设计了以下三个核心组件：

1. 欧氏到流形非对称三元量化器 (E2M-ATQ)

该模块旨在最小化权重三元化过程中的层输出误差。作者提出了一种两阶段优化过程：

• 欧氏初始化：首先基于欧氏空间进行初始化。
• 流形重定位：随后将参数重定位到流形空间。 通过这种非对称三元量化器，E2M-ATQ 提高了层输出的可校准性，使得在权重三元化的情况下仍能保持较好的层输出一致性。

2. 克罗内克正交三模态重塑 (KOTMS)

这是 TWLA 的核心创新之一。论文认为，单峰的权重分布不适合三元码本，而三模态分布（tri-modal distribution）则更匹配。

• 权重重塑：KOTMS 应用了一种具有克罗内克结构的正交旋转（Kronecker-structured orthogonal rotation），将权重重塑为三元友好的三模态分布。
• 激活异常值抑制：利用正交等价性，相同的旋转矩阵被共享用于激活值的混合。这种共享的正交变换在统计上缩小了激活异常值，稳定了动态范围，从而有利于低比特激活量化。
• 分布变化：论文在图 2 中对比了应用 TWLA 前后的分布。应用前，激活异常值阻碍低比特量化，权重为单峰分布；应用后，激活被平滑，权重转变为对称的三模态形式。

3. 层间感知激活混合精度 (ILA-AMP)

由于共享正交变换带来的量化增益在不同层之间是固有的异质的，某些增益较小的层可能成为低比特激活量化的瓶颈。为了解决这个问题，作者提出了 ILA-AMP：

• 层间交互成本：该策略首次在比特分配中明确引入了相邻层之间的二阶交互成本。
• 联合优化：它将每层的敏感性与由分布偏移引起的层间耦合统一到一个二次代理目标函数中。
• 防止级联效应：通过显式地将共享正交矩阵引起的激活增益层间差异纳入优化，ILA-AMP 防止了由少数弱层触发的误差级联，确保了整体系统的稳定性。

实验结果与分析

论文在多个 LLaMA 系列模型上进行了广泛实验，验证了 TWLA 在 W1.58A4 配置下的有效性。

• 精度保持：如图 1 所示，在七个零样本任务上，TWLA 在权重和激活同时量化的情况下（W1.58A4）保持了鲁棒性。相比之下，其他方法（如 QuaRot, PB-LLM 等）在 4 位激活量化下出现了显著的性能下降。TWLA 在同等内存成本下，实现了比仅权重量化更好的精度 - 效率权衡。
• 对比 SOTA：实验表明，TWLA 优于现有的最先进 2 位和亚 2 位方法，将 LLMs 推向了 PTQ 领域的 W1.58A4 制度。
• 推理加速：由于实现了端到端的低比特权重和激活，TWLA 能够交付显著的推理加速，解决了以往方法因保留高精度激活而限制加速的问题。

总结与展望

综上所述，作者提出的 TWLA 框架通过 E2M-ATQ、KOTMS 和 ILA-AMP 三个模块的协同工作，成功解决了 LLM 三元化中的权重分布不匹配和激活长尾问题。论文认为，TWLA 不仅实现了 1.58 位权重和 4 位激活的联合压缩，还通过无重新训练的后训练量化方式降低了部署门槛。

该工作的主要贡献在于：

1. 识别了 LLM 在 W1.58A4 下的关键瓶颈（三元量化器可校准性、权重分布不匹配、激活长尾）。
2. 提出了无需重训练的 PTQ 框架，实现了三元权重与低比特激活的联合支持。
3. 通过实验证明了在 PTQ 设置下，TWLA 能够超越现有的 2 位及亚 2 位方法，为高效 LLM 部署提供了新的解决方案。