PaperlayerView exact revision on arXiv

Source-linked AI summary

DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections

Xiaocong Yang, Yuren Hao, ChengXiang Zhai

arXiv:2610.19625v1cs.CLcs.AIcs.LG

TL;DR

DAGFormer addresses the lack of global planning in dynamic cross-layer routing by using a jointly trained predictor to plan connections across the model. It improves on the vanilla Transformer across model sizes and carries gains to evaluation, instruction finetuning, and structured pruning.

  • Problem

    Existing dynamic routing sets connections locally, so cross-layer coordination cannot be planned explicitly from the input and the topology cannot be inspected as an object.

    Source: §1 · p. 1

  • Method

    A jointly trained topology predictor reads input tokens to plan a global graph of weighted cross-layer connections, which the Transformer backbone then follows.

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

  • Results

    DAGFormer lowers loss at every size from 75M to 1B, beats four routing baselines at every size but 150M, carries gains to evaluation and instruction finetuning, and degrades least under structured pruning in five of six settings.

    Source: §1 · p. 2, §8 · pp. 8–9

  • Takeaways & Limitations

    The learned graph accounts for most of the gain and is an object that can be inspected and intervened on.

    Source: §6 · p. 7, §8 · p. 9

    The largest model is below its compute-optimal training budget, most pretraining runs use a single seed, and trained models rely on per-token refinement rather than a fixed topology at inference.

    Source: §6 · p. 7, Limitations · p. 9

Abstract

from arXiv · show

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.

1 Introduction

Existing dynamic cross-layer architectures route locally, leaving no module to plan or inspect the topology as a whole. DAGFormer uses an input-token-based predictor to plan global routing, and its matched-model results show gains over the vanilla Transformer and most routing baselines.

1 Introduction

  • The residual stream shares one running sum, so a late layer sees an early result only within a sum that intermediate layers have also modified.This makes which earlier computation a layer can read fixed by design.Source: §1 · p. 1
  • Unlike local routers that assemble connections layer by layer, DAGFormer predicts the global routing topology from input tokens before the backbone processes them.Each attention head and layer gains direct connections to earlier outputs, alongside the residual connections.Source: §1 · p. 1
  • DAGFormer is better than the matched vanilla Transformer at every size and equal compute, and beats four routing baselines at every size but 150M.A 1B DAGFormer also beats a 1B vanilla Transformer trained on twice as many tokens.Source: §1 · p. 2
  • The paper studies matched OLMo-2 pairs from 75M to 1B, compares four routing baselines on the same pipeline, and evaluates instruction finetuning and structured pruning.The pruning study examines degradation without the predictor re-planning the topology.Source: §1 · p. 2

2 DAGFormer

DAGFormer pairs a Transformer backbone with a causal topology predictor that plans weighted cross-layer connections for the whole model. The backbone runs along that plan, with each layer able to refine its own planned connections.

2 DAGFormer

  • A small topology predictor plans a weighted directed acyclic graph assigning each attention head and layer sources among earlier outputs; local refinement adjusts each layer’s planned rows.The predictor plans across the whole depth at once, while each layer’s refinement uses that layer’s input.Source: §2 · p. 2

2.1 Setup

  • In the vanilla setup, all 3H+1 reads in a layer see the same vector, whereas routed layers give each read its own weights over earlier outputs.Those weights form the graph for the token and fill four triangular matrices, with one row per read and one column per source.Source: §2.1 · p. 3

2.2 Topology predictor

  • The causal topology predictor reads the input sequence and produces a layer adjacency matrix with rows for each head’s query, key, and value reads, plus the residual.Its planned graph uses input tokens; local refinement lets a layer correct the plan using its own input.Source: §2.2 · p. 3

2.3 Routed layers

  • Each routed layer forms its query, key, value, and residual from weighted sums of earlier outputs; other Transformer modules remain unchanged.At initialization, each row selects the immediately preceding output, recovering the vanilla layer apart from V normalization where enabled.Source: §2.3 · p. 3

3 Experiments

DAGFormer is evaluated in matched OLMo-2 experiments across data, model sizes, training budgets, and routing baselines. It improves over the vanilla Transformer on language-model loss, downstream tasks, and instruction-following evaluations, with comparisons including routing cost.

