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

Venue: ICLR_cc_2024_Workshop_Re-Align
Published: 2024-02-08 21:03:30

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.

Main paper: 
 introduction: 
 Biological brains are spatially organized, containing category-selective areas (Kanwisher, 2010) , broad feature maps that tile individual cortical areas (Konkle & Oliva, 2012; Bao et al., 2020) and the cortex more broadly (Huth et al., 2012; 2016; Margulies et al., 2016) , and large-scale distributed networks (Yeo et al., 2011; Braga et al., 2020) . Particularly within brain regions, this spatial organization is one way in which the human brain, a vastly complex "black box", is more naïvely interpretable than modern deep neural networks (DNNs), whose units have functional properties organized without simple spatial priors. Recent work in computational neuroscience has bridged this gap in DNNs trained for vision, demonstrating that local smoothness or wiring cost minimization objectives can be incorporated into DNNs to encourage the development of smooth functional organization of responses, which can then be easily visualized in 2D (Lee et al., 2020a; Blauch et al., 2022; Keller & Welling, 2021; Doshi & Konkle, 2021; Margalit et al., 2023; Lu et al., 2023) , building upon classic approaches (Kohonen, 1982; Jacobs & Jordan, 1992) . In addition to simulating topographic properties within regions, topographic vision models have also explained the hierarchical organization of topographic information from earlier to later visual areas (Margalit et al., 2023; Lu et al., 2023) . One topographic vision model has even demonstrated the emergence of spatial clusters corresponding to ventral, lateral, and dorsal streams of the visual system (Finzi et al., 2021) . Collectively, topographic vision models are helping to unify a computational understanding of the functional organization of the visual system. However, topographical priors have not yet been built into models of linguistic processing, despite tremendous progress in the development of natural language processing (NLP) models and their application in cognitive science and neuroscience Wilcox et al. (2020) ; Gauthier et al. (2020) ; Schrimpf et al. (2021) ; Caucheteux & King (2022) ; Goldstein et al. (2022a) ; Tuckute et al. (2023) . In NLP, Transformer language models (LMs) have undoubtedly established themselves as the leading architecture for language tasks (Vaswani et al., 2017; Radford et al., 2018; Brown et al., 2020; OpenAI, 2023) , displaying human-like language understanding and generation for the first time. In cognitive science and neuroscience, these LMs have emerged as the most quantitatively accurate models of human language processing. They generate probabilities of upcoming words that explain reading behavior of humans Wilcox et al. (2020) ; Merkx & Frank (2021) ; Shain et al. (2022) , and their internal activations can explain the neural signals of humans reading or listening to naturalistic sentences or stories at the granularity of fMRI voxels and intracranial recordings Schrimpf et al. (2021) ; Goldstein et al. (2022b) ; Tang et al. (2023) . Despite the success of these LMs, they remain difficult to interpret, and incomplete as models of brain function.In the current work, our aim is to bridge these gaps by inducing a topographic organization of features within the Transformer architecture. We employ local-connectivity based approaches inspired by recent topographic vision models (Blauch et al., 2022; Keller & Welling, 2021) to the language domain, asking whether we can obtain topographic organization of linguistic representations within a Transformer architecture via spatial constraints. To do so, we introduce two computational motifs spatial querying and spatial reweighting -to the self-attention layer, which encourage the development of topographic organization in separate components of the self-attention layer. We call Transformer models employing these constraints Topoformers. We show that we can scale these topographic motifs to a large BERT Topoformer model trained with a masked language modeling objective, and that topographic organization develops within each hierarchical layer of the network, without significantly compromising task performance. We interpret this topography using a novel suite of 8 semantic and syntactic tests. Last, we demonstrate that the topographic representations of the Topoformer can be aligned with the topographic representations of the human functionallydefined language network in multiple subjects. In summary, our work demonstrates for the first time that Transformer models can be trained to exhibit topographic organization similar to the human brain, and paves the way for further interpretability work leveraging spatial priors. 
