DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections

Xiaocong Yang; Yuren Hao; ChengXiang Zhai

cs.CLOct 14, 2026

DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections

Authors: Xiaocong Yang, Yuren Hao, ChengXiang Zhai

Abstract

We propose DAGFormer, a Transformer language model whose cross-layer connections are planned by a jointly trained topology predictor. Beyond the residual connections in a standard Transformer, DAGFormer adds direct connections from earlier layers to later ones, and the predictor controls the strength of every connection based on input tokens. Unlike existing cross-layer routing architectures, which set each layer's connections separately, DAGFormer plans the routing of the whole model at once. On matched OLMo-2 pairs from 75M to 1B parameters, DAGFormer lowers WikiText-2 perplexity by 17% to 25% over the vanilla Transformer, and a 1B DAGFormer trained on 5B tokens outperforms the vanilla model trained on 10B. Against four routing baselines with the same data and budget, it reaches the lowest perplexity at 75M, 300M and 1B. The gain carries over to 13 commonsense, science, reading and math evaluation tasks (2.0 to 2.5 points higher mean accuracy at 300M and 1B) and to instruction finetuning (63% to 70% AlpacaEval win rate against the vanilla Transformer, leading every routing baseline), and under structured pruning DAGFormer degrades the least in five of six settings.

Figures & tables

Figure 1: DAGFormer. For token

x_{t}

the predictor outputs the planned graph as four triangular matrices

\hat{A}^{Q}

\hat{A}^{K}

\hat{A}^{V}

and

\hat{A}^{R}

that collect the planned weights

\hat{\boldsymbol{\alpha}}^{(l)}

of Eq. 1 : one column per earlier output and, per layer, one row per head for

Q

K

and

V

(the weights

\hat{\alpha}^{(l)}_{q,h,j}

) and a single row for

R

(the weights

\hat{\alpha}^{(l)}_{r,j}

). Only the edges into layer 5, its outlined rows, are drawn, as rings on its heads (value in front; key, query and residual behind) with edge width proportional to the magnitude of the entry; every routed layer reads along its own rows in the same way, after the auxiliary local refinement of Eq. 2 (not drawn for demonstration). All weights shown are real values from a trained DAGFormer on training data.

	baselines (WikiText-2 NLL $\downarrow$ )
size	vanilla	DenseFormer	hyper-conn.	mHC	MUDDFormer	DAGFormer	gap to vanilla
75M	4.997	5.027	4.964	4.913	4.974	4.711	+0.29
150M	4.268	4.246	4.220	4.209	4.041	4.085	+0.18
300M	3.808	3.745	3.692	3.652	3.649	3.544	+0.26
1B, 5B tokens	3.569	3.513	3.492	3.438	3.423	3.307	+0.26
1B, 10B tokens	3.446	–	–	–	–	3.219	+0.23

Table 1: Held-out WikiText-2 NLL (nats per token); every model of a row shares recipe, data, token budget and seed. Gap to vanilla: vanilla

-

DAGFormer. Bottom: measured prefill time per token relative to the vanilla Transformer (a forward pass over 8 sequences of 1024 tokens); training cost in Appendix A . Bold: best per row; underlined: second best (not marked at 1B, 10B tokens, with two models). –: not run (the baselines are trained at 5B tokens only due to computing budget limit).

		mean accuracy (%) per kind of task $\uparrow$
size	model	commonsense (5)	science (4)	reading (2)	math (2)	Mean	mean rank $\downarrow$
300M	vanilla	39.8	36.5	38.1	11.7	34.2	4.15
	DenseFormer	38.9	36.5	40.8	12.7	34.4	4.23
	hyper-conn.	38.5	38.5	38.3	12.3	34.4	3.69
	mHC	39.2	38.8	41.8	12.9	35.4	2.54
	MUDDFormer	39.7	38.6	40.9	11.5	35.2	3.08

