PaperlayerView exact revision on arXiv

Source-linked AI summary

When Can Attention Heads Be Statically Defined?

Weixian Waylon Li, Yintao Tai, Marcio Fonseca, Shay B. Cohen

arXiv:2609.34650v1cs.LGcs.CL

TL;DR

Some attention heads learn similar patterns across inputs, so recomputing their query-key scores and softmax may be unnecessary. The paper introduces SAF, which selects heads for fixed-pattern replacement and represents those patterns compactly. At 25% replacement, updates are faster at both reported model scales with modest perplexity increases.

  • Problem

    Some heads depend mainly on positions rather than token content, but determining which heads and patterns to replace and whether replacement improves efficiency remains unclear.

    Source: §1 · p. 1

  • Method

    SAF selects attention heads and replaces their weights with fixed causal patterns, represented by absolute-position and relative-distance preferences.

    Source: §1 · pp. 1–2

  • Results

    At 25% replacement, updates are 1.056× faster at 124M parameters and 4K context with a 0.77 ± 0.07% perplexity increase, and 1.068× faster at 1B and 8K with a 0.51% increase.

    Source: §1 · p. 2

  • Takeaways & Limitations

    SAF models also accelerate long-input finetuning and causal prefill, and the adapted 124M model generalises to more key-value pairs than ordinary attention and pruning controls.

    Source: §1 · p. 2, §5.3 · p. 8

Abstract

from arXiv · show

Some attention heads learn similar patterns across inputs. Reusing these patterns could reduce training cost by avoiding repeated query-key score computation and softmax. Through controlled pretraining comparisons, we identify Selective Attention Freezing (SAF), which selects heads with low attention-pattern variance and replaces their attention weights with fitted post-softmax means halfway through training. We represent these fixed patterns with absolute-position and relative-distance preferences, reducing storage from quadratic to linear in sequence length. A fused kernel reconstructs the patterns and executes ordinary-attention and replaced heads together. At matched training-token budgets, replacing 25% of attention heads gives 1.056x faster post-replacement optimiser updates at 124M parameters and 4K context, with a 0.77% perplexity increase. At 1B and 8K context, post-replacement updates are 1.068x faster on four GPUs including communication, with a 0.51% perplexity increase. The resulting models also accelerate long-input finetuning and causal prefill. After associative-recall adaptation, the 124M model with 25% replacement generalises to more key-value pairs at a fixed length better than ordinary attention and two pruning controls.

1 INTRODUCTION

The introduction asks when fixed causal attention patterns can replace input-dependent attention and improve efficiency. It presents Selective Attention Freezing (SAF), compact pattern storage, and faster updates at a modest perplexity cost.

1 INTRODUCTION

  • Because some heads depend mainly on query and key positions, recomputing their attention can reproduce the same positional bias at computational cost.Ordinary attention computes input-dependent weights for every head and input, including query-key interactions and softmax.Source: §1 · p. 1
  • The paper asks which heads and patterns to replace, when to replace them, and whether mixed layers can save memory and wall-clock time with modern kernels.Its framing highlights that existing results leave these choices unclear.Source: §1 · p. 1
  • SAF replaces selected attention matrices with fixed causal patterns while retaining trainable value and output projections.The hybrid layer combines ordinary query-key attention with fixed-pattern heads.Source: §1 · pp. 1–2
  • Two length-T vectors encode absolute-key-position and relative-distance preferences, replacing quadratic pattern storage with linear storage.A fused kernel reconstructs fixed patterns and executes them alongside ordinary attention.Source: §1 · p. 2
  • At 25% replacement, updates are 1.056× faster at 124M parameters and 4K context with a 0.77 ± 0.07% perplexity increase; at 1B and 8K, they are 1.068× faster with a 0.51% increase.The four-GPU 1B result includes communication.Source: §1 · p. 2

2 RELATED WORK

Related work connects fixed-pattern replacement to studies of input-independent attention, head heterogeneity, and hardware-aware execution. The paper distinguishes its approach by retaining token mixing of current input values while combining ordinary and fixed-pattern heads in one kernel.

