Softmax 重参数化用于输出头量化

TL;DR 针对大词表输出头量化导致的预测扭曲，提出 softmax 重参数化：量化前减去词表行均值的标量倍数，一维搜索选系数最小化 KL 散度，使 Phi-4-mini 的 W4 量化 KL 从 0.936 降至 0.256，且不增加推理开销。

问题

问题背景

当前小型语言模型推理优化重点正从线性层转向输出头，因为大词汇表使 output head 权重矩阵成为推理延迟和内存带宽的主要瓶颈，尤其在 batch-one 生成场景。

现有方法局限

传统后训练量化（RTN、GPTQ、AW-MSE）直接对输出头权重做低比特近似，但 softmax 对 logit 误差高度敏感：量化引入的逐行误差会放大 top-k 预测顺序扰动，导致 KL 散度飙升（如 Phi-4-mini 上 W4 AW-MSE KL 达 0.936）。已有的 mean-centering 利用 softmax 平移不变性，但仅固定减去所有行均值，没有考虑不同量化器误差结构的差异，无法自适应选择最优等价变换。此外，现有量化方案通常需要额外推理操作（如 per-channel scaling 的反量化乘法），破坏了 packed W4 推理的部署效率。

为什么这个问题难/重要

难点在于输出头参数的等效变换空间很大，但需要在保持全精度 softmax 分布不变的前提下，找到对量化误差最鲁棒的表示。量化误差与 logits 的 Fisher 信息加权相关，单纯最小化重构误差（MSE）并不等价于最小化分布偏移。业界对低延迟、低内存的小语言模型部署需求强烈，输出头量化若失真会直接损害生成质量，而高精度保留输出头又抵消了推理加速收益。

行业类比

类似在设备端语音识别或输入法预测中量化大词汇表 softmax 层：既要压缩权重减少内存带宽，又要确保 top-k 候选不因量化噪声而错位。

核心洞察

• 利用 softmax 对 logits 平移的不变性，通过重参数化选择更利于量化的 logits 表示，而不改变模型输出分布。与已有工作要么固定均值中心化、要么只优化量化器（如 GPTQ 重建）不同，该方法把参数空间的等价变换作为量化前预处理，用一维验证 KL 搜索最佳系数，且全精度 softmax 分布保持不变。对 shift-compatible 输出头，无需额外推理操作，可直接嵌入现有量化流程，为输出头量化提供了新的自由度。
• 该方法挑战了以重建误差最小化作为量化质量代理的假设。实验表明，softmax 重参数化后，即使 logits 重建误差增大，Fisher 加权代价降低，预测 KL 失真反而改善。这说明对输出头，直接优化预测分布（KL）比优化 logits 重建（MSE 或 Hessian 加权 MSE）更有效，因为 softmax 对某些方向的误差不敏感。这一发现提示输出头量化应重新考虑优化目标，以分布失真的直接度量为导向设计量化方法。

方法

方法详解

Softmax Reparameterization 是一种输出头量化前的后训练重参数化方法，旨在通过调整输出头权重分布来降低量化误差。

输入：预训练语言模型的输出头权重矩阵 ( W \in \mathbb{R}^{V \times D} )（( V ) 为词表大小，( D ) 为隐藏维度），以及一个验证集用于评估量化后的输出分布。

关键模块：