methods: 
 In this study, we propose two approaches for enforcing topographic organization in a Transformer layer. Both methods rely on the use of local communication to introduce spatial constraints that encourage the formation of spatially organized linguistic representations.We begin with the standard self-attention operation used by Vaswani et al. (2017) . In this formulation, every token embedding is projected onto a set of queries, keys, and values, and the query of a given token is associated with a corresponding key of all other tokens. Spatial querying works by associating a local pool of queries with a given key. The locality is parameterized with a width parameter r SQ determining the fraction of units in a given key's circular receptive field (RF). For simplicity, we examine the case of a simple non-weighted sum of queries. This is achieved by inserting a binary intermediate matrix M ∈ R d×d , where d is the embedding dimension, and the columns of M determine the spatial pool of queries associating with a given key. Essentially, this makes it such that the dot product attention between a given pair of tokens is not between individual queries and keys, but local pools of queries and individual keys. This biases the representations of queries to be locally smooth, and the representations of keys to have a spatial correspondence with the queries. The local pooling of spatial querying can be visualized with a simple example, assuming a model dimension of d = 3, and 2 tokens. EQUATION= Q 1,1 + Q 1,2 Q 1,2 + Q 1,3 Q 1,1 + Q 1,3 Q 2,1 + Q 2,2 Q 2,2 + Q 2,3 Q 2,1 + Q 2,3 K 1,1 K 2,1 K 1,2 K 2,2 K 1,3 K 2,3We can see that, instead of the rows of the matrix multiplication containing individual queries, they now contain summed local pools of queries. This is the essence of spatial querying. The full selfattention equation with spatial querying (SQ) is given as follows:Attention SQ (Q, K, V ) = softmax QM K T √ d k V W O (2)For simplicity of visualization, we use a single attention head in the Topoformer implementation, but we retain the outer reweighting matrix W O used in multi-head attention Vaswani et al. (2017) .Our motivation for using single-head attention is to ensure that the dominant functional organization occurs within a head rather than across heads; without further constraints than Eq 2, organization across heads would be non-topographic and thus complicate interpretability and visualizations. Although the dimensionality of the model could be of any size, it's convenient in our implementation for it to be a perfect square, such that it can be reshaped into a √ d × √ d grid for visualization purposes. While we work with square grids, theoretically, any 1D, 2D, or 3D arrangements of units could be used to define the spatial position of units. For a visual explanation, Figure 1 compares the differences between standard operations within a self-attention block, and their spatial counterparts.Spatial querying only imposes a topographic relationship between queries and keys to encourage the development of topographic organization of the values and self-attention outputs (hereafter fc out), we convert the outer reweighting matrix W O to a locally connected layer W O local . By using locally connectivity in W O local , we encourage the model to learn more localized feature representations in the values and attention outputs (analogous to what spatial querying does for the queries and keys). We parameterize local connectivity using a width parameter r SR that determines the fraction of units within a given unit's circular receptive field (RF). Our locally connected linear layer W O local ∈ R d×d is situated in our network as follows:Attention SQR (Q, K, V ) = softmax QM K T √ d k V W O local (3)Preliminary experiments demonstrated the need to use large positive weights to fully encourage the development of topographic organization. Thus, we initializeW O = |W O i * 10|, where W 0 i is a standardly initialized PyTorch linear layer. This operation, denoted as spatial reweighting, has the effect of enhancing local correlations, commonly viewed as a hallmark of topographic organization (Lee et al., 2020b; Blauch et al., 2022; Margalit et al., 2023) . These excitatory feedforward connections mimic the dominant role of excitatory pyramidal neurons in between-area cortical communication in biological brains (Laszlo & Plaut, 2012; Blauch et al., 2022) . 
