🤖 machine learning

BCJR-QAT: A Differentiable Relaxation of Trellis-Coded Weight Quantization

本文介绍了 BCJR-QAT，这是一种对格码权重量化的可微分松弛方法，它用温度控制的 BCJR 和积算法替代了不可微分的维特比 argmax，从而实现了端到端的量化感知训练，并在大语言模型上经验性地超越了当前最先进的后训练量化前沿。

原作者： Venugopalan Iyengar

发布于 2026-05-12

📖 1 分钟阅读☕ 轻松阅读

原作者： Venugopalan Iyengar

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象你拥有一个庞大的图书库（即大型语言模型），你希望将其缩小，以便能放入普通笔记本电脑或手机中。为此，你需要压缩库中的“单词”（即权重）。

本文介绍了一种名为BCJR-QAT的新方法，可将这些图书进一步压缩至每个单词仅2 比特（这就像将一张高清照片压缩成一张微小且颗粒感十足的缩略图）。

以下是他们如何实现这一目标的简略故事，使用了简单的类比：

1. 问题：“单行道”陷阱

此前，缩小这些模型的最佳方法是一种称为QTIP的技术。可以将 QTIP 想象为一位非常聪明的图书管理员，他查看一页文本后说道：“好的，我将用字典中最近的简单短语来替换这个复杂句子。”

然而，这种方法存在局限。一旦图书管理员做出选择，他就无法回头更改决定以使整本书听起来更好。这就像一条“单行道”。如果图书管理员在早期选了一个略有偏差的短语，整本书都会受到影响，而他无法修正，因为该决定是“硬性”且最终的。

为了解决这个问题，研究人员通常使用训练（QAT），让图书管理员练习并调整其选择。但这里有个关键问题：图书管理员的决策过程涉及一个复杂的迷宫（称为“格状图”）。为了找到穿过迷宫的最佳路径，他们使用一种称为Viterbi的规则，这就像是一个“选择绝对最佳选项”的按钮。

问题在于：你无法让计算机通过一个“选择最佳”的按钮来学习，因为该按钮在数学上是一个死胡同。一旦按下该按钮，计算机就不知道如何微调其选择以变得更好；它只会瞬间跳转到最近的选项。这就像试图只允许汽车瞬间跳入下一条车道来驾驶，而无法进行平滑转弯。

2. 解决方案：“软”决策（BCJR）

作者 Venugopalan Iyengar 发明了一种新方法，让图书管理员得以练习。他们引入了一个称为**温度（ $T$ ）**的概念，而不是立即强制做出硬性的“选择最佳”决定。

高温（热）：想象图书管理员非常放松且优柔寡断。他们不是只选择一个短语，而是同时考虑许多短语，并为每个短语分配一个“概率”。这就像在雾天，你能看到多条路径，而不仅仅是一条。这种“软”视角在数学上是平滑的，意味着计算机可以轻松计算出如何微调选择以改善整本书。
低温（冷）：随着训练结束，图书管理员变得“冷静”且更加果断。雾气散去，他们瞬间锁定到唯一的最佳短语（即硬决策）。

这种方法称为BCJR，它将“单行道”变成了一条平滑、可行驶的道路。计算机现在可以确切地学习如何调整权重，以使整本书听起来更好，而不仅仅是改善单个句子。

3. “过冲”错误

作者发现了这种“温度”工作机制中一个令人惊讶的怪癖。

在传统物理学中，当你试图寻找最佳解决方案时，通常是从非常热（即非常优柔寡断）的状态开始，以探索所有可能性，然后缓慢冷却。

本文的发现：对于这种特定类型的压缩，起始“过热”是一场灾难。如果图书管理员在开始时过于优柔寡断，他们就会 wander 进入一个糟糕的街区（一个“更差的 Voronoi 盆地”）并被困在那里。等到他们冷却下来并试图寻找最佳路径时，他们已经迷失在糟糕的区域，无法返回。
解决方法：他们发现，从中等温度（不太热，也不太冷）开始效果最好。这就像带着清晰的地图开始徒步，而不是在雾中盲目 wandering。通过跳过“超级热”阶段，他们避免了迷路，并找到了更好的解决方案。

