Initialisation Determines the Basin: Efficient Codebook Optimisation for Extreme LLM Quantization

1. 研究背景与核心问题 (Problem)

背景：
随着大语言模型（LLM）在边缘设备（如移动端、消费级 GPU）上的部署需求增加，模型权重压缩变得至关重要。虽然 4-bit 量化已接近无损，但在 2-bit 精度（每个参数仅 4 种可能值）的极端压缩场景下，传统的量化方法往往表现不佳。

核心问题：

• 灾难性失败： 在 2-bit 精度下，即使使用广泛的搜索（Beam Search）和微调（PV-tuning），基于自由形式加性量化（Free-form Additive Quantization，如 AQLM）的方法仍经常发生灾难性的性能下降（困惑度 Perplexity 飙升）。
• 错误的归因： 业界通常认为性能瓶颈在于搜索空间不足或微调不够，因此倾向于增加 Beam 宽度或训练轮数。
• 真正的瓶颈： 本文指出，主导瓶颈是 码本初始化（Codebook Initialisation）。在极端压缩下，贪婪的串行初始化（Greedy Sequential Initialisation）会将模型置于一个糟糕的优化盆地（Optimisation Basin）中，后续的搜索和微调难以跳出该盆地。

关键指标：表示比率 (Representational Ratio, $\rho$)
作者定义了一个关键指标 $\rho = N / (KM)$，其中 $N$ 是权重组数量，$K$ 是码本条目数（通常 256），$M$ 是加性码本数量。

• $\rho < 1$ (过完备)： 可表示的点多于权重组，初始化误差可被吸收。
• $\rho > 1$ (欠完备)： 权重组竞争有限的码本容量。在 2-bit 场景下（如 Llama 3.2 3B），$\rho \approx 18$，此时初始化至关重要。

---

2. 方法论 (Methodology)

为了解决初始化问题，作者提出了 OA-EM (Output-Aware Expectation-Maximisation) 算法。

2.1 理论分析：贪婪初始化的缺陷

• 贪婪串行拟合的局限性： 传统方法先拟合第一个码本 $C_1$，再拟合残差 $C_2$。这种方法忽略了码本间的联合结构。
• 误差耦合： 当 $\rho \gg 1$ 时，第二个码本不仅要表示真实的权重结构，还要补偿第一个码本的错误。由于 2-bit 下剩余码本组合数极少（$K=256$），贪婪选择导致的初始位移（Displacement）无法被后续码本修正，导致误差累积。

2.2 OA-EM 算法设计

OA-EM 旨在通过迭代优化，直接最小化输出重建误差，而非简单的权重空间距离。

• 核心机制： 在 AQLM 的残差框架内，使用 Hessian 加权的马氏距离 (Hessian-weighted Mahalanobis Distance) 替代欧氏距离。
• 步骤：
  1. M-step (质心优化)： 固定分配，优化码本质心以最小化 Hessian 加权的重建误差。
     $$L_{EM} = \frac{1}{N} \sum e_i^\top H_i e_i$$
     其中 $H_i$ 是基于校准激活计算的阻尼块对角 Hessian 近似（类似 GPTQ 中的二阶信息），$e_i$ 是重建误差。
  2. E-step (硬重分配)： 固定质心，根据马氏距离将权重组重新分配给最近的质心。
• 优势： Hessian 加权使得算法能够识别对输出敏感度高的权重组，优先为这些组分配码本资源，从而减少关键位置的初始位移 $\delta$。

---

3. 主要贡献 (Key Contributions)

