# Sparks of cognitive flexibility: self-guided context inference for flexible stimulus-response mapping by attentional routing

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

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 15 of its 45 references that are in the map.
- First posted on arXiv 21 Feb 2025.

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

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Model-Agnostic Meta-Learning for Fast Adaptation of Deep Networks](https://mlsanity.com/paper/1703.03400/) (2017)
- [Overcoming catastrophic forgetting in neural networks](https://mlsanity.com/paper/1612.00796/) (2016)
- [Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification](https://mlsanity.com/paper/1502.01852/) (2015)
- [Gradient Episodic Memory for Continual Learning](https://mlsanity.com/paper/1706.08840/) (2017)
- [On First-Order Meta-Learning Algorithms](https://mlsanity.com/paper/1803.02999/) (2018)
- [Disentangling by Factorising](https://mlsanity.com/paper/1802.05983/) (2018)
- [Task-Free Continual Learning](https://mlsanity.com/paper/1812.03596/) (2018)
- [Continual Learning and Catastrophic Forgetting](https://mlsanity.com/paper/2403.05175/) (2024)
- [Superposition of many models into one](https://mlsanity.com/paper/1902.05522/) (2019)
- [Thalamus: a brain-inspired algorithm for biologically-plausible continual learning and disentangled representations](https://mlsanity.com/paper/2205.11713/) (2022)
- [Thalamocortical contribution to solving credit assignment in neural systems](https://mlsanity.com/paper/2104.01474/) (2021)
- [Characterising representation dynamics in recurrent neural networks for object recognition](https://mlsanity.com/paper/2308.12435/) (2023)
- [Flexible task abstractions emerge in linear networks with fast and bounded units](https://mlsanity.com/paper/2411.03840/) (2024)
- [Modulation of early visual processing alleviates capacity limits in solving multiple tasks](https://mlsanity.com/paper/1907.12309/) (2019)

## Abstract (arXiv)

Flexible cognition demands discovering hidden rules to quickly adapt stimulus-response mappings. Standard neural networks struggle in such tasks requiring rapid, context-driven remapping. Recently, Hummos (2023) introduced a fast-and-slow learning algorithm to mitigate this shortcoming, but its scalability to complex, image-computable tasks was unclear. Here, we propose the Wisconsin Neural Network (WiNN), which extends Hummos' fast-and-slow learning to image-computable tasks demanding flexible rule-based behavior. WiNN employs a pretrained convolutional neural network for vision, coupled with an adjustable "context state" that guides attention to relevant features. If WiNN produces an incorrect response, it first iteratively updates its context state to refocus attention on task-relevant cues, then performs minimal parameter updates to attention and readout layers. This strategy preserves generalizable representations in the sensory and attention networks, reducing catastrophic forgetting. We evaluate WiNN on an image-based extension of the Wisconsin Card Sorting Task, revealing several markers of cognitive flexibility: (i) WiNN autonomously infers underlying rules, (ii) requires fewer examples to do so than control models reliant on large-scale parameter updates, (iii) can perform context-based rule inference solely via context-state adjustments-further enhanced by slow updates of attention and readout parameters, and (iv) generalizes to unseen compositional rules through context-state updates alone. By blending fast context inference with targeted attentional guidance, WiNN achieves "sparks" of flexibility. This approach offers a path toward context-sensitive models that retain knowledge while rapidly adapting to complex, rule-based tasks.