# Small transformer architectures for task switching

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

## Key facts

- Cited by 1 papers in the map (data through 14 Oct 2026).
- Builds on 28 of its 29 references that are in the map.
- First posted on arXiv 6 Aug 2025.

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

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [Scaling Laws for Neural Language Models](https://mlsanity.com/paper/2001.08361/) (2020)
- [Mamba: Linear-Time Sequence Modeling with Selective State Spaces](https://mlsanity.com/paper/2312.00752/) (2023)
- [Lost in the Middle: How Language Models Use Long Contexts](https://mlsanity.com/paper/2307.03172/) (2023)
- [Training Compute-Optimal Large Language Models](https://mlsanity.com/paper/2203.15556/) (2022)
- [Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention](https://mlsanity.com/paper/2006.16236/) (2020)
- [Linformer: Self-Attention with Linear Complexity](https://mlsanity.com/paper/2006.04768/) (2020)
- [Train Short, Test Long: Attention with Linear Biases Enables Input Length Extrapolation](https://mlsanity.com/paper/2108.12409/) (2021)
- [Toy Models of Superposition](https://mlsanity.com/paper/2209.10652/) (2022)
- [A Survey on Multi-Task Learning](https://mlsanity.com/paper/1707.08114/) (2017)
- [Long Range Arena: A Benchmark for Efficient Transformers](https://mlsanity.com/paper/2011.04006/) (2020)
- [A Comprehensive Overview of Large Language Models](https://mlsanity.com/paper/2307.06435/) (2023)
- [Neural Networks and the Chomsky Hierarchy](https://mlsanity.com/paper/2207.02098/) (2022)
- [Birth of a Transformer: A Memory Viewpoint](https://mlsanity.com/paper/2306.00802/) (2023)
- [What Formal Languages Can Transformers Express? A Survey](https://mlsanity.com/paper/2311.00208/) (2023)
- [Unveiling Transformers with LEGO: a synthetic reasoning task](https://mlsanity.com/paper/2206.04301/) (2022)
- [Reset-Free Reinforcement Learning via Multi-Task Learning: Learning Dexterous Manipulation Behaviors without Human Intervention](https://mlsanity.com/paper/2104.11203/) (2021)
- [Flowformer: Linearizing Transformers with Conservation Flows](https://mlsanity.com/paper/2202.06258/) (2022)
- [Multi-Task Learning in Natural Language Processing: An Overview](https://mlsanity.com/paper/2109.09138/) (2021)
- [Dynamic Movement Primitives in Robotics: A Tutorial Survey](https://mlsanity.com/paper/2102.03861/) (2021)
- [RoboHive: A Unified Framework for Robot Learning](https://mlsanity.com/paper/2310.06828/) (2023)
- [Scaling Laws for a Multi-Agent Reinforcement Learning Model](https://mlsanity.com/paper/2210.00849/) (2022)
- [Scaling Laws Under the Microscope: Predicting Transformer Performance from Small Scale Experiments](https://mlsanity.com/paper/2202.06387/) (2022)
- [Scaling Laws for Linear Complexity Language Models](https://mlsanity.com/paper/2406.16690/) (2024)
- [Multi-Task Learning with Summary Statistics](https://mlsanity.com/paper/2307.02388/) (2023)
- [AlphaZero Neural Scaling and Zipf's Law: a Tale of Board Games and Power Laws](https://mlsanity.com/paper/2412.11979/) (2024)
- [Kick control: using the attracting states arising within the sensorimotor loop of self-organized robots as motor primitives](https://mlsanity.com/paper/1806.09321/) (2018)
- [Reorganizing attention-space geometry with expressive attention](https://mlsanity.com/paper/2407.18601/) (2024)

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

- [Distance generalization in transformers: why bother with positional encoding?](https://mlsanity.com/paper/2609.11913/) (2026)

## Abstract (arXiv)

The rapid progress seen in terms of large-scale generative AI is largely based on the attention mechanism. It is conversely non-trivial to conceive small-scale applications for which attention-based architectures outperform traditional approaches, such as multi-layer perceptrons or recurrent networks. We examine this problem in the context of 'task switching'. In this framework models work on ongoing token sequences with the current task being determined by stochastically interspersed control tokens. We show that standard transformers cannot solve a basic task switching reference model based on finite domain arithmetics which contains subtasks dedicated to increment / addition / reverse copy / context (IARC). We show that transformers, long short-term memory recurrent networks (LSTM), and plain multi-layer perceptrons (MLPs) achieve similar, but only modest prediction accuracies. We enlarge our comparative study by including an extension of the standard transformer architecture to its non-translational invariant counterpart, the cisformer, and an alternative attention mechanism, extensive attention. A combination of the latter is found to be the only model able to achieve considerable performance levels, of around 95%. Our results indicate that the workings of attention can be understood better, and even improved, when comparing qualitatively different formulations in task-switching settings.