Self-Attention-Based Contextual Modulation Improves Neural System Identification

TL;DR

This paper demonstrates that self-attention mechanisms enhance neural response predictions in visual cortex models, outperforming traditional CNNs by better capturing contextual modulation and top feature preferences.

Contribution

It introduces self-attention as a novel mechanism to improve modeling of contextual effects in neural systems, showing its effectiveness over standard CNN approaches.

Findings

Self-attention improves tuning curve correlation.

Self-attention enhances peak tuning accuracy.

Local receptive fields are crucial for overall tuning, surround information is key for peak tuning.

Abstract

Convolutional neural networks (CNNs) have been shown to be state-of-the-art models for visual cortical neurons. Cortical neurons in the primary visual cortex are sensitive to contextual information mediated by extensive horizontal and feedback connections. Standard CNNs integrate global contextual information to model contextual modulation via two mechanisms: successive convolutions and a fully connected readout layer. In this paper, we find that self-attention (SA), an implementation of non-local network mechanisms, can improve neural response predictions over parameter-matched CNNs in two key metrics: tuning curve correlation and peak tuning. We introduce peak tuning as a metric to evaluate a model's ability to capture a neuron's top feature preference. We factorize networks to assess each context mechanism, revealing that information in the local receptive field is most important for modeling overall tuning, but surround information is critically necessary for characterizing the tuning peak. We find that self-attention can replace posterior spatial-integration convolutions when learned incrementally, and is further enhanced in the presence of a fully connected readout layer, suggesting that the two context mechanisms are complementary. Finally, we find that decomposing receptive field learning and contextual modulation learning in an incremental manner may be an effective and robust mechanism for learning surround-center interactions.