Efficient Reasoning at Fixed Test-Time Cost via Length-Aware Attention Priors and Gain-Aware Training
Not on the map. Why?
Abstract (arXiv)
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.
Key facts. No citing papers in the map yet (data through 14 Oct 2026). Builds on 9 of its 25 references that are in the map. First posted on arXiv 10 Mar 2026.
Builds on
9 of 25 references.
References · 9
- 2017Attention Is All You Needcited by 60,361
- 2016SGDR: Stochastic Gradient Descent with Warm Restartscited by 5,045
- 2017On Calibration of Modern Neural Networkscited by 4,960
- 2021RoFormer: Enhanced Transformer with Rotary Position Embeddingcited by 4,703
- 2017Outrageously Large Neural Networks: The Sparsely-Gated Mixture-of-Experts Layercited by 3,509
- 2013Sinkhorn Distances: Lightspeed Computation of Optimal Transportation Distancescited by 2,316
- 2020GShard: Scaling Giant Models with Conditional Computation and Automatic Shardingcited by 1,631
- 2018Self-Attention with Relative Position Representationscited by 1,120
- 2021Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolationcited by 1,103
Built on by
Not cited yet.
Similar, not cited
Close in wording, no citation either way
- 2026When2Think: Learning Difficulty-Aware Length Control for Efficient Hybrid Reasoning Models
- 2025PERK: Long-Context Reasoning as Test-Time Learning
- 2026On-Policy Supervised Fine-Tuning for Efficient Reasoning
- 2026Efficient LLM Reasoning via Variational Posterior Guidance with Efficiency Awareness
- 2025Virtual Parameter Sharpening: Dynamic Low-Rank Perturbations for Inference-Time Reasoning Enhancement
- 2026SageBwd: A Trainable Low-bit Attention
- 2026LongFlow: Efficient KV Cache Compression for Reasoning Models
- 2025MiniMax-M1: Scaling Test-Time Compute Efficiently with Lightning Attention