2 RELATED WORK

  • Prior work studies fixed positional patterns, input-independent encoders and token mixing, input-averaged attention, relaxed sequence dependence, and the expressivity of constrained attention.SAF builds on input-averaged attention for selective replacement during causal language-model training.Source: §2 · p. 2
  • Pruning and component-importance methods address head heterogeneity, while retrieval and sparse-attention methods assign different functions or patterns to heads.Unlike pruning, fixed-pattern replacement retains token mixing of the current input’s values.Source: §2 · p. 2
  • The fused kernel combines ordinary attention and fixed-pattern token mixing, removing score computation and softmax for replaced heads.The pretraining path also omits replaced query and key projections, with projection savings depending on key sharing.Source: §2 · p. 2

3 REPLACING ATTENTION WITH A FIXED PATTERN

SAF calibrates attention patterns, ranks heads for replacement, and resumes training with fixed causal weights for selected heads. It fits compact patterns from positional preferences and uses fused execution to combine replaced and ordinary-attention heads.

3 REPLACING ATTENTION WITH A FIXED PATTERN

  • After ordinary-attention training, SAF calibrates a checkpoint, selects heads, constructs fixed patterns, and resumes training with the same model weights and optimiser state.Value and output projections continue training while selected attention weights remain fixed.Source: §3 · p. 3, §3.3 · p. 4

3.1 CALIBRATION AND HEAD SELECTION

  • Heads are ranked globally by increasing attention-variance score, which measures reconstruction error across calibration inputs rather than functional importance.A compact-fit tolerance can cause candidate heads to be skipped until enough fits pass or candidates run out.Source: §3.1 · p. 3

3.2 FIXED-PATTERN CONSTRUCTION AND FITTING

  • Candidate fixed patterns include post-softmax means, averaged-logit patterns, Gaussian and Dirichlet samples, and structured-random controls.Sampled patterns are drawn once at replacement and retained thereafter.Source: §3.2 · pp. 3–4
  • Compact patterns encode absolute-key-position and relative-distance preferences in length-T vectors, reducing storage from O(T^2) to O(T) per head.The two marginals are sufficient statistics for fitting; the reported row-averaged KL fitting tolerance is 0.2 nats per row.Source: §3.2 · p. 4

3.3 REPLACEMENT SCHEDULES AND FUSED EXECUTION

  • One-time, gradual, validation-budget, and train-loss-triggered schedules vary when and how selected heads are replaced.Gradual replacement installs the same final heads in groups, while the budgeted and triggered schedules use validation-perplexity cost and training-loss plateaux, respectively.Source: §3.3 · p. 4
  • The fused kernel runs ordinary-attention and replaced heads in one launch per layer, reconstructing fixed-pattern tiles in registers and multiplying them with value tiles.Replaced heads avoid score computation, softmax, and dense pattern reads; the multi-head pretraining path also omits their query and key projections.Source: §3.3 · pp. 4–5

4 EXPERIMENTAL SETUP

The experiments compare SAF with matched-token and matched-time controls across model sizes, contexts, and downstream uses. The setup specifies separate quality and runtime measurements, plus longer-context and four-GPU evaluations.

4 EXPERIMENTAL SETUP

  • The 124M decoder is trained for 5,000 updates totaling 2.4576B tokens, with replacement at the midpoint unless timing is varied.The study compares learned absolute positions with RoPE at 4K context and uses disjoint calibration and reporting data.Source: §4 · p. 5
  • Three-seed controls compare mean replacement with pruning and random head selection at matched tokens, and compare continuations at matched time.Random selection matches per-layer head counts; Gate-Taylor pruning removes globally low-scoring heads.Source: §4 · p. 5
  • Quality is evaluated with held-out perplexity and three-seed finetuning, while update timings include backward computation and optimiser steps but exclude calibration and fitting.Quality and runtime use separate runs in the 124M study; systems measurements use GH200 GPUs.Source: §4 · p. 5
  • Longer-context runs retain the 2.4576B-token budget and support square causal-prefill benchmarks, while the 1B, 8K comparison uses four GPUs and includes communication in update timing.The Qwen3-4B evaluation also tests five zero-shot tasks without further training.Source: §4 · pp. 5–6

5 RESULTS

SAF’s results favor low-variance head selection and post-softmax mean patterns, with one-time midpoint replacement balancing perplexity and update speed. Tests report benefits across training scales and use stages, plus improved fixed-length associative recall after adaptation.

