大词表会显著增加小语言模型的输出头推理成本。我们提出一种后训练方法——Softmax 重参数化,在量化前搜索功能等价的输出头。该方法对每一行输出减去词表行均值的标量倍数,并通过验证集 KL 散度选取系数。对于线性‑softmax 头,这种平移在全精度下保持预测不变,无需重新训练解码器;对非线性 logits 路径,可通过秩一校正扩展该构造。实验在七个输出头和三种量化器上进行,W4 量化的收益在基线量化导致预测严重失真的情况下最大:在 XGLM 上使用 RTN 时测试 KL 降低 93%,在 Phi、BLOOM、BLOOMZ 上使用激活加权 MSE 时降低 73%‑77%。基线误差本就低的头部变化不大;在 W2 作为压缩压力测试时,收益更为广泛。Phi 上的收益在更强的 GPTQ 校准下仍然存在,独立的未触及 holdout 数据同样复现了 Phi 与 BLOOM 的改进。冻结的 WikiText‑选取系数在 C4 与 OpenWebMath 上无需重新调优即可迁移。对 Phi 的残差分析表明,尽管整体 logits 误差增大,所选代表能够降低高概率输出的误差并减少其 Fisher 加权代价。对于兼容平移的头部,平移本身不增加推理运算。将解码器保持在 BF16 时,打包的 W4 Phi 输出头将单批次生成延迟降低 10.8%,且重参数化保持了该加速。
点评