DAGFormer: Global Topology Prediction for Transformer Cross-Layer Connections
AuthorsXiaocong Yang, Yuren Hao, ChengXiang Zhai
AffiliationsAI Interpretability @ Illinois · University of Illinois Urbana-Champaign · {xy51, yurenh2, czhai}@illinois.edu
DAGFormer lets a Transformer learn how information should flow across its layers, improving language-model performance across several evaluations.
Key results
Reduction versus the vanilla Transformer across 75M to 1B model sizes.
Percentage-point improvement on 13 evaluation tasks at 300M and 1B.
Instruction-tuned DAGFormer versus size-matched vanilla models.
The 1B model trained on 5B tokens outperformed the vanilla 1B model trained on 10B tokens at roughly comparable compute.
DAGFormer FLOPs per token relative to the vanilla Transformer across model sizes.
What the paper found
DAGFormer changes Transformer cross-layer wiring with a shared causal topology predictor: before the backbone processes a token, it predicts a signed weighted directed acyclic graph across all layers. Each attention head builds its queries, keys, and values from mixtures of the embedding and earlier layer outputs, while each layer also mixes its residual input; a small local refinement adjusts the plan using the layer’s hidden state. Trained jointly with an OLMo-2 backbone on Dolma v1.7, DAGFormer achieves 17-25% lower WikiText-2 perplexity than a vanilla Transformer from 75M to 1B parameters. It has the lowest loss among routing baselines at 75M, 300M, and 1B, though MUDDFormer is slightly better at 150M. Across 13 commonsense, science, reading, and math tasks, mean accuracy improves by 2.0-2.5 percentage points at 300M and 1B; instruction-tuned models reach 63-70% AlpacaEval win rates against size-matched vanilla models. A 1B DAGFormer trained on 5B tokens outperforms a vanilla model trained on 10B tokens at roughly comparable compute, though routing increases FLOPs to 1.18-2.13 times the vanilla model. Experiments used NVIDIA A100 GPUs; the mHC baseline used DeepSeek’s PyTorch reference implementation, and Qwen3-Coder-30B-A3B-Instruct judged MT-Bench comparisons.
Original 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.
Read the original paperMore in Transformers
Browse all 48 papers →A Conditional No-Hair Theorem for Transformers: How Language Models Forget Their Prompt
Shuyang Xiang
The paper finds that a transformer’s direct memory of its prompt fades as it generates text, while the text itself can carry that memory forward.
Towards Looped Models Done Right, Part I: Topology, Input Injection, Recurrent-State Design
Benhao Huang, Chufan Shi, Junlin Chen, Shicheng Wen, Zhengzhong Liu, Eric Xing, Xuezhe Ma
The paper finds that a carefully designed looped Transformer can outperform comparable models while using far fewer resident parameters.
Decoding the Functional Roles of Register and High-Norm Patch Tokens in Vision Transformers
Neel Varma, Andrew Rufail, Dipika Khullar, Vasu Sharma
This study shows that ViT register tokens carry crucial high-level semantics, while high-norm patch tokens mainly encode lower-level visual structure and contribute little to overall representations.