5.1 WHICH FIXED PATTERNS AND HEADS SHOULD BE USED?

  • The post-softmax mean gives the lowest or similar perplexity across fixed patterns, while variance selection beats forward KL after adaptation and random selection nearly doubles the 25% penalty.Retaining token mixing with mean patterns improves on same-head pruning at matched tokens, though Gate-Taylor pruning has lower perplexity under both reported budgets.Source: §5.1 · p. 6

5.2 HOW MANY HEADS SHOULD BE REPLACED, AND WHEN?

  • At 25% replacement, perplexity increases by 0.768 ± 0.067%; at 50%, it increases by 2.492 ± 0.083%, making 25% the preferred quality trade-off.One-time midpoint replacement is robust across a broad range of training stages; gradual and automatic schedules do not improve the quality–speed trade-off.Source: §5.2 · pp. 6–7

5.3 DOES THE TRADE-OFF CARRY ACROSS MODEL USE STAGES?

  • SAF updates are faster and use less peak memory at 8K and 16K, while the perplexity penalty decreases from 4K to 16K without retuning.Long-input QuALITY finetuning reaches 1.118× faster updates at 50% replacement; 124M 16K causal prefill is 1.09–1.12× faster at 25% and 1.20–1.24× at 50%.Source: §5.3 · pp. 8–9
  • At 64 key-value pairs, SAF reaches 54.4% MQAR accuracy versus 26.2–29.0% for ordinary attention and pruning, despite adaptation on eight pairs at 512 tokens.At fixed input length, increasing key-value pairs tests generalisation to more associations rather than longer sequences.Source: §5.3 · p. 8, Figure 5 · p. 8
  • At 10% replacement, Qwen3-4B zero-shot accuracy changes range from −1.13 to +0.33 percentage points, averaging −0.37 points versus ordinary attention.Variance selection has higher accuracy than forward KL on all tested tasks at 10%, 20%, and 30% replacement.Source: §5.3 · p. 9, Figure 6 · p. 9

5.4 DOES SAF TRANSFER TO A LARGER MODEL?

  • At 1B, 25% replacement increases perplexity by 0.507% and changes mean finetuning accuracy by less than 0.72 percentage points per task.At 50%, perplexity increases by 2.027%; mean finetuning accuracy has no consistent advantage over Gate-Taylor pruning.Source: §5.4 · p. 9, Table 3 · p. 9

6 CONCLUSION

The conclusion presents Selective Attention Freezing (SAF) as a recipe for using fixed causal patterns while retaining useful token mixing. The accompanying notation and calibration details define the attention setup and when its statistics are measured.

6 CONCLUSION

  • SAF combines variance-based head selection and fitted mean patterns with compact storage and fused execution, while retaining token mixing through trainable value and output projections.At 25% replacement, perplexity increases remain below 1% at matched tokens, with post-replacement updates 1.056× faster at 124M and 4K, and 1.068× faster at 1B and 8K on four GPUs.Source: §6 · p. 9

A.1 ATTENTION NOTATION AND CALIBRATION

  • The notation defines L attention layers and H query heads per layer, with causal attention lower triangular and row-stochastic.The causal softmax masks future keys and normalizes each query row.Source: §A.1 · p. 15
  • Attention statistics are measured at the replacement checkpoint in evaluation mode with weights fixed, using calibration sequences separate from final reporting.The training-data sampler remains unchanged.Source: §A.1 · p. 15

A.2 TRAINING PROCEDURE

The method calibrates attention once, selects and fits eligible heads, and installs fixed patterns while ordinary attention and the remaining model parameters continue training. The accompanying analysis connects the post-softmax mean and attention variance to a reconstruction objective, while distinguishing that objective from language-model loss.

A.2 TRAINING PROCEDURE

  • At the replacement update, SAF measures attention on calibration data, orders heads by increasing variance, and installs fitted patterns only when their compact-fit error meets the threshold.Calibration leaves weights and training-data order unchanged; candidate fits do not receive intervening model updates, and the requested replacement rate is reached only if enough fits pass.Source: §A.2 · p. 15

A.3 ALTERNATIVE SELECTOR AND SCHEDULE DEFINITIONS

  • The default schedule replaces heads once near the training midpoint; alternatives install groups gradually or use validation-budget and training-loss triggers.Adaptive schedules change replacement timing and acceptance policy rather than recalibrating on every training update.Source: §A.2 · p. 16, §A.3 · p. 17

