Figure 1: The depth read is an attention, so it should be multi-head. Left: attention residuals (19) feed every sublayer a learned mixture of the depth history (embedding and every earlier attention/MLP output). Zoom left: one routing site is exactly single-head attention over that history—one query, one depth map shared by all d channels (Eq. 2). Zoom right: MHAR gives each of H subspaces its own map (Eq. 3), parameter-free; H=1 recovers attention residuals exactly.
Table 1: From-scratch validation loss (↓, tail-mean of the last eleven evaluations, steps 15–20k; §3) at the tuned peak learning rate 1×10−3; Δ is relative to the same-rate baseline; CEGFLOPs is the compute-equivalent gain of MHAR over the baseline (see text). MHAR is best at every scale; single-head routing regresses below baseline at 350M and 1B.
Scale
d/L/KV
Baseline
Hyper-conn.
Single-head
MHAR
CEGFLOPs
100M
512/12/4
3.336
3.322 (−0.014)
3.297 (−0.039)
3.287 (−0.049)
1.27×
350M
1024/24/8
3.290
3.260 (−0.030)
3.345 (+0.055)
3.210 (−0.080)
1.49×
1B
1280/36/8
3.175
3.172 (−0.002)
3.315 (+0.140)
3.111 (−0.063)
1.38×
Figure 2: What each method lets the current sublayer read across depth. Each column: forward stack (top) and the depth-read weight matrix (bottom; entry [i,c] is how strongly channel c reads depth row i; illustrative). Standard pre-norm: a one-hot read of the previous state. Attention residuals: one shared query reads all earlier sources—every column identical (rank-1). MHAR (ours): each subspace reads depth independently (block-wise columns).
Table 2: Zero-shot downstream evaluation at 100M, 350M, and 1B (single run). Each model is evaluated at its training context (100M: seq. 2048; 350M/1B: seq. 1024), so perplexity is comparable only within a scale. MHAR transfers its gain to held-out perplexity and LAMBADA at every scale, and HellaSwag at 350M (tied at 1B, marginally lower at 100M). 100M is the earlier 6×10−4 batch; 350M and 1B are the 5×10−4 batch.
Scale
Method
WikiText-2 PPL ↓
LAMBADA ↑
HellaSwag ↑
100M
Baseline
48.6
10.4%
33.0%
100M
MHAR
43.8
13.4%
32.5%
350M
Baseline
53.9
9.6%
33.5%
350M
MHAR
46.0
13.4%
35.5%
1B
Baseline
49.6
12.4%
34.0%
1B
MHAR
40.1
15.0%
34.0%
Figure 4: The H heads carry genuinely different depth-links. Each head’s token-averaged deviation w¯h−w¯ from the head-consensus routing at the first ten routing sites of the trained 1B MHAR model (dark cells: source not yet available). Deviations reach ±0.28 vs. 0.067 under a matched-norm random query, replicate on disjoint evaluation text (r=0.77), and are near-uncorrelated across heads.
Table 3: Downstream accuracy (↑) after 8B mid-training on anneal_pt_v3 (final checkpoints, EMA weights). The control and MHAR columns are schedule-matched: identical LR schedule, batch, data order, and ≈10 B-token budget.
Task
Base
Plain CPT
+MHAR
General
MMLU
0.554
0.643
0.645
GPQA
0.266
0.315
0.346
Math & Coding
GSM8K
0.190
0.470
0.502
MATH
0.053
0.191
0.191
HumanEval
0.122
0.409
0.415
MBPP
0.148
0.386
0.392
Figure 5: The conversion is exactly identity and introduces no optimization shock. Train loss of the schedule-matched 8B mid-training pair (identical seed and data order). Left: the two curves coincide. Right: the per-step |Δloss| starts at ∼10−5, rises only as the zero-initialized gate opens, and saturates 36× below step-to-step batch noise (dotted).
Table 4: Training speed and memory. Baseline Transformer vs. MHAR (torch.compiled reference kernels) vs. MHAR with our fused routing kernels, at the three model settings of Table 1. Throughput is relative to the same-scale baseline; MHAR is parameter-matched (+0.02%); single-head attention residuals are cost-identical to MHAR by construction. Protocol, absolute medians, and the full breakdown in Appendix F (Table 8).
Throughput (vs. base)
Peak mem. (GB)
Method
100M
350M
1B
100M
350M
1B
Baseline Transformer
1.00×
1.00×
1.00×
41.5
19.4
19.0
MHAR
0.54×
0.32×
0.23×
52.4
40.1
47.5
MHAR + fused kernels (ours)
0.88×
0.71×
0.55×
42.0
20.0
20.1
Figure 6: The bypass wants more heads. Under-trained regime (100M, 1×10−3, 5K steps): routing-head (H) × KV grid of validation loss (lower/yellow better; gold stars: per-KV optimum; red dashed line: the H= KV diagonal). The boxed KV= 1 row is the one supra-noise signal (see text).
Table 5: Routing-operation speedup, isolating the kernels from Table 4’s end-to-end numbers: wall-clock of all routing calls of one microbatch, forward+backward, bf16, single H100; the 8B column is the mid-training delta variant. End-to-end gains are smaller because routing is only part of a training step (batch settings and the eager-reference comparison in Appendix F).
Routing per microbatch (ms)
100M
350M
1B
8B (delta)
torch.compile
30.2
53.1
153.2
651
Fused Triton (ours)
5.7
11.4
42.4
323
Speedup
5.3×
4.7×
3.6×
2.0×
Figure 8: Validation loss of MHAR at 100M as a function of routing heads H (x) and KV heads (y). Lower is better. Gold stars mark the best H per KV; the red dashed line is the head-matched diagonal H= KV. At 100M the KV= 4 and KV= 8 optima sit on the H= KV diagonal. The 100M architecture of Table 1, trained at 5×10−4.
Table 6: Direct probe of the trained single-head routing queries (no new training; S=4 slices). Width-disagreement is the mean KL(as∥α) between each query slice’s preferred depth distribution and the shared α; a matched-norm random query gives the null, and the learned excess (trained−null) is the genuine learned signal. The middle column is a width-isolating control (d512→d768 at fixed L12/kv4/N): the learned disagreement rises while the null and collinearity stay flat, so width itself drives it. Across scale the learned excess grows monotonically 2.6× (→0.606) and is the primary, empirically-supported factor; source collinearity (mean pairwise cosine of the N=2L+1 sources)—a secondary, hypothesized factor we do not cleanly isolate—is non-monotonic (1.9× from 100M to 350M, then lower at 1B, partly a dimension effect), so the disagreement factor is the clean cross-scale signal. Caveat: slices are sub-parts of one trained query (a proxy for independent heads) rather than independently trained routers.
100M
124M (width ctrl)
350M
1B
Probe quantity (S=4)
d512/L12/kv4
d768/L12/kv4
d1024/L24/kv8
d1280/L36/kv8
Width-disagreement KL (trained)
0.273
0.311
0.570
0.701
random-query null
0.038
0.030
0.058
0.095
learned excess (trained−null)
0.235
0.281
0.512
0.606
Source collinearity (mean cos)
0.060
0.063
0.114
0.087
N sources (=2L+1)
25
25
49
73
Figure 11: Training loss at 100M (EMA-smoothed): MHAR stays below the standard baseline throughout training (inset: tail zoom, steps 12k–20k). The two runs are identical except for the routing mechanism (same node, software, data order, and global batch).
Table 7: Head-alignment ablation (validation loss ↓, final-step eval): routing over arbitrary subspaces (MHAR) vs. aligned to KV head groups (MHAR-HW), with matched head count. Aligning to attention heads shows no measurable benefit over arbitrary subspace routing; all differences fall within the ±0.07 eval-noise floor (single-seed, final-step). Separate matched-run batch (§3); read the within-table Δ.
Scale
MHAR
MHAR-HW
Δ(hw−mh)
100M
3.330
3.333
+0.003
350M
3.225
3.223
−0.002
1B
3.213
3.229
+0.016
Table 8: Training compute and memory at the three model settings of Table 1 (single H100; per-GPU microbatch as in training: batch 8 / sequence 2048 at 100M, 4/1024 at 350M, 2/1024 at 1B; identical node and software for all cells; 120-step runs on a verified-healthy GPU). Each cell is the median of three steady-state windows (within ±1%); peak memory reproduces to the decimal across sessions, and the 350M compiled/fused cells reproduce in a second independent session. Throughput is relative to the same-scale baseline (absolute baseline medians: 113.9k / 45.0k / 19.7k tokens/s). Parameters are matched to the baseline up to the O(d) routing queries (+0.02%). †Single-head attention residuals are cost-identical to MHAR by construction in every row—the head split is a parameter-, FLOP-, and memory-free reshape (§2)—so they are not listed separately.
100M
350M
1B
Method†
thr.
mem. (GB)
thr.
mem. (GB)
thr.
mem. (GB)
Baseline (additive residual)
1.00×
41.5
1.00×
19.4
1.00×
19.0
MHAR, reference kernels
0.31×
68.8
0.16×
70.9
OOM (>80 GB)
MHAR, torch.compile
0.54×
52.4
0.32×
40.1
0.23×
47.5
MHAR, fused Triton (ours)
0.88×
42.0
0.71×
20.0
0.55×
20.1
Table 9: Width-control (tail-mean validation loss, steps 15–20k). Holding depth (L=12) and KV (4) fixed and only widening d→d768: single-head routing’s edge over the baseline erodes (−→−0.033) while MHAR holds (−→−0.053), so MHAR’s advantage over single-head more than doubles (−→−0.020). This is the loss-level counterpart of the width-isolated disagreement in Table 6 (which rises →0.281 over the same widening). Single-seed; the single-head deficit shift (0.008) is near the per-seed noise floor, so the robust signal is the widening MHAR−single gap.
Width
Baseline
Single-head
MHAR
MHAR−single
100M (d512)
3.462
3.421 (−0.041)
3.412 (−0.050)
−0.009
124M (d768)
3.220
3.187 (−0.033)
3.167 (−0.053)
−0.020
Table 10: Best-vs-best learning rate: each method at its own optimum over peak LR ∈{1,5}×10−4, 10−3 (validation loss ↓, tail-mean of the last eleven evaluations; Δ vs. the same-scale best baseline; superscript is the selected LR). MHAR wins at every scale even when every method is given its own optimal learning rate; single-head, even at its kinder 5×10−4, still regresses at 350M and 1B.
Scale
Baseline
Hyper-conn.
Single-head
MHAR
100M
3.336 10−3
3.322 10−3 (−0.014)
3.297 10−3 (−0.039)
3.287 10−3 (−0.049)
350M
3.290 10−3
3.260 10−3 (−0.030)
3.328 5e–4 (+0.038)
3.210 10−3 (−0.080)
1B
3.175 10−3
3.172 10−3 (−0.002)
3.279 5e–4 (+0.105)
3.111 10−3 (−0.063)
Table 11: Seed robustness: paired per-seed delta vs. baseline (validation loss ↓, tail-mean of the last eleven evaluations; three seeds per method per scale, shared 5×10−4 recipe; ± is the per-seed standard deviation). Seed s fixes the data order across methods, so the comparison is paired. MHAR improves at every scale with |Δ|/SE≥10; single-head crosses from help to harm and is the highest-variance method.
Scale
Hyper-conn.
Single-head
MHAR
100M
−0.031±0.003
−0.043±0.003
−0.045±0.007
350M
−0.028±0.009
−0.005±0.015
−0.090±0.004
1B
−0.016±0.015
+0.093±0.049
−0.071±0.002
Table 12: Composition of the anneal_pt_v3 mid-training corpus: share of text bytes per source, measured over a uniform random sample of 48 of the 2,048 shards (187 GB of text; group rows sum their constituents, so rounded columns may differ in the last digit). The corpus is English-only; FinePDFs is restricted to its top three of twenty quality bins (verified from the per-document final_bucket metadata).
Transformers propagate information across depth through a single additive residual stream: every sublayer reads only the most recent state. Attention residuals relax this by letting each sublayer attend, through a learned softmax. However, that read uses a single query shared across the entire width, so every feature subspace must read the depth history through one distribution. The cost of this forced compromise grows with how much the subspaces disagree about which layers to read, and disagree
![Figure 2: What each method lets the current sublayer read across depth. Each column: forward stack (top) and the depth-read weight matrix (bottom; entry [i,c] is how strongly channel c reads depth row i; illustrative). Standard pre-norm: a one-hot read of the previous state. Attention residuals: one shared query reads all earlier sources—every column identical (rank-1). MHAR (ours): each subspace reads depth independently (block-wise columns).](https://media.metallab.ai/papers/2607.27230/f1.svg)