4. 结果：小型计算机的胜利

该团队在两种不同类型的模型上测试了这种方法：

“代理”测试（OLMoE）：他们尝试仅通过观察模型各层重建自身数据的程度来优化压缩（就像检查复印件是否像原件一样）。结果：失败了。模型的表现实际上比旧方法更差。这让他们认识到，仅仅制作“更好的复印件”并不意味着这本书读起来更好。
“真实”测试（Llama-3.2）：他们优化模型以更好地阅读和理解文本（使用一种称为“蒸馏”的方法，即由一个聪明的教师模型指导学生）。结果：成功了！
- 在模型的特定层上，他们的新方法以微小但显著的幅度击败了旧的最佳方法。
- 当他们将此方法应用于多层时，改进以“超加性”的方式累积（整体大于部分之和）。

5. 引擎内部

执行这些数学运算通常非常缓慢，且需要昂贵的超级计算机。作者还构建了一个专门的“引擎”（一个Triton 内核），能够在单个消费级显卡（如 RTX 4080）上极快地运行这些复杂数学运算。他们将其速度提高了6.5 倍，超过了标准方法，证明了这种先进技术无需数据中心即可运行。

总结

本文提出了一种将 AI 模型缩小以适应消费级设备的新方法。

技巧：他们用一种“软”的、可学习的决策过程取代了僵硬的、不可学习的决策过程，并使其逐渐硬化。
洞察：你不应该让学习过程起始得“太热”（太随机）；适度的起始可以防止模型迷失。
成果：当正确训练以理解文本（而不仅仅是复制数据）时，该方法产生的结果优于之前的最先进技巧，并且能在普通游戏 PC 上快速运行。

作者已发布其代码和训练好的模型，以便其他人尝试使用。

技术摘要：BCJR-QAT

问题陈述

在消费级硬件上为大型语言模型（LLM）实现每权重 2 比特（bpw）或更低的量化是一个关键瓶颈。虽然像 QTIP 这样的训练后量化（PTQ）方法利用网格编码量化（TCQ）为 2 比特量化建立了强大的前沿，但它们受限于"PTQ 天花板”。要突破这一限制，需要量化感知训练（QAT）。然而，标准的 QAT 技术在网格码上失效，因为维特比（Viterbi）解码算法涉及对指数级路径进行不可微的 argmax 操作。对网格应用直通估计器（STE）会破坏相邻权重决策之间的耦合，而这正是网格编码相对于标量量化的主要优势。

方法论：BCJR-QAT

本文介绍了BCJR-QAT，这是一种网格编码量化的可微松弛方法，它用基于BCJR（Bahl-Cocke-Jelinek-Raviv）前向 - 后向算法的软期望取代了硬维特比 argmax。

核心机制

统计力学类比：将网格视为等同于 1D 伊辛（Ising）类自旋链的因子图。维特比 argmax 被视为路径上玻尔兹曼分布在零温极限（ $T \to 0$ ）下的结果。
软码字：在有限温度 $T > 0$ 下，该方法利用 BCJR 算法计算每个状态的精确边缘后验概率。量化输出是发射码字的玻尔兹曼期望：
$\hat{w}_t(T) = \langle c(s_t) \rangle_{p_T} = \sum_s c(s) p_T(s_t = s | w)$
该软码字关于权重原像和网格参数是完全可微的。
退火：在训练过程中，温度 $T$ 从初始值 $T_0$ 退火至接近零的 $T_{end}$ 。随着 $T \to 0$ ，软码字收敛至硬维特比输出，有效地将连续优化“结晶”为离散解。

关键技术革新

