Abstract
Spatial functional organization is a hallmark of biological brains: neuronsare arranged topographically according to their response properties, atmultiple scales. In contrast, representations within most machine learningmodels lack spatial biases, instead manifesting as disorganized vector spacesthat are difficult to visualize and interpret. Here, we propose a novel form ofself-attention that turns Transformers into "Topoformers" with topographicorganization. We introduce spatial querying - where keys and queries arearranged on 2D grids, and local pools of queries are associated with a givenkey - and spatial reweighting, where we convert the standard fully connectedlayer of self-attention into a locally connected layer. We first demonstratethe feasibility of our approach by training a 1-layer Topoformer on a sentimentclassification task. Training with spatial querying encourages topographicorganization in the queries and keys, and spatial reweighting separatelyencourages topographic organization in the values and self-attention outputs.We then apply the Topoformer motifs at scale, training a BERT architecture witha masked language modeling objective. We find that the topographic variantperforms on par with a non-topographic control model on NLP benchmarks, yetproduces interpretable topographic organization as evaluated via eightlinguistic test suites. Finally, analyzing an fMRI dataset of human brainresponses to a large set of naturalistic sentences, we demonstrate alignmentbetween low-dimensional topographic variability in the Topoformer model andhuman brain language network. Scaling up Topoformers further holds promise forgreater interpretability in NLP research, and for more accurate models of theorganization of linguistic information in the human brain.