TopoLM: brain-like spatio-functional organization in a topographic language model (2410.11516v3)

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 LLM 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 the spatio-functional organization of a 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.

• The paper introduces a novel topographic transformer that integrates a 2D spatial loss to form semantically organized clusters.
• The model mirrors neural clustering in the human cortex, replicating verb/noun and concrete/abstract word selectivity observed in neuroscience studies.
• Experimental benchmarks like GLUE and Brain-Score demonstrate TopoLM’s competitive performance and enhanced cognitive alignment.

TopoLM: A Topographic Approach to LLMing

In the presented paper, Rathi et al. explore the spatial organization of neuronal clusters in the human language system by introducing a novel LLM called TopoLM. This model adapts the transformer architecture to integrate a two-dimensional spatial representation, mirroring the spatial clustering observed in the human cortex. Through a combined objective incorporating both next-token prediction and spatial smoothness loss, TopoLM effectively organizes its representations into semantically interpretable clusters analogous to those found in brain activity.

Model Architecture and Implementation

TopoLM is an adaptation of the conventional transformer architecture, distinctively embedding model units into a two-dimensional grid structure. The spatial correlation loss introduced in TopoLM optimizes the arrangement of these units by minimizing wiring costs. As a result, neurons with similar response profiles naturally form clusters. This architecture facilitates brain-like spatio-functional organization without the need for explicit brain data during training, relying solely on natural text input.

Experimental Validation

The researchers validated TopoLM's performance through various neuroscience-inspired metrics:

1. Core Language System Selectivity: TopoLM successfully replicates the clustering of language selective regions in the brain. The model's response profiles align with known brain data, demonstrating coherent clustering across different linguistic stimuli.
2. Verb-Noun Clustering: Using a paradigm from previous empirical studies, TopoLM was evaluated for its ability to model verb- and noun-selective regions. The simulation results reflected brain data with significant spatial clustering for these categories.
3. Concrete vs. Abstract Word Selectivity: TopoLM also shows selective clustering specific to concrete word stimuli, replicating patterns observed in neuroimaging studies. Notably, abstract words elicit less clustering, aligning with empirical findings and supporting the model's cognitive validity.

Benchmarking Performance

TopoLM was put to the test against several benchmarks to evaluate its functional alignment with human cognition:

• BLiMP served to ascertain the model's linguistic proficiency through minimal pairs, revealing slightly hindered performance compared to non-topographic transformers.
• GLUE benchmarks tested downstream task capabilities, where TopoLM marginally outperformed the baseline due to the potential regularizing effect of spatial loss during fine-tuning.
• Brain-Score assessments compared neural alignment, showing competitive performance and, in some cases, superior brain alignment on specific tasks.

Implications and Future Directions

The introduction of TopoLM exemplifies an innovative approach to LLMing, extending beyond mere functional similarities with the brain to incorporate spatial coherence in its architecture. This development posits a unified explanation for cortical organization principles that transcend both visual and linguistic domains.

From a theoretical standpoint, TopoLM enriches our understanding of the spatial basis in cognitive processing. Practically, the model paves the way for future enhancements in AI systems that require deeper cognitive and neural alignment. Moreover, TopoLM's capability to predict new clustering patterns could inspire targeted experimental designs in neuroscience, facilitating the discovery of as yet unidentified linguistic patterns in the human brain.

TopoLM represents a significant advancement towards understanding the computational underpinnings of language processing. Its architecture serves as a promising framework for future research into integrated spatial-functional modeling, potentially unveiling new avenues in AI and cognitive neuroscience.