training a 1-layer topoformer on a supervised task: 
 We begin by training 1-layer single-head encoder-only Topoformer (bidirectional attention) on the IMDB sentiment analysis dataset Maas et al. (2011) , which classifies movie reviews as having a positive or negative overall sentiment. The local connectivity in the spatial querying and spatial reweighting operations are controlled through a hyperparameter r SR that sets the radius of spatial receptive fields (RFs). We investigated the effect of different RF sizes in the Topoformer-SQR model, finding that smaller r SQ values yield better accuracy and topography, while the network is more robust to the r SR hyperparameter (see Appendix A.9 for details). In the following, we report results for an RF size of 0.3 for the SQ model and 0.1 for the SQR model. We set d = 400.We first describe the results of a model using only spatial querying (Topoformer-SQ). Following training, Topoformer-SQ achieved an accuracy of 0.81 on the IMDB sentiment test set. In comparison, an identical 1-layer Transformer model without spatial querying achieved an accuracy of 0.83. To probe its topographic organization, we conducted selectivity analyses and principal component analysis (PCA) to investigate the unit activation patterns in the different layers, as shown in Figure 2 (see Appendix A.1, A.3 for details). Our selectivity analysis was designed to contrast the response magnitudes for positive and negative sentiment sentences. As expected, this analysis revealed a topographic organization in the keys and query sublayers, but not in the values of fully-connected layers. We next performed PCA to assess generic forms of topographic variability in the model representations. We found that the weights of the first two principal components (PCs) exhibited a smooth topography in the keys and queries, with the second PC spatially aligned to the selectivity for both representations. This demonstrates that the network has learned to organize its dominant modes of variability spatially, a hallmark of topographic functional organization.We next trained a model incorporating both spatial querying and reweighting (Topoformer-SQR). This model achieved a test set accuracy of 0.75 on the IMBD sentiment test set, slightly lower than the Topoformer-SQ model. We performed identical probing analyses to those performed in the previous section, highlighting the results for the values and fully-connected (fc out) representations in Figure 2B . We found that the Topoformer-SQR exhibited more pronounced topographic organization in the values and fc out layers compared to Topoformer-SQ. This suggests that the local connectivity matrix W O local (see Figure 1 A., and equation 3) successfully enforced a topographic correspondence between the values and attention outputs, as predicted. 1 : Comparison of GLUE performance between multi-head and single-head non-topographic BERT control models and Topoformer-BERT, each trained with the Cramming procedure (Geiping & Goldstein, 2022) .We next scaled up the Topoformer motifs to train a BERT model using a Masked Language Modeling objective (Topoformer-BERT). We followed the training paradigm introduced by (Geiping & Goldstein, 2022) . We trained a 16-layer BERT model on the Bookcorpus-Wikipedia dataset (Zhu et al., 2015) for 12 hours (see Appendix A.2 for more details). To provide a control for our Topoformer-BERT model, we trained a standard, non-topographic single-head BERT model with identical parameters and training procedure as our Topoformer-BERT (besides the lack of topographical motifs, Appendix A.7). To evaluate the models' performance on natural language tasks, we followed the General Language Understanding Evaluation (GLUE) benchmark (Wang et al., 2019) procedure as described in (Geiping & Goldstein, 2022) , testing each model on all tasks besides WNLI as in (Devlin et al., 2019) . Critically, we observed that the task performance of Topoformer-BERT on the GLUE Benchmark was similar to that of the non-topographic model counterpart, suggesting that our added spatial constraints were not significantly hindering task performance Table 1 . Having established that Topoformer-BERT is capable of performing linguistic tasks, we move on to characterizing the topographic organization in Topoformer-BERT. First, we systematically quantified the topography of each the 16 Topoformer-BERT layers using a statistic that relates the degree of correlation with the distance between pairs of units (Appendix A.3.2, Equation 4). A high value of the statistic t g indicates that nearby units tend to be more correlated in their response pattern across sequences than distant units. We plot the mean statistic tg over distance thresholds for all layers in Figure 3A , and the distance-threshold-specific statistic t g,d for layer 15 (Figure 3B . In general, the keys and queries have the greatest degree of topographic organization, and the values show the weakest organization. Nevertheless, each is consistently above 0, driven by very local decay in correlation, as seen in the analysis of t g,d across different maximum distances (Figure 3 .Second, we took a step towards interpreting the emergent topographical structure in Topoformer-BERT. Specifically, we evaluated the selectivity of the unit activations to a set of eight test suites targeting different linguistic properties. All eight test suites consisted of 76 sentences each, and were either based on carefully designed minimal pair sentences based on prior work (Gauthier et We performed selectivity analyses for these eight test suites (Figure 4 ). These analyses intuitively ask whether a given unit shows a preference for a particular contrast (e.g., animate versus inanimate sentences). We evinced a strong topographic organization according to broad semantic categories (top row, Figure 4 ), both in terms of significant topographic selectivity, as well as significant decodability of condition from the distributed pattern of activities. Intriguingly, the selectivity patterns were different across contrasts, implying that semantic distinctions are represented in topographic activity pattern differences across categories. It is important to note that despite the strongly significant selectivity, the mean activity patterns were highly similar across categories within each contrast (Appendix A.9.3, Figure 12 ): rather than indicating contrasting hot spots of activation for animate and inanimate content, for example, the rank order of unit activities tends to be similar Each panel shows the selectivity of Topoformer-BERT layer 15 (keys), for a given contrast. Each test suite contains two contrasting conditions each with a set of sentences; unit activities are computed as the mean over tokens for each sentence, and the conditions are contrasted with a t-test. We plot the selectivity significance value (see Appendix A.3), where s = 2 corresponds to positive selectivity with p = 0.01, and s = -2 corresponds to negative selectivity with the same significance level. The first row contains sentences with natural variability, whereas the bottom row contains results from constructed minimal pairs differing in only one word across conditions. To ensure visibility of effects regardless of size, we used different statistic ranges for plotting of each row: s = 10 for the top row, and s = 2 for the bottom row.across sentences, and selectivity across conditions is indicative of small yet distinct deviations from the dominant pattern. This is not unique to the Topoformer (Appendix A.9.3, Figure 12E .), however, the Topoformer allows a uniquely intuitive visualization of selectivity patterns in 2D, aiding interpretability.Second, we turned to more controlled test suites constructed using pairs of minimally different sentences in line with prior work in psycholinguistics and natural language processing (e.g., Linzen et al. 2016; Warstadt et al. 2020) . As expected, the effects were lower (bottom row, Figure 4 ) relative to sentences only matched on length and surprisal (top row). We evidenced weak topographic selectivities to sentences with correct syntactic agreement versus those that do not, for example. The weak selectivities were also reflected in the fact that the condition of interest could not be decoded accurately from the patterns of key activation.In summary, we quantified the extent of generic topographical organization across all sublayers across the full Topoformer-BERT model, and honed in on selectivities of the topographic organization of the final layer. We analyzed eight different linguistic properties, finding strong effects for naturalistic semantic dimensions: animacy, concreteness, and visuomotor properties. Future work should investigate the organization of finer-grained semantic dimensions, as well as more extensive tests for syntactic knowledge.To assess the topographic organization of language in the human brain, we recorded brain responses using event-related fMRI from N=5 participants (4 female, native English speakers) during a sentence reading task. Participants read 1,000 6-word, corpus-extracted sentences that were selected to maximize semantic and stylistic diversity (see A.4). Following standard preprocessing, we used a set of five language masks ("parcels") that denote brain regions within which most or all individuals in prior studies Fedorenko et al. (2010) ; Lipkin et al. (2022) showed activity for an extensively validated language localizer contrast between reading of sentences and non-word strings (Fedorenko et al., 2010) . For each participant, within these anatomical parcels, we then computed individual functionally-defined regions by comparing responses to sentences and non-words, and taking all voxels with at least weak preferences for sentences (t > 1). We then restricted our analyses to these voxels, henceforth the "language network". To determine that the language network exhibits spatial smoothness, as in the model, we computed the generic topographic statistic t g (Equation 4) on unsmoothed brain responses within the functionally-defined language network of each participant, splitting the network into 5 spatial subregions (see Appendix A.4). We compared this statistic to a null distribution, using shuffled brain responses, finding that the t g value for each cluster fell outside this null distribution, indicating significant decay in unit response correlations with distance.To determine whether the topography of the human language network is linguistically meaningful and corresponding to that of Topoformer-BERT, we performed representational alignment using partial least squares singular value decomposition (PLS-SVD). are the i-th aligned component scores.Given the spatial organization of both brain and Topoformer responses, we can visualize the SVD weights of individual brain and model components W (i)x and W (i) y , respectively, reshaped into their native spatial format. To compute the alignment of components, we perform a cross-validated analysis that ensures generalization, where SVD is computed using 80% of the sentences, and the scores are computed for the remaining 20% of the data. These scores can then be correlated across brain and model, for each dimension, to determine their alignment.Figure 5A plots example alignments between the first three brain and model components, using the first participant and the Topoformer-BERT layer 15 (final layer, zero-indexed) keys representation. We see that the first two components are strongly aligned, as well as strongly topographically organized in both model and brain spaces. The third component is not aligned, despite being spatially organized in each representational space. Figure 5B repeats this analysis for all participants and sublayers, using layer 15 again. In general, the first two components were significantly aligned for each sublayer, whereas later components were less likely to be aligned. This result demonstrates that the low-dimensional variability can be aligned in the topographic representations of the human language network and Topoformer language model. The fact that we used functionally-defined language regions suggests that there is spatial functional organization even within this relatively functionally homogeneous brain network (e.g., Blank & Fedorenko (2020); Fedorenko et al. (2020) ), rather than simply across different functional networks with heterogeneous response profiles, similar to the organization that emerges in the Topoformer model.To determine the specificity of this alignment, we performed an identical analysis using a control network and untrained Topoformer-BERT variant (Appendix A.6). Alignment, as well as voxel encoding model prediction, was significantly greater between the trained Topoformer-BERT and language network compared to a non-language control network and an untrained model, highlighting the linguistic nature of the alignment. 
