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.
Figures & tables
Hardware
Batch × Seq
Steps/epoch
Throughput
Epochs
Total steps
Wall-clock
GH200 (single)
48×512
73
32.5 it/s
110
8030
1:11:46
Table 1: Compute disclosure for WT2 (raw-v1, GPT-2 BPE).
Run
Extras
Context mix
s/it (mean)
p50
p95
n
all_but_no_swaselect
align, guardian, cg
512:1.0
25.90
25.71
26.27
119
full_run
align, guardian, swa , cg
768:1.0
32.79
32.62
33.06
110
hyperparam_tuning_raw_run
guardian, swa
–
30.58
30.61
30.81
50
context_game_presinkhorn_train_run
guardian, cg
{256,512,1024}
168.99
166.86
180.83
100
Table 2: Training step latency with roughly fixed tokens_per_batch. cg = context-game.
Context (tokens)
Val CE ↓
PPL ↓
Rel. change vs 512
512
5.4547
≈233.9
–
768
5.2461
≈189.8
−3.8% CE, −18.8% PPL
Table 3: Context-length gains on WT2 (raw-v1, GPT-2 BPE). Lower is better.
Raw WT2 baseline (Table 1, sparsity 0.0). [ Syed et al., 2023 ]
Appendix
Table 6: Reference language-modeling numbers reported elsewhere on WikiText-2. Not directly comparable to our protocol (different training data/schedules/tokenizers or pretrained checkpoints). Lower is better.