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

大词表使输出头成为小型语言模型推理成本的重要组成部分。我们提出 Softmax 重参数化，这是一种训练后方法，可在量化前选择一个功能等价的输出头。该方法从每个输出行中减去词表行均值的一个标量倍数，并分别针对 RTN、激活加权 MSE 和全 Hessian GPTQ 通过验证 KL 选择系数。这个一维搜索包含原始头和固定均值中心化，保留全精度 Softmax 分布，并使训练好的解码器保持不变；秩一校正可处理非线性 logit 路径，例如 soft-capping。在七个头上，W4 的收益集中在基线量化显著扭曲预测的地方：在 Phi-4-mini 上，AW-MSE KL 从 0.936 降至 0.256。这些收益在更强的 GPTQ 校准下仍然存在，并且与精确的逐通道缩放和仿射量化保持互补。在四个头和三种 W4 量化器上，使用 WikiText 选择并冻结的系数也能迁移到 C4 和 OpenWebMath，在冻结系数不同于 $1$ 的全部 18 次比较中均优于均值中心化，并在其余 6 次中与之相当。在 W2 这一压缩压力测试设置下，收益扩展到几乎整个模型--量化器矩阵。匹配残差分析表明，保真度的提升可以伴随更大的 logit 重构误差，同时降低残差的 Fisher 加权成本。对于与平移兼容的头，重参数化不增加推理操作，并保留打包 W4 执行：在解码器保持 BF16 的情况下，量化 Phi 输出头相对于 BF16 头基线将 batch-one 生成延迟降低了 10.8%。