3.1 Data

  • Evaluation uses held-out WikiText-2, MathInstruct, GSM8K, and balanced Dolma sets, plus 13 downstream tasks and instruction-following evaluations after finetuning.The instruction-following comparisons include held-out instruction-data likelihood, AlpacaEval, and MT-Bench.Source: §3.1 · p. 4, §3.2 · p. 5

3.2 Setup and baselines

  • Pretraining uses Dolma v1.7, with matched vanilla and DAGFormer pairs sharing backbone, tokenizer, data order, token budget, and optimizer; routing cost is included in compute comparisons.The study uses 75M, 150M, 300M, and 1B backbones; four cross-layer routing baselines are trained with the same recipe, data, tokens, and seed as the corresponding pair.Source: §3.2 · pp. 4–5, §3.1 · p. 4

3.3 Main results

  • 0.18 to 0.29 nats lower WikiText-2 loss than the vanilla Transformer at every size, with 17% to 25% lower perplexity; it has the lowest loss among routing baselines except at 150M.At 75M, 300M, and 1B it is 0.1 to 0.2 nats below the strongest routing baseline; at 150M, MUDDFormer is 0.04 nats lower. Figure 2 also compares loss against parameters, compute including routing cost, and training progress.Source: §3.3 · p. 6, Figure 2 · p. 6
  • 2.0 to 2.5 points higher mean accuracy than the vanilla Transformer across 13 evaluation tasks at the two tested sizes, and 0.7 to 0.8 points above the strongest routing baseline.DAGFormer leads on commonsense, science, and reading tasks; mHC is 0.1 points ahead on math, where MathQA and GSM8K are close to chance.Source: Table 2 · p. 5, §3.3 · p. 6
  • 0.63 to 0.70 AlpacaEval win rate against the vanilla Transformer, above every routing baseline; DAGFormer also has the lowest held-out instruction-data negative log-likelihood across sizes and both recipes.On MT-Bench, DAGFormer wins 0.56 to 0.69 of comparisons against the vanilla Transformer, with the 95% interval excluding 0.5 in seven of ten comparisons.Source: Table 3 · p. 5, §3.3 · p. 6

4 Analysis

Scaling fits show a larger compute advantage for DAGFormer at larger tested sizes, while controls indicate that parameter count alone does not explain most of the gain there. Ablations attribute most of the improvement to the shared topology predictor, with local refinement completing it.

4.1 Scaling with compute and parameters

  • 1.7 to 2.0× more training compute is needed by the vanilla Transformer to reach DAGFormer’s loss at 300M and 1B, versus 1.2 to 1.4× at 75M and 150M.The fitted compute exponents are 0.30 for DAGFormer and 0.27 for vanilla; the exponents differ by about one standard error, and nothing beyond the largest run is tested. At about the same compute, the 1B DAGFormer trained on 5B tokens beats the 1B vanilla model trained on 10B tokens on WikiText-2, MathInstruct, and GSM8K.Source: §4.1 · pp. 6–7
  • 70 to 89% of the gap at 75M is closed by giving vanilla Transformers DAGFormer’s routing parameters as extra width, but only 4 to 42% at 150M; DAGFormer remains ahead of every control on every held-out set.Extra parameters explain most of the gain at the smallest size and less than half at the larger one.Source: §4.1 · p. 7

4.2 Where the gain comes from

  • About 90% of DAGFormer’s gain over vanilla is recovered by the predictor alone, and adding local refinement gives the best result on every held-out set.Independent per-layer predictors reading the same tokens recover under 80%, despite having 7M more parameters; the refinement costs under 0.5% of forward multiply-accumulates.Source: §4.2 · p. 7

5 Structured pruning

Under structured pruning during MathInstruct finetuning, DAGFormer generally loses less than the compared models, including with its predictor frozen. Its advantage is associated with the sources available to each head, rather than routers replanning during pruning.