discussion: 
 Here, we introduced the first topographically organized Transformer language models, "Topoformers". Across small and large models, we found that these spatial querying and reweighting operations produced topographic organization in Topoformer models trained on natural language processing tasks. This organization was revealed with specific hypotheses by contrasting different linguistic properties, as well as generically through PCA. Finally, analyzing brain responses to a large number of sentences in the human language network, we uncovered topographic variability with low-dimensional alignment to that found in the Topoformer-BERT model.Introducing topography into language models may improve interpretability in NLP. We took some initial steps with our suite of tests, but the interpretability problem is far from solved. One issue is that of "polysemanticity," whereby units are involved in the representation of several distinct concepts (Bricken et al., 2023) . Despite strong semantic selectivity, we found that Topoformer-BERT's activations were highly overlapping across categories, similar to non-topographic models. While our 2D visualizations aided interpretability of selectivity, efforts to improve "monosemanticity" or to encourage disentangled representations (Higgins et al., 2021) may prove fruitful in yielding even more interpretable topographic organizaton when combined with the Topoformer motifs. Additionally, topographically constrained sparse autoencoders might allow for greater interpretability of entangled representations in Topoformers (Cunningham et al., 2023; Bricken et al., 2023) .Introducing topography is also necessary to imrpove the biological realism of models of language processing in the human cortex, and understand how biological constraints (e.g. wiring cost) shape the emergence and organization of the language network. Future work should aim to scale the approach towards foundation-model level. One critical insight is that scale will not only improve the performance of these models, but also improve their brain predictivity (Schrimpf et al., 2021) . In parallel, a greater focus on biological plausibility may prove fruitful for basic neuroscientific investigation of the language network (Jain et al., 2023) . This work marks the beginning of topographic modeling of language processing. We hope that other researchers will be persuaded to embrace topography in language models, and push the development and use of Topoformers along several new directions. organization (Kiani et al., 2007; Lee et al., 2020b) . Given the Pearson correlation matrix R i,j = r p (a i , a j ) -where a i gives the activity vector of unit i over sentences, r p (x, y) gives the Pearson correlation of x and y, and r s (x, y) gives the Spearman rank correlation of x and y -and the pairwise Euclidean distances D i,j computed in volumetric space, along with a maximum distance over which to compute the statistic d, we compute the generic statistic as follows:t g,d = r s (-R i,j , D i,j ) ∀ i, j : D i,j < d (4)Because topographic organization can exist at multiple scales, with long-range correlation violating the locally distance-decaying correlation, we may wish to compute t g,d at a range of maximum distances, as in Figure 3 . Thus, for those analyses, we computed a vector t g = {t g,d0 , ..., t g,dn } over a linearly spaced range maximum distance values d. We summarized this vector by taking the mean tg = 1 n i t i g (Figure 3A ), and plotting the vector against d (Figure 3B ). When no maximum distance is used, we refer to the statistic as simply t g , as used in the brain analyses.We recorded brain responses using fMRI from N=5 participants during a sentence reading task. The participants were neurotypical native speakers of English (4 female), aged 21 to 30 (mean 25; std 3.5), all right-handed. Participants read 1,000 6-word, corpus-extracted sentences that were selected to maximize semantic and stylistic diversity. Participants completed two scanning sessions where each session consisted of 10 runs of the sentence reading experiment (sentences presented on the screen one at a time for 2s with an inter-stimulus interval of 4s, 50 sentences per run) along with additional tasks. Participants were exposed to the same set of 1,000 sentences (no repetitions), but in fully randomized order. Structural and functional data were collected on the whole-body, 3 Tesla, Siemens Prisma scanner with a 32-channel head coil. T1-weighted, Magnetization Prepared RApid Gradient Echo (MP-RAGE) structural images were collected in 176 sagittal slices with 1 mm isotropic voxels (TR = 2,530 ms, TE = 3.48 ms, TI = 1100 ms, flip = 8 degrees). Functional, blood oxygenation level dependent (BOLD) were acquired using an SMS EPI sequence (with a 90 degree flip angle and using a slice acceleration factor of 2), with the following acquisition parameters: fiftytwo 2 mm thick near-axial slices acquired in the interleaved order (with 10% distance factor) 2 mm × 2 mm in-plane resolution, FoV in the phase encoding (A ≪ P) direction 208 mm and matrix size 104 × 104, TR = 2,000 ms and TE = 30 ms, and partial Fourier of 7/8. All participants gave informed written consent in accordance with the requirements of an institutional review board.A.4.2 DATA PREPROCESSING AND FIRST-LEVEL MODELING fMRI data were preprocessed using SPM12 (release 7487), and custom CONN/MATLAB scripts. Each participant's functional and structural data were converted from DICOM to NIfTI format. All functional scans were coregistered and resampled using B-spline interpolation to the first scan of the first session. Potential outlier scans were identified from the resulting participant-motion estimates as well as from BOLD signal indicators using default thresholds in CONN preprocessing pipeline (5 standard deviations above the mean in global BOLD signal change, or framewise displacement values above 0.9 mm; Nieto-Castañón (2020)). Functional and structural data were independently normalized into a common space (the Montreal Neurological Institute [MNI] template; IXI549Space) using SPM12 unified segmentation and normalization procedure Ashburner & Friston (2005) with a reference functional image computed as the mean functional data after realignment across all timepoints omitting outlier scans. The output data were resampled to a common bounding box between MNI-space coordinates (-90, -126, -72) and (90, 90, 108) , using 2 mm isotropic voxels and 4th order spline interpolation for the functional data, and 1 mm isotropic voxels and trilinear interpolation for the structural data. Last, the functional data were smoothed spatially using spatial convolution with a 4 mm FWHM Gaussian kernel or no smoothing (we investigated both to ensure that are claims about topographic organization could not be explained by smoothing). A General Linear Model (GLM) was used to estimate the beta weights that represent the blood oxygenation level dependent (BOLD) response amplitude evoked by each individual sentence trial using GLMsingle Prince et al. (2022) (fixation was modeled implicitly, such that all timepoints that did not correspond to one of the conditions (sentences) were assumed to correspond to a fixation period). Within the GLMsingle framework, the HRF which provided the best fit to the data was identified for each voxel (based on the amount of variance explained). Data were modeled using 5 noise regressors and a ridge regression fraction of 0.05. The 'sessionindicator' option in GLMsingle was used to specify how different input runs were grouped into sessions. By default, GLMsingle returns beta weights in units of percent signal change by dividing by the mean signal intensity observed at each voxel and multiplying by 100. Hence, the beta weight for each voxel can be interpreted as a change in BOLD signal for a given sentence trial relative to the fixation baseline. After first-level modeling, the voxels within a set of 5 masks ("parcels") were extracted. These parcels were derived from n=220 independent participants using a Group-Constrained Subject-Specific (GSS; Julian et al. (2012) based on an extensively validated language localizer contrast between reading of sentences and non-word strings Fedorenko et al. (2010) ; Mahowald & Fedorenko (2016) ; Lipkin et al. (2022) . These parcels delineate the expected gross locations of language-selective brain regions but are sufficiently large to encompass individual variability. The parcels are in the left hemisphere, three frontal parcels (inferior frontal gyrus [IFG] , its orbital portion [IFGorb] , and middle frontal gyrus [MFG] ) and two temporal ones (anterior temporal [AntTemp], posterior temporal [PostTemp] ). The mean number of voxels in these five parcels were (they differ slightly across participants because of lack of spatial coverage in the functional acquisition sequence): IFG=743 (SD=0); IFGorb=364 (SD=13.4); MFG=462 (SD=0); AntTemp=1623 (SD=4.1); PostTemp=2948 (SD=0). The parcels are available for download at ANONYMIZED. As a control brain region, we extracted voxels within a set of motor-and supplementary motor areas in the left hemisphere. Specifically, we used the Glasser parcellation Glasser et al. (2016) to extract responses within 5 motor parcels: 1, 2, 3a, 3b, and 4. Moreover, we identified a set of supplementary regions using the grouping category "Paracentral lobular and mid-cingulate cortex" in Tait et al. (2021) which consisted of 8 additional parcels: 24dd, 24dv, 6mp, 6ma, SCEF, 5m, 5L, and 5mv, yielding a total of 13 Glasser parcels of interest. The mean number of voxels in these 13 parcels were 3753 (SD=363.2).A.5 QUANTIFYING BRAIN TOPOGRAPHY Here, we quantified the spatial smoothness of brain representations using the generic topographic statistic t g , in each language ROI and a control ROI (see A.4 for detail on ROI definition). Due to differences in sizes across ROIs, we use the version computed without a maximum distance, over the full range of voxel pairs in each ROI. Statistics are shown in Figure 6A . To determine the significance of these statistic values, which were low in some cases, a null permutation analysis was performed by shuffling the voxel responses of each region 100 times with respect to their positions, and computing t g . This allowed us to construct a null distribution against which to compare real t g values. As can be seen in Figure 6 , the variability of this null distribution is very small. All of the brain data fell outside the 95th percentile, indicating that each brain region -including the control region -exhibited significant spatial smoothness in responses to natural sentences. A.6 BRAIN-MODEL CONTROL ANALYSESIn the main text, we performed a PLS-SVD analysis to assess the alignment of human language network and model components. Here, we repeated this analysis considering two important controls: i) motor-related control brain regions, and ii) untrained Topoformer-BERT. First, we confirmed that the alignment with Topoformer-BERT (Figure 7A ) is stronger in the language network than in a control brain network (motor-related regions, see Appendix A.4 for details; Figure 7B ). Particularly, whereas significant alignment was seen for the first two components in the language network, this was not true for most sublayers when comparing to the control brain network (exception: weak significant alignment of values with component 0, and keys with component 1) with substantially reduced alignment overall. Second, we analyzed an untrained version of the Topoformer-BERT model. We found that this untrained model demonstrated some alignment with the first (0th) component, however this was highly variable across participants and not statistically significant (Figure 7C .). When comparing to the control brain network, the untrained model showed very little alignment (Figure 7D ). Thus, the alignment was particularly strong between the trained Topoformer-BERT model and language network.A.6.1 ENCODING MODEL ANALYSES Since our PLS-SVD alignment analysis is less common for comparing representations in the neuroscience literature than other analyses, we opted to perform an additional encoding model analysis for greater comparability to prior work.The encoding model analysis approach is used to predict a given voxel's response based on a representational space. We employed ridge regression, or L2-regularized least squares regression. Formally, let us consider the model embedding space X ∈ R n×d , where each of n rows is a ddimensional vector consisting of the model's representation of a given sentence. A given voxel's responses are given as y ∈ R n . For each voxel, we can formulate our regularized regression problem as finding a vector ŵ ∈ R d such thatEQUATIONwhere ∥ • ∥ 2 is the Euclidean norm. The λ multiplier is a hyperparameter that specifies the relative weight of the regularization term in the loss. This value is chosen using leave-one-sentence-out cross-validation using 80% of the total data, as summarized belowEQUATIONwhere X (k) and y (k) are held out cross-validation data at the k-th fold. This cross-validation is performed efficiently using the scikit-learn function RidgeCV (Pedregosa et al., 2011) . The remaining 20% of the data is used to test the generalization of the encoding model, by predicting held-out voxel responses. We report the correlation of held-out voxel responses with predictions, and take the average over all voxels from all participants.We perform this encoding model analysis for both trained and untrained Topoformer-BERT models, for voxels in both language and motor-related, non-linguistic control brain networks. The results are seen in Figure 8 . We see that the trained Topoformer provides superior neural fits, and that language network voxels are predicted better than the control network. This confirms the previous results found using the PLS-SVD alignment approach. We note that the relatively low performance of the encoding model is partly attributed to the fact that we included a large subset of voxels. The results obtained via PLS-SVD in Figures 5 and 7 indicate that the shared variability across voxels aids the alignment between brain and model. Finally, we note that although the untrained Topoformer-BERT had lower predictivity performance than the trained counterpart, it was still above zero, in line with prior findings Schrimpf et al. (2021) ; Caucheteux & King (2022) ; Pasquiou et al. (2022) . A.9 RECEPTIVE FIELD (RF) SIZE ANALYSIS A.9.1 EFFECT ON PERFORMANCE In this section, we discuss our methodology for choosing a receptive field (RF) size. First, we train the 1-layer (Topoformer-SQR) across a sweep of different RF sizes for both spatial querying (r SQ ) and spatial reweighting (r SR ). As seen in 10, we can see a clear inverse relationship between the r SQ (R 2 = 0.8612), in line with the idea that larger spatial querying pools reduce the amount of information due to the local averaging of queries within the pool. However, we found that the size of spatial reweighting r SR has no measurable effect on model performance r SR (R 2 = 0.0137). This deviates from the result for spatial querying because there is no averaging performed in spatial reweighting. Future work could explore a weighted version of spatial querying that would be expected to show less of a decrement in performance with increasing RF size. Importantly, substantial topographic organization can be seen with a very small SQ value, such that performance decrements are minimal. Using the models from the previous section, we plotted a matrix grid containing the weights of the second PC across all RF values. For brevity, we visualized the the keys and fc out sublayer representations using the same procedure mentioned in Figure 2 , so as to determine the effects of RF size on both spatial querying and reweighting. As seen in Figure 11 , having smaller r S Q values results in much stronger topographic organization in the keys, in line with the better performance; larger values of r SQ were associated with poor performance and minimal organization, suggesting that the model struggled to learn representations in the presence of large averaging pools. However, also similar to the performance results, the spatial reweighting RF r SR had minimal effect on the topography, with topographic organization developing in the fc out sublayer across a large range of RF sizes. 
 
 Appendix:  