Table 2: Base models on the evaluation tasks of Section 3.2 : mean accuracy (%) over the commonsense (HellaSwag, PIQA, WinoGrande, CommonsenseQA, Social IQa), science (ARC-Easy/Challenge, OpenBookQA, SciQ), reading (BoolQ, LAMBADA) and math tasks (MathQA, GSM8K; close to chance at these sizes), and over all 13 (Mean). Mean rank: rank within the block, averaged over the 13 tasks (1 = best). Bold: best per block; underlined: second best per kind of task (neither marked nor ranked in the two-model block). 1B (5B), 1B (10B): pretrained on 5B and 10B tokens. Per-task accuracies in Table 8 .

		+Alpaca		+SmolTalk
size	model	NLL $\downarrow$	win $\uparrow$	NLL $\downarrow$	win $\uparrow$
300M	vanilla	2.392	–	1.760	–
	DenseFormer	2.342	0.55	1.722	0.57
	hyper-conn.	2.298	0.59	1.684	0.58
	mHC	2.277	0.58	1.664	0.62
	MUDDFormer	2.256	0.63	1.639	0.64

Table 3: Instruction following after finetuning. NLL: held-out negative log-likelihood on the instruction data. Win: AlpacaEval ( Li et al., 2023 ) win rate against the vanilla Transformer of the same size and recipe, both answer orders judged by gpt-5.6-sol , ties half, 249 of the 805 prompts. 1B (5B), 1B (10B): the 1B models pretrained on 5B and 10B tokens.

Figure 2: Scaling of DAGFormer (solid) against the matched vanilla OLMo-2 Transformer (dashed). (a) Final WikiText-2 loss against backbone parameters, with power laws whose irreducible loss is shared by the two families (coefficients in the panel); the 1B runs, trained on fewer tokens per parameter, appear only in (b). (b) Loss against training FLOPs with the routing cost included, fitted likewise: the 1B DAGFormer at 5B tokens is below the 1B vanilla Transformer at 10B tokens at about the same compute. (c) Held-out loss along training for the four matched pairs: the gap opens early and persists.

			held-out NLL $\downarrow$
size	model	params	WikiText	Math	GSM8K	balanced
75M	vanilla	77M	4.997	4.271	4.242	5.167
	vanilla, hidden 656	107M	4.769	4.057	4.052	4.961
	vanilla, 13 layers	106M	4.871	4.215	4.174	5.094
	DAGFormer	106M	4.711	3.984	3.973	4.937
150M	vanilla	153M	4.268	3.599	3.608	4.471

Table 4: Capacity controls: vanilla Transformers given DAGFormer’s routing parameters as extra width or, at 75M, as extra depth, trained with the pair’s recipe, data, tokens and seed (12B’ corpus). Params: total, including the predictor and refinement. Math: MathInstruct; balanced: the balanced Dolma set. Bold: best per size and column.

	75M			150M			300M
	NLL $\downarrow$	damage (%) $\downarrow$		NLL $\downarrow$	damage (%) $\downarrow$		NLL $\downarrow$	damage (%) $\downarrow$
model	0%	50%	70%	0%	50%	70%	0%	50%	70%
vanilla	2.387	9.0	17.9	1.854	9.9	21.2	1.676	7.4	17.8
DenseFormer	2.383	9.5	18.7	1.830	9.5	20.7	1.641	7.9	18.5
hyper-conn.	2.306	9.5	19.5	1.804	9.7	21.4	1.604	7.7	17.8
mHC	2.328	9.0	20.1	1.813	9.5	21.0	1.585	8.4	18.4

Table 5: Pruning during domain finetuning on MathInstruct, every model of Table 1 from its own checkpoint (75M and 150M: 12B’ corpus; 300M: 21B): held-out NLL of the unpruned finetune (0%) and the damage (loss increase relative to it) at 50% and 70% backbone sparsity. Bold: best per column. All sparsities and the frozen-router controls: Table 23 .

Appendix figures & tables21 assets

Supplementary material from the paper’s appendix.

Appendix

	training: time per token relative to vanilla $\downarrow$						prefill
