# Unveiling Transformers with LEGO: a synthetic reasoning task

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

## Key facts

- Cited by 106 papers in the map (data through 14 Oct 2026).
- Builds on 48 of its 67 references that are in the map.
- First posted on arXiv 9 Jun 2022.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [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)
- [LoRA: Low-Rank Adaptation of Large Language Models](https://mlsanity.com/paper/2106.09685/) (2021)
- [Chain-of-Thought Prompting Elicits Reasoning in Large Language Models](https://mlsanity.com/paper/2201.11903/) (2022)
- [Swin Transformer: Hierarchical Vision Transformer using Shifted Windows](https://mlsanity.com/paper/2103.14030/) (2021)
- [Hierarchical Text-Conditional Image Generation with CLIP Latents](https://mlsanity.com/paper/2204.06125/) (2022)
- [Measuring Mathematical Problem Solving With the MATH Dataset](https://mlsanity.com/paper/2103.03874/) (2021)
- [PaLM: Scaling Language Modeling with Pathways](https://mlsanity.com/paper/2204.02311/) (2022)
- [SQuAD: 100,000+ Questions for Machine Comprehension of Text](https://mlsanity.com/paper/1606.05250/) (2016)
- [A ConvNet for the 2020s](https://mlsanity.com/paper/2201.03545/) (2022)
- [OPT: Open Pre-trained Transformer Language Models](https://mlsanity.com/paper/2205.01068/) (2022)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [ALBERT: A Lite BERT for Self-supervised Learning of Language Representations](https://mlsanity.com/paper/1909.11942/) (2019)
- [Xception: Deep Learning with Depthwise Separable Convolutions](https://mlsanity.com/paper/1610.02357/) (2016)
- [Decision Transformer: Reinforcement Learning via Sequence Modeling](https://mlsanity.com/paper/2106.01345/) (2021)
- [LaMDA: Language Models for Dialog Applications](https://mlsanity.com/paper/2201.08239/) (2022)
- [Scaling Language Models: Methods, Analysis & Insights from Training Gopher](https://mlsanity.com/paper/2112.11446/) (2021)
- [What Does BERT Look At? An Analysis of BERT's Attention](https://mlsanity.com/paper/1906.04341/) (2019)
- [Deep Networks with Stochastic Depth](https://mlsanity.com/paper/1603.09382/) (2016)
- [GLaM: Efficient Scaling of Language Models with Mixture-of-Experts](https://mlsanity.com/paper/2112.06905/) (2021)
- [Show Your Work: Scratchpads for Intermediate Computation with Language Models](https://mlsanity.com/paper/2112.00114/) (2021)
- [Right for the Wrong Reasons: Diagnosing Syntactic Heuristics in Natural Language Inference](https://mlsanity.com/paper/1902.01007/) (2019)
- [Analyzing Multi-Head Self-Attention: Specialized Heads Do the Heavy Lifting, the Rest Can Be Pruned](https://mlsanity.com/paper/1905.09418/) (2019)
- [A Primer in BERTology: What we know about how BERT works](https://mlsanity.com/paper/2002.12327/) (2020)
- [A Generalist Agent](https://mlsanity.com/paper/2205.06175/) (2022)
- [Annotation Artifacts in Natural Language Inference Data](https://mlsanity.com/paper/1803.02324/) (2018)
- [CoAtNet: Marrying Convolution and Attention for All Data Sizes](https://mlsanity.com/paper/2106.04803/) (2021)
- [Using DeepSpeed and Megatron to Train Megatron-Turing NLG 530B, A Large-Scale Generative Language Model](https://mlsanity.com/paper/2201.11990/) (2022)

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

- [Sparks of Artificial General Intelligence: Early experiments with GPT-4](https://mlsanity.com/paper/2303.12712/) (2023)
- [What Can Transformers Learn In-Context? A Case Study of Simple Function Classes](https://mlsanity.com/paper/2208.01066/) (2022)
- [Faith and Fate: Limits of Transformers on Compositionality](https://mlsanity.com/paper/2305.18654/) (2023)
- [Hyena Hierarchy: Towards Larger Convolutional Language Models](https://mlsanity.com/paper/2302.10866/) (2023)
- [The Impact of Positional Encoding on Length Generalization in Transformers](https://mlsanity.com/paper/2305.19466/) (2023)
- [Transformers as Statisticians: Provable In-Context Learning with In-Context Algorithm Selection](https://mlsanity.com/paper/2306.04637/) (2023)
- [Transformers Learn Shortcuts to Automata](https://mlsanity.com/paper/2210.10749/) (2022)
- [Neural Networks and the Chomsky Hierarchy](https://mlsanity.com/paper/2207.02098/) (2022)
- [Exploring Length Generalization in Large Language Models](https://mlsanity.com/paper/2207.04901/) (2022)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [Looped Transformers are Better at Learning Learning Algorithms](https://mlsanity.com/paper/2311.12424/) (2023)
- [LIFT: Language-Interfaced Fine-Tuning for Non-Language Machine Learning Tasks](https://mlsanity.com/paper/2206.06565/) (2022)
- [A Toy Model of Universality: Reverse Engineering How Networks Learn Group Operations](https://mlsanity.com/paper/2302.03025/) (2023)
- [Skill-it! A Data-Driven Skills Framework for Understanding and Training Language Models](https://mlsanity.com/paper/2307.14430/) (2023)
- [Linear attention is (maybe) all you need (to understand transformer optimization)](https://mlsanity.com/paper/2310.01082/) (2023)
- [How Do Transformers Learn Topic Structure: Towards a Mechanistic Understanding](https://mlsanity.com/paper/2303.04245/) (2023)
- [Transformers Can Achieve Length Generalization But Not Robustly](https://mlsanity.com/paper/2402.09371/) (2024)
- [Same Pre-training Loss, Better Downstream: Implicit Bias Matters for Language Models](https://mlsanity.com/paper/2210.14199/) (2022)
- [Transformers as Decision Makers: Provable In-Context Reinforcement Learning via Supervised Pretraining](https://mlsanity.com/paper/2310.08566/) (2023)
- [Transformers, parallel computation, and logarithmic depth](https://mlsanity.com/paper/2402.09268/) (2024)
- [Exposing Attention Glitches with Flip-Flop Language Modeling](https://mlsanity.com/paper/2306.00946/) (2023)
- [Generalization on the Unseen, Logic Reasoning and Degree Curriculum](https://mlsanity.com/paper/2301.13105/) (2023)
- [Understanding Transformer Reasoning Capabilities via Graph Algorithms](https://mlsanity.com/paper/2405.18512/) (2024)
- [Length Generalization in Arithmetic Transformers](https://mlsanity.com/paper/2306.15400/) (2023)
- [Positional Description Matters for Transformers Arithmetic](https://mlsanity.com/paper/2311.14737/) (2023)
- [Active-Dormant Attention Heads: Mechanistically Demystifying Extreme-Token Phenomena in LLMs](https://mlsanity.com/paper/2410.13835/) (2024)
- [Mimetic Initialization of Self-Attention Layers](https://mlsanity.com/paper/2305.09828/) (2023)
- [Learning threshold neurons via the "edge of stability"](https://mlsanity.com/paper/2212.07469/) (2022)
- [The Closeness of In-Context Learning and Weight Shifting for Softmax Regression](https://mlsanity.com/paper/2304.13276/) (2023)
- [What can a Single Attention Layer Learn? A Study Through the Random Features Lens](https://mlsanity.com/paper/2307.11353/) (2023)

## Abstract (arXiv)

We propose a synthetic reasoning task, LEGO (Learning Equality and Group Operations), that encapsulates the problem of following a chain of reasoning, and we study how the Transformer architectures learn this task. We pay special attention to data effects such as pretraining (on seemingly unrelated NLP tasks) and dataset composition (e.g., differing chain length at training and test time), as well as architectural variants such as weight-tied layers or adding convolutional components. We study how the trained models eventually succeed at the task, and in particular, we manage to understand some of the attention heads as well as how the information flows in the network. In particular, we have identified a novel association pattern that globally attends only to identical tokens. Based on these observations we propose a hypothesis that here pretraining helps for LEGO tasks due to certain structured attention patterns, and we experimentally verify this hypothesis. We also observe that in some data regime the trained transformer finds ``shortcut" solutions to follow the chain of reasoning, which impedes the model's robustness, and moreover we propose ways to prevent it. Motivated by our findings on structured attention patterns, we propose the LEGO attention module, a drop-in replacement for vanilla attention heads. This architectural change significantly reduces Flops and maintains or even improves the model's performance at large-scale pretraining.