TopoLM: brain-like spatio-functional organization in a topographic language model

Authors: Neil Rathi, Johannes Mehrer, Badr AlKhamissi, Taha Osama A Binhuraib, Nicholas Blauch, Martin Schrimpf
Primary area: applications to neuroscience & cognitive science
Keywords: language modeling, topography, fMRI, neuroscience

Abstract:
Neurons in the brain are spatially organized such that neighbors on tissue often exhibit similar response profiles. In the human language system, experimental studies have observed clusters for syntactic and semantic categories, but the mechanisms underlying this functional organization remain unclear. Here, building on work from the vision literature, we develop TopoLM, a transformer language model with an explicit two-dimensional spatial representation of model units. By combining a next-token prediction objective with a spatial smoothness loss, representations in this model assemble into clusters that correspond to semantically interpretable groupings of text and closely match the functional organization in the brain's language system. TopoLM successfully predicts the emergence of a spatially organized cortical language system as well as the organization of functional clusters selective for fine-grained linguistic features empirically observed in human cortex. Our results suggest that the functional organization of the human language system is driven by a unified spatial objective, and provide a functionally and spatially aligned model of language processing in the brain.Neurons in the brain are spatially organized such that neighbors on tissue often exhibit similar response profiles. In the human language system, experimental studies have observed clusters for syntactic and semantic categories, but the mechanisms underlying this functional organization remain unclear. Here, building on work from the vision literature, we develop TopoLM, a transformer language model with an explicit two-dimensional spatial representation of model units. By combining a next-token prediction objective with a spatial smoothness loss, representations in this model assemble into clusters that correspond to semantically interpretable groupings of text and closely match the functional organization in the brain's language system. TopoLM successfully predicts the emergence of a spatially organized cortical language system as well as the organization of functional clusters selective for fine-grained linguistic features empirically observed in human cortex. Our results suggest that the functional organization of the human language system is driven by a unified spatial objective, and provide a functionally and spatially aligned model of language processing in the brain.

Decision: accept (oral)
Reviewer scores: 8, 8, 8

Reviews and discussion:
<thread>
creation date:
1730029218670

thread_type:
official_review

<thread_object>
creator: reviewer_u3qv

<content>
summary:
The paper presents a language model whose training objective is based on a joint loss function: the standard next-word prediction loss and an additional spatial regularization term, which enforces a constraint where adjacent neurons exhibit more similar activation profiles than those of more distant neurons. The primary contributions of the paper are: (1) demonstrating that such a model can be successfully trained, resulting in an increase in local spatial correlations, (2) showing that its performance on standard benchmarks is comparable to models trained without the spatial regularization, and (3) providing evidence that the learned representations exhibit some degree of isomorphism with neural representations. For example, certain spatial clusters show stronger activation in response to well-formed sentences compared to nonsensical ones, and there appears to be a distinction between cluster-responses to nouns and verbs.

strengths:
The paper addresses a novel and important issue. There has been little work on training LLMs with topography constraints.  Successful training opens the door to interesting analyses, and potentially more flexible approaches for supervised learning of model-to-brain alignment. The writing is clear, and the analyses are interesting in that the authors apply well established statistical approaches used in the human neuroimaging literature to the model itself.

weaknesses:
**P.S  following the discussion period: I believe the more critical perspective that was adopted by the authors during the review process improves the soundness of the paper. I have modified my evaluation.**

1. The spatial loss term is a key element in the work and should be better motivated. 
1a. As phrased in the paper, the loss term constrains neurons that are closer to each other (according to a predefined spatial arrangement) to have more strongly  correlated activation profiles than those farther apart. However, this introduces a global, long-range constraint, because activation profiles of neurons that are far apart are not independent, but constrained to differ. To enforce a strictly local connectivity constraint, the authors could have employed a formulation that emulates Gaussian smoothing or another local averaging kernel, which would only consider the immediate neighborhood of each neuron. Reading the authors' motivation for the spatial regularization term on page 3, I expected a local spatial smoothing constraint rather than the more global constraint that was implemented. They write: "This loss function measures spatial smoothness, which serves as an efficiently computable proxy for neural wiring length: neurons located close to one another should have similar response profiles—i.e., their activations should be correlated (Lee et al., 2020)". In practice, for computational efficiency, the authors do not transform an entire layer to a grid-topology, but select 5 smaller areas; however, the point still holds.