size	vanilla (tok/s)	DenseFormer	hyper-conn.	mHC	MUDDFormer	DAGFormer	vanilla (tok/s)
75M	108.8K	1.02	1.36	1.63	1.17	1.44	851K
150M	79.7K	1.03	1.50	1.66	1.21	1.57	474K
300M	49.1K	1.03	1.59	1.68	1.28	1.77	245K
1B	18.3K	1.03	1.56	1.51	1.19	1.65	76K

Appendix

Table 6: Measured training cost: time per token of a forward and backward pass relative to the vanilla Transformer of the same size (every model compiled, bf16 weights, 2 sequences of 1024 tokens, one A100-40GB, median of 20 iterations, fastest of two runs; DAGFormer in the faster of its two computation orders), with the vanilla Transformer’s throughput in tokens per second for training and prefill. Prefill ratios are in Table 1 .

300M, 21B corpus
		MC acc.	bits per byte $\downarrow$
model	stage	$\uparrow$	GSM8K	WikiText
vanilla	base	0.369	1.458	1.128
DAGFormer	base	0.391	1.344	1.057
vanilla	+Alpaca	0.371	1.452	1.184
DAGFormer	+Alpaca	0.390	1.295	1.108

Appendix

Table 7: Downstream evaluation before (base) and after each instruction-finetuning recipe, summarized (Table 9 gives per-task accuracy). MC acc.: mean accuracy over the twelve tasks (eleven multiple-choice and LAMBADA); GSM8K and WikiText: bits per byte on the gold GSM8K solutions and on WikiText (lower is better). One pretraining run per model. Repeating the 300M instruction finetuning with a second seed moves the 12-task mean by at most 0.001 and WikiText bpb by less than 0.001 (BoolQ alone by up to 0.02), and the routed-minus-vanilla gaps reproduce within 0.003, so the instruction-finetuning stage adds little noise; pretraining-seed variance is not measured. Shared-slice 75M and 150M pairs show the same pattern (Appendix C ).

		accuracy (%) $\uparrow$
size	model	ARC-e	ARC-c	HSwag	PIQA	WinoG	OBQA	SciQ	BoolQ	LAMB	CSQA	SIQA	MathQA	GSM8K	Mean
300M	vanilla	38.1	20.6	27.4	61.2	52.2	24.6	62.5	56.7	19.5	20.5	37.7	22.4	1.0	34.2
	DenseFormer	37.8	21.2	27.3	60.8	51.1	22.4	64.5	60.8	20.7	19.7	35.9	22.9	2.5	34.4
	hyper-conn.	39.6	22.5	27.8	60.8	48.6	25.8	65.9	55.0	21.6	19.7	35.8	21.6	3.0	34.4
	mHC	38.7	22.9	27.6	61.2	50.0	26.2	67.4	61.3	22.2	19.0	38.0	23.3	2.5	35.4
	MUDDFormer	38.3	22.6	28.5	60.8	52.2	25.0	68.4	58.7	23.0	19.7	37.3	22.0	1.0	35.2

Appendix

Table 8: Per-task accuracy (%) of the base models in Table 2 , normalized for ARC, HellaSwag, PIQA, OpenBookQA and MathQA, GSM8K 3-shot exact match with flexible numeric extraction on the first 200 test problems; Mean: over the 13 tasks. Between the vanilla Transformer and DAGFormer: the baselines of Table 1 . Bold: best per block. 1B (5B), 1B (10B): the 1B models pretrained on 5B and 10B tokens.

			accuracy (%) $\uparrow$
size	stage	model	ARC-e	ARC-c	HSwag	PIQA	WinoG	OBQA	SciQ	BoolQ	LAMB	CSQA	SIQA	MathQA	GSM8K	Mean
300M	base	vanilla	38.1	20.6	27.4	61.2	52.2	24.6	62.5	56.7	19.5	20.5	37.7	22.4	1.0	34.2
		DenseFormer	37.8	21.2	27.3	60.8	51.1	22.4	64.5	60.8	20.7	19.7	35.9	22.9	2.5	34.4
		MUDDFormer	38.3	22.6	28.5	60.8	52.2	25.0	68.4	58.7	23.0	19.7	37.3	22.0	1.0	35.2
		DAGFormer	38.8	22.1	28.2	61.7	52.3	24.4	73.7	60.4	26.5	19.6	37.8	23.1	2.5	36.2
	+Alpaca	vanilla	37.1	21.8	27.7	60.6	53.7	25.6	59.3	59.1	20.8	20.0	37.4	22.6	2.5	34.5