1. 提出表示比率 $\rho$： 证明了 $\rho$ 是预测量化对初始化敏感度的关键指标。当 $\rho > 1$ 时，初始化质量成为决定性因素。
2. 提出 OA-EM 算法： 一种即插即用的初始化方法，利用输出感知（Output-Aware）的 EM 算法和 Hessian 加权距离，显著改善了码本几何结构。
3. 揭示“盆地持久性” (Basin Persistence)： 通过实验证明，初始化的差异会导致模型收敛到不同的优化盆地。即使经过强大的端到端微调（PV-tuning），由 OA-EM 初始化的模型依然优于贪婪初始化的模型，且这种优势在多种架构和压缩率下均存在。
4. 帕累托优势 (Pareto Dominance)： OA-EM 在相同的计算预算下实现了更低的困惑度，或者在更短的时间内达到相同甚至更好的性能。

---

4. 实验结果 (Results)

实验在 Llama 3.2 3B, Llama 3.1 8B, Qwen 2.5 3B 三个模型上进行，主要关注 2-bit 和 3-bit 压缩。

4.1 预微调 (Pre-PV) 表现

• 2-bit (Llama 3.2 3B, $\rho \approx 18$)：
  • 贪婪初始化 + Beam=4：WikiText-2 困惑度高达 352.39（灾难性失败）。
  • 贪婪初始化 + Beam=16（增加 2.8 倍时间）：困惑度降至 46.01，仍远差于 FP16 (7.28)。
  • OA-EM + Beam=4： 困惑度仅为 16.82，且耗时更短。
  • 结论： 增加搜索预算无法弥补糟糕的初始化，OA-EM 在低 Beam 宽度下即取得优异效果。

4.2 后微调 (Post-PV) 表现

• 盆地持久性： 经过 PV-tuning 后，所有模型的困惑度都大幅下降，但 OA-EM 始终优于贪婪初始化。
  • 例如 Llama 3.2 3B (2-bit)：贪婪 (Beam=8) 最终 PPL 为 11.76，而 OA-EM (Beam=8) 为 11.53。
  • 即使在贪婪初始化表现最好的配置下，OA-EM 依然能进一步降低 PPL。
• 非对称缩放： 增加 Beam 宽度对贪婪基线有时甚至有害（如从 8 增加到 16，PPL 从 11.76 升至 12.01），而对 OA-EM 则保持稳定或微升。这表明两者处于不同的优化轨迹。
• 下游任务： 在 ARC, HellaSwag 等零样本任务中，OA-EM 在大多数配置下取得了更高的准确率。

4.3 跨架构与压缩率

• 3-bit ($\rho \approx 0.07$)： 虽然瓶颈较轻，OA-EM 依然能带来 0.12 PPL 的提升，且优势在微调后依然存在。
• 不同模型： 在 Llama 8B 和 Qwen 3B 上均观察到类似趋势，尽管 Qwen 由于权重分布较平滑，差距较小，但 OA-EM 依然保持领先。

---

5. 意义与启示 (Significance)

1. 重新定义优化瓶颈： 在极端量化（特别是 2-bit）中，初始化质量比搜索强度更重要。盲目增加 Beam 宽度或微调轮数可能无法解决根本问题，甚至造成资源浪费。
2. 优化几何的重要性： 研究揭示了压缩模型空间中的优化几何特性。初始化的选择决定了模型进入哪个“优化盆地”，而后续的微调往往只能在当前盆地内寻找局部最优，难以跨越到全局最优。
3. 边缘部署的实用性：
   • OA-EM 允许在极低的计算成本（如 Beam=4）下获得高质量模型。
   • 对于 ARM CPU、微控制器等资源受限设备，OA-EM 能在不增加推理延迟（保持 O(1) 查找表解码）的前提下，显著提升模型可用性。
4. 通用性框架： 提出的 $\rho$ 指标和盆地持久性现象不仅适用于 AQLM，也为未来结合结构化与自由形式量化的混合方法提供了理论指导。

总结： 该论文证明了在极端 LLM 量化中，通过改进初始化策略（使用 OA-EM），可以突破贪婪初始化带来的性能天花板，以极低的额外计算成本实现显著的模型质量提升，是边缘设备部署大模型的关键技术突破。