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

AI summary

Overview

• Research area: Natural language processing and computational neuroscience — specifically, inductive biases for Transformer language models and the alignment of those models with human brain language networks.
• Technical level: Advanced (requires familiarity with self-attention, query/key/value projections, BERT, masked language modeling, fMRI decoding, and dimensionality-reduction methods such as PCA and PLS-SVD).
• Scope: The paper introduces two spatial constraints on self-attention — spatial querying and spatial reweighting — that produce topographic (spatially smooth, 2D-visualizable) organization in Transformer language models, evaluates them on sentiment classification, GLUE, and eight linguistic test suites, and aligns their topographic variability with fMRI responses in the human language network.

What This Paper Is About

Biological brains organize neurons spatially according to what they respond to, which makes brain activity relatively easy to visualize and interpret. Standard machine learning models have no such spatial bias: their internal representations are disorganized vector spaces that are hard to visualize. The paper's goal is to build spatial constraints directly into the Transformer's self-attention layer so that language models develop brain-like topographic organization without losing task performance, and then test whether that organization lines up with human brain responses to sentences.

Key Contributions

1. Two new self-attention motifs. The paper introduces spatial querying, where keys and queries are arranged on 2D grids and a local pool of queries is associated with a given key (via a binary intermediate matrix M, replacing QK^T with QMK^T), and spatial reweighting, where the standard fully connected output layer W^O is replaced with a locally connected layer W^O_local.

2. Feasibility demonstration at small scale. A 1-layer, single-head Topoformer is trained on IMDB sentiment classification, showing that spatial querying produces topography in the queries and keys, and spatial reweighting separately produces topography in the values and self-attention outputs (fc_out).

3. Scaling to a BERT architecture. The motifs are applied to a 16-layer BERT model trained with a masked language modeling objective (Topoformer-BERT), performing comparably to a non-topographic single-head control on the GLUE benchmark while showing topographic organization in every evaluated sublayer.

4. Brain alignment. Using fMRI data from N=5 participants reading 1,000 six-word sentences, the authors show that low-dimensional topographic variability in Topoformer-BERT aligns with that of the human functionally-defined language network.

Main Findings

• Topoformer-SQ performance: The 1-layer model with spatial querying alone reached 0.81 accuracy on the IMDB sentiment test set, versus 0.83 for an identical 1-layer Transformer without spatial querying.

• Topoformer-SQR performance: Adding spatial reweighting lowered test accuracy to 0.75 on the IMDB sentiment test set.

• Receptive field effects: Smaller r_SQ values yielded better accuracy and stronger topography, with a clear inverse relationship between r_SQ and IMDB accuracy (R² = 0.8612); r_SR had no measurable effect on performance (R² = 0.0137) and minimal effect on topography.

• Component-specific topography: Spatial querying encouraged topography in queries and keys only; spatial reweighting encouraged topography in the values and self-attention outputs. In Topoformer-BERT, keys and queries showed the greatest degree of topographic organization and values the weakest, though all were consistently above 0.

• Comparable GLUE results: Topoformer-BERT scored 80.1/80.1 on MNLI, 90.9 on SST-2, 75.1 on STSB, 51.2 on RTE, 86.6 on QNLI, 86.0 on QQP, 81.5 on MRPC, 46.3 on CoLA, and 75.31 overall on GLUE — similar to the non-topographic single-head control (76.9 overall) and the multi-head model (78.1 overall).

• Strong semantic selectivity, weaker syntactic selectivity: Across eight test suites of 76 sentences each, layer 15 of Topoformer-BERT showed strong topographic selectivity for broad semantic contrasts (animacy, concreteness, visuomotor properties) and weaker effects for minimally paired syntactic contrasts. Decoding for Licensing was approximately .18, with a Ridge classifier reaching 31% (interpreted as overfitting); chance level was 50%.

• Overlapping activations across categories: Despite significant selectivity, mean activity patterns were highly similar across conditions within a contrast — a property shared with a non-topographic BERT control and an off-the-shelf all-miniLM-L6-v2 model.

• Brain alignment: The first two components of the PLS-SVD alignment were significantly aligned between the human language network and each Topoformer-BERT sublayer, while later components were less likely to be aligned. The third component was spatially organized in each space but not aligned.

• Specificity controls: Alignment and voxel encoding model prediction were significantly greater between the trained Topoformer-BERT and the language network than for non-linguistic control regions and an untrained Topoformer variant. The authors note that alignment itself does not depend on using a topographic rather than standard transformer, but the Topoformer uniquely allows spatial visualization of the aligned component.

• Topographic organization in the brain: The generic topography statistic t_g for each of the five language parcels fell outside a null distribution computed from shuffled brain responses, indicating significant decay in unit response correlations with distance.

Methodology in Plain English

