DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections

AI summary

Overview

Research area: Natural Language Processing / large language model architecture, specifically cross-layer connectivity and learned routing inside Transformer language models.

Technical level: Advanced — it assumes familiarity with residual streams, attention heads, multi-head projections, pretraining scaling laws, structured pruning, and instruction tuning.

Scope: The paper introduces DAGFormer, a Transformer variant in which a small separately trained network predicts the entire cross-layer connection graph from the input tokens before the backbone runs, and it validates that design with a matched scaling study from 75M to 1B parameters.

What This Paper Is About

In a standard Transformer, every attention head and MLP reads the same running residual sum, so which earlier computation a late layer can see is fixed by design. Recent architectures add learned, input-dependent cross-layer connections, but each of them routes locally: every layer carries its own router that reads only the hidden state entering that layer, so the overall topology is assembled greedily one layer at a time. DAGFormer replaces this with a single predictor that reads only the input tokens and plans the connections for all layers at once, turning the routing topology into one global object that can be inspected, frozen, averaged, or swapped.

Key Contributions

1. DAGFormer architecture. A Transformer built on an OLMo-2 backbone in which every attention head forms its query, key and value from weighted mixtures of the embedding and all earlier layer outputs (and the residual input is such a mixture too). The mixing weights are signed, non-softmax-normalized coefficients over a weighted directed acyclic graph. A separate causal-encoder topology predictor with its own randomly initialized 256-dimensional token table predicts all of these weights from the token ids alone and is pretrained jointly with the backbone from scratch. A zero-initialized auxiliary local refinement lets each layer correct its own planned rows from its layer input.

2. Matched scaling study from 75M to 1B parameters. Every comparison uses paired models with identical backbone, tokenizer, data order, token budget and optimizer, with DenseFormer, dynamic hyper-connections (four streams), mHC and MUDDFormer trained on the same pipeline, and with routing cost charged in every compute comparison. The study is carried through instruction finetuning and AlpacaEval/MT-Bench evaluation.

3. Ablations of the routing source. Experiments isolating where the mixing weights come from (predictor alone, refinement, independent per-layer predictors, static weights, per-position mean of the predicted graph), plus capacity controls giving vanilla Transformers the routing parameters as extra width or depth.

4. Structured-pruning study. Every model is finetuned on MathInstruct with attention heads and MLP channels permanently removed on a gradual cubic schedule, with routers frozen in control runs, showing that DAGFormer degrades the least in five of six settings without re-planning.

Main Findings

• Lower language-modeling loss at every size except one. Against the matched vanilla OLMo-2 Transformer, DAGFormer is 0.18 to 0.29 nats lower on held-out WikiText-2 NLL (17% to 25% lower perplexity) at 75M, 150M, 300M and 1B. Concretely: 4.711 vs 4.997 (75M), 4.085 vs 4.268 (150M), 3.544 vs 3.808 (300M), 3.307 vs 3.569 (1B, 5B tokens), and 3.219 vs 3.446 (1B, 10B tokens).

• It beats the routing baselines at every size but 150M. DAGFormer has the lowest NLL at 75M (4.711), 300M (3.544) and 1B at 5B tokens (3.307). At 150M, MUDDFormer is lower (4.041 vs 4.085), a 0.04-nat gap the authors flag as a single-seed result.

• A smaller-budget model beats a larger-budget vanilla model. The 1B DAGFormer trained on 5B tokens has lower loss (3.307) than the 1B vanilla Transformer trained on 10B tokens (3.446) at roughly the same training compute.

• Gains carry to 13 downstream tasks. The evaluation set covers commonsense (5 tasks), science (4), reading (2) and math (2). Mean accuracy is 2.0 to 2.5 points above vanilla and 0.7 to 0.8 points above the strongest routing baseline: 36.2 vs 34.2 at 300M, 38.1 vs 35.7 at 1B (5B tokens), and 39.1 vs 36.6 at 1B (10B tokens). DAGFormer also has the best mean rank at both tested sizes (2.38 at 300M, 2.08 at 1B with 5B tokens). On math, mHC leads by 0.1 points at 300M. At 1B with 10B tokens DAGFormer is ahead of vanilla on all 13 tasks.