5 Structured pruning

  • DAGFormer loses the least of all models in five of six pruning settings; with 50% of heads and channels removed, it degrades less than the vanilla Transformer and every routing baseline at every size.At 70% removal, it degrades the least at 150M and 300M; at 75M several models lose less, although DAGFormer’s pruned loss remains lowest.Source: §5 · p. 7
  • Freezing DAGFormer’s predictor during pruning barely changes the result, and random unit removal causes both model families to lose much more while shrinking the gap at 75M.The reported advantage lies in the sources each head can draw on, not routing replanning.Source: §5 · p. 7

6 Discussion

The discussion distinguishes learning a topology from using token-specific routing at inference and identifies boundaries in scale and interpretability. DAGFormer’s graph can be manipulated as an object, but its meaning for individual inputs remains unresolved.

Token dependence

  • Most of DAGFormer’s gain lies in the learned graph rather than replanning it for every token; at inference, models still rely on the refinement’s per-token corrections.Static weights trained directly recover little of the gain, and post-hoc graph sparsification is costly.Source: §6 · p. 7

Shared predictor and per-layer routers

  • Whether the shared predictor keeps its lead at larger scale is untested; DAGFormer’s loss is lower than hyper-connections, mHC, and MUDDFormer at every size but one.The exception is MUDDFormer at 150M; the comparison uses one seed per model.Source: §6 · pp. 7–8

Inspecting the learned graph

  • The graph is read only in aggregate, and whether it is meaningful to a human for a given input remains open.Per-token variation in trained models lives mostly in the refinement; a predictor-only model whose graph varies in nameable ways is proposed as a possible direction.Source: §6 · p. 8

7 Related work

Cross-layer architectures range from static access to earlier outputs to input-dependent mixing, while conditional-computation methods instead choose how much computation an input receives. DAGFormer’s distinction is to keep computation fixed and choose which earlier results each layer reads.

7 Related work

  • DenseNet connects each layer to earlier layers, while transparent attention, ELMo, DenseFormer and LAuReL learn static weights over layer outputs.The passage places these approaches in the history of cross-layer connectivity.Source: §7 · p. 8
  • Hyper-connections and mHC use hidden-state-dependent residual mixing, while MUDDFormer generates per-position query, key, value and residual weights from each layer’s hidden state.DAGFormer shares MUDDFormer’s decomposition but plans weights for all layers with one token-based predictor, then refines each layer’s rows per head.Source: §7 · p. 8
  • Unlike methods that choose how much computation an input receives, DAGFormer keeps computation fixed and chooses which earlier results each layer reads.The passage identifies BlockDrop as the closest planned-topology analogue: one policy network selects image blocks to skip.Source: §7 · p. 8

8 Conclusion

DAGFormer plans cross-layer connections with one input-token predictor trained jointly with the backbone. It improves across model sizes and several evaluation settings, and its learned graph can be inspected and intervened on.

8 Conclusion

  • One predictor over the input tokens plans the cross-layer connections and is trained jointly with the backbone.The separate predictor makes the learned graph an object that can be inspected and intervened on.Source: §8 · pp. 8–9
  • From 75M to 1B parameters, DAGFormer lowers loss at every size and equal compute, and beats four routing baselines at every size but 150M.Its gains also carry over to evaluation tasks and instruction finetuning.Source: §8 · p. 8
  • DAGFormer degrades the least under structured pruning in five of six settings.The conclusion attributes most of the gain to the learned graph.Source: §8 · p. 9

Limitations

The reported scale, seed, compute, comparison, evaluation, and interpretability evidence has specific boundaries. In particular, the largest model is below its compute-optimal training budget, and whether gains persist beyond 1B parameters remains open.

Scale and seeds

  • The 1B model is trained on 5 to 10B tokens, below its compute-optimal budget, and every pretraining run except the 300M pair uses a single seed.Repeating the 300M pair changes held-out NLL by up to 0.08 nats per token, or at most 0.03 on WikiText-2, so smaller differences are described as noise; pruning cells are also single runs.Source: Limitations · p. 9

Cost

  • Routing costs 1.18 to 2.13× the vanilla Transformer’s FLOPs per token, and trained models rely on per-token refinement corrections.As a result, they cannot run with a fixed topology without retraining.Source: Limitations · p. 9

