Cortico-subcortical multi-head self-attention as a substrate for cognitive performance

Abstract

The neocortex is central to mammalian cognition, yet a computational framework that is both biologically constrained and capable of performing complex cognitive tasks remains missing. Here we show that cortico-thalamic circuits are well suited to implement multi-head self- and cross-attention, the mechanism underlying the cognitive abilities of transformer networks. We propose that layer 2/3 pyramidal cells maintain a recurrent key-value memory, while layer 5 pyramidal cells decode the memory retrieved by an incoming query. The computation of keys, values and queries maps onto core and matrix thalamo-cortical projections, distributed across the micro- and macro-columns of a cortical area. One cortical area forms an attention head, and cortex a multi-head self-attention network. The same thalamo-cortical microcircuit also calculates sensory prediction errors guiding gradient-based synaptic plasticity. A reward-prediction error gates via basal ganglia the cortical output and the re-activation of hippocampal memories. The trained network aligns with human intracranial recordings during speech perception. Overall, the suggested cortico-subcortical attention circuit may represent a substrate for the cognitive capacity of mammals.