Start here
About 60 secondsThe idea, in a minute.
- The question
- Does headwise projected routing improve causal language modeling and structured sequence tasks enough to justify its training, inference, and memory overhead?
- What we found
- Three-seed byte-level WikiText tests found mean test improvements of 0.0293 BPB at depth 8 and 0.0207 BPB at depth 12, with consistent routing ablation effects. A longer paired five-task A40 study found causal quality gains on eligible Dyck-2 and TinyStories runs, but both missed batch-one latency limits; Dyck-2 also failed training-time and peak-memory limits.
- Why it matters
- Quality, causal use, and cost are separate gates. These results show why a favorable loss curve alone is insufficient to establish an efficiency or deployment advantage.
The local language model is a small byte-level MPS proxy; the RunPod study has one seed and eager CUDA implementation. Neither establishes performance on a standard large subword model or a fused production kernel.
A concrete example
Ablation confirms HCRR use, while cost guards still fail
In the longer RunPod paired screen, Dyck-2 and TinyStories candidates were run only after baselines met the registered eligibility band.
- 01
Pair
Baseline and candidate used seed 101, matched task-specific steps and tokens, data order, dimensions, optimizer, BF16 precision, and evaluation.
- 02
Measure quality
Dyck-2 validation exact match rose from 0.813477 to 0.908203; TinyStories validation NLL fell by 0.041537.
- 03
Disable routes
Routing ablation worsened Dyck-2 validation exact match by 0.060547 and TinyStories NLL by 0.579032, confirming causal use.
- 04
Apply cost gates
Dyck-2 latency reached 1.577x baseline and memory 2.609x; TinyStories latency reached 1.542x. Both failed the 1.35x latency guard.
Causal quality gains are real in these paired runs, but they did not satisfy the implementation efficiency rule.
Longer HCRR paired summary (in download archive)Read the formal abstract
Can a model improve by deciding which earlier layer updates each attention head should reuse? Small WikiText experiments suggest that headwise routing can lower held-out loss at two depths. But a longer five-task study shows the practical tension: the quality-positive candidates were slower at batch one, and Dyck-2 also used more training time and peak memory than allowed by the registered gates.
The evidence supports continued implementation work, not a broad efficiency claim. In a paired three-seed Apple M3 Pro proxy, headwise projected routing improved mean test loss by 0.0293 bits per byte at depth 8 and 0.0207 at depth 12. In a separate one-seed A40 screen with longer training, Dyck-2 and TinyStories both showed causal quality improvements after baseline eligibility, yet neither passed latency limits. The larger depth and task tests therefore sharpen the question from “does routing help?” to “can its quality gain fit within a useful systems budget?”
01 Research question
Does headwise projected routing improve causal language modeling and structured sequence tasks enough to justify its training, inference, and memory overhead?
02 Experiment setup
The full procedure, model settings, datasets, and evaluation choices are recorded here.
Explore the complete methodology
HCRR adds independent weights over non-adjacent normalized residual updates for each attention head, followed by a head-local projection. The local proxy used byte-level WikiText-2, sequence length 64, width 64, four heads, and MLP width 128. Baseline and candidate began with exact initialization parity and shared paired seeds and ordered sampled streams. At depth 8, three seeds (101, 202, 303) trained for 2,000 steps, batch 16, or 2.048 million tokens per model. At depth 12, the same token budget and architecture except depth were used. Both were FP32 on MPS. Validation/test used the first 131,072 bytes of each split. Routing-disabled evaluation measured causal contribution.
The longer RunPod screen tested five tasks with width 128, depth 8, four heads, MLP width 256, BF16 CUDA, seed 101, and task-specific steps: 2,000 for MQAR, ListOps, permutation, and TinyStories; 600 for Dyck-2. Data, ordered examples, optimizer, tokens, and within-task training budgets were paired. Candidates ran only when a synthetic baseline reached 55–90% validation exact match. Of five baseline tasks, only Dyck-2 and TinyStories were eligible for candidate runs. Frozen gates included parameter ratio 1.005x, training time 1.25x, batch-one latency 1.35x, and peak allocation 1.25x, along with task-specific quality and ablation requirements.
The longer screen used Python 3.12.3 and PyTorch 2.8.0+cu128 on an NVIDIA A40. Reproduction command for the paired Dyck-2 candidate, from the repository root, is `python3 projects/learned-depth-routing/five_task_screen.py --task dyck2 --model headwise_projected --seed 101 --steps 600 --batch-size 128 --eval-examples 2048 --eval-batch-size 256 --sequence-length 128 --width 128 --depth 8 --heads 4 --mlp-width 256 --learning-rate 0.001 --training-seconds 140 --device cuda --data-dir data/tinystories-hcrr-v1 --output <new-run-directory>/result.json`; baseline uses the same command with `--model baseline`. The screen was launched through `scripts/sprout-measure --require-idle-gpu` with a 180-second process timeout and 10-second kill grace. Preserve the original immutable result paths when reproducing.
The local three-seed confirmation can be reproduced with `python3 projects/learned-depth-routing/depth_routing_lm.py --depths 8 --steps 2000 --batch-size 16 --seeds 101 202 303 --models baseline headwise_projected --output artifacts/mac/2026-09-02-headwise-depth-routing-confirmation/results.json`; depth transfer changes `--depths 8` to `--depths 12` and uses the depth-12 result path. This local run used Python 3.12.13, PyTorch 2.13.0, macOS MPS FP32. These local results are distinct from the one-seed RunPod A40 task study.
03 What happened?
Three-seed byte-level WikiText tests found mean test improvements of 0.0293 BPB at depth 8 and 0.0207 BPB at depth 12, with consistent routing ablation effects. A longer paired five-task A40 study found causal quality gains on eligible Dyck-2 and TinyStories runs, but both missed batch-one latency limits; Dyck-2 also failed training-time and peak-memory limits.
Validation exact match (fraction) · scale shown from 0 to 1
How to read thisThe paired candidate gained 9.47 percentage points in validation exact match while taking 1.33994x the baseline training time, above the 1.25x gate.
The limitOne seed and one eligible synthetic task; HCRR also exceeded latency and peak-memory guards. The chart covers Dyck-2 only.
Paired RunPod A40 seed-101 baseline and HCRR candidate, each trained for 600 steps. Score is validation exact-match fraction (higher is better); training time is seconds.
Full results and data table
| Model | Validation exact match (fraction) | Training time (s) |
|---|---|---|
| Baseline | 0.8134765625 | 18.46658294799272 |
| HCRR headwise projected | 0.908203125 | 24.744115907000378 |
At depth 8, baseline and HCRR means were 2.9800 and 2.9500 validation BPB, and 2.9855 and 2.9562 test BPB. All three paired seeds improved test BPB; the validation improvements were 0.0198, 0.0399, and 0.0304. Disabling routing raised validation BPB by 0.2137–0.2681. HCRR added 0.44% parameters; mean training time was 1.145x baseline and eager batch-one latency 1.300x. At depth 12, mean validation BPB was 2.8734 versus 2.8948, and test BPB 2.8750 versus 2.8956. Test improvements were 0.0062 worse for seed 101, then gains of 0.0383 and 0.0299 for seeds 202 and 303; the mean test gain was 0.0207. Routing-disabled penalties ranged from 0.3036 to 0.4991 BPB. Mean training-time and latency ratios were 1.187x and 1.301x.
In the longer paired screen, Dyck-2 baseline validation exact match was 0.813477; HCRR reached 0.908203 (+0.094727), while validation answer NLL fell by 0.067495. Test exact match reached 0.911133. Its routing-disabled validation exact match was 0.847656, with a 0.032141 NLL penalty. However, training time was 1.339940x, batch-one latency 1.577067x, and peak CUDA allocation 2.608784x baseline, failing three cost guards. TinyStories validation NLL improved by 0.041537 to 2.957975 and test NLL by 0.037349 to 2.989672; routing ablation worsened validation NLL by 0.579032. Training time was 1.045256x and memory 1.000053x, but latency was 1.541564x, above the 1.35x guard. MQAR, ListOps, and permutation remained below the baseline eligibility band, so no candidate results exist for those tasks. Suite decision: investigate.
04 What does this mean?
What this suggests
- Per-head selection and projection can improve a small byte-level causal model at two tested depths.
- Ablation effects show the learned routes affect outcomes on the confirmed WikiText runs and the eligible longer-screen tasks.
- The current eager implementation has a material inference cost that blocks the suite promotion rule.
What we still don’t know
- Whether the gains persist in larger standard-tokenizer language models and across multiple seeds on the five-task screen.
- Whether fusion and activation-allocation changes can reduce latency and memory while preserving exact candidate outputs.
- Whether routing remains useful at fixed compute-to-quality rather than at fixed token budgets.
Read the full discussion
The local WikiText results establish a repeatable quality signal across paired seeds and two depths, with causal ablations showing that the learned routes are used. Their cost remains visible even after caching normalized updates and replacing generic einsums: inference latency is roughly 30% higher in eager execution. Depth 12 has greater seed variance, and one seed slightly regresses on test loss.
The longer screen adds a more demanding cost lens but has only one seed. Its quality-positive Dyck-2 and TinyStories candidates do not pass the suite gates because both exceed latency limits; Dyck-2 also exceeds training and memory guards. The three ineligible synthetic tasks are negative evidence about baseline readiness, not evidence of candidate quality. The next registered step is an implementation-only diagnostic that freezes the accepted checkpoints/configurations, profiles activation allocation, batches or fuses routing, and requires exact output parity before remeasuring cost.
05 Limitations
The local language model is a small byte-level MPS proxy; the RunPod study has one seed and eager CUDA implementation. Neither establishes performance on a standard large subword model or a fused production kernel.
Read all study limitations
The local study uses a tiny byte-level model, short fixed token budgets, one device family, and eager MPS timings. The longer study uses one seed, a specific RunPod A40 environment, task-specific horizons, and eager CUDA execution. Its two candidate tasks do not establish suite-wide gains. No fused production kernel, convergence-at-fixed-compute comparison, cached decoding, or large standard-tokenizer model was evaluated. Checkpoints are absent from the selected local depth confirmation and RunPod longer candidate/baseline directories; result records and system samples exist but are not weights.
06 References & materials
Keep the complete record
Take the research with you.
The full write-up, source code, and recorded experiment materials.
Download everything together or choose individual files.
The five-task script imports its sibling depth_routing_lm.py; both are included. Runtime dependencies are Python 3.12.3, NumPy, and PyTorch 2.8.0+cu128 for the RunPod CUDA study. The paired Dyck-2 records, environment JSON, and system samples are preserved under both longer-run directories; the code hash is recorded in each result. No trained HCRR checkpoints are present in those selected directories; result JSON and system samples are measurements, not weights. The TinyStories directory is an input dependency for other tasks, but Dyck-2 generation is synthetic and does not require it.