Appendix

Table 9: The evaluation tasks of Table 2 at every stage: the base models and the same models after each instruction-finetuning recipe. The 13-task mean falls after SmolTalk for every model, vanilla and routed alike, while the same models fit the instruction data better (Table 3 ) and their bits per byte on the gold GSM8K solutions improve (1.344 to 1.097 for the 300M DAGFormer, 1.242 to 1.038 at 1B): instruction finetuning moves the output distribution away from zero-shot multiple choice, and the loss concentrates in LAMBADA, SciQ and BoolQ. This is why the body compares instruction-finetuned models on instruction following instead. Bold: best per stage.

SFT	DAGFormer vs. vanilla	win rate $\uparrow$ [CI]	W/L
Alpaca	300M vs. 300M, seed 1	0.57 [0.50, 0.64]	16/10
Alpaca	300M vs. 300M, seed 2	0.61 [0.54, 0.67]	19/8
Alpaca	1B@5B vs. 1B@5B	0.58 [0.51, 0.66]	26/12
Alpaca	1B@5B vs. 1B@10B	0.56 [0.48, 0.64]	23/17
Alpaca	1B@10B vs. 1B@10B	0.56 [0.48, 0.63]	19/14
SmolTalk	300M vs. 300M, seed 1	0.69 [0.63, 0.75]	27/2

Appendix

Table 10: Judged chat quality on the 80 MT-Bench first turns. Win rate of DAGFormer against the vanilla Transformer, judged by Qwen3-Coder-30B-A3B-Instruct ( Yang et al., 2025 ; Qwen Team, 2025 ) with the MT-Bench pairwise prompts (with the reference answer for math, reasoning and coding), with ties counted as half, averaged over both answer orders, with a 95% bootstrap interval; W/L counts prompts where both orders agree. “Seed 1/2” are two independent instruction-finetuning runs of both models from the same pretrained checkpoints. The judge prefers DAGFormer in every comparison; the interval excludes 0.5 in seven of ten.

		+Alpaca		+SmolTalk
size	routed model	win rate	W/T/L	win rate	W/T/L
300M	DAGFormer	0.65 † [0.60, 0.70]	124/69/55	0.70 † [0.65, 0.75]	138/69/41
	MUDDFormer	0.63 † [0.58, 0.68]	117/78/53	0.64 † [0.59, 0.69]	120/73/55
	DenseFormer	0.55 † [0.50, 0.61]	91/84/73	0.57 † [0.52, 0.63]	101/81/66
	hyper-conn.	0.59 † [0.54, 0.64]	109/80/60	0.58 † [0.53, 0.64]	106/76/66
	mHC	0.58 † [0.53, 0.63]	100/89/60	0.62 † [0.57, 0.67]	118/73/57

Appendix

Table 11: Every judged AlpacaEval comparison, with the intervals and counts behind the win rates of Table 3 : AlpacaEval ( Li et al., 2023 ) win rate of each routed model against the vanilla Transformer of the same size and recipe, on a stratified 249-prompt subset of the 805 (sampling seed 1234). Every pair is judged in both answer orders by gpt-5.6-sol (2026-10-09); the win rate counts wins plus half the ties, averaged over the two orders and then over prompts, with a 95% bootstrap interval. W/T/L counts a prompt as a win or a loss only when both orders agree, so disagreements become ties.

\dagger

: the interval excludes 0.5. Last row: DAGFormer trained on 5B tokens against the vanilla Transformer trained on 10B, at about equal compute.