Reviews: review: title: title: The paper shows potential in tackling an important problem to representational alignment, in the context of natural language processing.
main_review: main_review: This work is concerned with the lack of spatial constraints in machine learning models, in particular for the transformer architecture, which hinders comparisons with the brain. The paper aims to implement such spatial constraints in the transformer, by making changes that encourage topographic organization to the self-attention layer. The more biologically plausible topographic structure may closer resemble structure found in the brain, and potentially be easier to interpret.

The paper aims to provide evidence using natural language experiments that 
 - These changes result in the transformer learning a topographic structure.
 - The emergent topographic structure has hints of increased intepretability.
 - The structure shares some topographic alignment with brain responses.


Strong points:

The approach of adding topographic organization to transformers comes across as important and relevant in the context of representational alignment between machine learning models and the brain, for the following two reasons: 

It seems likely that the existence of spatial constraints will affect the representational structure that is formed in the learning process. Thus, in order make comparisons between machine learning models and the brain in the area of natural language processing, we should consider using a spatially constrained architecture, which the leading NLP architecture, the transformer, currently is not.
 
Additionally, it is plausible that topographic organization aids in interpretability. If so, this may help in crafting comparisons between different representational structures.

The methods used to implement topographic organization in the transformer architecture make sense and are relatively straightforward. It is clear why they may result in topographic organization and convincing experimental evidence is provided for this as well.

