通过长度感知注意力先验和收益感知训练实现固定测试成本下的高效推理

我们研究在紧张计算资源下的高效推理。我们探讨如何在不增加测试时间成本的情况下进行结构化、正确的决策。我们引入两个仅用于训练的组件到小型和中型 Transformer 模型中，这些组件也可迁移到更广泛的可微优化器。首先，构建一种长度感知注意力先验，通过模糊场景位置对齐实现，这产生一个归一化的预 softmax 偏置，像结构化正则器一样引导注意力，同时不添加新的推理参数。其次，一个最小化收益感知控制器，名为 Guardian，仅在验证改进值得时轻微调整注意力锐度，遵循非凸优化的双时间尺度策略梯度视角。在推理时禁用该控制器。KL 视角表明，softmax(z) 加上 log π 作为带 KL 正则化的 MAP 估计，将该先验置于原则目标中。我们在 WikiText 2 上严格保持计算平衡的情况下，减少了验证交叉熵，同时匹配基线的延迟和内存。推理时，我们添加一个预计算、缓存的偏置 B 作为每个注意力头的单个加法偏置。控制器不运行。实际上，这只是在每个注意力头上添加一个缓存偏置，几乎没有可测量的 p50 延迟变化。我们的结果表明，长度感知先验和后期阶段收益控制在保持测试时间成本基本不变的同时，尤其是在长序列、嘈杂的逻辑分布中，能保留稀缺的改进。