1.b In addition, the spatial-loss term introduces certain computational characteristics that should be considered in greater depth. Given its properties, it may encourage learning of decorrelated features, which could be advantageous for downstream transfer learning and fine-tuning.  In fact, this principle is used in a regularization technique introduced by Zhu et al (2023; https://arxiv.org/abs/2306.13292) to learn a potentially richer feature set. On the other hand, the spatial regularization term also pushes in the opposite direction: it enforces strong correlations between nearby units, potentially reducing the model’s effective capacity. In the extreme case, it can force two adjacent units to have almost identical activation patterns. This could limit the model’s ability to learn distinct features, which is a fundamental goal in neural network training. The authors briefly touch on this issue in the results section (p. 9), where they mention that the spatial regularization may help prevent overfitting, but this point asks for a more thorough treatment. 

2. The effects of the spatial loss on neuron-to-neuron similarity should be quantitatively evaluated
It is not clear how and whether TopoLM training changes the covariance structure of neural firing. It is important to rule out the possibility that non topographic models produce similar unit-to-unit correlation patterns compared to TopoLM. Otherwise, it would mean that training with spatial loss is not necessary for obtaining the resulting covariance structure.  This could be done by computing, for each neuron, a histogram of correlation values and averaging these histograms across all neurons.  If both TopoLM and standard models produce similar correlation distributions, it would suggest that typical neural networks naturally develop a graded connectivity structure similar to what TopoLM enforces, but this structure may manifest in a shuffled form in the absence of an explicit spatial arrangement.

3. Some of the spatial results appear determined or constrained by the training objective and other procedures. 
The finding that training produces stronger spatial smoothness and clustering seems to me to be a direct consequence of the spatial constraint in the training objective, which inherently encourages the formation of locally correlated clusters with weak correlations between remote clusters. Essentially, it implements a result similar to spatial data smoothing. Similarly, the fact that a simulation of fMRI readout-sampling (which applies further local averaging of unit activations) produces an even smoother data map is expected, as this is a direct result of the averaging imposed by the readout kernel.  Consequently, these properties do not emerge organically from the model’s spatial-loss training or fmri-like sampling, but instead reflect the successful implementation of the externally imposed constraints and workflow.  Because of this, the documentation of spatial clusters in statistical analyses is not surprising. For example, the authors write (p. 7) that “Contrasting model activations to verb and noun stimuli yields clear verb- and noun-selective clusters across layers in our topographic model”, but given the smoothing constraint applied, it would be very likely that the results would manifest in the form of clusters, as nearby units are constrained to be non-independent.  If these results are rephrased to be presented as a validation or sanity check that would make sense, but currently they are phrased as if they are a novel result. 

4. Some of the language-activation results appear dependent on the definition of language regions
The authors define a language system in TopoLM as the set of units that respond more strongly to grammatically valid sentences than to strings of non-words, which is analogous to the use of a functional localizer task in fMRI studies.  Then, for the network units identified by this condition, they evaluate responses to well-formed sentences and different types of language and language-like stimuli including strings of scrambled non-words or stings of phonotactically plausible non-words (jabberwocky sentences).  However, the aforementioned test used to define the localizer is not entirely independent of the subsequent evaluations. While it does not fully determine them, the localizer identifies units that are likely sensitive to word meaning and syntactic relationships, which are properties not found in the other stimulus types. Therefore, it is not entirely surprising that in the areas identified by the localizer, the strongest responses were to sentences, as compared to unconnected words, non-words, or Jabberwocky sentences.

It would have been more convincing if these results were found without the use of a localizer, using the spatial distribution alone: if the local clusters formed in the network are functionally meaningful, it would be important to see what information is encoded in those clusters.  If the strongest clusters naturally show a gradient of responses (e.g., from sentences to unconnected non-words), this would suggest that such distinctions are inherently emphasized during learning. Clusters can be identified via regional homogeneity (as done for fMRI) by computing correlation in response profiles across neurons.

The last two points also bear on the author’s ultimate conclusion, where they write, “Our results suggest that the functional organization of the human language system is driven by a unified spatial objective..”  I don’t think this follows from the results,  because I don’t see a basis for the analogy between the explicit spatial objective implemented in the paper (which produces some of the results) and those factors that ultimately lead to local correlations in the brain. Specifically, in the brain it is probably *unlikely* that there exists an objective that constraints nearby areas to fire similarly. This more likely emerges from cytoarchitecture constraints which produce different gradients of correlation (sometimes within the same neural population; https://pmc.ncbi.nlm.nih.gov/articles/PMC3412441/). 

5. General comment: I do not see a general weakness in the fact that benchmarks are not consistently improved by TopoLM. But what is lacking are more clear statements about whether TopLM develops different representations than non-Topo architectures.  This can be quantified by comparing representation-dissimilarity matrices, CCA, CKA or other methods for quantifying similarity of representations.  Of course, a comparison of these against RDMs computed from brain areas could also be helpful.

questions:
1.  I did not understand the argument about smoothing reducing wiring length. From my understanding, the TDANN network, even after training with the spatial regularization, retains the same dimensionality (weight matrix) as a non-spatial network. What then is the basis for the claim that this minimizes wiring length? (This refers to p. 3: “The TDANN’s spatial smoothness, as implemented with an additional loss term, is an indirect but efficient approach to minimizing local wiring-length.”) Additionally, from a functional perspective, in vision, hierarchical processing may reduce the need for long-range connectivity, as much processing occurs in localized circuits (even though reciprocal connections still exist). But in language processing, long-range connectivity seems essential. For instance, world knowledge can immediately influence lexical processing, and predictions can impact activity in primary and secondary auditory cortices, indicating that long-range connections are necessary in this domain.

2. page 6, the statement that a successful alignment between brain and model is accompanied by “these clusters have consistent response profiles”; the meaning of this phrase was unclear.

3. The analysis is focusing on univariate methods, where preference for sentence or word class is identified by differences in single-unit activation. But it is well known that the brain also represents information using multivariate (population) codes. What if a set of units, each of which shows similar activation for two classes of linguistic stimuli, can still be used to decode stimulus class through a population code?  From the perspective of the current work, these units will be defined as being outside the identified ‘selectivity’ clusters (and hence not considered ‘core’ to language). However, to the extent that a linear probe indicates that this population-code contains relevant information, there is no a priori reason to exclude them as being less relevant, essential or core to language.  For example, the null effect reported for sentences vs. unconnected words (p.6; results section) could still be consistent with the possibility that this set of units can differentiate between the two classes if a classifier trained and tested on their activations using linear probing. In short, the authors should better motivate why the operationalize ‘centrality’ to language based on univariate tests rather than population codes. 

4. For the benchmarks (p. 9) it seemed that fine-tuning involved retraining the entire model, rather than freezing the trained model, using it as feature extractor, and using a linear head for classification.  To evaluate the quality of features produced (rather than the model’s capacity to learn a new task) it would be useful to treat TopoLM and the control architecture as feature extractors, and apply the same procedure for learning a new task.

References

Zhu, J., Evtimova, K., Chen, Y., Shwartz-Ziv, R., & LeCun, Y. (2023). Variance-covariance regularization improves representation learning. arXiv preprint arXiv:2306.13292.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730538247100

thread_type:
official_review

<thread_object>
creator: reviewer_tbeb

<content>
summary:
The study introduces TopoLM, a model based on the transformer architecture, aimed at exploring the spatial functional organization of human language systems. This model optimizes autoregressive language modeling by incorporating two-dimensional spatial encoding, along with task and spatial loss, which encourages response smoothness among adjacent units. Experiments demonstrate that the spatial functional organization generated by TopoLM closely aligns with the neuronal clusters found in the brain's language system. Although TopoLM scores slightly lower on certain behavioral tasks compared to non-topological baselines, its performance is comparable on other downstream tasks and brain alignment benchmarks. This indicates that TopoLM successfully captures the topological structure involved in language processing when trained on natural text, providing new insights into the functional organization of language systems.

strengths:
TopoLM introduces a novel approach to understanding the spatial organization of language processing in the human brain.  By integrating two-dimensional spatial encoding into a transformer architecture, the model effectively captures the topological structure of language, which has not been extensively explored in existing literature.  This innovative perspective opens new avenues for research in language modeling and cognitive neuroscience.  Meanwhile, this research has significant implications for both artificial intelligence and neuroscience.  By aligning computational models more closely with the biological structures of the brain, TopoLM paves the way for more effective natural language processing applications.  Furthermore, it contributes to a deeper understanding of the neural mechanisms underlying language, potentially informing both theoretical frameworks and practical applications in cognitive science and related fields.

weaknesses:
While the paper presents a promising theoretical framework, the empirical evaluation of TopoLM is somewhat limited.  The experiments primarily focus on a small set of tasks, which may not fully capture the model's robustness across diverse linguistic contexts.  Meanwhile, this paper does not sufficiently explore or compare TopoLM with other existing topological models in language processing.  Including a comparative analysis with relevant models could strengthen the argument for TopoLM's effectiveness and novelty.  Highlighting specific improvements or unique aspects of TopoLM in relation to these models would enhance the discussion around its contributions to the field.

questions:
Could you elaborate on the specific architectural choices made in TopoLM?  For instance, what motivated the selection of particular layers or activation functions, and how do these choices contribute to the model's performance compared to traditional language models?

Can you provide more insight into the training process for TopoLM?  Specifically, what datasets were used, and what were the criteria for their selection?  Additionally, how did you address potential issues like overfitting during training?

In your discussion of the biological implications of TopoLM, could you provide references or data from neuroimaging studies that support the connection between the topological aspects you incorporate and the way the human brain processes language?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730593234830

thread_type:
official_review

<thread_object>
creator: reviewer_3qer

<content>
summary:
This work introduces TopoLM, a transformer language model incorporating information about the spatial arrangement of the units to model brain-like activity in the processing of language. In particular, the spatial information is integrated by adding a spatial smoothing term to the classical autoregressive loss of decoder-only transformers. By performing functional localization experiments, the authors identify clusters of artificial neurons in TopoLM whose activations are language-selective. The authors highlight the similarity between TopoLM clusters of activations and response in the brain by conducting experiments on verb-noun and concrete-abstract selectivity.

strengths:
The authors contextualized and described with clarity the strategy introduced to incorporate spatial information of activations in transformers for natural language processing. Both the training details and experiments are very well explained and organized.

The work, inspired by previous efforts in vision models, introduces for the first time (to the reviewer knowledge) explicit spatial information in the loss of transformers trained on natural language.

weaknesses:
Given the strong inspiration from human neural activity studies of the work, the work would benefit from more extended discussion and analyses related to some architectural choices (e.g. random permutation of spatial position between layers and role of multi-head attention).

questions:
Some further experiments and discussions could further strengthen the work:
- How does the choice of randomly permuting positions at each layer influence the representations? It would be extremely interesting to consider different types of mapping between spatial positions at successive layers.
- The authors stress the role of considering architectures with multi-head self-attention. Is this type of architecture expected to better model human neural activity? 
- In the NLP community, it has been observed a hierarchical organization of syntactic and semantic information in hidden layers of deep learning models (e.g. Belinkov, ACL, 2017; Manning, PNAS, 2019). Do the TopoLM clusters of activation exhibit some level of hierarchy through the layers of the model?
- In case it is feasible, comparing TopoLM representations/clusters also to Topoformer would further strengthen the quality of experiments.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732426040061

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
General response #1

comment:
We are very grateful to all 3 reviewers for their critical and very constructive evaluation of our work on TopoLM, a topographic language model showing hallmarks of language processing in the human brain. With this first response, we are addressing a number of important points raised by the reviewers, while other points will be addressed as soon as the associated additional analyses are completed.

Multiple reviewers suggested a comparison against another topographic baseline model. The version of our paper being evaluated here contains comparisons of TopoLM with its non-topographic equivalent (“non-topographic”), in which the spatial correlation loss term is switched off during training. We are currently working on evaluating the only other existing topographic language model (Topoformer BERT; BinHuraib 2024, https://openreview.net/forum?id=R6AA1NZhLd) on the main datasets we presented in the paper (Fedorenko 2024, fig. 2; Hauptman 2024, fig. 3; Moseley 2014, fig. 4) and will add the results in our response as soon as possible. 

We address the other points raised by the reviewers in individual responses below.
</content>
</thread_object>

</thread>
<thread>
creation date:
1733220517274

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Summary of reviewer feedback and request for area chair consideration

comment:
Dear area chairs,

please find below a summary of the points raised by the reviewers and how we responded to them. 

## Response to reviewer 3QeR

In our latest response to reviewer 3QeR (https://openreview.net/forum?id=aWXnKanInf&noteId=yz1b8Rc7DQ), we have summarized the  points we addressed:
- Evaluation of Topoformer BERT as a baseline topographic language model (same point requested by the other 2 reviewers)
- Effect of random permutations of unit locations on topographic organization
- The role of multi- vs. single-head attention 

As a result, reviewer 3QeR raised their score from 6 to 8. 

## Response to reviewer TBeb

In our latest responses to reviewer TBeb (https://openreview.net/forum?id=aWXnKanInf&noteId=v48FMPue7b), we have summarized the  points we addressed:
- Evaluation of Topoformer BERT as a baseline topographic language model (same point requested by the other 2 reviewers)
- Detailed architectural and training decisions
- Topography in TopoLM and the brain

Since we believe we have addressed the main points reviewer TBeb had raised, we have asked for a reassessment of their evaluation of our paper; we have not received a response to our request at the time of writing. 

## Response to reviewer U3qv

Here, we are summarizing the  points we addressed in multiple responses to reviewer U3qv:
- **Evaluation of Topoformer BERT** as a baseline topographic language model (same point requested by the other 2 reviewers)
- **Objective function and biological implementation:** Relation between the spatial correlation loss and reduced wiring length
- **Discussion on whether cluster formation are surprising**
- **Clarification** with regard to the second condition under which we see the spatio-functional alignment between brain and model
- **Statistical analysis:** Uni- vs. multivariate 
- **Implementation of spatial correlation loss**
- **Potential recursiveness of the core language system results**
- **Clarification** of claim on unified spatial objective in the conclusion
- **Relation between our TDANN spatial correlation loss and factors ultimately leading to local correlations in the brain**
- **Brain alignment** analysis using representational similarity instead of linear predictivity. 

As a result, reviewer U3qv raised their score from 6 to 8. 

We hope that the area chair considers this context in their evaluation of the TopoLM paper. 

Thank you!
</content>
</thread_object>

</thread>


Metareview:
metareview:
This paper modifies the standard autoregressive setting by adding a regularizing term that encourages nearby parts the model to behave similarly, with the motivation that this mimics known brain function. All reviewers agree that the idea is interesting and that the paper is ably executed.