1. 行均值计算：对权重矩阵 ( W ) 的每一行（对应一个词汇）计算均值 ( \mu_i = \frac{1}{D} \sum_j W_{ij} )。
2. 重参数化：引入标量系数 ( c )，构造新权重 ( W' = W - c \cdot \mu_i \cdot \mathbf{1}^T )，即对每一行减去该行均值乘以 ( c ) 的常数。该操作改变了权重的数值分布，但保持模型结构不变。
3. 量化与评估：使用目标量化器（如 RTN、AW-MSE、GPTQ）对 ( W' ) 进行量化，并在验证集上计算量化模型与原始浮点模型的 KL 散度（衡量 softmax 输出分布的差异）。
4. 一维系数搜索：在实数范围内（通常为 ( c \in [0, 2] ) 或更广）搜索最优 ( c )，最小化 KL 散度。搜索空间天然包含两个特例：( c=0 )（原始头）和 ( c=1 )（固定均值中心化）。

非线性 logit 路径处理：对于带有 soft-capping 等非线性变换的模型，直接减去行均值会破坏等价性，因此需要额外的 rank-one 校正 来保持原始 logits 的功能等效。

输出：在最优 ( c ) 下量化后的输出头权重，该权重在推理时与原始浮点输出头在功能上等效（量化后分布接近），且完全不影响解码器部分。

与同类方法的差异点：与固定均值中心化（( c=1 )）或直接量化相比，该方法通过数据驱动的方式选择系数 ( c )，能够自适应不同量化器的误差特性，从而在更低位宽下保持更优的预测保真度。

实验

实验设计

论文在 7 个输出头 上评估 softmax reparameterization,覆盖 RTN / AW-MSE / GPTQ 三种量化器,通过验证集 KL 散度 一维搜索标量系数,与固定 均值中心化 对比。跨域迁移测试使用 WikiText 选出的系数,评估在 C4 与 OpenWebMath 上的表现;另设 W2 压力测试与 Phi 输出头的打包推理延迟测量。

关键发现

• 在 Phi-4-mini 上,AW-MSE KL 从 0.936 降至 0.256,收益集中于基线量化严重扭曲预测的头部。
• 跨 4 个头 与 3 种 W4 量化器,冻结的 WikiText 系数在 18/24 个比较中优于均值中心化,其余持平。
• W2 压力测试下收益扩展到几乎整个模型-量化器矩阵。
• 对 shift-compatible 头,不增加推理操作;Phi 输出头量化使 batch-one 生成延迟降低 10.8%(解码器保持 BF16)。

与基线对比的深度解读

Softmax reparameterization 通过一维标量搜索保持全精度 softmax 分布,与固定均值中心化不同,能自适应选择最优偏移;秩一修正解决非线性 logit 路径(如 soft-capping),弥补了仅中心化的局限。匹配残差分析显示,方法可容忍更大 logit 重构误差却降低 Fisher 加权成本,说明其优化目标与下游保真度更一致。工程上,该法不增加推理开销,兼容打包 W4 执行,且与精确 per-channel scaling 和 affine quantization 互补,为小语言模型输出头量化提供了简洁有效的后训练策略。

行业影响

落地场景

Softmax reparameterization 面向大词汇量语言模型的输出头量化，尤其适合端侧部署与实时推理场景。典型应用包括：

• 多语言翻译服务：词汇表常达数十万，输出头占据主要推理成本。例如在 Phi-4-mini 上，AW-MSE 量化后 KL 散度从 0.936 降至 0.256，显著减少量化引起的预测失真。
• 电商搜索与推荐系统：基于小语言模型的 query 理解与商品描述生成模型，可借助该方法在边缘设备上保持低延迟与高保真度。
• 内容平台辅助创作：轻量级生成模型嵌入移动端应用时，输出头量化可降低内存占用，同时避免生成质量下降。

商业价值

该方法直接降低推理成本：论文报告在 BF16 解码器保持不变时，量化 Phi 输出头使 batch-one 生成延迟降低 10.8%。对于按 token 计费或大规模推理服务，等效于降低硬件需求和能耗。

同时，由于重新参数化保持全精度 softmax 分布不变，量化后的模型保真度更高，用户体验不会因量化而明显降级，减少因模型质量下降导致的用户流失或业务风险。

与现有产品/工作流的接口

作为后训练方法，可以无缝插入现有量化管线：

1. 在应用 RTN、AW-MSE 或 GPTQ 等量化器之前，对输出头权重执行一维系数搜索，得到最优移位。
2. 对于 shift-compatible 的 logit 路径，量化后的推理无需额外操作，可直接使用打包的 W4 kernel，无需修改推理代码或导出格式。

因此，可以将其封装为 Hugging Face Transformers、ONNX Runtime 或 TensorRT 量化工具的可选 pass，或作为模型导出前的自动化预处理步骤。对非线性 logit 路径（如 soft-capping），只需额外存储一个 rank-one 校正向量，接口改动很小。

局限

• **适用场景受限于 softmax 平移不变性**。该方法依赖输出头 logits 的全局平移不影响 softmax 概率分布这一性质。对于引入非线性 logit 路径的模型（如 soft-capping、温度缩放或其他归一化层），需要额外的**秩一校正**，但并非所有非线性变换都能精确保持等价。论文只验证了 soft-capping 这一种情况，对于更复杂的后处理（如 top-k/top-p 截断在推理时的动态行为）未作讨论。此外，rank-one correction 会引入额外存储和计算，可能部分抵消推理延迟收益，尤其在小批量或单 token 生成场景下。
• **系数选择需要额外校准数据与调参步骤**。该方法通过验证集 KL 散度搜索最优标量系数，这引入了新的超参数和校准流程。虽然论文展示了 WikiText 选定系数在 C4 和 OpenWebMath 上的跨域迁移能力，但在更极端的分布偏移（如领域外数据、多语言、代码生成）下系数稳定性缺乏验证。此外，每次模型更新或部署到新任务时都可能需要重新搜索系数，增加了工程复杂度和持续集成成本。
• **仅改善输出头量化误差，无法覆盖其他层**。论文聚焦输出头这一特定组件，而实际端到端推理中，注意力层、前馈层等同样存在量化误差积累。在 W4 量化下，如果其他层的量化误差占主导，输出头的改善可能被掩盖；论文在 W2 压力测试中观察到收益几乎遍及整个模型-量化器矩阵，但实际部署很少使用 W2 精度。此外，方法属于后训练量化范畴，不适用于量化感知训练流程，限制了在训练阶段就优化权重分布的可能性。