Comparisons

  • MUDDFormer beats DAGFormer at 150M and loses at other sizes, but each comparison uses a single seed, so loss alone does not settle it.The 1B baselines use 5B tokens only; parameter-matched vanilla controls exist at 75M and 150M, and only the 75M pair was swept over learning rates.Source: Limitations · p. 9

Evaluation

  • Models of this size are near chance on harder benchmarks, and one LLM judge scored the chat and AlpacaEval comparisons.Source: Limitations · p. 9

Interpretability and scale

  • The graph is intervened on only as a whole and read only in aggregate; naming its input-dependent part is left to future work.Source: Limitations · p. 9
  • Whether DAGFormer’s gains persist beyond 1B parameters remains an open question.Source: Limitations · p. 9

A Compute accounting

The compute accounting includes routing operations and compares two exact computation orders with different FLOP and memory costs. The appendix also documents held-out data, evaluation details, baseline setup, and matched-pair results.

A Compute accounting

  • 1.18× to 2.13× the vanilla Transformer’s FLOPs per token are required by project-then-mix, with source projections dominating and the predictor under 2% of total cost.Training FLOPs are estimated as 6× forward MACs times tokens, including unembedding, attention-context, and routing terms.Source: A Compute accounting · p. 12

A.1 Two exact computation orders

  • 1.02 to 1.04× vanilla FLOPs are possible with mix-then-project, but it is memory-bound; projecting first is faster at inference through 300M, while mixing first is faster at 1B.Both orders give the same logits up to 10^-4 in fp32 at every size; the modular variant costs 1.33 to 2.30×.Source: §A.1 · p. 13

B Held-out sets

  • The 21B-token held-out tail consists only of templated instruction and Wikipedia text, so it is used only for comparisons within that corpus.Cross-corpus numbers use WikiText-2, MathInstruct, and GSM8K.Source: B Held-out sets · p. 13

C Additional tables

  • The appendix provides complete downstream and chat evaluation stages, post-hoc baselines and learning rates, training hyperparameters, held-out results, routing-source ablations, and pruning results.Source: C Additional tables · p. 13

C.1 Downstream and chat per stage

  • Repeating 300M instruction finetuning with a second seed moves the 12-task mean by at most 0.001 and routed-minus-vanilla gaps by at most 0.003.SmolTalk lowers the multiple-choice mean by 0.7 to 1.4 points for every model; absolute judge scores are about 2 of 10.Source: §C.1 · p. 13

C.2 Post-hoc baselines and learning rates

  • Hyper-connection baselines use one module per attention and MLP sub-layer; the authors use an open implementation because the paper has no official code release.mHC uses DeepSeek’s official reference operations with four streams, 0.01 gating initialization, and 20 Sinkhorn-Knopp iterations.Source: §C.2 · pp. 13–14

C.3 Training hyperparameters

  • Table 15 lists the architecture, token budget, and optimization settings for every pretraining run and the instruction-tuning recipe.Routed baselines use the settings of the same-size pair.Source: §C.3 · p. 14

C.4 Matched pairs on all held-out sets

  • DAGFormer has lower loss in every matched-pair cell across the four held-out sets, while a second seed changes loss by up to 0.08 nats and at most 0.03 on WikiText-2.Source: §C.4 · p. 14

C.5 Gap along training

Across training, DAGFormer maintains a held-out-cache advantage, and downstream comparisons favor it across evaluation tasks and instruction tuning. The routing analyses indicate that the predictor finds a useful static graph, while local refinement carries more token-specific variation.

C.5 Gap along training

  • At 1B, the vanilla-minus-DAGFormer held-out-cache gap is +0.110 nats at 9.96B tokens; it opens early, peaks, then narrows slowly.The 1B entry is the 5B-token run.Source: §C.5 · p. 14

C.6 Router-locality ablation on all metrics

  • At 75M, the four router-locality arms’ multiple-choice means lie within 0.02 of each other across the reported sets.The arms are DAGFormer, shared predictor alone, per-layer predictors, a static table, and the vanilla Transformer.Source: §C.6 · p. 14

C.7 Shared-slice 75M, 150M and 300M pairs downstream

  • Across the shared-slice pairs, DAGFormer has higher multiple-choice means and lower bits per byte on GSM8K solutions and WikiText at every size and stage.The comparison covers evaluation before and after instruction tuning.Source: §C.7 · p. 14