Results hinting at increased interpretability are interesting and show potential. Similarly, results for the alignment displayed in between topoformers and the brain hold promise.

Experimental detail were fairly comprehensive, and for the most part relevant control experiments and additional experimental details were present, or could be found in the appendix.


Weak points:

This work could benefit from further experimental work. The addition of either

 - an experiment more explicitly demonstrating how topographically organized transformers can be seen as more interpretable
 - an experiment providing more evidence of alignment between the brain and the topoformer

would add a lot, as the experiments currently presented still feel a little limited in scope to tackle the questions raised. However, the topoformer approach appears sound and holds potential.

Although methods were mostly well-detailed, there were a few occasions were some additional clarity could be added.

For instance, in section 2.2 the process of converting the outer reweighting matrix $W^O$ to a locally connected matrix $W_{local}^O$ is not made very formally explicit. An example or an explicit formula may add some clarity.

Another example would be in figure 3, where it would be nice to see a typical value of the topographic statistic for the non-topographic control model. This gives us a frame of reference to understand the degree of topography displayed here in topoformer-BERT, as this topographic statistic I believe is not common knowledge.



I recommend this work to be accepted, the main reason being that I find it convincing that the lack of topography in transformers is important for representational alignment in context of natural language processing, and should be further discussed. The approach taken here makes sense, the methods are sound, and there is potential for interesting results.