B UNDERSTANDING THE SELECTED RECIPE

  • The controlled study selects the post-softmax mean as the fixed pattern and attention variance as the head-selection score.Under the fixed-matrix reconstruction objective, the mean is the best fixed pattern and variance measures how accurately each head can be approximated by one.Source: B UNDERSTANDING THE SELECTED RECIPE · p. 17

B.1 FIXED-PATTERN RECONSTRUCTION

  • Squared distance correlates 0.950 with perplexity increase, versus 0.625 for DKL(A∥P) and 0.900 for squared Hellinger distance and total variation.These are descriptive stratified Spearman correlations, with five patterns compared within each replacement rate; the post-softmax mean has zero squared distance by construction.Source: Table 4 · p. 18, §B.1 · p. 18
  • The mean attention matrix minimizes expected reconstruction error, whose minimum is population attention variance per matrix entry; the same objective therefore guides pattern choice and head selection.Calibration estimates the population variance used for selection.Source: §B.1 · pp. 17–18
  • The reconstruction argument does not directly establish language-model-loss outcomes, which also depend on values, output projections, head interactions, and further training.The empirical mean estimates the population mean, and input-independent sampling adds variation without conditioning it on the current input.Source: §B.1 · p. 18
  • Forward- and reverse-KL objectives select different fixed summaries: the post-softmax mean and the sharp mean, respectively.When attention is a softmax of scores, the sharp mean is softmax of the expected scores.Source: §B.1 · p. 18

B.2 SELECTING HEADS FOR CONTINUED PRETRAINING

In matched settings, attention variance yields lower final perplexity increases than row-wise forward KL, while comparisons of one-head costs and final jointly replaced models show that local diagnostics do not always predict continued-training outcomes. Selection patterns and replacement timing also shape the interpretation of these results.

B.2 SELECTING HEADS FOR CONTINUED PRETRAINING

  • In preliminary one-run 124M, 1K-context comparisons, final perplexities are 26.558 for Q/K-gradient selection, 26.568 for variance, 26.627 for rank sum, 26.607 for gradient veto, and 26.382 for ordinary attention.These results use dense post-softmax means and 25% midpoint replacement, separate from compact-pattern comparisons.Source: §B.2 · p. 18
  • Residual cosine and variance yield immediate PPL increases of 4.136% and 5.092%, but final PPL increases of 0.792% and 0.665% after 2,500 continuation updates.The selected 25% sets share 23 of 36 heads; the final residual-cosine-minus-variance difference is 0.126% (95% interval: 0.102–0.150%).Source: Table 5 · p. 19, Table 6 · p. 19, §B.2 · p. 19
  • Attention variance gives the lower final perplexity increase in all six matched settings.The comparisons use complete head sets selected at the midpoint and the same 2,500 continuation updates.Source: Table 6 · p. 19, §B.2 · p. 19
  • At 25% replacement, 30–31 heads overlap between quarter-training and midpoint, versus 12.6–13.4 expected for independent selections with matching per-layer counts.In three-seed 4K controls, the first four layers contain 21–24 of the 36 selected heads.Source: §B.2 · p. 19
  • Selected heads have higher mean normalized entropy than layer-matched unselected heads, 0.837 versus 0.678, and less mass on the most recent 64 tokens, 19.0% versus 40.7%.Only 0.023% of their attention falls on the first token on average.Source: Figure 7 · p. 20, §B.2 · p. 20
  • Variance correlates 0.723 with one-head replacement cost; its ranking correlation with final ordering is 0.869–0.904 at quarter-training and 0.945–0.966 at midpoint.The one-head cost correlation remains 0.543 after controlling for layer and entropy.Source: §B.2 · p. 20

B.3 ADAPTATION AFTER REPLACEMENT

  • The immediate replacement cost rises from 2.72% at τ = 0.25 to 7.99% at τ = 1; two hundred updates after replacement, smoothed train-NLL gaps are 0.0061 nats at τ = 0.25 and 0.0123 at τ = 0.75.By the final update, the respective gaps are 0.0062 and 0.0086 nats, consistent with continued training offsetting part of the replacement cost.Source: Figure 8 · p. 21, §B.3 · p. 21

C.1 CONTROLLED PRETRAINING RESULTS

This section reports controlled attention-pattern diagnostics and replacement-versus-pruning comparisons, including matched-time perplexity results and associative-recall evaluations across pair counts and sequence lengths.

