# Flexible task abstractions emerge in linear networks with fast and bounded units

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

## Key facts

- Cited by 2 papers in the map (data through 14 Oct 2026).
- Builds on 18 of its 48 references that are in the map.
- First posted on arXiv 6 Nov 2024.

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

- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Fashion-MNIST: a Novel Image Dataset for Benchmarking Machine Learning Algorithms](https://mlsanity.com/paper/1708.07747/) (2017)
- [Neural Tangent Kernel: Convergence and Generalization in Neural Networks](https://mlsanity.com/paper/1806.07572/) (2018)
- [Neural Module Networks](https://mlsanity.com/paper/1511.02799/) (2015)
- [Recurrent Independent Mechanisms](https://mlsanity.com/paper/1909.10893/) (2019)
- [A mathematical theory of semantic development in deep neural networks](https://mlsanity.com/paper/1810.10531/) (2018)
- [Neural Networks as Kernel Learners: The Silent Alignment Effect](https://mlsanity.com/paper/2111.00034/) (2021)
- [Alleviating catastrophic forgetting using context-dependent gating and synaptic stabilization](https://mlsanity.com/paper/1802.01569/) (2018)
- [Modular Networks: Learning to Decompose Neural Computation](https://mlsanity.com/paper/1811.05249/) (2018)
- [Learning to Combine Top-Down and Bottom-Up Signals in Recurrent Neural Networks with Attention over Modules](https://mlsanity.com/paper/2006.16981/) (2020)
- [The Neural Race Reduction: Dynamics of Abstraction in Gated Networks](https://mlsanity.com/paper/2207.10430/) (2022)
- [The Onset of Variance-Limited Behavior for Networks in the Lazy and Rich Regimes](https://mlsanity.com/paper/2212.12147/) (2022)
- [Discovering modular solutions that generalize compositionally](https://mlsanity.com/paper/2312.15001/) (2023)
- [Improving Expert Specialization in Mixture of Experts](https://mlsanity.com/paper/2302.14703/) (2023)
- [Why Do Animals Need Shaping? A Theory of Task Composition and Curriculum Learning](https://mlsanity.com/paper/2402.18361/) (2024)
- [Learning, Planning, and Control in a Monolithic Neural Event Inference Architecture](https://mlsanity.com/paper/1809.07412/) (2018)
- [Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations](https://mlsanity.com/paper/2205.11713/) (2022)
- [Gradient-based inference of abstract task representations for generalization in neural networks](https://mlsanity.com/paper/2407.17356/) (2024)

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

- [Separating the what and how of compositional computation to enable reuse and continual learning](https://mlsanity.com/paper/2510.20709/) (2025)
- [Sparks of cognitive flexibility: self-guided context inference for flexible stimulus-response mapping by attentional routing](https://mlsanity.com/paper/2502.15634/) (2025)

## Abstract (arXiv)

Animals survive in dynamic environments changing at arbitrary timescales, but such data distribution shifts are a challenge to neural networks. To adapt to change, neural systems may change a large number of parameters, which is a slow process involving forgetting past information. In contrast, animals leverage distribution changes to segment their stream of experience into tasks and associate them with internal task abstracts. Animals can then respond flexibly by selecting the appropriate task abstraction. However, how such flexible task abstractions may arise in neural systems remains unknown. Here, we analyze a linear gated network where the weights and gates are jointly optimized via gradient descent, but with neuron-like constraints on the gates including a faster timescale, nonnegativity, and bounded activity. We observe that the weights self-organize into modules specialized for tasks or sub-tasks encountered, while the gates layer forms unique representations that switch the appropriate weight modules (task abstractions). We analytically reduce the learning dynamics to an effective eigenspace, revealing a virtuous cycle: fast adapting gates drive weight specialization by protecting previous knowledge, while weight specialization in turn increases the update rate of the gating layer. Task switching in the gating layer accelerates as a function of curriculum block size and task training, mirroring key findings in cognitive neuroscience. We show that the discovered task abstractions support generalization through both task and subtask composition, and we extend our findings to a non-linear network switching between two tasks. Overall, our work offers a theory of cognitive flexibility in animals as arising from joint gradient descent on synaptic and neural gating in a neural network architecture.