可微性：BCJR 前向 - 后向递归（复杂度 $O(L \cdot S^2)$ ）提供了通过网格的精确梯度，保留了标量 STE 方法所丢失的状态间耦合。
Triton 内核：作者实现了一个融合的 Triton 内核，将 gather、log-sum-exp 和加法操作合并为单次启动。这将前向/后向传递从数百次内核启动减少为一次，在保持 FP32 数值一致性的同时，相比参考的 PyTorch 自动梯度实现实现了6.57 倍的加速。
调度修正（Skip-High-T）：本文诊断了标准模拟退火调度（从 $T_0=1.0$ 开始）的一个关键失效模式。在高温下，软码字梯度变得“模糊”且无信息量，导致优化器在冷却阶段能够恢复之前，就漂移到了更差的 Voronoi 盆地。提出的修复方案是跳过高温阶段，在梯度已有信息量的较低温度下初始化（例如 $T_0=0.3$ ）。

实验结果

本文在两种架构上评估了 BCJR-QAT：OLMoE-1B-7B（MoE）和Llama-3.2-1B（Dense）。

1. OLMoE-1B-7B（负面结果：代理差距）

设置：每层贪婪 QAT，优化重建均方误差（MSE）（一种代理目标）。
结果：与 QTIP-PTQ 相比，BCJR-QAT 成功将每层重建 MSE 降低了 17–20%。然而，这种改进并未转化为最终任务性能。与 QTIP-PTQ 基线相比，模型在 WikiText-2 上的困惑度（PPL）退化了约 1.3。
结论：这证实了存在“代理差距”，即在 2 bpw 下，最小化局部重建误差并不能保证更好的最终任务性能。

2. Llama-3.2-1B（正面结果：端到端 KL）

设置：单层 QAT，针对 FP16 教师模型优化端到端前向 KL 蒸馏。
结果：
- 调度过冲：使用朴素的高温调度（ $T_0=1.0$ ）导致性能略差于 PTQ（PPL 增加 +0.005 至 +0.022）。
- Skip-High-T 修复：使用修正后的调度（ $T_0=0.3$ ）消除了过冲。在第 4 层，模型实现了10.1347 PPL，比 QTIP-PTQ 基线（10.2189）改善了**−0.084 PPL**。
- 超可加性：当组合训练过的层（第 4 层和第 8 层）时，联合增益（−0.077 PPL）超过了单独增益之和（−0.062 PPL），展示了 BCJR 训练出的码字相互增强的“合作盈余”。

3. 理论分析

漂移预算理论：作者推导了逃离 PTQ Voronoi 盆地的条件：累积权重漂移（ $\eta \cdot N \cdot g_{max}$ ）必须超过盆地半径（ $r_{Voronoi} \approx 10^{-3}$ ）。实验证实，不满足此条件的调度无法逃离 PTQ 解，而超过此条件的调度可以找到更好的解，前提是目标和调度正确。

意义与主张

本文声称有三个主要贡献：

算法层面：BCJR-QAT 是首个使网格编码量化原生可微的方法，实现了在不损失路径耦合优势的情况下对 TCQ 进行 QAT。
工程层面：融合的 Triton 内核使其在消费级硬件（单张 12GB GPU）上计算可行，消除了研究领域的推理部署障碍。
实证洞察：
- 它指出每层重建 MSE 是 2 比特 QAT 的糟糕代理，而端到端 KL 蒸馏是有效的。
- 它挑战了离散松弛模拟退火的传统智慧，表明对于 BCJR-QAT，从“热”开始（ $T=1.0$ ）是有害的，因为高温梯度以噪声为主导。最佳策略是在梯度已有信息量的温度下开始。

本文结论指出，虽然单层结果令人鼓舞，但 BCJR-QAT 的全部潜力可能在于跨所有层的联合端到端训练，这项任务因发布的内核而变得经济可行，但对于参数超过 10 亿（1B）的模型，仍需要云规模硬件（H100/H200）。作者明确表示，本文未执行这种全规模的联合训练。