Some additional questions:

The motivation to restrict to a single head is clear, but I wonder if this is not a limitation, as in practice often many heads are needed to capture different dependencies. Will this limit the topoformer when scaling to larger and more complex language datasets? Perhaps you could elaborate on why this is/isn't important, or how to remedy it.

For the PLS-SVD comparison between the topoformer and the fMRI data there is a control experiment with an untrained topoformer. Perhaps a control experiment with the trained non-topographic transformer can also be valuable, to determine the importance of the topographic constraints introduced here in the topoformer architecture to the alignment.

I find the (somewhat large) differences in selectivity patterns for the different constraints very interesting. It would be interesting to know where they are coming from, e.g. is there some property of the data in the test suites that is correlated to the type of pattern you get?


Minor comments:

The text in some plots is pretty small and hard to read.
strength_weakness: strength_weakness: The approach of adding topographic constraints to the transformer architecture seems like a promising path towards studying alignment between language processing areas in the brain and transformer architectures. The paper takes steps in this direction I think it is worth additional exploration and discussion.
title: title: Topoformer: brain-like topographic organization in Transformer language models through spatial querying and reweighting
main_review: main_review: **Summary:** 

DNNs typically have no topographic organization in the layers or spatial priors while these types of organization are key to neural computations. Successful NN models are not often interpretable. The authors provide a method of providing topographic organization in the layers of a Transformer network, forming so-called ‘Topoformers’. 
This is done through two adjustments. 

*Spatial querying (SQ)*: A binary intermediate matrix $M$ is inserted into the attention computation, associating local pools of queries with a single key. Locality is parameterized with parameter $r_{SQ}$. Representations of keys will have spatial correspondence with queries. 

*Spatial reweighting (SR)*: Converts the fully-connected layer into a locally connected layer by using a locally connected re-weighting matrix $W^O_{local}$ with parameter $r_{SR}$. Large positive weights were used to encourage topographic organization. 

Topoformer networks perform similar to, often slightly worse than non-topographic control models on NLP benchmarks, while retaining interpretable topographic organization. Using fMRI analysis of human brain responses to naturalistic sentences, they show alignment between the variability in the Topoformer and the language network in the brain.

*Experiment 1*: 1-layer single-head encoder-only Topoformer on the IMDB sentiment analysis dataset. Transformer without SQ or SR had accuracy of 0.83. Topoformer with only SQ on achiveved 0.81, with SR achieved 0.75. Topoformer-SQ produced topography in keys and queries,  which Topoformer-SQR also had topography in the values and output layers.