• Gains carry to instruction following. After finetuning, DAGFormer has the lowest held-out instruction NLL at every size under both recipes (0.15 to 0.19 nats below vanilla, 0.04 to 0.06 below the strongest routing baseline). Its AlpacaEval win rate against the matched vanilla model is 0.63 to 0.70, above every routing baseline (0.46 to 0.64). On 80 MT-Bench first turns judged pairwise, it wins 0.56 to 0.69 of comparisons, with the 95% interval excluding 0.5 in seven of ten.

• The predictor, not per-token re-planning, carries most of the gain. In a 75M routing-source ablation, the predictor alone — which never sees a hidden state — recovers about 90% of DAGFormer's gain over vanilla; adding the refinement gives the best result on every held-out set. Independent per-layer predictors reading the same tokens recover under 80% while adding 7M more parameters. Replacing the predictor output with its per-position mean costs almost nothing, implying the learned graph itself, not its input dependence, is where the benefit sits.

• Extra parameters explain only part of the gain. Vanilla Transformers given DAGFormer's routing parameters as extra width close 70 to 89% of the gap at 75M but only 4 to 42% at 150M; extra depth closes less.

• Best degradation under structured pruning. With 50% of heads and channels removed, DAGFormer degrades less than vanilla and every routing baseline at every size, and stays about as good as or better than the unpruned vanilla model. At 70% sparsity it degrades least at 150M and 300M; at 75M three baselines lose less, though DAGFormer's pruned loss is still the lowest. MUDDFormer loses the most at every size and sparsity. Freezing the predictor or the baselines' routers during pruning barely changes results.

• Scaling fits favor the architecture, with caveats. Fitting L = E + A·N^−α and L = E + K·C^−γ with a shared irreducible loss gives larger exponents for DAGFormer (0.56 vs 0.48 for parameters, 0.30 vs 0.27 for compute). The compute multiplier is 1.2 to 1.4× at 75M and 150M but 1.7 to 2.0× at 300M and 1B, extrapolating to 1.9× at 10^20 FLOPs and 2.3× at 10^21 FLOPs — though the exponents differ by about one standard error in each fit.

• Routing costs more compute than most baselines. DAGFormer adds 1.18× the vanilla FLOPs per token at 75M rising to 2.13× at 1B in the training order, and takes 1.5 to 1.9× the prefill time (the predictor itself is 1 to 3% of that). Training time is 1.44 to 1.77× vanilla. This is more than MUDDFormer (1.4 to 1.5× prefill) and, from 300M, mHC (1.5 to 1.8×), but less than hyper-connections (1.9 to 2.1×); DenseFormer adds nothing measurable. Peak inference memory is 1.9 vs 1.8 GB at 75M and 9.1 vs 4.2 GB at 1B.

Methodology in Plain English

Start with a normal OLMo-2 Transformer. Then, before the backbone processes a token, hand the token ids to a small separate encoder — two causal Transformer layers, four heads, its own 256-dimensional embedding table, a shared 512-dimensional trunk — and let it output a set of signed numbers. Those numbers say, for every layer, every attention head and every earlier layer output, how much of that earlier output to mix in when building the query, the key, the value, and the residual. Because the weights are signed and not softmax-normalized, a head can amplify one source or subtract one from another.

The predictor's output layer starts at zero with a bias that selects only the immediately preceding layer, so at initialization the model is exactly the vanilla Transformer and training starts from a familiar point. The predictor and backbone train jointly from scratch on the language-modeling loss. Each layer also has a tiny correction network that adjusts its own rows of the graph using what the backbone has actually computed, initialized to zero so it starts inert.

