Paper: Topoformer: brain-like topographic organization in Transformer language models through spatial querying and reweighting (ICLR-2024)

Reviewer comment (Theory):
The approach to link the topological structure of the fMRI data to that of the Topoformer model was not well explained and is questionable. These include Section 3.3, Fig 1B, and Fig 6.

Author 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.

Reviewer comment (Presentation):
What does Fig 1b plot exactly? What do the different colors represent? The description in the caption was insufficient for understanding what’s really going on. Same questions apply for Fig 6.

Author response:
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.

Reviewer comment (Evaluation):
How to interpret the statistics in Fig 4? The calculation of the Typograph statistic is based on Equation 4, but there is a lack of explanation regarding how to interpret it and why this is a good statistic to use.

Author response:
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).

Reviewer comment (Writing):
In Section 3.3, the authors stated “we then visualized the weights in the cortical MNI surface space (Gao et al., 2015a)”. What does the “weights” mean here?

Author response:
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.

Reviewer comment (Theory):
When connecting the data with the model, the authors say “Because we have already shown the spatial smoothness of these components, if we see smoothness in the correlation of these components with voxels/units in the target space, we can infer that there is a correspondence between the two topographies.” This is not obvious. The argument should be unpacked, and clarification is needed on how to quantify the strength of the correspondence.

Author response:
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.

Reviewer comment (Experiments):
The authors performed a control analysis in Appendix A 6.1, i.e., “We sanity-checked that a control model did not show any topographic correspondence with the brain organization (see Appendix A.6.1).” It is unclear why this is the appropriate control because Fig 9 is not about the mapping between the PCs and single units for the control model & brain response. A figure analogous to Fig 5, showing the Brain-PC correlation map under the control model, would be expected.

Author response:
- 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.

Reviewer comment (Presentation):
The appendix is poorly formatted and appears almost incomplete. There are equations (such as (4)) which are important to the text but are left with undefined symbols. Furthermore, some figures (such as the Brain-PC correlation with a non-topographic control) appear to be missing (or is this Figure 9?).

Author response:
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.

Reviewer comment (Experiments):
No baseline for IMDB accuracy performance without topographic organization.

Author response:
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.”

Reviewer comment (Experiments):
Only a single attention head is used on all experiments, limiting the comparison with most state of the art transformer architectures. (To be clear, I think the Topoformer-BERT also only uses a single head but could not find this in the text).

Author response:
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.

Reviewer comment (Presentation):
Too many important aspects are relegated to the Appendix making the text challenging to read. The authors should improve the formatting of the appendix (ensure headings for sub-sections are in the correct location) and potentially move some figures (such as Figure 9, and equation 4) to the main text.

Author response:
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.

Reviewer comment (Evaluation):
The topographic organization for the values and fc_out appears quite weak (despite the topographic metric). This makes one either question the topographic metric, or question the usefulness of this model for organizing the actual output of self-attention. It would be helpful if the authors could comment on this. (To be clear, despite this weak organization of values, I still think there is significant value in the strong organization of keys and queries and therefore do not think this should be held strongly against the quality of this paper).

Author response:
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.

Reviewer comment (Experiments):
It would be helpful if the authors included additional controls comparing non-topographic models with brain PCs for the correlation plots of Section 3.3. Without these, the strengths of the conclusions drawn from this correlation analysis are significantly lowered.

Author response:
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).

Reviewer comment (Evaluation):
Is there any intuition why in Figure 1, you see that the Keys and Queries show strong selectivity for negative sentiment but not positive? While the values appear to show only very strong selectivity (very little intermediary), but for both classes? Is this somehow a result of the different mechanisms used to induce topographic organization? (Bilinear vs. Matrix-vector product?). Similarly, is there a reason why these components appear to have patchier organization in the BERT setting as well?

Author response:
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.

Reviewer comment (Evaluation):
In the appendix you say SQR was less stable and required larger batch size, do you have intuition for why this might be?

Author response:
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.

Reviewer comment (Evaluation):
Would it be possible to again compute the topography metric (Equation 4) on the correlation maps between the model and Brain PCs? Would this make sense as another quantitative metric of the alignment of the two representations?

Author response:
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.

Reviewer comment (Evaluation):
Do the PCs found for the human participants correspond to any semantic categories of the text?

Author response:
The PCs are quite difficult to interpret in both the brain and the model, however their alignment is intriguing. We find this interpretability problem fascinating and challenging. We made some headway on this in the model, with our new suite of selectivity tests. We find substantial selectivity for several semantic distinctions, including animacy and concreteness. This is more difficult in the brain, given we have less abiility to run controlled experiments. However, we are excited by the possibility of further work using computational models to drive insights that predict particular selectivities in the brain, and believe the Topoformer will be helpful in this vein.

