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

We appreciate the reviewer's generally positive assessment of our work, and hope to clarify many points in our response.

In response to your one listed weakness: we have substantially revised the brain modeling component of the work, to make alignment between brain and model topographies clearer. We now use a technique called partial least squares singular value decomposition (PLSSVD) that projects brain and model embeddings into a joint space, by performing SVD on the cross-covariance matrix. This allows us to recover shared dimensions, and evaluate the correspondence of brain and model projections along these dimensions using held-out sentences. This allowed us to discover alignment between the representations, and reveal the topographic nature of the shared dimensions. We have also added multiple control analyses, discussed in our comments to the other reviewers, that we believe help clarify the contribution of our work.

Answers to questions:

1. Figure 1b has been removed, but generally, we believe we have vastly improved the visual presentation of our results, which we hope can help clarify the reviewer's understanding of our work.
2. We apologize for the missing information. We have added this information to the text. The generic topographic statistic is a measure of the degree to which response correlations decay with distance, over some maximum distance. In the model, we compute this for a range of distances and take the mean, as well as plotting the statistic over all distances for layer 15 (Figure 3).
3. By weights, we meant the weights of a given principal component. We have now substantially revised our presentation of the results to focus more on the PLS-SVD results, but the same idea holds, only now the weights are supervised for alignment across brain and model representations.
4. We thank you for raising this point, which was also raised by another reviewer. Now, rather than relying on heuristics such as this one, we have taken substantial efforts to demonstrate the alignment of brain and model topographies, with the PLS-SVD approach, as well as targeted control analyses using a different brain network and an untrained Topoformer-BERT variant. We find that the alignment between trained Topoformer and language networks is substantially stronger than that including either the control network or untrained variant; however, the untrained variant also shows a degree of low-dimensional alignment.

• A related issue regarding the comparison of topographies between two different systems is that, without an explicit theory of the global organization, and a careful consideration of the shape of the cortical sheet, one cannot simply compare topographies one-to-one. Thus, our PLS-SVD approach is a useful alternative approach to demonstrate alignment and visualize the individual topographies of each representation in their native spatial layout. Tackling the global layout is a long-term problem, but is of great interest for future work.
• We apologize for the confusion. That figure (now Figure 9, A.7) merely demonstrates the lack of topographic organization in a control model without the topographic constraint. Moreover, we hope that the various controls already discussed take care of any remaining concerns raised in this comment.

Thank you again for your insightful comments and questions. We hope our responses clarify the work, and persuade you of its importance and rigor.

While we are disappointed by the reviewer's negative overall view of our paper, we greatly appreciate the many important critiques, and believe that our answer's should help the reviewer see our paper more favorably. Let us address the concerns in turn:

1. We appreciate the reviewer’s careful assessment of our writing. Regarding the first “falsehood”, the reviewer misunderstands our point: indeed, convnets exhibit and exploit spatial structure, but it is of a different form (explicitly spatial information in the stimulus) than the spatial structure explored in this work (spatial organization of learned features). A more similar form of “spatial structure” to that exploited by convnets would be the sequence position, but this is not the target of our work. Rather, we are interested in exploring the spatial organization of featural information, not immediately available in the input. Beyond this, we have given the paper a comprehensive edit and hope that the reviewer will find the text less problematic and more factual and descriptive, backed up by several further analyses.

2. Equation 4 indeed was in the appendix, and we apologize for not explaining it more clearly. This has been fixed. Additionally, we have added an intuitive explanation about this statistic: “A high value of the statistic indicates that nearby units tend to be more correlated in their response pattern across sequences than distant units.”.

3. We appreciate the insightful comments. We have re-analyzed the data and now use a version without smoothing, and have moved the functionally localized language network results (previously in the appendix) to the main paper, to quell any concerns about individual localization or false smoothness. The smooth organization of the brain data is still very clear and statistically significant, as now confirmed by a permutation testing analysis. To be clear, none of the statistics that are computed are affected by the surface projection, which is merely a visualization tool. Lastly, we have moved towards lateral surface views, as we agree that they are substantially easier for the average reader to interpret, and more clearly demonstrate the topographic organization.

Para below 3. We thank the reviewer for correctly pointing out this poorly framed claim. We have now added additional control analyses to make the point of greater topographic alignment more clear.

