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

TL;DR

This paper introduces Topoformer, a novel Transformer variant with topographic organization inspired by biological brains, achieved through spatial querying and reweighting, leading to more interpretable models and alignment with human brain data.

Contribution

The paper proposes a new self-attention mechanism that induces topographic organization in Transformers, enhancing interpretability and biological plausibility.

Findings

Topoformer achieves topographic organization in NLP models.

Topoformer performs comparably to standard models on NLP benchmarks.

Topographic features in Topoformer align with human brain language responses.

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.