Asymmetric Attention Heads: Structured Head-Wise Context Allocation for Transformer Attention
arXiv:2608.192032026-08-21
트랜스포머의 모든 어텐션 헤드가 같은 길이의 문맥을 볼 필요는 없다는 실험
표준 트랜스포머는 모든 어텐션 헤드에게 똑같이 긴 문맥 범위를 보게 하지만, 저자는 헤드마다 필요한 문맥 길이가 다를 수 있다고 보고 헤드별로 문맥 창 크기를 다르게 배분하는 AAH-v3라는 방법을 제안했다. 4096 토큰, 시드 0(난수 고정값 0) 설정의 실험에서 여러 AAH 변형이 순수 풀 어텐션보다 검증 손실이 낮게 나왔다. 다만 이 결과가 실제 GPU 연산량 감소로 이어진다는 증거는 아직 없다고 저자는 분명히 밝힌다.
무엇을 했나
- 어텐션 헤드마다 특징 통계(EMA로 평활화한 쿼리/키/값 활성값과 이전 어텐션 진단치)를 뽑아 계층적으로 그룹을 만들고, 각 그룹에 512/1024/2048/4096 토큰 중 하나의 인과적(미래를 보지 않는) 로컬 윈도우를 할당한다
- 형제 그룹끼리 따로 점수를 매기던 기존 방식 대신 두 그룹을 함께 비교해 점수를 매기는 '조인트 시블링 스코어링'을 도입했고, 자식 그룹은 부모 그룹보다 큰 윈도우를 선택할 수 없게 제약을 걸었다
- 선택된 윈도우 결정은 미분 불가능한 하드 선택(argmax)이며, 언어모델 학습 손실만으로 트랜스포머 가중치를 학습하고 컨트롤러 자체는 시드로 고정된 통계 기반 정책으로 작동한다
- 4096토큰, 1B(10억 파라미터) 규모, 시드 0 조건에서 여러 AAH 국소 배분 변형이 풀 어텐션보다 낮은 검증 손실을 기록했고, 3000스텝 짧은 예산 실험에서는 고정/로컬 방식도 적응형 계층 구조와 경쟁할 수 있음을 보였다
- Attention Coverage Ratio(ACR, 선택된 윈도우 비율 지표)를 라우팅 진단 지표로만 보고했고, Nsight(GPU 연산 실측 도구) 기반 측정에서는 실제 GPU 연산량(FLOPs) 감소가 확인되지 않았다고 명시했다
Table 1: Relationship between ACR and measured FLOPs Ratios. ACR is a selected-window routing diagnostic. Measured FLOPs Ratios require matched profiler counters and are not inferred from selected windows.| Metric | What it measures | Source in this paper | Claim supported |
|---|
| ACR | selected attention-window coverage under rectangular query-by-window accounting | AAH policy logs and final branch-usage frequencies | head-wise context-allocation structure |
| Measured Attention FLOPs Ratio | profiler GPU FP operations inside attention ranges divided by the matched pure-attention baseline | Nsight / profiler suite only | hardware-measured attention operation count, when available |
| Measured Total FLOPs Ratio | profiler GPU FP operations for the full profiled forward or step divided by the matched pure-attention baseline | Nsight / profiler suite only | hardware-measured end-to-end operation count, when available |
Table 2: Main 1B / 4096-token protocol fields and W&B row-selection rule. Additional artifact-packaging metadata is listed in Appendix A.5.| Protocol field | Value used for the main 1B / 4096-token suite |
|---|
| Principal regimes | Full attention, Grouping off, Full adaptive, Shallow freeze, and Deep practical reuse. |
| Intended controlled fields | Tokenizer, data pipeline, optimizer schedule, batch construction, model width/depth, seed, and Transformer output interface are held fixed by local experiment configuration across the five paper-facing regimes. |
| Exported metadata | The W&B export records per-regime run names, config paths and hashes, run seeds, checkpoint step, and git commit IDs. Appendix A.5 lists the remaining artifact-packaging fields to include in a public release bundle. |
| Precision | bf16 in the W&B run configuration. |
| Batch size | 1. |
| Context length | 4096 tokens. |
| Candidate windows | [512,1024,2048,4096]. |
| Control/EMA settings | Control updates every five steps; feature EMA coefficient 0.9; final resolution EMA coefficient 0.15. |
| Checkpoint rule | Save checkpoints at optimizer steps 1000, 5000, and 10000; paper tables use the 10000-step final checkpoint. |
| W&B row-selection rule | For Full attention rows, set ACR to 1.0000 by definition when AAH-specific ACR fields are absent. For AAH rows, select completed seed-0 rows with final checkpoint step 10000 and non-null validation, ACR or branch-usage data, throughput, and peak allocated memory fields when available; ignore stale summaries that conflict with wandb_results_new/. |
| Dense backend status | The main tables report ACR as a routing diagnostic, not as a FLOPs Ratio. The dense-masked backend has not established skipped query–key work; therefore these rows omit backend work-accounting and measured FLOPs Ratio columns. |
Table 3: Configuration differences among the main 1B / 4096-token regimes. Performance metrics are reported separately in Tables 4 and 6.| Method | Grouping | Hierarchy | Joint scorer |
|---|
| Full attention | - | - | - |
| Grouping off | off | - | off |
| Full adaptive | feature-derived adaptive | [2,2,2,2] adaptive | wide joint |
| Shallow freeze | feature-derived frozen | [2] | wide joint |
| Deep practical reuse | cached level-0 | [2,2,2,2] reuse | wide joint |
Table 4: Main 1B / 4096-token training comparison for the Full attention reference and principal AAH-v3 execution regimes. ACR is the selected attention-coverage policy proxy. The dense-masked path does not establish skipped backend work, so backend work accounting is omitted. Memory is peak allocated GPU memory from perf/gpu_alloc_max_mb/1024. Bold indicates the best value in each metric column.| Method | Val. loss | ACR | Token/s | Memory (GB) |
|---|
| Full attention | 6.5672 | 1.0000 | 4447 | 35.18 |
| Grouping off | 6.5655 | 0.5814 | 3653 | 36.00 |
| Full adaptive | 6.5590 | 0.3066 | 3319 | 36.78 |
| Shallow freeze | 6.5367 | 0.3724 | 3479 | 36.77 |
| Deep practical reuse | 6.5549 | 0.2891 | 3381 | 36.62 |
Table 5: Phase 1 3000-step quality/structure screening ablation at 4096-token context. Lower validation loss is better. These short-budget rows test whether head-window assignment structure and routing stability matter; they are not a replacement for matched 10000-step paper-grade confirmation.| Phase 1 row | Val. loss | Conservative interpretation |
|---|
| Pure baseline | 7.3226 | uniform full-span reference for the screening budget |
| Shallow control interval 10 | 7.2775 | best screening row; stable shallow routing is useful in this budget |
| Shallow freeze | 7.2920 | feature-derived shallow structure improves over pure baseline |
| Fixed random grouping | 7.2831 | strong control; hierarchy adaptivity alone is not proven causal |
| Fixed 1024 | 7.2834 | strong fixed-window control; local context allocation itself is competitive |
| Shallow shuffle-post-select | 7.3074 | shuffling selected assignments weakens the shallow result |
| Full adaptive shuffle-post-select | 7.3404 | shuffling selected assignments is worse than the pure baseline |
Table 6: Main 1B / 4096-token inference comparison for the Full attention reference and principal AAH-v3 execution regimes. Inference ACR is reconstructed from final inference branch-usage frequencies as ∑WfWW/T with T=4096. The dense-masked path does not establish skipped backend work; no measured FLOPs Ratio is reported without matched profiler GPU FP-operation totals. Bold indicates the best value in each metric column.| Method | Val. loss | Val. ppl | ACR | Token/s |
|---|
| Full attention | 6.5598 | 706.16 | 1.0000 | 14425 |
| Grouping off | 6.5580 | 704.88 | 0.5793 | 4882 |
| Full adaptive | 6.5546 | 702.50 | 0.2484 | 4318 |
| Shallow freeze | 6.5293 | 684.93 | 0.3400 | 6983 |
| Deep practical reuse | 6.5514 | 700.24 | 0.3616 | 4595 |
Table 7: Systems diagnostics for the measured-GPU-FLOPs claim boundary. ACR is a routing diagnostic; measured FLOPs claims require Nsight-derived GPU floating-point-operation counters.| Evidence source | Current result | Claim boundary |
|---|
| Backend 4096 realized-attention suite | AAH rows expose selected-window ACR and backend span diagnostics under FlexAttention or FlashAttention execution paths. | Useful for routing/backend checks, but not a measured FLOPs metric. |
| Backend 4096 Nsight summary | FlashAttention AAH rows are about 1.59x–1.61x the pure FlashAttention measured GPU FLOPs. | Negative evidence for a GPU-FLOPs-reduction claim in the current backend implementation. |
| PRO6000 FLOPs-lab lower-overhead probes | Later probes approach but do not beat pure FlashAttention or pure dense MHA. | Overhead can be reduced, but measured GPU-FLOPs savings are not established. |
| Dense-framework FLOPs probes | Dense-framework variants approach a ratio of 1.0 but remain above the pure dense baseline. | Dense-framework probes also do not establish measured GPU-FLOPs savings. |
Table 8: Planned appendix diagnostic runs for the fixed 1B / 4096 / 10000-step / seed-0 protocol. These rows define intended mechanism checks, not completed main-text evidence.| Appendix run | Variable changed |
|---|
| appendix_4096_control_off | disable controller decisions and use no adaptive control |
| appendix_4096_fixed_random_grouping | replace feature-derived grouping with fixed random grouping |
| appendix_4096_freeze_after_warmup_passthrough | freeze or pass through topology after warmup |
| appendix_4096_independent_scoring | use an independent fixed scorer instead of joint sibling scoring |
| appendix_4096_no_parent_constraint | disable parent index constraint |
| appendix_4096_no_feature_ema | disable feature EMA smoothing |
| appendix_4096_simple_policy | planned deterministic entropy/norm-based or random-window policy baseline |
Table 9: Appendix-only internal pretrained compatibility smoke test on capped deterministic Qwen3-4B-Base subsets at 4096-token context. These are not official full benchmark scores and should not be cited as benchmark results or used for external model comparisons. Best score for each benchmark is bolded, with ties bolded.| Benchmark | Full attention | Full adaptive | Deep practical reuse | Shallow freeze | Grouping off |
|---|
| Language |
| MMLU | 70.9 | 71.1 | 70.9 | 70.7 | 70.9 |
| MMLU-Pro | 28.5 | 27.7 | 27.7 | 27.7 | 28.5 |
| GPQA-Diamond | 34.8 | 34.3 | 34.3 | 33.8 | 35.4 |
| ARC-Challenge | 86.3 | 86.3 | 86.3 | 86.3 | 86.3 |
| HellaSwag | 44.3 | 44.3 | 44.3 | 44.3 | 44.5 |
| TriviaQA | 37.5 | 37.5 | 37.5 | 37.5 | 37.5 |
| C-Eval | 71.5 | 71.5 | 71.5 | 71.5 | 71.9 |
| Math and code |
| GSM8K | 13.3 | 13.3 | 13.3 | 14.8 | 13.3 |
| HumanEval | 34.4 | 34.4 | 34.4 | 34.4 | 34.4 |
| MBPP | 3.1 | 3.1 | 3.1 | 3.1 | 3.1 |
Table 10: Separate 1B-suite routing-diagnostic context for the AAH regimes used in the Qwen3 compatibility check. ACR is taken from the 1B final-checkpoint inference comparison in Table 6; it is not a Qwen3 downstream-run hardware or FLOPs measurement. Max absolute delta is the largest benchmark-score change, in percentage points, relative to the Full attention row in Table 9.| Regime | 1B inference ACR | Max |Δ| vs. Full attention (pp) |
|---|
| Full attention | 1.0000 | 0.0 |
| Full adaptive | 0.2484 | 0.8 |
| Deep practical reuse | 0.3616 | 0.8 |
| Shallow freeze | 0.3400 | 1.5 |
| Grouping off | 0.5793 | 0.6 |
왜 중요한가
이 연구는 어텐션 헤드가 모두 같은 문맥 길이를 볼 필요가 없다는 구조적 가정을 실증적으로 검증해, 향후 효율적인 트랜스포머 설계에 참고할 근거를 제공한다. 다만 검증 손실 개선과 실제 하드웨어 연산량 절감을 명확히 구분해 보고한 점은, 비슷한 주장을 하는 다른 연구를 평가할 때도 유의해야 할 기준을 제시한다.
이 논문의 용어
- 멀티헤드 어텐션(MHA) · 트랜스포머가 문장 내 단어 간 관계를 여러 개의 병렬 '헤드'로 나눠 계산하는 핵심 구조
- 인과적 로컬 윈도우 · 현재 위치 기준으로 과거 방향으로만, 정해진 길이만큼만 참조하도록 제한한 문맥 범위
- EMA(지수이동평균) · 최근 값에 더 큰 가중치를 주며 시간에 따라 값을 부드럽게 평활화하는 방법
- ACR(Attention Coverage Ratio) · 실제 선택된 어텐션 윈도우가 전체 문맥에서 차지하는 비율을 나타내는 진단 지표, 실측 연산량과는 다름
- FLOPs · 부동소수점 연산 횟수로, GPU가 실제로 수행한 계산량을 재는 단위
논문 원문 초록 (영문)
Standard multi-head attention (MHA) gives every head the same full causal context span, although heads can serve different contextual roles. Some heads may rely mainly on nearby lexical or syntactic context, while others may depend on longer-range relations such as entity interactions, discourse links, or state changes. We present Asymmetric Attention Heads (AAH), a head-wise context- allocation framework that treats context length as an explicit per-head or per-group allocation variable. AAH groups heads using feature-derived statistics, organizes these groups hierarchically, and assigns causal local windows while preserving the standard flat MHA output interface. In 4096- token seed-0 experiments, several AAH-style local-allocation variants achieve lower validation loss than pure full attention. Short-budget ablations show that stable local allocation and head-window assignment structure matter, while fixed/local controls can be competitive with adaptive hierarchy. We interpret AAH as a structured head-wise context-allocation mechanism for quality and analysis, with Attention Coverage Ratio (ACR) reported as a selected-window routing diagnostic
저자 · Zimu Zhao
arXiv에서 원문 보기