run	size	WikiText-2	MathInstruct	GSM8K	balanced
vanilla	75M	4.997	4.271	4.242	5.167
vanilla, hidden 656	75M	4.769	4.057	4.052	4.961
vanilla, 13 layers	75M	4.871	4.215	4.174	5.094
DenseFormer	75M	5.027	4.303	4.249	5.196
hyper-connections (static)	75M	5.042	4.301	4.246	5.191
hyper-connections (DHC $\times$ 4)	75M	4.964	4.202	4.157	5.114

Appendix

Table 12: Post-hoc baselines of Table 1 on every common held-out set (NLL, nats per token). The parameter-matched vanilla Transformers carry DAGFormer’s routing parameters as extra width (hidden 656 and 864; 107M and 182M parameters) or, at 75M, as extra depth (13 layers, 106M); hyper-connections are the static 4-stream variant. All runs of a size share recipe, data, tokens and seed; the 75M and 150M runs use the 12B’ corpus and the 300M and 1B runs the 21B corpus.

75M, WikiText-2 NLL $\downarrow$	2.5e-4	5e-4	1e-3	2e-3
vanilla	5.878	4.997	4.654	4.563
DAGFormer	5.558	4.711	4.413	4.329
gap	0.32	0.29	0.24	0.23

Appendix

Table 13: Learning-rate sweep at 75M (12B’ corpus, 1.58B tokens; the predictor’s learning rate is 0.6

\times

the backbone’s). Both models improve up to the largest rate; the matched-rate gap moves from 0.32 to 0.23 nats and the best-against-best gap is 0.23. The same holds on MathInstruct, GSM8K and the balanced set (best-against-best narrows the 5e-4 gap by 5 to 19%; 30% on the Dolma tail).

routing source (75M)	routing	held-out NLL $\downarrow$				tasks
	params	own	WikiText	Math	GSM8K	acc. $\uparrow$
none (vanilla)	0	4.213	6.593	5.373	5.404	25.7
static table, no input	500	4.171	6.529	5.243	5.279	n/a
per-layer predictors	35.6M	4.058	6.256	5.025	5.148	26.0
predictor only	28.7M	4.031	6.207	4.981	5.116	26.2
predictor + refinement	29.1M	4.014	6.163	4.932	5.080	26.5

Appendix

Table 14: Routing-source ablation: one 75M backbone, granularity and 1.57B tokens of a 1.7B-token Dolma slice, varying only the source of

\boldsymbol{\alpha}

. Static table: one learned weight per read and source, no input; per-layer predictors: an encoder over the token ids per layer, nothing shared; predictor + refinement is DAGFormer. Own: the run’s held-out Dolma tail; Math: MathInstruct; tasks acc.: the 13-task mean of Table 2 before instruction finetuning (the static table was not run on the tasks).

size	75M	150M	300M	1B
layers / width / heads	6 / 512 / 8	8 / 768 / 12	12 / 1024 / 16	16 / 2048 / 16
MLP width	2048	3072	4096	8192
backbone params (M)	76.6	152.6	304.1	1279.4
routing params (M)	29.1	30.0	32.3	36.7
steps	3000	6000	12000	9540 (+9540)
tokens (B)	1.57	3.15	6.29	5.0 (10.0)

Appendix

Table 15: Pretraining configurations. Common to all: 1024-token sequences, 512 sequences (524K tokens) per step, AdamW

(0.9,0.95)

, weight decay 0.1, gradient clipping 1.0, linear decay to zero, bf16, tied embeddings, vocabulary 100,352 (the OLMo-2 tokenizer). Predictor: 256-d token and position embeddings, 2 causal encoder layers with 4 heads, trunk width 512, local refinement width 128. The 1B runs were continued from 5B to 10B tokens with the schedule re-stretched and the learning rate re-warmed over 500 steps. Instruction tuning: global batch 64 sequences of 1024 tokens, cosine learning rate

2{\times}10^{-4}

