Code for the paper Routing Absorption in Sparse Attention: Why Random Gates Are Hard to Beat by Keston Aquino-Michaels (2026).
Paper: arXiv:2603.02227
When sparse attention is trained end-to-end, the model's Q/K/V projections co-adapt to whatever mask is imposed, absorbing the routing signal until learned gates perform little better than frozen random gates. We call this routing absorption and present four independent lines of evidence:
- Soft gating converges to nearly the same perplexity whether the gate is learned or random (48.73 ± 0.60 vs 49.83 ± 0.04 over 3 seeds)
- Hard top-k gating receives exactly zero gradient through the mask
- A gate distilled onto co-adapted Q/K/V achieves high F1 but catastrophic perplexity (601.6 vs 48.6)
- Stochastic mask randomization fails to prevent co-adaptation (78.2 ppl vs 37.3 baseline)
Post-hoc distillation sidesteps absorption entirely: freeze the model, train a tiny bilinear gate with KL divergence for 1K steps, and deploy sparse. This achieves near-oracle performance at all sparsity levels (99.9% efficiency at 87.5% sparsity on Qwen3-1.7B).
pip install -r requirements.txt
All experiments run on Modal for cloud GPU access. You'll need a Modal account and modal token new to authenticate.
31M model (WikiText-103, GPT-2 tokenizer):
modal run modal_app/data_prep.py
Qwen3-1.7B (WikiText-103, Qwen3 tokenizer + model download):
modal run modal_app/qwen3_setup.py
This downloads Qwen/Qwen3-1.7B-Base from HuggingFace and tokenizes WikiText-103 into 512-token chunks.
Run 6 training jobs (3 seeds ×ばつ {learned, random}) in parallel:
modal run modal_app/diagnostic_experiments.py --experiment soft-gating-seeds
This trains 50K steps per seed (42, 137, 256) on L4 GPUs and produces mean ± std results (learned: 48.73 ± 0.60, random: 49.83 ± 0.04).
The dense baseline (37.32 ppl) can be trained separately:
modal run modal_app/train.py --config configs/phase1/wikitext_dense.yaml --name dense --gpu l4
Distill gates on the dense checkpoint vs the soft-gated checkpoint:
modal run modal_app/diagnostic_experiments.py --experiment distill-topk
This freezes each checkpoint, trains gate projections with BCE for 1K steps, then deploys with hard top-k and measures perplexity.
Train with stochastic random masks, then deploy:
# Train with stochastic masks modal run modal_app/train.py --config configs/phase2/stochastic_k64.yaml --name stochastic --gpu l4 # Evaluate deployment modal run modal_app/diagnostic_experiments.py --experiment stochastic-deploy
modal run modal_app/diagnostic_experiments.py --experiment kl-sweep
Trains a bilinear gate with KL divergence on the frozen dense checkpoint and evaluates at k={32, 64, 128, 256}.
modal run modal_app/diagnostic_experiments.py --experiment gate-only-seeds
# KL distillation at d_gate=32 (the main result) modal run modal_app/qwen3_distill.py --experiment kl-sweep # BCE distillation at d_gate={32, 64, 128} modal run modal_app/qwen3_distill.py --experiment distill-sweep modal run modal_app/qwen3_distill.py --experiment gate-sweep
modal run modal_app/qwen3_absorption.py --experiment single-layer
Freezes all 28 layers except layer 14, adds soft gates, trains for 5K steps. Compares learned vs random vs no-gate.
modal run modal_app/qwen3_distill.py --experiment absorption-ablation
Sweeps unfrozen layers {0, 2, 4, 8} to measure how co-adaptation capacity drives absorption.
The 31M numbers come from the KL sweep (Table 4). The Qwen3 numbers:
# Dense eval + oracle sweep modal run modal_app/qwen3_distill.py --experiment dense-eval modal run modal_app/qwen3_distill.py --experiment oracle-sweep # KL distillation modal run modal_app/qwen3_distill.py --experiment kl-sweep
Runs 7 benchmarks (HellaSwag, ARC-Easy, ARC-Challenge, WinoGrande, PIQA, BoolQ, OpenBookQA) at k={64, 128, 256}:
modal run modal_app/qwen3_eval.py
After generating the analysis JSON files (saved to analyses/ by the scripts above), regenerate figures:
python papers/generate_figures.py
Outputs PDF and PNG to papers/figures/.
src/
models/
gated_attention.py # GatedSparseAttention: bilinear gate + top-k masking
transformer.py # 31M pre-norm transformer with gated attention
config.py # Pydantic model/training/data configs
training/
trainer.py # Training loop with gate diagnostics
losses.py # CE + sparsity + approximation losses
data/
wikitext.py # WikiText-103 dataset loading
modal_app/
train.py # 31M model training (dense, soft, hard, stochastic)
diagnostic_experiments.py # 31M ablations (oracle, distill, gate-only, SVD)
qwen3_distill.py # Qwen3 distillation (KL + BCE) and benchmarks
qwen3_absorption.py # Qwen3 single-layer absorption test
qwen3_eval.py # Downstream task evaluation (lm-eval-harness)
qwen3_setup.py # Download Qwen3 model + prepare data
data_prep.py # WikiText-103 preprocessing (GPT-2 tokenizer)
configs/
phase1/ # Dense baseline and soft gating configs
phase2/ # Hard top-k sweeps, stochastic masking
papers/
generate_figures.py # All publication figures
figures/ # Output PDFs and PNGs
tests/ # Unit tests for attention and transformer modules
The gate adds per-head projections W_gq, W_gk projecting to d_gate dimensions:
G(x) = (x @ W_gq) @ (x @ W_gk)^T / sqrt(d_gate)
With d_gate=32 this adds 393K parameters (1.3% of the 31M model). At deployment, only the top-k entries per query are retained.
KL distillation trains this gate against the model's own soft attention distribution (not a binary mask), preserving full ranking information. This is the critical difference vs BCE distillation, which loses ranking and fails at high sparsity.
The full experimental suite costs ~125ドル on rented A100 GPUs via Modal. Individual experiments:
| Experiment | GPU | Time |
|---|---|---|
| 31M training (50K steps) | L4 | ~1 hour |
| 31M diagnostics | T4 | ~30 min |
| Qwen3 KL distillation | A100-80GB | ~2 hours |
| Qwen3 absorption ablation | L4 | ~4 hours |
| Qwen3 downstream eval | A100-80GB | ~3 hours |
@article{aquino2026routing, title={Routing Absorption in Sparse Attention: Why Random Gates Are Hard to Beat}, author={Aquino-Michaels, Keston}, journal={arXiv preprint arXiv:2603.02227}, year={2026}, url={https://arxiv.org/abs/2603.02227} }