# The Neural Data Router: Adaptive Control Flow in Transformers Improves Systematic Generalization

arXiv: https://arxiv.org/abs/2110.07732 · HTML: https://mlsanity.com/paper/2110.07732/

## Key facts

- Cited by 62 papers in the map (data through 14 Oct 2026).
- Builds on 45 of its 61 references that are in the map.
- First posted on arXiv 14 Oct 2021.

## Builds on (most cited first, 30 of 45)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Deep Residual Learning for Image Recognition](https://mlsanity.com/paper/1512.03385/) (2015)
- [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://mlsanity.com/paper/1810.04805/) (2018)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://mlsanity.com/paper/1901.02860/) (2019)
- [GLU Variants Improve Transformer](https://mlsanity.com/paper/2002.05202/) (2020)
- [Self-Attention with Relative Position Representations](https://mlsanity.com/paper/1803.02155/) (2018)
- [Language Modeling with Gated Convolutional Networks](https://mlsanity.com/paper/1612.08083/) (2016)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [Universal Transformers](https://mlsanity.com/paper/1807.03819/) (2018)
- [Adaptive Computation Time for Recurrent Neural Networks](https://mlsanity.com/paper/1603.08983/) (2016)
- [Nystr"omformer: A Nystr"om-Based Algorithm for Approximating Self-Attention](https://mlsanity.com/paper/2102.03902/) (2021)
- [Measuring Compositional Generalization: A Comprehensive Method on Realistic Data](https://mlsanity.com/paper/1912.09713/) (2019)
- [Stabilizing Transformers for Reinforcement Learning](https://mlsanity.com/paper/1910.06764/) (2019)
- [Compositionality decomposed: how do neural networks generalise?](https://mlsanity.com/paper/1908.08351/) (2019)
- [On the Binding Problem in Artificial Neural Networks](https://mlsanity.com/paper/2012.05208/) (2020)
- [Thinking Like Transformers](https://mlsanity.com/paper/2106.06981/) (2021)
- [Training Very Deep Networks](https://mlsanity.com/paper/1507.06228/) (2015)
- [PonderNet: Learning to Ponder](https://mlsanity.com/paper/2107.05407/) (2021)
- [Compositional Generalization and Natural Language Variation: Can a Semantic Parsing Approach Handle Both?](https://mlsanity.com/paper/2010.12725/) (2020)
- [Character-Aware Neural Language Models](https://mlsanity.com/paper/1508.06615/) (2015)
- [ListOps: A Diagnostic Dataset for Latent Tree Learning](https://mlsanity.com/paper/1804.06028/) (2018)
- [The Devil is in the Detail: Simple Tricks Improve Systematic Generalization of Transformers](https://mlsanity.com/paper/2108.12284/) (2021)
- [Compositional generalization through meta sequence-to-sequence learning](https://mlsanity.com/paper/1906.05381/) (2019)
- [Language Modeling with Deep Transformers](https://mlsanity.com/paper/1905.04226/) (2019)
- [Compositional Generalization in Semantic Parsing: Pre-training vs. Specialized Architectures](https://mlsanity.com/paper/2007.08970/) (2020)
- [Going Beyond Linear Transformers with Recurrent Fast Weight Programmers](https://mlsanity.com/paper/2106.06295/) (2021)
- [Compositional Generalization via Neural-Symbolic Stack Machines](https://mlsanity.com/paper/2008.06662/) (2020)
- [Making Transformers Solve Compositional Tasks](https://mlsanity.com/paper/2108.04378/) (2021)

## Built on by (most cited first, 30 of 62)

- [Randomized Positional Encodings Boost Length Generalization of Transformers](https://mlsanity.com/paper/2305.16843/) (2023)
- [Block-Recurrent Transformers](https://mlsanity.com/paper/2203.07852/) (2022)
- [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022)
- [Linear algebra with transformers](https://mlsanity.com/paper/2112.01898/) (2021)
- [A Generalist Neural Algorithmic Learner](https://mlsanity.com/paper/2209.11142/) (2022)
- [MoEUT: Mixture-of-Experts Universal Transformers](https://mlsanity.com/paper/2405.16039/) (2024)
- [Positional Description Matters for Transformers Arithmetic](https://mlsanity.com/paper/2311.14737/) (2023)
- [Forgetting Transformer: Softmax Attention with a Forget Gate](https://mlsanity.com/paper/2503.02130/) (2025)
- [A Modern Self-Referential Weight Matrix That Learns to Modify Itself](https://mlsanity.com/paper/2202.05780/) (2022)
- [How Far Can Transformers Reason? The Globality Barrier and Inductive Scratchpad](https://mlsanity.com/paper/2406.06467/) (2024)
- [SwitchHead: Accelerating Transformers with Mixture-of-Experts Attention](https://mlsanity.com/paper/2312.07987/) (2023)
- [PaTH Attention: Position Encoding via Accumulating Householder Transformations](https://mlsanity.com/paper/2505.16381/) (2025)
- [SALSA: Attacking Lattice Cryptography with Transformers](https://mlsanity.com/paper/2207.04785/) (2022)
- [V2PE: Improving Multimodal Long-Context Capability of Vision-Language Models with Variable Visual Position Encoding](https://mlsanity.com/paper/2412.09616/) (2024)
- [Approximating Two-Layer Feedforward Networks for Efficient Transformers](https://mlsanity.com/paper/2310.10837/) (2023)
- [When can transformers reason with abstract symbols?](https://mlsanity.com/paper/2310.09753/) (2023)
- [Compositional Capabilities of Autoregressive Transformers: A Study on Synthetic, Interpretable Tasks](https://mlsanity.com/paper/2311.12997/) (2023)
- [Complexity Control Facilitates Reasoning-Based Compositional Generalization in Transformers](https://mlsanity.com/paper/2501.08537/) (2025)
- [Systematic Generalization and Emergent Structures in Transformers Trained on Structured Tasks](https://mlsanity.com/paper/2210.00400/) (2022)
- [ModuleFormer: Modularity Emerges from Mixture-of-Experts](https://mlsanity.com/paper/2306.04640/) (2023)
- [Initialization is Critical to Whether Transformers Fit Composite Functions by Reasoning or Memorizing](https://mlsanity.com/paper/2405.05409/) (2024)
- [Transformer Working Memory Enables Regular Language Reasoning and Natural Language Length Extrapolation](https://mlsanity.com/paper/2305.03796/) (2023)
- [Scaling Stick-Breaking Attention: An Efficient Implementation and In-depth Study](https://mlsanity.com/paper/2410.17980/) (2024)
- [An Analysis for Reasoning Bias of Language Models with Small Initialization](https://mlsanity.com/paper/2502.04375/) (2025)
- [CTL++: Evaluating Generalization on Never-Seen Compositional Patterns of Known Functions, and Compatibility of Neural Representations](https://mlsanity.com/paper/2210.06350/) (2022)
- [Hyperloop Transformers](https://mlsanity.com/paper/2604.21254/) (2026)
- [SALSA PICANTE: a machine learning attack on LWE with binary secrets](https://mlsanity.com/paper/2303.04178/) (2023)
- [From Interpolation to Extrapolation: Complete Length Generalization for Arithmetic Transformers](https://mlsanity.com/paper/2310.11984/) (2023)
- [Learning to Skip the Middle Layers of Transformers](https://mlsanity.com/paper/2506.21103/) (2025)
- [Decoupling the "What" and "Where" With Polar Coordinate Positional Embeddings](https://mlsanity.com/paper/2509.10534/) (2025)

## Abstract (arXiv)

Despite progress across a broad range of applications, Transformers have limited success in systematic generalization. The situation is especially frustrating in the case of algorithmic tasks, where they often fail to find intuitive solutions that route relevant information to the right node/operation at the right time in the grid represented by Transformer columns. To facilitate the learning of useful control flow, we propose two modifications to the Transformer architecture, copy gate and geometric attention. Our novel Neural Data Router (NDR) achieves 100% length generalization accuracy on the classic compositional table lookup task, as well as near-perfect accuracy on the simple arithmetic task and a new variant of ListOps testing for generalization across computational depths. NDR's attention and gating patterns tend to be interpretable as an intuitive form of neural routing. Our code is public.