(

10^{-4}

at 1B), assistant-only loss, ChatML format. The 1B hyper-connections model was pretrained and instruction-tuned on four H100 GPUs, with 8 sequences per micro-batch in pretraining against 2 for the other 1B models on A100s; the 1B mHC and DenseFormer instruction tuning used 2 sequences per micro-batch instead of 4 to fit in memory. Micro-batching changes neither the global batch nor the sequences in each step.

size	corpus	WikiText-2		MathInstruct		GSM8K		balanced Dolma
		vanilla	DAGFormer	vanilla	DAGFormer	vanilla	DAGFormer	vanilla	DAGFormer
75M	12B	4.980	4.775	4.267	4.053	4.216	3.998	5.175	4.992
75M	12B’	4.997	4.711	4.271	3.984	4.242	3.973	5.167	4.937
150M	12B	4.297	4.100	3.630	3.424	3.626	3.465	4.514	4.334
150M	12B’	4.268	4.085	3.599	3.383	3.608	3.402	4.471	4.305
300M	21B	3.808	3.544	2.742	2.565	3.232	2.937	3.832	3.606

Appendix

Table 16: Held-out NLL (nats per token) of every matched pair on the four cross-corpus sets (final checkpoints): WikiText-2, MathInstruct, GSM8K and a source-balanced Dolma set (64 windows from each of 14 Dolma sources, drawn from files that no training corpus consumed; books, Wikipedia and open-web-math could not be represented). DAGFormer is lower in every cell; “seed 2” is the 300M pair repeated with a second pretraining seed.

routing source	own cache	WikiText-2	MathInstruct	GSM8K	MC w/o BoolQ
					base / Alpaca / SmolTalk
none (vanilla)	4.213	6.593	5.373	5.404	0.269 / 0.278 / 0.293
static learned table (no token input)	4.176	6.529	5.243	5.279	-
per-layer predictors (unshared)	4.058	6.256	5.025	5.148	0.271 / 0.290 / 0.304
shared predictor only	4.031	6.207	4.981	5.116	0.274 / 0.289 / 0.301
shared + local (DAGFormer)	4.014	6.163	4.932	5.080	0.276 / 0.290 / 0.302

Appendix

Table 17: The 75M router-locality arms (1.57B tokens of a 1.7B-token Dolma slice, 4 GPUs each) on every held-out set, and the 11-task multiple-choice mean (without BoolQ, whose majority-class flips move the 12-task mean by up to 20 points at this size) before and after instruction tuning. The modular variant, which also routes the intra-layer edge, is best on loss at 75M; its routing cost is 1.33

\times

the vanilla Transformer here and it is not better than FourWay routing at 150M or 300M (Appendix H ).

size	model	stage	MC acc. (12)	GSM8K acc.	GSM8K bpb	WikiText bpb
75M	vanilla	base	0.304	0.010	1.840	1.393
75M	vanilla	+Alpaca	0.309	0.010	1.801	1.459
75M	vanilla	+SmolTalk	0.309	0.015	1.613	1.595
75M	DAGFormer	base	0.316	0.025	1.759	1.347
75M	DAGFormer	+Alpaca	0.319	0.015	1.733	1.417
75M	DAGFormer	+SmolTalk	0.336	0.010	1.502	1.550

Appendix

Table 18: The original 12B-slice pairs before and after instruction tuning (the 300M DAGFormer of this slice stopped at 9000 of 12000 steps, so its margin is a lower bound). MC acc.: mean accuracy over the 12 multiple-choice tasks; GSM8K acc.: 3-shot exact match on the first 200 test problems; bpb: bits per byte of the GSM8K solutions and of WikiText.

size	setting	vanilla	DAGFormer	peak post-prune vanilla	peak post-prune DAGFormer
75M	unpruned finetune	2.388	2.223	-	-
75M	30% heads+channels	2.478	2.298	2.749	2.578
75M	50% heads+channels	2.590	2.391	2.841	2.604
75M	70% heads+channels	2.829	2.614	3.199	2.913
75M	50%, random importance	3.029	2.935	4.513	4.729
75M	33% whole blocks	2.610	2.547	4.760	4.377