*Experiment 2*: BERT model using a Masked Language Modeling objective on the Bookcorpus-Wikipedia dataset. Using the GLUE benchmark, the Topoformer performance was similar to the non-topographic counterpart. Topographic organization was quantified using the generic topography $t_g$ statistic, which averages correlation over a range of scales. To interpret the organization, they evaluated the selectivity of activations to eight test suites containing linguistic or semantic properties such as Intactness, Animacy etc. While the unit activations showed selectivity towards contrasts, these were different across contrasts, implying that distinctions between the properties were represented in the topographic activity pattern differences. Selectivity was based on distinctions from the overall topographic activity pattern. With minimally different sentences, the effects were lower and the Topoformer showed weak selectivity to sentences with correct syntactic agreement.

*Experiment 3*: Recorded brain responses from 5 participants reading 1000, 6-word sentences, some of which were non-word strings. Voxels with preference for sentences were chosen as the ‘language network’. Computed $t_g$ on unsmoothed BOLD responses in these voxels, which indicated topographic organization.

Using PLS-SVD, they performed representational alignment between the language network response and Topoformer-BERT layers. In general, the first two components were strongly aligned for each sublayer. Low-dimensional variability between activity in the brain’s language model and the Topoformer can thus be aligned. 

**Comments:**

Overall, explicitly introducing topography into an NN model is a good way of increasing interpretability. However, as the authors found out and discuss, this does not guarantee interpretability as there can still be significant variability and polysemanticity even with topographic organization. It is unfortunate that the Topoformers showed no marked performance improvements – however this may be different in tasks that explicitly demand spatial reasoning and separation. The analysis and alignment to human fMRI data is a challenging but well-explained and well-executed work that greatly increases the applicability and impact of this paper. The authors have done an admirable job of explaining their methodology and rationale in the appendix.

It is not clear why the $t_g$ statistic was not employed for the first test on the IMDB sentiment analysis dataset. I believe this should be done, as it is good to quantify the organization of the layers in this analysis.

Similarly, it is not clear why the SQ and SQR versions were not tested for Topoformer-BERT. 

*Sec 3.1*: “The local connectivity in the spatial querying and spatial reweighting operations are controlled through a hyperparameter $r_{SR}$ that sets the radius of spatial receptive fields (RFs).” I believe this should be hyperparameters $r_{SQ}$ and $r_{SR}$.

*Figure 3*: It may be a lot clearer if the statistics in A are shown not as a bar graph, but as datapoints like in B.

General comment for all figures: please increase the font size in the axes labels and especially the legends.
strength_weakness: strength_weakness: The paper is an excellent fit to the workshop as it directly deals with creating representations in a machine learning model and aligning it with biological represenations. The work is complex, yet clearly explained and well-analyzed.
title: title: A solid start to creating more brain-like and interpretable LMs but the insights presently obtained from this are a little lacking
main_review: main_review: In this work the authors seek to embed Language models with topographical organization similar to the type of topological organization seen in many areas of the brain. To achieve this the authors add two constraints to the traditional attention layer of their transformer architecture: a spatial querying term that pools together local queries within some prespecified receptive field and spatial reweighting that locally weights the attention values. The authors very clearly show that these two constraints leads to topographical activity in their Topoformer network and show that it only slightly reduces performance. Additionally the authors design 8 datasets specifically designed to test sensitivity to distinct linguistic properties. Surprisingly, they show that the constraints they added simply to impose topography onto the network also lead to interpretable functional organization. This finding is noteworthy, but I do wish that it was followed with ablation experiments to test how modular their network is as well. Additionally, here the topography seems an inconvenience to the network (as performance always decreases when adding the two constraints) but I imagine that there are situations where it would actually benefit the network as well. For example, the induced topography may make the model more robust to noise in the token as both the queries and values are being pooled now. This would potentially help in answering the question of if the topography we see in the brain is simply a by product of biological constraints that evolution never found a way to relax or if these constraints endow the brain with a functional advantage as well.

 After showing that they can successfully induce topographical structure into their transformer,  the authors ask how similar their model behaves to the brain. To do this they collect fMRI data from 6 participants and compare model response to fMRI. Their analyses argue that Topoformer is aligned with human data but to show this they develop a novel method making it difficult to interpret the magnitude of the alignment. Here it would have been nice to see what typical alignment scores are between the different subjects and maybe using this as an upper bound. Importantly, the authors do perform a negative control comparing alignment between Topoformer and motor-related activity and show that there is no alignment: suggesting that their metric is not so powerful that it always finds alignment. Additionally, the authors do also calculate alignment using more traditional encoding metrics, but I found the results here puzzling. Specifically, while they found that the topoformer had an above zero prediction score: their largest r was 0.03 this is almost an order of magnitude lower than alignments reported in previous literature that had a minimum alignment of around 0.2. Here again it would help to have a ceiling for how high these alignments could possibly be in their dataset. Further, it’s unclear how much the induced topography aids in model alignment. The authors should run a control using a performance matched BERT model to see whether their topography actually increases brain alignment more than just having more accurate models does.

All that aside, the work presented here as is already very promising and presented very clearly, but has the potential to be much better with some more careful comparisons to the brain.
strength_weakness: strength_weakness: From an ML perspective this paper is quite strong. They develop two constraints to the traditional attention layer that leads to topographical activity and functional selectivity in a LM. To my knowledge this is the first time anyone has ever demonstrated either of these two points.