初始化决定优化盆地：面向极低比特大语言模型量化的高效码本优化方法 Initialisation Determines the Basin: Efficient Codebook Optimisation for Extreme LLM Quantization

OA-EM 的完整流程嵌入 AQLM 的层内量化循环，对每一层独立执行。步骤 1 是输入准备：给定当前层权重矩阵 $W \in \mathbb{R}^{d_\text{out} \times d_\text{in}}$、校准激活 $X$ 和分组大小 $g=8$，把 $W$ 切分为 $N=d_\text{out} d_\text{in}/g$ 个权重组 $w_i \in \mathbb{R}^g$，并预算 Hessian $\mathcal{H}_i = X_i^\top X_i + \lambda I$（阻尼 $\lambda = 0.01 \cdot \text{diag}(\mathcal{H})$）。步骤 2 是 k-means 热启动：用欧氏残差 k-means 得到初始码本 $\{c_{m,k}^0\}$ 与初始赋值 $\{b_i^0\}$，这与标准 AQLM 完全一致。步骤 3 是 OA-EM 主体循环：对 $r=1,\dots,R=3$ 轮，每轮先做 M 步——固定赋值，把每本码字的每个码字 $c_{m,k}$ 用 Adam 在损失 $\mathcal{L}_{EM}=\frac{1}{N}\sum_i e_i^\top \mathcal{H}_i e_i$（其中 $e_i = w_i - \sum_m c_{m,b_{i,m}}$）上做 $S=100$ 步更新，余弦退火 $\eta \to 0.1\eta$；再做 E 步——固定码字，对每个权重组 $w_i$ 按 Mahalanobis 距离 $\arg\min_k (w_i - c_{m,k})^\top \mathcal{H}_i (w_i - c_{m,k})$ 重新选码字 $b_{i,m}$，保持顺序拟合结构。步骤 4 是下游接力：把 OA-EM 输出的码本与赋值交给 AQLM 的束搜索（$b \in \{4,8,16\}$, $e \in \{5,100\}$）做组合优化，再交给 PV-Tuning 做端到端微调（Adam, $\eta_{ft}=3\times10^{-4}$, batch 32, 10K 样本, 5 epoch）。整个过程不修改推理阶段的 LUT 反量化路径，因而对 ARM CPU、单片机等边缘硬件完全透明。

本文的技术新颖性体现在三个层面。理论层面，Proposition 1 给出贪心次优性的显式误差界 $\varepsilon_g - \varepsilon^* = \underbrace{\|\delta\|^2}_{\text{直接代价}} + \underbrace{2\langle \delta, c_{2,j_g} - r^*\rangle}_{\text{耦合项}} + \underbrace{\|r^* - c_{2,j_g}\|^2 - \varepsilon^*}_{\text{残差不匹配} \ge 0}$，把「贪心顺序分配为什么会失败」从经验观察转化为可分解、可量化的几何解释，并由此自然导出表征比率 $\rho$ 的判别准则。算法层面，OA-EM 是首个把 Hessian 加权 Mahalanobis 距离同时引入码本质心更新（Adam 优化）和码字赋值（E 步）的自由形式加性量化初始化器，区别于 GPTVQ 的「单码本 EM」和 LSQ++ 的「信息检索场景 AQ」。实证层面，作者系统性地把「pre-PV 困惑度差距」「post-PV 困惑度差距」「束宽响应曲线」「领域距离梯度」四类证据拼成一幅完整的盆地持久性图谱，论证了「PV-Tuning 在盆地内精修而非跨盆地跳跃」这一反直觉结论。

独立审视论文可识别出几个值得改进的薄弱环节。第一，OA-EM 的超参（$R=3$ 轮、$S=100$ 步 Adam、$\eta=10^{-4}$、$\lambda=0.01\cdot\text{diag}(\mathcal{H})$）源自经验，作者未做消融实验也未说明是否对所有模型一致——尤其是在 8B 模型上 Hessian 块更大、激活校准仅 128 条序列的情况下，是否需要调整 $R$、$S$ 仍未知，改进方向是补充系统性的超参敏感性分析。第二，Hessian 近似 $\mathcal{H}_i = X_i^\top X_i + \lambda I$ 用的是块对角形式，忽略了权重组之间的耦合，而 AQLM 本身的 beam search 正是在补偿这种耦合，理论上更精确的全 Hessian 近似可能进一步缩小初始化差距，改进方向是引入 Fisher 信息矩阵或低秩 Hessian 近似并比较代价收益。第三，论文对 Qwen 2.5 3B 上「下游准确率 OA-EM 微负 0.3pp」的解释仅停留在「权重分布更平滑」的口语化论述，缺乏直方图、谱分析或 outlier 比率等定量证据，改进方向是补充对不同架构权重分布的统计刻画。第四，OA-EM 仍沿用「先 $C_1$ 再 $C_2$」的顺序拟合结构，理论上若能用更复杂的联合 EM（如同时优化 $C_1, C_2$）打破顺序假设，可能进一步提升质量，改进方向是探索多码本联合 OA-EM。

作者在 Limitations 与 Conclusion 中明确点出两个延伸方向：一是「将 $\rho$ 框架与盆地持久性分析迁移到未来结合结构化与自由形式的混合量化器」——例如先用 GLVQ 风格的格码本初始化，再用 OA-EM 精调，或在 QuIP# 的 Hadamard 不相干处理之后接 OA-EM；二是把 OA-EM 推广到 1-bit 或亚 2-bit 极限制压缩（如 1.5 bpp 或三元量化），因为更低的位宽会进一步推高 $\rho$，理论上盆地效应会更强。基于论文成果还可进一步延伸：(1) 把 Hessian 信息融入 PV-Tuning 自身，使微调阶段也具备「输出感知」能力，与初始化阶段的 OA-EM 形成端到端一致性；(2) 探索 OA-EM 在多模态大模型（如 LLaVA、Qwen-VL）上的效果，因为视觉适配器权重分布通常更稀疏、更不平滑，可能为 $\rho$ 框架提供新的边界案例；(3) 把表征比率 $\rho$ 拓展为「面向任务」的版本——例如按各权重组的下游损失贡献加权得到 $\rho_{\text{task}}$，用于指导特定任务（数学、代码）的量化策略。

论文公开了代码仓库（论文脚注标注的 GitHub 链接 https://github.com/kenno94-IK/aqlm-oaem），且在 AQLM 框架内做最小改动即可替换初始化器，复现门槛较低。实验所用校准数据是 C4 数据集的 128 条长度 4096 的序列（与 GPTQ、AWQ 一致），评价使用 LM Evaluation Harness 框架，所有任务的标准做法有据可查。算力方面，Llama 3.2 3B 与 Qwen 2.5 3B 在单卡 A100 80GB 上完成，Llama 3.1 8B 的 PV-Tuning 需单卡 B200 192GB；整体一轮主实验大约 6–17 小时，普通研究组可承受。超参选择（$R=3, S=100, \eta=10^{-4}, \lambda=0.01\cdot\text{diag}(\mathcal{H})$）作者有明确描述；唯一的不确定项是 Adam 内部的具体实现细节（如 betas、eps、weight decay），以及并行化的具体策略——这些通常需要直接读代码补齐。总体而言复现难度为「中等」，瓶颈不在算力而在把 OA-EM 的 EM 主循环正确接入 AQLM 的层内流水线。