Appendix

Table 19: Final held-out MathInstruct NLL after 2000 finetuning steps for every pruning run of the shared-slice pairs, and the highest NLL observed immediately after a pruning update (before recovery). Whole-block pruning removes entire attention or MLP blocks; the 75M DAGFormer at half the blocks is the one setting where it loses (Section 5 ). With random unit selection both families lose far more than with Taylor importance and the DAGFormer edge shrinks at 75M.

model	set	intact	predictor $\to$ position mean	predictor $\to$ global mean	local $\to$ mean	both $\to$ mean
75M (locality arm, both)	wikitext2	6.164	6.161	6.162	6.582	6.595
75M (locality arm, both)	mathinstruct	4.933	4.928	4.934	5.267	5.272
75M (locality arm, both)	gsm8k	5.080	5.075	5.076	5.353	5.353
75M (locality arm, both)	own	4.009	4.012	4.038	4.465	4.502
300M	wikitext2	3.544	3.547	3.604	4.251	4.297
300M	mathinstruct	2.565	2.556	2.600	3.065	3.065

Appendix

Table 20: Held-out NLL when one routing channel is replaced by its mean over 128 held-out windows of the model’s own corpus (‘own’ / ‘dolma21b’ is that corpus’s cache). The predictor’s per-token output is replaceable at every size and on every set; the local refinement channel’s is not.

SFT	comparison	writing	roleplay	reasoning	math	coding	extraction	stem	humanities
Alpaca	300M vs. 300M	0.72	0.57	0.47	0.40	0.60	0.55	0.75	0.47
Alpaca	300M mod. vs. 300M	0.80	0.53	0.42	0.50	0.50	0.72	0.55	0.68
Alpaca	1B@5B vs. 1B@5B	0.55	0.60	0.42	0.45	0.42	0.62	0.75	0.85
Alpaca	1B@5B vs. 1B@10B	0.42	0.82	0.50	0.28	0.53	0.65	0.53	0.75
Alpaca	1B@10B vs. 1B@10B	0.40	0.75	0.42	0.33	0.65	0.65	0.53	0.75
SmolTalk	300M vs. 300M	0.82	0.82	0.53	0.80	0.65	0.57	0.60	0.72

Appendix

Table 21: Per-category win rate of DAGFormer (10 prompts per category, so

\pm 0.15

is noise). “300M mod.” is the modular variant.

model	mean score (1–10)	answer tokens	ended on end-of-turn
300M vanilla / Alpaca	1.91	109	0.74
300M DAGFormer / Alpaca	2.15	109	0.72
300M modular / Alpaca	1.95	114	0.71
1B vanilla 5B / Alpaca	2.06	103	0.74
1B DAGFormer 5B / Alpaca	2.14	122	0.75
1B vanilla 10B / Alpaca	2.19	111	0.71

Appendix

Table 22: Single-answer judge scores (MT-Bench single-answer prompt, same judge), mean answer length and the fraction of answers that stopped on the end-of-turn token rather than the 256-token cap. Scores sit near 2 of 10 for every model at this scale; the pairwise win rate is the informative number.

		held-out NLL $\downarrow$ at sparsity				NLL at 50%, routers frozen
size	model	0%	30%	50%	70%	all routers	predictor only
75M	vanilla	2.387	2.475	2.602	2.815	–	–
	DenseFormer	2.383	2.482	2.610	2.829	2.640	–
	hyper-conn.	2.306	2.400	2.526	2.756	2.517	–
	mHC	2.328	2.408	2.537	2.796	2.549	–
	MUDDFormer	2.362	2.461	2.612	2.889	2.618	–

Appendix

Table 23: All matched-corpus pruning runs of Table 5 : held-out MathInstruct NLL after 2000 finetuning steps of 32K tokens (units removed between steps 200 and 1400, then 600 recovery steps) at 0 to 70% sparsity, and at 50% with every router frozen during finetuning (DenseFormer’s depth weights, the hyper-connection modules, MUDDFormer’s routers, DAGFormer’s predictor and local refinement) or, for DAGFormer, only the predictor. Bold: best per size and column.

		held-out NLL $\downarrow$				damage (%)		frozen
