# Chain and Causal Attention for Efficient Entity Tracking

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

## Key facts

- Cited by 8 papers in the map (data through 14 Oct 2026).
- Builds on 27 of its 37 references that are in the map.
- First posted on arXiv 7 Oct 2024.

## Builds on (most cited first, 27 of 27)

- [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)
- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Language Models are Few-Shot Learners](https://mlsanity.com/paper/2005.14165/) (2020)
- [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://mlsanity.com/paper/1910.10683/) (2019)
- [RoBERTa: A Robustly Optimized BERT Pretraining Approach](https://mlsanity.com/paper/1907.11692/) (2019)
- [Distilling the Knowledge in a Neural Network](https://mlsanity.com/paper/1503.02531/) (2015)
- [BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension](https://mlsanity.com/paper/1910.13461/) (2019)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Switch Transformers: Scaling to Trillion Parameter Models with Simple and Efficient Sparsity](https://mlsanity.com/paper/2101.03961/) (2021)
- [XLNet: Generalized Autoregressive Pretraining for Language Understanding](https://mlsanity.com/paper/1906.08237/) (2019)
- [DeBERTa: Decoding-enhanced BERT with Disentangled Attention](https://mlsanity.com/paper/2006.03654/) (2020)
- [Deep Equilibrium Models](https://mlsanity.com/paper/1909.01377/) (2019)
- [Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task](https://mlsanity.com/paper/2210.13382/) (2022)
- [Estimating the Carbon Footprint of BLOOM, a 176B Parameter Language Model](https://mlsanity.com/paper/2211.02001/) (2022)
- [Implicit Representations of Meaning in Neural Language Models](https://mlsanity.com/paper/2106.00737/) (2021)
- [Fine-Tuning Enhances Existing Mechanisms: A Case Study on Entity Tracking](https://mlsanity.com/paper/2402.14811/) (2024)
- [Chess as a Testbed for Language Model State Tracking](https://mlsanity.com/paper/2102.13249/) (2021)
- [WIQA: A dataset for "What if..." reasoning over procedural text](https://mlsanity.com/paper/1909.04739/) (2019)
- [Multi-hop Attention Graph Neural Network](https://mlsanity.com/paper/2009.14332/) (2020)
- [The Materials Science Procedural Text Corpus: Annotating Materials Synthesis Procedures with Shallow Semantic Structures](https://mlsanity.com/paper/1905.06939/) (2019)
- [Subformer: Exploring Weight Sharing for Parameter Efficiency in Generative Transformers](https://mlsanity.com/paper/2101.00234/) (2021)
- [Simpler Context-Dependent Logical Forms via Model Projections](https://mlsanity.com/paper/1606.05378/) (2016)
- [Code Pretraining Improves Entity Tracking Abilities of Language Models](https://mlsanity.com/paper/2405.21068/) (2024)
- [Causal Reasoning of Entities and Events in Procedural Texts](https://mlsanity.com/paper/2301.10896/) (2023)
- [OpenPI2.0: An Improved Dataset for Entity Tracking in Texts](https://mlsanity.com/paper/2305.14603/) (2023)
- [ResidualTransformer: Residual Low-Rank Learning with Weight-Sharing for Transformer Layers](https://mlsanity.com/paper/2310.02489/) (2023)

## Built on by (most cited first, 8 of 8)

- [PaTH Attention: Position Encoding via Accumulating Householder Transformations](https://mlsanity.com/paper/2505.16381/) (2025)
- [Understanding Transformer from the Perspective of Associative Memory](https://mlsanity.com/paper/2505.19488/) (2025)
- [Lost in Transmission: When and Why LLMs Fail to Reason Globally](https://mlsanity.com/paper/2505.08140/) (2025)
- [The Topological Trouble With Transformers](https://mlsanity.com/paper/2604.17121/) (2026)
- [KisMATH: Do LLMs Have Knowledge of Implicit Structures in Mathematical Reasoning?](https://mlsanity.com/paper/2507.11408/) (2025)
- [MET-Bench: Multimodal Entity Tracking for Evaluating the Limitations of Vision-Language and Reasoning Models](https://mlsanity.com/paper/2502.10886/) (2025)
- [Structured-Sparse Attention for Entity Tracking with Subquadratic Sequence Complexity](https://mlsanity.com/paper/2605.22476/) (2026)
- [Trading Complexity for Expressivity Through Structured Generalized Linear Token Mixing](https://mlsanity.com/paper/2605.31367/) (2026)

## Abstract (arXiv)

This paper investigates the limitations of transformers for entity-tracking tasks in large language models. We identify a theoretical constraint, showing that transformers require at least log2 (n+1) layers to handle entity tracking with n state changes. To address this issue, we propose an efficient and frugal enhancement to the standard attention mechanism, enabling it to manage long-term dependencies more efficiently. By considering attention as an adjacency matrix, our model can track entity states with a single layer. Empirical results demonstrate significant improvements in entity tracking datasets while keeping competitive performance on standard natural language modeling. Our modified attention allows us to achieve the same performance with drastically fewer layers. Additionally, our enhanced mechanism reveals structured internal representations of attention. Extensive experiments on both toy and complex datasets validate our approach. Our contributions include theoretical insights, an improved attention mechanism, and empirical validation.