Every layer runs its fused query/key/value projection once over the stacked earlier outputs and then mixes in head space, which is cheaper than it sounds because projection and mixing commute (the projections have no bias). The authors verify the two computation orders give the same logits to within 10^−4 in fp32.

For evaluation, they train matched pairs using identical recipes, data order, budgets and seed (42), at 75M (6 layers, width 512, 8 heads), 150M (8, 768, 12), 300M (12, 1024, 16) and 1B (16, 2048, 16) on Dolma v1.7. The 75M/150M/300M runs train about 21 tokens per parameter; the 1B pair trains 5B tokens and is continued to 10B with a re-warmed schedule. Routing adds 29M to 37M parameters, 25.7M of which is the predictor's token embedding. They compare against DenseFormer, four-stream dynamic hyper-connections, mHC and MUDDFormer, then measure WikiText-2, MathInstruct, GSM8K-solution and balanced-Dolma held-out NLL, twelve lm-evaluation-harness tasks plus GSM8K, and instruction-tuned NLL plus AlpacaEval win rates judged by an LLM.

Why This Matters

Impact on research. The paper reframes cross-layer routing from a per-layer, greedy decision into a global prediction problem that a single module solves before the backbone runs. That makes the topology a first-class object that can be averaged, frozen, patched and audited — the authors use exactly these interventions to locate where the gain comes from. It also gives a clean, matched scaling comparison across four routing designs that are usually evaluated on incomparable setups, and it charges routing cost in every compute comparison rather than reporting FLOP-free wins.

Real-world applications:

• Training more capable small models under fixed compute budgets, where a 1B model at 5B tokens matches a vanilla model trained on twice the tokens.
• Deploying compact models after aggressive structured pruning to 50% backbone sparsity, where DAGFormer loses the least and stays at or above the unpruned vanilla model's quality.
• Cheaper instruction tuning of small assistants: DAGFormer models win 0.63 to 0.70 of AlpacaEval comparisons against their vanilla counterparts of the same size and recipe.
• Auditing or steering an inference pass: because the predicted graph is computed from the tokens before the backbone runs, its output is a natural intervention point for interpretability work.

Industry relevance. The result matters most for teams that must trade quality against serving cost. DAGFormer's quality-per-token advantage is real, but so is its price: 1.5 to 1.9× vanilla prefill time and 9.1 vs 4.2 GB peak memory at 1B, plus dependence on the refinement's per-token corrections, meaning trained models cannot simply be frozen into a fixed topology without retraining. The paper also shows asymmetric cost across computation orders (project-then-mix is faster up to 300M, mix-then-project at 1B), which is directly actionable for kernel-level engineering.

Future Directions

• Distilling the learned graph into a fixed topology that a dense kernel can execute, since post-hoc graph sparsification is costly and trained models rely on the refinement's per-token corrections.
• Testing whether the shared predictor keeps its lead at larger scale, in particular whether it can plan routing that no single layer can see; nothing beyond 1B parameters is tested.
• Making the input-dependent part of the graph human-nameable. The authors read the graph only in aggregate, and the per-token variation in trained models lives mostly in the refinement; they suggest a larger predictor-only DAGFormer whose graph varies by task or domain, and routing among specialized modules, as a route to a readable topology.
• Settling the MUDDFormer comparison and seed noise. MUDDFormer beats DAGFormer at 150M and loses at the other sizes, each a single seed; repeating the 300M pair shows seed variation of up to 0.08 nats on GSM8K solutions (at most 0.03 on WikiText-2), so smaller differences are noise. Only the 75M pair was swept over learning rates.

Target Audience

Architecture researchers and graduate students working on Transformer variants, cross-layer connectivity and conditional computation; practitioners who train small language models from scratch and need to weigh quality against FLOPs, prefill latency and memory; compression researchers interested in how learned connectivity interacts with structured pruning; and interpretability researchers looking for a routing signal that is computed from the input before the backbone runs. The paper is not beginner-friendly — it assumes comfort with residual streams, attention projections, scaling-law fitting and pruning schedules.

Authors’ 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.