size	model	0%	30%	50%	70%	at 50%	at 70%	NLL at 50%
75M	vanilla	2.388	2.478	2.590	2.829	8.4	18.5	–
	DAGFormer	2.223	2.298	2.391	2.614	7.5	17.6	2.388
150M	vanilla	1.883	1.957	2.071	2.278	10.0	21.0	–
	DAGFormer	1.743	1.807	1.899	2.091	8.9	20.0	1.899
300M	vanilla	1.608	1.665	1.762	1.967	9.6	22.3	–

Appendix

Table 24: The earlier pruning study on the original 12B-slice pairs (75M to 300M) and a 150M comparison with MUDDFormer on a common corpus (150M ∗ ): held-out NLL, the performance damage relative to the same model’s unpruned finetune, and the NLL at 50% with the predictor (DAGFormer) or every router (MUDDFormer) frozen. Bold: best per block and column.

Figure 3: Pruning during domain finetuning on MathInstruct. (a) Held-out NLL relative to the same model’s unpruned finetune with 30/50/70% of heads and MLP channels removed, the original 12B-slice pairs at three sizes: DAGFormer loses less at every budget, more so at higher sparsity. (b) 150M, common corpus: DAGFormer against MUDDFormer and the vanilla Transformer, equal unpruned; DAGFormer ahead by 0.05/0.07/0.10 nats at 30/50/70%; frozen routers (crosses) change neither.

Figure 4: The modular variant against DAGFormer (FourWay routing) and the vanilla OLMo-2 Transformer, as in Figure 2 a and b: (a) final WikiText-2 loss against backbone parameters, (b) against training FLOPs with the routing cost included. On the same corpus, modular routing is 0.03 nats behind FourWay at 75M and 150M and 0.08 behind at 300M, and per FLOP it does not repay its extra cost above 75M.

Transformer

Transformers are neural architectures that use attention to model relationships among sequence elements in language, vision, and scientific data.

See all 536 papers

Directional Routing in Transformers
Mar 16, 2026Kevin TaylorAttention Routing Transformer
Adaptive Computation Depth via Learned Token Routing in Transformers
Apr 18, 2026Ahmed Abdelmuniem Abdalla MohammedEfficient Transformer Inference Transformer
Independent Researcher
WhiteMatter: All-to-All Cross-Layer Connections via KV Source Mixing
Aug 19, 2026Wenbo Zhang, Xiang RenEfficient Transformer Inference Cross-Layer Attention
University of Southern California

Mar 16, 2026J/K move · Enter open · S save

Directional Routing in Transformers

Kevin Taylor

Attention Routing Transformer LLM Routing Transformer Attention

Open paper AskPDF arXiv

We introduce directional routing, a lightweight mechanism that gives each transformer attention head learned suppression directions controlled by a shared router, at 3.9% parameter cost. We train a 433M-parameter model alongside an identical baseline in a single run, then trace the resulting circuits through mechanistic interpretability. Routing becomes the model's dominant computational pathway. Disabling it collapses factual recall to near-zero probability across all 8 test prompts and drops induction accuracy from 93.4% to 0.0%. Knocking out individual attention heads has negligible effect: the primary mover head's removal actually increases target probability, and induction heads retain 98.6% accuracy without their strongest member. The coordination mechanism is irreplaceable; the components it coordinates are not. The model also self-organizes, without explicit pressure, into two regimes: domain-adaptive routing in early layers and fixed syntactic pruning in late layers, where the least-varying layer is the most critical (+42.6 PPL when disabled). Routing reduces perplexity 31-56% relative to the baseline, though downstream multiple-choice benchmarks do not yet reflect these gains.

Abstract

Transformer

Explore similar work

Directional Routing in Transformers

Adaptive Computation Depth via Learned Token Routing in Transformers

WhiteMatter: All-to-All Cross-Layer Connections via KV Source Mixing