C.2 GAUSSIAN AND DIRICHLET DISTRIBUTION DIAGNOSTICS

  • 1.44 and 2.58 normalized energy distances for Gaussian and Dirichlet sampling exceeded 1.10 for whole-row resampling, and neither parametric family reproduced every head’s attention distribution.These diagnostics used concentration scale 4; visible clusters in the binned-profile visualization do not establish a mixture model or a uniquely suitable distribution.Source: §C.2 · pp. 24–25, Figure 9 · p. 25, Figure 10 · p. 25

C.3.1 MATCHED-TOKEN PRETRAINING

  • 0.772 ± 0.021% was the relative perplexity change for the early 25% replacement setting in the matched-token controls.The controls compared SAF with pruning and random head selection over three seeds at matched token budgets.Source: §C.3.1 · p. 26, Table 12 · p. 26
  • 3.30 ± 0.17 seconds at 25% replacement and 5.40 ± 0.10 seconds at 50% were the fitting-and-installation times in the 124M controls.Ranking took less than 0.012 seconds in every run; total intervention time also included calibration, ranking, and the KL check.Source: §C.3.1 · p. 26, Table 13 · p. 27

C.3.2 PATTERN CONTENT AND TRAINABILITY

  • −0.010 ± 0.006% was the paired relative perplexity change from learning patterns, while update and continuation times increased by 11.9 ± 1.2% and 12.4 ± 1.5%, respectively.In fixed-length pair-count tests, frozen means had higher mean accuracy throughout, while trainable means were slightly better in length-only tests.Source: §C.3.2 · p. 27, Table 14 · p. 27

C.3.3 MATCHED-TIME PRETRAINING

  • 0.786 ± 0.082 and 0.866 ± 0.154 percentage points were the excess perplexity penalties for random selection over variance selection at 25% and 50% replacement, respectively.The ordering was the same across all three seeds under matched-time budgets.Source: §C.3.3 · p. 28, Table 15 · p. 28

C.3.4 LATER REPLACEMENT AND PRUNING-SPECIFIC SELECTION

  • 22.683 was the perplexity after training a reference-derived variance-pruned architecture from initialization, versus 23.049 after pruning the same architecture following 5,000 updates.In seed 1337, the from-initialization run took 248.26 minutes, excluding the cost of obtaining the reference-derived layout; full-path times were 286.86 minutes for ordinary attention, 282.87 for SAF, 275.43 for variance pruning, and 275.93 for gate-Taylor pruning. Finetuning differences for SAF minus gate-Taylor were +0.765 ± 0.894 points on SST-2 and −0.051 ± 0.795 on BoolQ, with reversals across pretraining seeds.Source: §C.3.4 · p. 29, Table 16 · p. 29

C.3.5 DOWNSTREAM FINETUNING

  • Downstream comparisons selected an ordinary-attention model’s epoch count on an internal development split, then evaluated alternatives at that same epoch count and on the same ordered training examples.SST-2 and BoolQ used batch size 16, length caps of 128 and 512 tokens, and learning rate 2 × 10−5; downstream updates were matched within each seed and task.Source: §C.3.5 · p. 29

C.3.6 ASSOCIATIVE RECALL: ADAPTATION AND IMMEDIATE REPLACEMENT

  • The associative-recall evaluation adapted models for 1,500 updates on 512-token sequences with eight key-value pairs, then tested 8 to 64 pairs without further training or head selection.The comparison varied pretraining seeds with task seed 1337 fixed, and separately varied task-training seeds with pretraining seed 1337 fixed; the groups were not pooled as independent repetitions.Source: §C.3.6 · p. 29
  • The mean accuracy advantage over variance pruning rises from 5.65 points at 16 pairs to 25.39 points at 64; SAF beats both pruning variants across seeds at 24–64 pairs.At 16 pairs, seed 1338 does not share the mean advantage; a separate task-seed experiment confirms the ordering at 32 pairs but does not test additional pair counts. Length-only tests instead favor pruning, distinguishing generalisation to more associations from generalisation to longer inputs.Source: §C.3.6 · p. 30
  • Immediate replacement on an MQAR-trained model does not reproduce adapted SAF’s consistent 32-pair advantage, and the ordering changes during recovery updates.The immediate-replacement study calibrates at native 4K context and then trains all arms on the same 512-token, eight-pair stream. Because intervention stage and calibration domain differ between studies, they distinguish measured behaviors without isolating a single cause.Source: §C.3.6 · p. 30