The authors take the ordinary self-attention computation and change two things about how units are arranged and connected. Units are placed on a 2D square grid (embedding dimension d = s², where s is the grid side length), so every unit has a spatial position.

First, in spatial querying, instead of matching each query to each key one-to-one, they match a small local pool of queries (defined by a circular receptive field with width r_SQ) to a single key. This is implemented with a binary matrix M inserted between Q and K, so the dot product is between a pooled query and a key. This smooths query representations locally and gives keys a spatial correspondence with queries.

Second, in spatial reweighting, the output projection W^O is turned into a locally connected layer whose connectivity is controlled by a receptive field width r_SR. The authors also found in preliminary experiments that large positive weights were needed to fully encourage topography, so they initialize the locally connected layer as the absolute value of a standard initialization multiplied by 10. This mimics the excitatory feedforward connections of biological brains and enhances local correlations.

They use a single attention head to keep organization within a head rather than across heads. For the feasibility study they train a 1-layer encoder-only Topoformer on IMDB with cross-entropy loss, 20 epochs, an Adam optimizer with learning rate 0.001, and batch sizes of 128 (SQ) and 256 (SQR), using d = 400 and RF sizes of 0.3 (SQ) and 0.1 (SQR). To interpret topography, they run selectivity analyses (two-tailed t-tests contrasting sentence conditions, reported as −s·log10(p)) and PCA, reshaping principal component weights onto the grid.

For the larger model, they train a 16-layer BERT on Bookcorpus-Wikipedia for 12 hours using the Cramming procedure, evaluate on GLUE (all tasks except WNLI), quantify topography with a statistic t_g relating unit correlation to unit distance (averaged over nine maximum distance values into t̄_g), and test eight linguistic suites. Finally, they collect event-related fMRI from N=5 participants reading 1,000 6-word sentences, define language regions within five anatomical parcels as voxels with t > 1 for sentences over non-words, and align brain and model representations using PLS-SVD computed on 80% of sentences and cross-validated on the remaining 20%.

Why This Matters

The paper is the first to demonstrate that Transformer language models can be trained to exhibit topographic organization similar to the human brain, opening a route to models that are both accurate at language tasks and more visually interpretable. It also provides evidence that the human language network contains spatial functional organization within a relatively functionally homogeneous network, rather than only across distinct networks.

Potential real-world applications:

• Model interpretability tooling: 2D topographic visualizations of selectivity could make it easier for researchers and auditors to inspect what a language model's units encode.
• Interpretable NLP systems: Spatial priors could support building models whose internal structure is easier to diagnose and debug in deployed language applications.
• Cognitive and clinical neuroscience: Topoformer models can generate testable hypotheses about the spatial layout of linguistic information in the human brain, where experiments are costly and slow.
• General brain modeling: Because Transformers are domain-general, the same motifs could be applied to model functional organization in non-language brain areas that Transformers can simulate.

Industry relevance: The work is directly relevant to teams working on language model interpretability and mechanistic analysis. It shows that spatial constraints can be added at scale (a 16-layer BERT trained in 12 hours) with GLUE scores close to a non-topographic control (75.31 vs 76.9 for the single-head control), suggesting a workable trade-off for organizations that prioritize transparency. The authors note support from Novus Technologies for computing resources and released training code at https://github.com/TahaBinhuraib/topoformer (fMRI analysis code available on request).

Future Directions

• Learned pooling and finer-grained interpretation: The authors suggest converting the binary pooling matrix M into a learnable locally connected matrix, and investigating finer-grained semantic dimensions and more extensive tests of syntactic knowledge.
• Monosemanticity and disentanglement: Addressing polysemanticity through extraction of monosemantic features or objectives that encourage disentangled representations could yield more interpretable topography in both models and brain hypotheses.
• Better brain data and broader brain coverage: The fMRI data were acquired in a rapid event-related design without sentence repetitions, making it noise-prone; future work should use datasets with multiple repetitions over specific linguistic content. The authors also propose extending analyses beyond the functionally-defined language network to larger-scale semantic organization across the whole brain.
• Scaling up and generalizing: The paper states that scaling Topoformers further holds promise for greater interpretability in NLP research and more accurate models of linguistic organization in the brain, and that topographic transformers may be useful for modeling brain organization more broadly.

Target Audience

This paper is best suited to researchers in NLP interpretability, computational neuroscience, and cognitive science of language — particularly those working on brain-language model alignment, neural encoding models, or inductive biases in Transformer architectures. It will also interest engineers building interpretable language systems, and readers already comfortable with self-attention mechanics, masked language modeling, and fMRI analysis pipelines. Readers without a background in Transformer internals or neuroimaging statistics will find the methods sections demanding, though the core idea of arranging units on a 2D grid and constraining connections locally is conveyed clearly in the introduction and Figure 1.

Authors’ 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.