트랜스포머가 예전 층의 정보를 다시 읽어올 때, 질문을 하나만 던지지 말고 특징별로 여러 개 던지게 했더니 더 좋아졌다
arXiv:2607.272302026-07-21
Multi-Head Attention Residuals
트랜스포머가 예전 층의 정보를 다시 읽어올 때, 질문을 하나만 던지지 말고 특징별로 여러 개 던지게 했더니 더 좋아졌다
표준 트랜스포머는 바로 이전 층의 결과만 더해서 다음 층으로 넘기는데, '어텐션 잔차(attention residuals)'라는 기존 방법은 지금까지 지나온 모든 층의 출력을 학습된 하나의 질문으로 훑어보게 해준다. 이 논문은 그 질문이 전체 폭(width)에 대해 딱 하나뿐이라 모든 특징이 같은 비율로 과거를 읽어야 한다는 점을 문제로 보고, 질문을 여러 개(헤드)로 나눠 각 특징 묶음이 독립적으로 과거를 읽게 하는 MHAR(Multi-Head Attention Residuals)를 제안한다. 처음부터 학습시킨 100M/350M/1B 모델 모두에서 표준 트랜스포머보다 검증 손실이 낮아졌고, 8B 모델에 이어서 학습(mid-training)시켰을 때도 GSM8K와 GPQA 점수가 올랐다.
METAL MEDIA 해설 도표
잔차 읽기 방식의 차이
증거 상태측정 결과가 보고됨
표준 잔차각 층이 바로 이전 층의 출력만 더해서 다음 층에 전달한다
어텐션 잔차(H=1)지금까지의 모든 층 출력을 하나의 공유 질문으로 훑어보고, 모든 특징이 같은 비율로 과거를 읽는다
MHAR(다중 헤드)특징을 H개 그룹으로 나눠 각 그룹이 독립적인 질문으로 과거 층들을 따로 읽는다, 파라미터 추가는 없다
헤드 수 = KV 헤드 수별도 튜닝 없이 거의 최적의 성능을 내는 기본 설정
8B 이어 학습초기 게이트를 0으로 두어 기존 모델을 그대로 유지한 채 서서히 라우팅을 열어 성능 손상 없이 추가 학습한다
METAL MEDIA이 원문을 바탕으로 재구성한 해설 도표이며, 논문 저자의 원문 figure가 아닙니다.
무엇을 했나
기존 방식인 '어텐션 잔차'는 모든 특징 차원이 하나의 공유된 질문(query)으로 과거 층의 출력을 읽는데, 이 논문은 이를 H개의 하위 그룹(헤드)으로 나눠 각 그룹이 독립적으로 자기만의 질문으로 과거를 읽게 하는 MHAR를 제안했다.
이 변경은 새 파라미터나 추가 연산이 거의 들지 않는 '재구성(reshape)'이며, 헤드 수 H=1이면 기존 어텐션 잔차와 완전히 동일해진다.
FineWeb-Edu 데이터로 100M, 350M, 1B 규모 모델을 처음부터 학습시켜 표준 트랜스포머, 하이퍼커넥션, 단일 헤드 어텐션 잔차와 비교했다.
헤드 수를 KV(키-값) 헤드 수와 같게 맞추는 것이 별도 튜닝 없이도 거의 최적인 기본값임을 확인했고, 실제로 학습된 질문들을 직접 분석해 특징 묶음들이 서로 다른 층을 선호한다는 것도 확인했다.
8B 모델에서는 기존 학습을 건드리지 않고 안전하게 이 방식을 얹을 수 있는 변환 방법을 만들고, 학습 속도를 크게 높이는 전용 커널(Triton)도 구현했다.
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).
![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)