C.3.7 TRAINING AND PREFILL COSTS

The appendix compares training-update costs, causal prefill, and downstream finetuning for replaced models. Replacement can accelerate some workloads, but gains depend on workload and replacement schedule; pruning saves more update time, with no quality data for its long-context timing models.

C.3.7 TRAINING AND PREFILL COSTS

  • Both pruning and fixed-pattern replacement reduce update time at 4K, 8K, and 16K, with pruning giving larger savings.The timing comparison uses matched head sets; the 8K and 16K pruned models are timing-only, without quality evaluation after adaptation.Source: §C.3.7 · pp. 30–31
  • At B = 1, mean replacement is slower in causal prefill: 0.933× for token-matched and 0.877× for time-matched groups.These measurements do not include token-by-token KV-cache decoding.Source: §C.3.7 · pp. 31–32

C.4 FINETUNING THE 124M CHECKPOINTS

  • One-time replacement during finetuning does not improve accuracy or update time; repeated train-loss-triggered replacement is especially costly on BoolQ.Post-training replacement avoids training-time kernel overhead but provides no adaptation after replacement.Source: §C.4 · p. 32
  • Zero-shot evaluation finds no Holm-corrected significant change after 25% replacement, but does not establish performance equivalence.The tasks are HellaSwag, PIQA, ARC-Easy, SST-2, and BoolQ.Source: §C.4 · p. 32

C.5 1B TRAINING AND EVALUATION

This section evaluates 1B models trained with replacement at the midpoint and compares their post-replacement update costs and downstream finetuning. The downstream comparison with pruning reports small paired accuracy differences, while finetuning times are observed rather than controlled paired speedups.

C.5 1B TRAINING AND EVALUATION

  • The 1B configurations share the update-18,756 midpoint checkpoint and select 128 or 256 of 512 heads by attention variance, fitting post-softmax means.Calibration uses 16 sequences of length 8,192; final perplexity uses 512 sequential 8K windows from reporting data disjoint from calibration.Source: §C.5 · p. 33
  • Four-GPU post-replacement update timing includes gradient all-reduce, while the single-GPU benchmark excludes distributed communication.The four-GPU measurements use four paired runs; the single-GPU measurements use fresh-process pairs and a matched per-update token workload.Source: §C.5 · p. 33, Table 23 · p. 34
  • The pretrained masks, patterns, and pruned layouts are inherited unchanged during finetuning, with no additional selection or fitting.Task epoch counts are selected using an internal development split for the ordinary-attention model, then shared across alternatives.Source: §C.5 · p. 34
  • The paired accuracy differences between SAF 25% and gate-Taylor pruning are −0.08 ± 0.04, −0.19 ± 2.07, and −0.06 ± 0.75 percentage points on SST-2, BoolQ, and QuALITY, respectively.The three finetuning repetitions measure task-training variation, not equivalence between methods or variation across pretraining seeds.Source: §C.5 · p. 34, Table 24 · p. 34
  • The finetuning-time table reports observed training times, not controlled paired speedups.Runs use different GPU allocations and some reuse previously completed ordinary-attention references.Source: §C.5 · p. 34, Table 24 · p. 34

C.6.1 QWEN3-4B ZERO-SHOT TASK ACCURACY

The Qwen3-4B evaluation compares head selectors, replacement accuracy, and prefill execution, then documents compact pattern storage and mixed-head kernel checks. Variance selection has higher accuracy in every tested task–rate comparison, while the implementation stores patterns linearly and avoids score computation and softmax for replaced heads.

C.6.1 QWEN3-4B ZERO-SHOT TASK ACCURACY

  • Variance selection gives higher accuracy in all fifteen task–rate comparisons, although the selected head sets overlap by only 31% at 10% and 43% at 30%.At 10%, its median per-input forward-KL score is 0.578 nats versus 0.296 for forward-KL selection; these observations distinguish the selected sets but do not establish which head functions cause accuracy differences.Source: §C.6.1 · p. 35, Table 25 · p. 35
  • At 20% replacement, forward-KL selection lowers SST-2 accuracy by approximately 25.6 percentage points, while variance selection’s loss is at most 3.21 points across tested rates.The strongest forward-KL loss coincides with a larger imbalance in predicted labels, but the aggregate counts do not establish the mechanism behind the non-monotone curve.Source: §C.6.1 · p. 35