• First, we applied a new analysis approach leveraging partial least squares singular value decomposition (PLS-SVD), where we can project the brain and model representations into a shared space (Figure 5). Using held-out data, we demonstrated that the low dimensional variability in these models could be substantially aligned. Visualizing the weights of these aligned components, we found that each was highly topographic. Regardless of any arguments about control regions or control models, this demonstrates significant generalizable alignment of low-dimensional topographic variability in the model and functional language network.
• Second, we added a control analysis of a separate brain region composed of motor and somatosensory strips, demonstrating weaker alignment of these voxels with TopoBERT, and worse encoding model prediction.
• Third, we added a control analysis of an untrained TopoBERT model, demonstrating weaker alignment of untrained vs. trained TopoBERT representations and brain representations beyond the first dimension, and poorer encoding model prediction (albeit above chance). Other studies have demonstrated significant brain predictivity of untrained transformer models, and our results corroborate this while demonstrating greater brain alignment with the trained model.
• Fourth, as mentioned, we computed the generic topography statistic for each language region and the control region, and demonstrated that while the values were somewhat low, they were significantly greater than the chance values estimated through a permutation of voxel responses with respect to positions. In summary, while it is true that a topographically organized system will always show topographic correlation structure with whatever it is being compared to, in many cases these correlations would be weak and insignificant. The stronger correlations between trained TopoBERT and the language network – both of which we demonstrate to exhibit topographic variation across natural language sentences – demonstrates an intriguing alignment of low dimensions. We do not want to make strong claims about the particular similarity of the spatial frequency, etc., of the emergent topographies at this stage, as optimizing for such properties is beyond the scope of this initial paper.

We will address the final paragraph and question in a follow up comment.

We thank the reviewer for an extremely positive review, noting the novelty and potential for impact of our work, while also providing several important criticisms that were easily addressed. Let us address the weaknesses in turn. While they were not numbered, we will number our responses (in order) for the sake of organization.

1. We apologize for the earlier poor formatting of our appendix. We have edited this to ensure greater readability, and sufficient information to understand each equation.

2. We apologize for the lack of control comparison for IMDB accuracy. We have hereby trained a non-topographic counterpart (identical to the topographical model, just without Spatial Querying or Spatial Reweighting), and this model achieved a performance of 0.83 on the IMDB test set. We have added this in the appropriate section: “In comparison, an identical 1-layer Transformer model without spatial querying achieved an accuracy of 0.83.”

3. We trained both a single head and multihead non-topographic model, and reported GLUE results in Table 1. While the multihead model performs the best, a single head does surprisingly well, and our topographic variant performs nearly as well as the non-topographic control model.

4. As stated before, we have improved the formatting of the appendix substantially. Unfortunately, we are heavily limited by space and have had to keep only the most important results in the main text, while trying to provide as much relevant information as we can to the reader. We hope that any limitations on readability are overweighed by the benefits of greater information and robustness of our results.

5. This is a great point. We have included these analyses in Appendix A.9.

6. The topographic organization of values and fc_out representations is a bit weaker, as you point out. We were also interested in this finding. The statistic has been updated slightly to look at the decay of response correlation in terms of the correlation of pairwise correlations and distances, as should be clearer now from Appendix 3.2. It still shows the same result. It appears to us that the fc_out layer in particular is a bit sparser in its activity --- more specifically, having few units with very high activations --- which makes the visualizations look a bit weaker. In any case, for space considerations, we have opted to focus primarily on the keys, while quantifying all of the layers for breadth of consideration. Many other formulations of spatial constraints are possible (Gaussian probability masks, wiring penalities rather than fixed local connectivities, etc.), and it is possible that other formulations would work better in the scaled up model. However, we note that the small-scale model did not exhibit this same issue, showing substantial organization in the values and fc_out representations, suggesting it is not a fundamental limitation of spatial reweighting.

7. This is an excellent point, and we have now added multiple controls. Moreover, we have now used a cross-validated approach to assess the alignment of brain and model topographies, which reduces the need for controls in the first place. Nevertheless, we added a control analysis of a separate brain region composed of motor and somatosensory strips, demonstrating weaker alignment of these voxels with TopoBERT, and worse encoding model prediction. We also added a control analysis of an untrained TopoBERT model, demonstrating weaker alignment of untrained vs. trained TopoBERT representations and brain representations beyond the first dimension, and poorer encoding model prediction (albeit above chance).