Reviewer comment (Writing):
Variables used in formulas are not introduced and formulas are hard to find and not explained. See, e.g., the first math items in 2.1. Further, see reference to a non-existent "Equation 4" (which may correspond to Eq4 found in the appendix, but the reader can't know this, because it is possible that enumeration restarts). On top of this, the equation from the appendix is stated without any context or description of what is what.

Author response:
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.”.

Reviewer comment (Evaluation):
The statement "Because we have already shown the spatial smoothness of these components, if we see smoothness in the correlation of these components with voxels/units in the target space, we can infer that there is a correspondence between the two topographies" is incorrect. One would see patterns and smoothness in the correlation of the topoformer pixels with many things, possibly even certain noise vectors (though with lower correlation values), but very likely with brain data from other regions. The latter would be a good test. It would be very useful to see the correlation drop when using, e.g., data from a visual area or a motor area. The currently important thing to read from the diagram is the relatively high correlation value, which shows that there is a correspondence between the transformer activations and brain activity. However, this is also true for unstructured transformers.

Author response:
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.

Reviewer comment (Experiments):
It is worth considering removing the brain analysis part and focusing on more detailed interpretability of the transformer model. For example, why is the focus on the final layer of the topoformer model? What do the other ones look like? Even if they exhibit the spatial property less, this would be interesting. A comprehensive study would have been more helpful instead of a link to brain activity that does not strongly corroborate similar topological information, but that only shows what several papers have already shown — a correspondence of these models to brain activity.

Author response:
Regarding the paragraph beginning "It is worth considering":
- No paper has quantified the fine-grained spatial topography of regions of the language system, so we believe that establishing the topographic organization of the language network is a critical motivation for the work presented here, and thus an important aspect of the work and the audience we hope to receive it. That being said, we agree that further comparisons are an interesting path for future work beyond the scope of a single conference paper. We hope that the reviewer agrees that our new analyses of the brain data are a worthy inclusion to the paper. 
- We looked at the final layer because we expect it to be the most context-aware layer, and thus limited our visualization to a single layer. The other layers look similar as suggested by the statistic. A more comprehensive analysis of all layers, particularly looking at the nature of their representations rather than just the presence of topography -- leveraging similar approaches used here, such as the new selectivity test suite -- is unfortunately not possible within the scope of the paper, but would be an excellent direction for future work.

Reviewer comment (Theory):
Why is locality enforced using a matrix M that forces averaging of the adjacent keys and queries? This a priori does not actually have to enforce similarity between the points that are connected, but it does seem that it is enough pressure to make the locations become similar to their average. A different approach would be to simply constrain each location to only be able to use information from a neighborhood. Was this approach considered at all?

Author response:
Answer to question:
- The use of a mask in spatial querying is a mathematically simple way of introducing locality over feature dimensions. In a standard self-attention layer, each key is associated with a single query, so there is no limiting beyond this that could be done. Rather, spatial querying expands the amount of information that a given key has access to, in terms of its local pool of queries. That being said, we acknowledge there are other possible ways of introducing locality, such as a locally-biased weight matrix that allows for learned spatial combinations of queries – rather than an average – to be associated with a given key. Lastly, and perhaps this is what you meant by “constraining each location to only be able to use information from a neighborhood"), locality could be introduced by making the query and key matrices locally connected. One advantage of our formulation of spatial querying compared to the latter approach is that it does not require spatial organization of its inputs, and allows the network to combine information from the full input embedding before organizing it again. However, we agree that these other options would be interesting to explore, albeit beyond the scope of this first paper on topographic transformers. Lastly, we agree that spatial querying encourages, rather than forces, the development of local similarity in queries and keys, since this allows for reasonable local averages of queries and keys.

Reviewer comment (Experiments):
How to choose the optimal number of tokens in local spatial querying?

Author response:
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.

Reviewer comment (Experiments):
With the introduction of local pooling of spatial queries, the parameter differences between Topoformer and Vanilla BERT are not provided.

Author response:
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.

Reviewer comment (Writing):
What does fc_out refer to? It is not defined or referenced anywhere in the entire paper except in Fig 4.

Author response:
d. fc\_out is the output of the self-attention layer. We hope that our updated Figure 1 helps to clarify this.

Reviewer comment (Evaluation):
The validation of the proposed Topographic Transformer model appears insufficient. Similar to BERTology studies, has the proposed Topoformer been assessed for its ability to capture the hierarchy of linguistic structure (such as early layers capturing surface features, intermediate layers capturing syntax, and later layers representing semantics)?

Author response:
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.

Reviewer comment (Evaluation):
Why does the Topoformer with a single head result in better accuracy scores on the GLUE benchmark compared to the multihead attention?

Author response:
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.

Reviewer comment (Reproducibility):
What is the complexity of self-attention in Topoformer after introducing the local spatial querying mechanism? Did the authors maintain the same local spatial querying across layers, or did they increase the local pool size with the depth of the layer?

Author response:
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.

Reviewer comment (Writing):
The clarity can be improved.

Author response:
We apologize for any lack of clarity and have significantly revised the manuscript to address this.

Reviewer comment (Writing):
Several typos: In Fig 4: fc_out, rest of the paper: fc-out, FC-Out.

Author response:
We thank you for catching this and have fixed these.

Reviewer comment (Evaluation):
What is the representational similarity between Topoformer and Vanilla BERT across layers? During fine-tuning of Topoformer, similar to fine-tuned BERT, are the last layers significantly affected?

Author response:
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.