C.6.2 GQA PREFILL MEASUREMENTS

  • The Qwen prefill benchmark compares fused GQA with Flash scaled dot-product attention at batch sizes 4 and 8 and lengths 2,048, 4,096, and 8,192.The measurements beyond the 2,048-token calibration length assess execution speed; the supplied passages do not report the speedup values.Source: §C.6.2 · p. 35, Table 26 · p. 35

D.1 COMPACT FITTING AND STORAGE

  • At T = 106, one causal fp32 matrix occupies approximately 2 TB even with only its lower triangle stored; for the 4K RoPE head set, compact buffers occupy 0.133 GB versus 8.060 GB dense.The dense total covers ten affected layers, including unused head slots; ideally packing only 36 selected matrices would use 2.416 GB.Source: §D.1 · p. 36

D.2 MIXED-HEAD EXECUTION AND NUMERICAL CHECKS

  • The fused kernel reconstructs replaced heads’ causal tiles in registers and multiplies them by value tiles, avoiding query-key scores, softmax, and dense pattern reads.Persistent state per replaced head uses O(T) storage, but bPV multiplication still costs O(T^2dh).Source: §D.2 · p. 36
  • The mixed-head implementation is checked against an eager dense reference across fp32 and BF16 cases, including gradient comparisons for ordinary and replaced heads.For replaced heads on the multi-head long-context path, query/key gradients are zero by construction; value gradients include both head types.Source: §D.2 · pp. 36–37, Table 28 · p. 37

D.3 CAUSAL-PREFILL BENCHMARKS

Checkpoint-backed causal-prefill benchmarks measure replacement across input lengths and batch sizes. At batch size 64, replacement remains above break-even from 512 to 16K tokens; the separately reported batch-one results are slower on median and highly variable.

D.3 CAUSAL-PREFILL BENCHMARKS

  • The checkpoint-backed benchmark compares a seed-1337 ordinary-attention 16K RoPE checkpoint with its 25%- and 50%-replaced counterparts.It sweeps input length at B = 64 and batch size at T = 4096; a separate B = 1 measurement is retained only as a diagnostic.Source: §D.3 · p. 37
  • At fixed batch size 64, causal-prefill speedups remain above break-even across 512–16K: 1.092–1.116× at 25% and 1.200–1.244× at 50% replacement.At fixed 4K length, the primary batch sweep uses sizes two through 64 and is similarly stable.Source: §D.3 · p. 37, Table 29 · p. 37
  • At T = 4096 and B = 1, median paired speedups are 0.969× and 0.972× for 25% and 50% replacement, with broad observed ranges.The batch-one point is dominated by launch overhead and varies substantially across six pairs.Source: §D.3 · p. 38
  • Separate forward-only measurements report 1.08× causal-prefill speedup for the 124M 16K shape at 25% replacement and 1.33× with full replacement.These measurements do not assess perplexity, token-by-token decoding, K-cache removal, or static query/key parameter repacking.Source: §D.3 · p. 38

E DATA, MODELS, AND LICENCES

The appendix inventories external datasets, pretrained models, and implementation assets used in the experiments, alongside licensing details. It notes that the controlled 124M pretraining models were trained from scratch.

E DATA, MODELS, AND LICENCES

  • Table 30 inventories external datasets and pretrained models used in reported experiments, including auxiliary zero-shot evaluations; 124M controlled-pretraining models were trained from scratch.Asset names link to access repositories, and licence entries link to upstream statements.Source: E DATA, MODELS, AND LICENCES · p. 38, Table 30 · p. 38
  • FineWeb-Edu’s database licence does not replace rights in underlying web content, QuALITY records licences per source article, and GLUE points to original dataset licences.GLUE’s loading code does not establish a licence for SST-2 review text.Source: E DATA, MODELS, AND LICENCES · p. 38
  • The implementation uses nanoGPT and Triton under MIT licences, GPT-2’s tiktoken tokeniser for 124M experiments, Qwen3-4B’s tokeniser for Qwen experiments, and Zoology’s Apache 2.0 generator for synthetic MQAR data.The 124M tokeniser is MIT-licensed.Source: E DATA, MODELS, AND LICENCES · p. 38

Related papers

Loading 2609.34650v1…