Attention weights accurately predict language representations in the brain (2022)

Abstract: In Transformer-based language models (LMs) the attention mechanism converts token embeddings into contextual embeddings that incorporate information from neighboring words. The resulting contextual hidden state embeddings have enabled highly accurate models of brain responses, suggesting that the attention mechanism constructs contextual embeddings that carry information reflected in language-related brain representations. However, it is unclear whether the attention weights that are used to integrate information across words are themselves related to language representations in the brain. To address this question we analyzed functional magnetic resonance imaging (fMRI) recordings of participants reading English language narratives. We provided the narrative text as input to two LMs (BERT and GPT-2) and extracted their corresponding attention weights. We then used encoding models to determine how well attention weights can predict recorded brain responses. We find that attention weights accurately predict brain responses in much of the frontal and temporal cortices. Our results suggest that the attention mechanism itself carries information that is reflected in brain representations. Moreover, these results indicate cortical areas in which context integration may occur.

## introduction the attention mechanism has enabled transformerbased language models (lms) to achieve state-ofthe-art performance on many nlp tasks @xcite @xcite . in these lms the attention mechanism iteratively computes weighted sums of the embeddings of neighboring tokens in order to transform token embeddings into contextual representations. the resulting contextual representations have been used to produce predictions of brain responses that are more accurate than predictions derived from (static) lexical embeddings or from rnn-based contextual embeddings @xcite @xcite . these findings show that the attention mechanism produces contextual hidden state embeddings that contain more information reflected in brain responses than lexical embeddings alone. however, prior work has not determined whether the mechanism by which hidden state embeddings are constructed (i.e., the weights computed by the attention mechanism) is itself reflected in brain representations. to address this question, we used voxelwise encoding models (vm), a powerful data-driven method for determining what information is represented in the brain @xcite . in encoding models, brain responses to stimuli are modeled as a linear combination of feature spaces that each represent properties of the stimuli. these encoding models are used to determine where each feature space is represented in the brain. encoding models have previously been used to determine how variables such as lexical and contextual semantics are represented in brain responses @xcite @xcite . however, previous work did not directly use attention weights in encoding models to study whether and how the attention weights are reflected in brain responses. understanding the relationship between attention weights and brain representations would provide insight into where information integration across context occurs in the brain, and whether attention weights contain information relevant for human language processing. in this work, we use attention weights in encoding models to determine whether weights computed by the attention mechanism of transformer-based lms are represented in brain responses. we modeled functional magnetic resonance (fmri) data recorded from six human participants while they read a set of english natural language narratives. we presented the same stimuli as input to bert and gpt-2 and extracted their corresponding attention weights. we investigated how well attention weights predict brain responses, compared to other stimulus features such as lexical and contextual semantic word embeddings. we found that, across many cortical areas that have previously been associated with language processing, attention weights better predict brain responses than lexical embeddings. in a subset of these cortical areas, attention weights better predict brain responses than contextual embeddings, suggesting that the mechanism of context integration (attention weights) in transformer-based lms reflects aspects of brain responses that are not captured by the resulting contextual embeddings. moreover, these findings suggest cortical areas that could be involved in the integration of information from context words. ## methods to investigate whether attention weights predict brain responses, we analyzed functional magnetic resonance imaging (fmri) recordings of brain responses to english language narratives (section 3.1). we used the voxelwise encoding modeling (vm) framework @xcite to determine where attention weights are reflected in the recorded brain responses (section 3.2). we identified how well attention weights can predict brain responses relative to lexical semantic features, and compared the performance of attention weights across layers and across two lms (bert and gpt-2). then, to study whether attention weights contain additional information to lm hidden states, we determined how well attention weights extracted from bert can predict brain responses relative to contextual embeddings extracted from the hidden states of bert. we further did a focused analysis on cortical regions that have previously been associated with language representations: the high-level auditory cortex (ac), the superior temporal sulcus (sts), broca's area and the superior ventral premotor speech area (spmv) @xcite @xcite . while the attention mechanism in transformer-based lms and the concept of attention in cognitive neuroscience are distinct, the two share a conceptual relationship (lindsay, 2020). therefore, we also inspected brain representations in areas of the human attention network: the middle temporal visual area (mt), the intraparietal sulcus (ips) and the frontal eye field (fef) @xcite . we visualize the results on cortical maps made with pycortex @xcite . ## results to study the role of the attention mechanism in explaining brain responses, we tested how well attention weights can predict brain responses to natural language stimuli. we analyzed fmri recordings that were taken while human participants read a set of english narratives. we provided the text of the stimulus narratives to a transformer-based language model (bert or gpt-2) and extracted attention weights from each attention head. we used voxelwise encoding models to determine how well these attention weights can predict brain responses compared to semantic features and low-level sensory features. in order to determine how well attention weights extracted from bert explain brain responses compared to lexical semantic embeddings and sensorylevel features, we estimated an encoding model that maps lexical embeddings, bert attention weights and low-level features to brain responses. in figure 2 , we show the prediction performance of this joint model for one representative participant. results for five additional participants are qualitatively similar (see section c in the appendix). voxel hue denotes the feature space that best predicted brain response in that voxel: green for attention weights, red for lexical semantic embeddings and blue for low-level features. voxel opacity denotes the splitprediction performance of the corresponding feature space. the split-prediction performance denotes the contribution of a specific feature space in explaining the variance in the brain response. voxels shown in gray were not significantly predicted (fdr corrected, @xmath0.05). attention weights better predict brain responses in voxels located across the temporal and prefrontal cortices that have previously been associated with language processing: high-level ac, sts, broca's area and spmv. these results suggest that, in many language-related cortical areas, brain representations reflect more information about integration across words, than the lexical representation of the words themselves. ## conclusion in this work, we examined how well attention weights from transformer-based lms predict brain responses to linguistic stimuli. we find that attention weights from bert and gpt-2 produce accurate predictions across many cortical areas that have previously been associated with language processing (high-level auditory cortex, superior temporal sulcus, broca's area, superior ventral premotor speech area). attention weights better predict brain responses than lexical embeddings in most of these areas. attention weights are also better at predicting brain responses than contextual embeddings in portions of the language related areas. these findings show that attention weights capture information that is not contained in lexical or contextual embeddings, and indicate cortical areas that could be involved in integrating context across words. moreover, attention weights from the middle layers of the lm, that have been previously identified as linguistically relevant, contribute strongest to the predictions. while our results showed consistency between unidirectional and bidirectional models in terms of how well attention weights predict brain responses, more work is needed to understand the impact of model directionality on cognitive fit. in the future, comparing the results of our method of computing attention weights with more cognitivelyinspired methods could provide more specific insights into how attention weights in transformerbased lms relate to processes in the human brain.