用轻量机制提升模型推理效率,不增加测试开销。
Efficient Reasoning at Fixed Test-Time Cost via Length-Aware Attention Priors and Gain-Aware Training
- 引入长度感知注意力先验,训练时加偏置引导注意力
- 验证提升时才激活控制模块,推理时无额外开销
- 适合资源受限场景,尤其长文本噪声环境
我们研究在严格计算约束下的高效推理。目标是在不增加测试时间成本的前提下实现结构化、正确的决策。通过仅在训练阶段引入两个组件,适用于小型和中型Transformer,并可推广至其他可微优化器。首先,基于模糊制度位置对齐(RPA)构建长度感知注意力先验,生成归一化的预softmax偏置,以结构化正则化方式引导注意力,且不增加推理参数。其次,设计一个极简的增益感知控制器Guardian,仅在验证性能提升时调整注意力锐度,遵循非凸优化的双时间尺度策略梯度视角,推理时禁用。从KL视角看,softmax(z) + log π 可视为带KL正则化的最大后验估计,为先验提供理论基础。在WikiText 2上,保持基线延迟与内存不变的情况下,降低验证交叉熵。推理时,仅需添加预先计算并缓存的偏置矩阵B,每头增加一次加法操作。实际中该开销可忽略,无显著中位数延迟变化。结果表明,长度感知先验与后期增益控制能有效保留稀缺改进,尤其在长跨度、高噪声的logit环境下,同时维持测试时成本基本不变。
原文摘要 · Abstract (English)
We study efficient reasoning under tight compute. We ask how to make structured, correct decisions without increasing test time cost. We add two training only components to small and medium Transformers that also transfer to broader differentiable optimizers. First, a length aware attention prior built via fuzzy regime position alignment, RPA, yields a normalized pre softmax bias that guides attention like a structured regularizer while adding no new inference parameters. Second, a minimal gain aware controller, Guardian, nudges attention sharpness only when validation improvements warrant it, following a two timescale policy gradient view of nonconvex optimization. It is disabled at inference. A KL perspective shows softmax of z plus log pi as MAP with KL regularization, grounding the prior in a principled objective. Under strict compute parity on WikiText 2, we reduce validation cross entropy while matching baseline latency and memory. At inference, we add a precomputed, cached prior B of T as a single additive bias per head. The controller does not run. In practice, this incurs negligible overhead, a cached bias add per head, with no measurable p50 latency shift. Our results suggest that length aware priors and late phase gain control preserve scarce improvements, especially in long span, noisy logit regimes, while keeping test time costs effectively unchanged.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。