C.8 Pruning: absolute losses, controls and immediate damage

  • The pruning-loss tables report final MathInstruct NLL, immediate post-update damage, and controls using random unit selection or whole-block pruning.The comparison follows 2,000 finetuning steps.Source: §C.8 · p. 14

C.9 Static substitution on all held-out sets

  • Replacing the predictor output with its per-position mean raises loss by at most 0.005 nats, while replacing the refinement output costs 0.27 to 0.75 nats.The substitutions are evaluated across held-out sets.Source: §C.9 · p. 14

C.10 Judged chat evaluation: categories and scores

  • MT-Bench results are broken down by category win rate and single-answer scores; repeated 300M instruction tuning reproduces routed-minus-vanilla downstream gaps within 0.003.Pretraining-seed variance was not measured.Source: §C.10 · p. 14, Table 7 · p. 15

D The predictor learns a static graph

  • A per-position mean for the predictor changes WikiText-2 loss by at most 0.005 nats at 300M and 1B, while a global mean costs 0.04 to 0.06 nats.Replacing the refinement output by its mean costs 0.4 to 0.7 nats, more than the full vanilla gap; swapping predictor outputs between windows changes loss by at most 0.007 nats from 75M to 300M.Source: D The predictor learns a static graph · p. 15
  • The predictor’s token input helps find a good static graph, while the local refinement carries per-token variation; a static table recovers only 17 to 23% of predictor-only gains.The paper connects the pre-output interventions to generative interpretability, while leaving human understandability as an open question.Source: D The predictor learns a static graph · p. 15

E Topology analysis

The topology analysis finds substantial use of skip-source connections, and pruning results favor DAGFormer in the reported comparisons. The modular routing variant generally trails FourWay, while the paper also records the study’s artifact, risk, and compute boundaries.

E Topology analysis

  • Across 128 held-out windows, skip sources hold most mixing mass; removing all skips raises WikiText-2 loss to 9 to 10 nats, so post-hoc sparsification needs retraining.Zeroing the quarter of entries with the smallest mean absolute weights costs 0.27 nats at 300M and 0.15 to 0.22 at 1B.Source: E Topology analysis · p. 15

G Why reading many layers may soften pruning

  • The multiple-path explanation for pruning robustness remains untested directly; DAGFormer’s advantage shrinks when random units, rather than least-important units, are removed.The paper suggests that reading multiple earlier layers may provide alternative routes to information.Source: G Why reading many layers may soften… · p. 16

H The modular variant

  • At 300M, FourWay beats modular routing on every held-out set (WikiText-2 3.544 vs 3.628); modular routing’s extra FLOPs are not repaid above 75M.The modular variant is 0.03 nats behind FourWay at 75M and 150M, and 0.08 behind at 300M.Source: H The modular variant · p. 16, Figure 4 · p. 18
  • Modular routing predicts intra-layer and intermodular edges, and wins one reported pruning setting: at 150M, half-block removal yields 2.130 vs 2.229 nats.Its routed sources are module outputs, and it adds routed MLP inputs and a routed final read-out.Source: H The modular variant · p. 16

I Artifacts, licenses and compute

  • The models use OLMo-2 initialized from scratch and Dolma v1.7 for pretraining; the study collects no new data and relies on Dolma’s filtering and deduplication.The listed artifacts are used for non-commercial research.Source: I Artifacts, licenses and compute · pp. 16–17
  • The models have at most 1.3B parameters, are research artifacts rather than deployment models, and may generate false, biased, or offensive text.The paper reports no risk specific to the architecture.Source: I Artifacts, licenses and compute · p. 17
  • Reported experiments used about 3,700 GPU-hours, including about 1,850 for matched-pair and modular-variant pretraining.The estimate is based on job records and training logs.Source: I Artifacts, licenses and compute · p. 17

F Pruning: all runs

  • DAGFormer loses less under pruning at every budget, increasingly at higher sparsity; on the common-corpus 150M pair it leads by 0.05/0.07/0.10 nats at 30/50/70%.Figure 3 reports pruning during MathInstruct domain finetuning.Source: Figure 3 · p. 18

Related papers

Loading 2610.19625v1…