Topoformer: brain-like topographic organization in Transformer language models through spatial querying and reweighting

TL;DR

Topoformer introduces brain-inspired topographic priors into Transformer language models via spatial querying and spatial reweighting, inducing local organization in queries/keys and values/outputs. The approach is validated first on a 1-layer encoder and then scaled to Topoformer-BERT, achieving comparable NLP task performance while revealing clear topographic structure. In parallel, the authors show that the model's topographic representations align with the human language network measured by fMRI, suggesting shared low-dimensional variability between artificial and biological language systems. This work advances interpretability in NLP by coupling spatial priors with representational analysis and brain-alignment evidence, with potential implications for modeling brain organization and improving linguistic representations.

Abstract

Spatial functional organization is a hallmark of biological brains: neurons are arranged topographically according to their response properties, at multiple scales. In contrast, representations within most machine learning models lack spatial biases, instead manifesting as disorganized vector spaces that are difficult to visualize and interpret. Here, we propose a novel form of self-attention that turns Transformers into "Topoformers" with topographic organization. We introduce spatial querying - where keys and queries are arranged on 2D grids, and local pools of queries are associated with a given key - and spatial reweighting, where we convert the standard fully connected layer of self-attention into a locally connected layer. We first demonstrate the feasibility of our approach by training a 1-layer Topoformer on a sentiment classification task. Training with spatial querying encourages topographic organization in the queries and keys, and spatial reweighting separately encourages topographic organization in the values and self-attention outputs. We then apply the Topoformer motifs at scale, training a BERT architecture with a masked language modeling objective. We find that the topographic variant performs on par with a non-topographic control model on NLP benchmarks, yet produces interpretable topographic organization as evaluated via eight linguistic test suites. Finally, analyzing an fMRI dataset of human brain responses to a large set of naturalistic sentences, we demonstrate alignment between low-dimensional topographic variability in the Topoformer model and human brain language network. Scaling up Topoformers further holds promise for greater interpretability in NLP research, and for more accurate models of the organization of linguistic information in the human brain.

Figures (18)

• Figure 1: Spatial querying and reweighting operations in the "Topoformer". The first row shows standard querying operations in the attention module of a single-head Transformer, and the second row shows the spatial counterparts used in the Topoformer. Standard querying associates a single query dimension of token i with a single key dimension of token j. In contrast, spatial querying associates a local pool of query dimensions with a given key dimension, through the intermediate local pooling matrix M. Standard (dense) reweighting applies a fully connected linear layer $W^O$ to the outputs of a single attention head (typically to combine the outputs of multiple attention heads). In our formulation, we use a locally connected layer $W^{O}_{\text{local}}$ in its place (spatial reweighting). While the figure illustrates querying for a pair of tokens and reweighting for a single token, when processing a full sequence, there is a 2D grid of the form shown here for each token. Each heatmap shows the loadings of the second PC of responses (top: control model, bottom: Topoformer-SQR model with spatial querying and reweighting).
• Figure 2: Topographic organization across sublayers with spatial querying and reweighting.A. Topoformer-SQ produces topography in the keys and queries, but not the values or self-attention outputs. Each column shows a different sublayer representation within a self-attention block (keys, queries, values, and fc_out). The representations were obtained by averaging across the tokens in each sentence from the IMDB sentiment classification test set maas-etal-2011-learning. The first row shows selectivity for positive vs. negative sentiment sentences. We plot the selectivity significance value (see Appendix \ref{['appendix:selectivity']}), where $s=2$ corresponds to positive selectivity with $p=0.01$, and $s=-2$ corresponds to negative selectivity with the same significance level. The second and third rows show the PC weights for the first and second components, respectively. B. Topoformer-SQR produces topography in the values and self-attention outputs. The format is the same as for panel A, but for brevity we show only the values and self-attention outputs as the keys and queries show a similar topographic organization from spatial querying.
• Figure 3: Topographic organization across all layers of Topoformer-BERT. A. We quantified the extent of topography in Topoformer-BERT via the generic topography statistic $\bar{t}_g$ (Equation \ref{['eq:generic_topostat']}) which intuitively relates the degree of correlation with the distance between pairs of units averaged over a range of scales (nine maximum distance values). B. The topography statistic $t_{g,d}$ for layer 15 computed at each of a range of scales (maximum distance value), highlighting the particularly strong local topography for queries and keys. C. Visualization of the first principal component (PC) weights for keys and fc_out sublayers.
• Figure 4: Selectivity-based interpretation of topographic organization in Topoformer-BERT. Each panel shows the selectivity of Topoformer-BERT layer 15 (keys), for a given contrast. Each test suite contains two contrasting conditions each with a set of sentences; unit activities are computed as the mean over tokens for each sentence, and the conditions are contrasted with a $t$-test. We plot the selectivity significance value, as in Figure \ref{['fig:small_scale']}. The first row contains sentences with natural variability, whereas the bottom row contains results from constructed minimal pairs differing in only one word across conditions. To ensure visibility of effects regardless of size, we used different statistic ranges for plotting of each row: $s=10$ for the top row, and $s=2$ for the bottom row.
• Figure 5: Alignment of topographic representations in the human language network and Topoformer-BERT model.A. Illustration of the PLS-SVD alignment approach for a single participant and model sublayer representation. B. Alignment quantified across all 10 components, and each sublayer of Topoformer-BERT layer 15. The alignment of components is computed as the correlation of respective cross-validated PLS-SVD component scores for brain and model representations. Error bars indicate 95% confidence intervals over 5 participants.