Answers to questions:

1. Interesting question! This is not found for all models, so we believe this was a random occurence for this one model. We believe the patchier representation in BERT may owe to the greater effective dimensionality required when representing the features of a much larger corpora of linguistic information.
2. Unfortunately we do not yet have a strong intuition for this, however it may be due to the positivity constraint on the feedforward connections, which is not present in the SQ model.
3. The topography metric requires multiple trials, so it cannot be computed on the model correlation maps. However, we have added a quantification of brain topography using this statistic, finding significant smoothness (albeit reduced now that we have eliminated the smoothing from our preprocessing pipeline). While the statistic is not perfect, it is one useful way of summarizing organization quantitatively across a number of spatial representations.

We will address the final question in a follow-up comment.

We appreciate the reviewers generally positive remarks on our paper, and are glad to address their concerns.

Regarding the weakness of interpretability. We have added a new suite of 8 tasks designed to probe interpretability. We find selectivity for several semantic attributes such as animacy and concreteness. This selectivity exists within generally highly similar mean activations (true of our model, and standard language models), but is simply visualized in our 2D topographic layout. While much future work remains to be done to interpret language models (with superposition being a major obstacle), we believe our work provides some important first steps and a useful architecture for future work that can benefit from its simple topographic priors.

We believe this also addresses the previous concern regarding evaluation on multiple tasks. Additionally, we already analyzed the Topoformer-BERT model on the GLUE benchmark suite, which contains several tasks (Table 1). We believe our evaluation of Topoformer-BERT is extensive and demonstrates similar performance to a non-topographic model. Beyond this, state-of-the-art performance is not our goal, as we do not have the resources to train massive language models. We believe that our simple architectural motifs should be scalable to much larger models that can be trained by groups with more resources.

We thank the reviewer for noting the strengths of our paper and providing some relevant critiques. Let us address the weaknesses in turn:

a. There is a misunderstanding in the understanding of spatial querying. Our spatial operations do not operate over tokens, but features for a given token. Regardless, we have added an additional analysis of the role of the RF sizes for spatial querying and reweighting, using the small-scale model. b. Topoformer-BERT and the non-topographic control model are identical except for the presence of spatial querying and reweighting. We now make this clearer in the text. c. Regarding the generic topography statistic, we have modified the statistic for greater robustness to scale, and discussed it more thoroughly both in the appendix and what is now Figure 3. d. fc_out is the output of the self-attention layer. We hope that our updated Figure 1 helps to clarify this. e. We selected 16 layers because we followed the protocol of the CRAMMING paper (Geiping et all., 2022). Throughout our paper, we depict 16 layers, however, they are zero-indexed. We have clarified this in the manuscript, thank you for catching this. 2. We have added substantial new assessments of topographic selectivity in Topoformer-BERT, please see Figure 4. Beyond this, we found that Topoformer-BERT performed generally on par with a non-topoformer control – as we have shown on a variety of benchmarks present within the GLUE composite benchmark – and thus we have no reason to believe that these factors you mention would be different. In general, we believe these aspects of interpretability are beyond the scope of the present paper, which is designed to introduce the algorithm, provide some intuitions about its performance, and make comparisons with the human brain’s language network. 3. This is only the case for a single metric (COLA) - overall, the multihead model does perform better. However, the Topoformer performs similarly, albeit somewhat worse, than the matched single-head non-topographic model. 4. The complexity is unchanged and is the same as the original attention mechanism: O(n2). We use a single RF value for spatial querying and reweighting across the network, but further exploration of structured variation as you suggest would be interesting for future work. 5. We apologize for any lack of clarity and have significantly revised the manuscript to address this. 6. We thank you for catching this and have fixed these.

Answers to questions:

1. We expect that the representational similarity between Topoformer-BERT and non-topographic variants is generally on par with that between randomly initialized variants of non-topographic BERT models. We make no claims about the topographic bias fundamentally altering the representations, only that they allow the representations to be visualized in 2D.

Thank you again for your insightful comments. Addressing them has improved our paper substantially. We hope to have convinced you of the rigor, quality, and contribution of our work.