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

TL;DR

This paper introduces WiNN, a neural network model that combines fast context inference with attentional routing to enable flexible, rule-based behavior in image-based tasks, reducing catastrophic forgetting and generalizing to new rules.

Contribution

WiNN extends Hummos' fast-and-slow learning to complex image tasks by integrating a pretrained vision network with a dynamic context state for flexible rule inference.

Findings

WiNN autonomously infers underlying rules in image tasks.

Requires fewer examples than control models for rule inference.

Generalizes to unseen rules through context-state updates.

Abstract

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.