Evidence from fMRI Supports a Two-Phase Abstraction Process in Language Models

Authors: richard j antonello
emily  cheng
enc  perf
Venue: NeurIPS_cc_2024_Workshop_UniReps
Published: 2024-09-20 20:16:49

Abstract: Research has repeatedly demonstrated that intermediate hidden states extracted from large language models predict measured brain response to natural language stimuli. Yet, very little is known about the representation properties that enable this high prediction performance. Why is it the intermediate layers, and not the output layers, that are most capable for this unique and highly general transfer task? In this work, we show that evidence from language encoding models in fMRI supports the existence of a two-phase abstraction process within LLMs. We use geometric methods to show that this abstraction process naturally arises over the course of training a language model and that the first "composition" phase of this abstraction process is compressed into fewer layers as training continues. Finally, we demonstrate a strong correspondence between layerwise encoding performance and the intrinsic dimensionality of representations from LLMs. We give initial evidence that this correspondence primarily derives from the inherent compositionality of LLMs and not their next-word prediction properties.

Main paper: 
 introduction: 
 How do brains and machines take low-level information, such as a collection of sounds or words, and compose it into the rich tapestry of ideas and concepts that can be expressed in natural language? This question of composition, or abstraction, is at the heart of most studies of human language comprehension. Recent work has shown that representations from large language models (LLMs) are able to successfully model human brain activity at varying spatial and temporal resolutions with only a linear transformation [1, 2, 3, 4, 5, 6, 7] . This has led to questions about the reason for this brain-model similarity. Do LLMs and brains possess similar representations because they have similar learning properties or objectives? [8, 9, 1] Or is the similarity merely a consequence of shared abstraction, the ability to represent features not derivable from the lexical properties of language alone? [10] In this work, we present new evidence that it is the abstractive, compositional properties of LLMs that drive predictivity between LLMs and brains. We do this by examining an underexplored and unexplained phenomenon of the similarity -the tendency for intermediate hidden layers of LLMs to be optimal for this linear transfer task. We show that an LLM layer's performance at predicting brain activity is strongly related to intrinsic dimensionality of that layer relative to other layers in the same network. Furthermore, we demonstrate that this relationship is itself an indicator that pretrained LLMs naturally split into an early abstraction, or composition, phase, and a later prediction, or extraction, phase, a result independently suggested in the LM interpretability literature [11, 12] . We suggest that it is the first abstraction phase, rather than the latter prediction phase, that primarily drives the observed correspondence between brains and LLMs. Neural encoding model training To train encoding models, we use the method described in [4] . High-level details of the method are summarized here. For each word in the stimulus set, activations were extracted by feeding that word and its immediate preceding context into the LLM. A sliding window was used to ensure each word received a minimum of 256 tokens of context. Activations were then downsampled using a Lanczos filter and FIR delays of 1,2,3 and 4 TRs were added to account for the hemodynamic lag in the BOLD signal. A linear projection from the downsampled, time-delayed features was trained using ridge regression. Encoding models were built using the OPT language model [14] (three sizes -125M, 1.3B, 13B) and the 6.9B parameter deduped Pythia language model [15] .  
dimensionality of neural manifolds: 
 To measure representational complexity, we compute the intrinsic dimensionality I d as well as the linear effective dimensionality d of activations at each layer. I d and d describe different geometric properties of the representations: while the former is the dimension of the representations' underlying (nonlinear) manifold, the latter describes the number of linear directions that explain their variance up to a threshold. We will use dimensionality to refer to both I d and d, specifying when necessary.We are interested in an LLM's behavior on a representative sample of natural language, so that the computed dimensionality is informative about the model's general linguistic processing. For all models, we compute the ID on N = 10000 20-word contexts randomly sampled from Pythia's training data, 1 The Pile [16] , over 5 random data samples. The model's tokenization scheme can produce sequences of variable length, so we aggregated representations at each layer by taking the last token representation in the model's residual stream [17] ; this yields one N × D matrix of activations per layer, D being the model's hidden dimension, or extrinsic dimension.Nonlinear ID estimation To compute I d , we apply the Generalized Ratios Intrinsic Dimension Estimator (GRIDE) [18] , an extension of the popular TwoNN estimator [19] to general scales. GRIDE operates on ratios µ i,2k,k := r i,2k /r i,k , where r i,j is the Euclidean distance between point i and its j th neighbor. Assuming local uniform density up to the 2k th neighbor, the ratiosµ i,2k,k follow a generalized Pareto distribution f µi,2k,k (µ) = I d (µ I d -1) k-1 B(k,k)µ I d (2k-1)+1, where B(•, •) is the beta function. The I d is then recovered by maximizing this likelihood over points i for several candidate scales k.Finally, in order to choose the proper I d , a scale analysis over k, which controls the neighborhood size, is necessary: if k is too small, the I d likely describes local noise, and if k is too large, the curvature of the manifold will produce a faulty estimate. Instead, it is recommended to choose a k for which the I d is stable [18] . We provide an example of such a scale analysis in Appendix B.Linear dimensionality estimation In addition to nonlinear I d , we computed linear effective dimensionality d two ways: using PCA with variance cutoff 0.99 [20] , and the Participation Ratio (PR), defined as ( i λ i ) 2 /( i λ 2 i ) [21] . By definition, I d -dimensional manifolds can be embedded in d ≥ I d dimensions, so we expect that d ≥ I d . To determine whether predictive coding explains brain-LLM representational similarity, on The Pile, we computed the next token's surprisal from intermediate layers using TunedLens [22] . TunedLens learns an affine mapping from an intermediate layer to the vocabulary space in order to predict the next token, indicating how much intermediate layers (linearly) represent next-token identity. See Appendix C for implementation details. 
results: 
 Layerwise encoding performance and representational dimensionality, linear and nonlinear, are highly correlated across brain areas involved in linguistic processing. Table 1 shows the correlation between encoding performance and dimensionality averaged over all voxels, and Figure 1a shows the correlation between average encoding performance and (normalized) I d for OPT models. The positive relationship, ρ = 0.85, between I d and encoding performance suggests that in trained models, the I d of layer activations captures linguistic feature complexity needed to support language comprehension. Figure 1b overlays, for OPT-1.3B, the encoding performance, I d , and next-token prediction loss computed from each layer. Encoding performance peaks at layer 17, which exactly marks a sharp downwards turn in prediction loss. While Cheng et al. [12] showed pre-I d -peak layers to extract syntactic and semantic features, our results also suggest a functional shift post-I d -peak to next-token prediction. The sharp transition from abstraction to prediction is observed across OPT model sizes, but it is more gradual for Pythia (see Appendix F.1). To further evidence a transition in layer function, we report inter-layer representational similarity via linear Centered Kernel Alignment [23] . In Figure 1d , where lighter is more similar, the I d peak approximately marks a point where preceding layers are no longer similar to following ones. Results hold across models, see Appendix F.3.Figure 1c shows, for one subject, the voxelwise correlation of I d with encoding performance across layers (dark red is better). Except for the primary auditory cortex, which processes low-level auditory information, encoding performance in brain areas thought to handle higher-level linguistic processing is well-predicted by I d . Results hold across subjects and models, see Appendix F.2.The relationship between encoding performance and I d arises nontrivially from learning. Figure 2 plots the encoding performance (left) and I d (right) across layers over the course of training for Pythia-6.9B (each curve is a different checkpoint). We confirm an existing result from the literature that the I d peak emerges and that I d generally grows for all layers over training (Figure 2 right) [12] . Furthermore, encoding performance and I d increase at similar rates over training, seen by similar positions of the checkpoint curves in the two plots. The two plots are globally correlated with ρ = 0.94. Lastly, the location of the I d peak (red dots, right), changes over training, eventually settling at the same layers for peak encoding performance (red dots, left). This rules out that the I d peak trivially reflects the Transformer architecture, e.g., layer index. 
discussion: 
 Recent studies of the properties of language encoding models have observed that the intermediate layers of LLMs, rather than the output layers, have the highest linear similarity to measured brain activity. This is true regardless of the scanning modality (be it fMRI [4] , ECoG [24] , or MEG [8] ), and regardless of the chosen LLM. Despite this very frequently observed trend, little research has been dedicated to explaining this phenomenon. Yet, an understanding of this trend would greatly benefit our understanding of both brains and LLMs, not least because layerwise differences in LLMs have highly useful epistemic properties. LLM layers are invariant to many confounding variableseach layer has seen the same data in the same order, has an identical architecture, was trained on the same loss term, and built using the same hyperparameters. Therefore, differences between layers can only arise either as a result of the compositional nature of the transition from earlier layers to later ones, or due to the "time pressure" exerted by the loss term on the final output layers.These competing pressures, to first build up the most comprehensive representation of the input text possible, and to then ultimately use this representation to resolve towards a distribution over predicted next word outputs, have opposite effects, as we demonstrate here. The composition effect leads to a increase in encoding performance and dimensionality, whereas the prediction effect narrows the dimensionality to the detriment of encoding. Furthermore, we observe that as models get larger and more thoroughly trained, the best layer for encoding slowly drifts to earlier in the model, perhaps suggesting a saturation effect for this initial compositional phase.What conclusions should we draw from this? Firstly, that it is not likely to be the autoregressive nature of language models that drives brain-model similarity [9, 1, 10] . As models get more potent at prediction, their most predictive and most descriptive layers drift apart. 2 Secondly, we can draw that the multi-phase abstraction process in LLMs that has been proposed independently by other authors [12, 11] is supported by evidence from the only other system known to effectively reason with complex language, the human brain. As the present work only tests two model families, it will be necessary to test more models for conclusions to hold in the general case.From a practical perspective, conclusions point to a potential new avenue for improving the performance of encoding models. If the spectral properties of different LLM layers can be measured and efficiently combined to produce a representation with higher I d than any individual layer, then we might expect that new representation to outperform any single layerwise encoding model coming from the same LLM. As linear layerwise encoding models reach their limit, such methods may be necessary to see further benefits. We parallelized all computation, and estimate the overall parallelized runtime, including preliminary experiments and failed runs to be around 6 days.Ridge regression was performed using compute nodes with 128 cores (2 AMD EPYC 7763 64-core processors) and 256GB of RAM. In total, roughly 1,000 node-hours of compute was expended for these models. Feature extraction for language models was performed on specialized GPU nodes similar to the AMD compute nodes but with 3 NVIDIA A100 40GB cards. Feature extraction required roughly 300 node-hours of compute on these GPU nodes. Pycortex [26] and Numpy [27] were used for flatmap visualization and figure generation.B.1 Nonlinear ID GRIDE operates on ratios µ i,2k,k := r i,2k /r i,k , where r i,j is the Euclidean distance between point i and its j th neighbor. Assuming local uniform density up to the 2k th neighbor, the ratiosµ i,2k,k follow a generalized Pareto distribution f µi,2k,k (µ) = I d (µ I d -1) k-1 B(k,k)µ I d (2k-1)+1, where B(•, •) is the beta function. The I d is then recovered by maximizing this likelihood over points i for several candidate scales k.Finally, in order to choose the proper I d , a scale analysis over k, which controls the neighborhood size, is necessary: if k is too small, the I d likely describes local noise, and if k is too large, the curvature of the manifold will produce a faulty estimate. Instead, it is recommended to choose a k for which the I d is stable [18] .For ID estimation using GRIDE, we reproduce the setup in Cheng et al. [12] . For each model, checkpoint, and layer, we perform a scale analysis. The intrinsic dimension of the manifold is sensitive to the scale, or neighborhood size, for which it is estimated [19, 18] . Figure B.1 shows an example, where the GRIDE scale k varies from 2 0 to 2 12 . As recommended in Denti et al. [18] , we choose a scale k corresponding in a range where the intrinsic dimension is stable, or plateaus, by visual inspection. For simplicity, we choose one scale k per model. In the particular example in Figure B .1, we choose k = 2 4 , where the derivative of the curve is closest to 0 for as many layers as possible. Scales chosen for all models are in Table B To compute the linear effective dimensionality of a mean-centered representation matrix X, we compute the eigenspectrumλ 1 ≥ λ 2 ≥ • • • λ D of its covariance matrix X ⊤ X ∈ R D×D .Then, the dimensionality given by Principal Component Analysis with a threshold of 0.99 isd P CA,0.99 (X) = min i∈1•••D i s.t. i j=1 λ j D j=1 λ j ≥ 0.99. (B.1)In words, this is the minimal number of principal components that explain at least 99% of the variance in X.We also compute the Participation Ratio (PR), a non-integer measure of linear dimensionality defined asd P R (X) = ( j λ j ) 2 /( j λ 2 j ). (B.2)The PR is designed to smoothly interpolate between 1 and D: one can verify that when λ i̸ =1 = 0, then d P R (X) = 1, and when data are isotropic, that is,λ i = λ j ∀i ̸ = j, then d P R (X) = D [21].We used the TunedLens [28] implementation by Ghandeharioun et al. [29] . TunedLens ascertains the amount of information linearly encoded in hidden layer t about the next token. To do so, an affine mapping is learned from the last-token hidden representation h t at layer t as follows:min At,bt D KL (f >t (h t ) || LayerNorm(A t h t + b t )W U ). (C.3)Here, A t ∈ R D×D , b t ∈ R D are the learnable parameters of the affine mapping. W U is the LM's unembedding matrix that maps the final layer to the vocabulary. Finally, f >t (h t ) is the layers of the LM f after layer t, producing the model's original distribution over the vocabulary. In the provided code [29] , TunedLens is implemented using a direct solver numpy.linalg.lstsq on N = 8000 randomly sampled sequences from The Pile dataset [16] , returning the least squares solution that minimizes the l 2 -norm between h t and last layer representation h T . Finally, we compute the nexttoken surprisals on a validation set of The Pile (N = 2000) from the TunedLens-modified hidden layers.MRI data were collected on a 3T Siemens Skyra scanner at The University of Texas at Austin Biomedical Imaging Center using a 64-channel Siemens volume coil. Functional scans were collected using a gradient echo EPI sequence with repetition time (TR) = 2.00 s, echo time (TE) = 30.8 ms, flip angle = 71°, multi-band factor (simultaneous multi-slice) = 2, voxel size = 2.6mm x 2.6mm x 2.6mm (slice thickness = 2.6mm), matrix size = 84x84, and field of view = 220 mm. Anatomical data were collected using a T1-weighted multi-echo MP-RAGE sequence with voxel size = 1mm x 1mm x 1mm.In addition to motion correction and coregistration [30] , low frequency voxel response drift was identified using a 2nd order Savitzky-Golay filter with a 120 second window and then subtracted from the signal. The mean response for each voxel was subtracted and the remaining response was scaled to have unit variance. Layer Num.SurprisalLayer Num. Similarity The training data for OPT are not publicly downloadable.We note that it is inherently challenging to draw affirmative conclusions about underlying mechanisms in the brain from from brain-model similarity alone[25,10]. 
 
 Appendix:  


Decision: In light of the positive reviewers' feedback and relevancy of the submission, we are pleased to accept this paper for presentation at UniReps 2024. We kindly ask the authors to incorporate the reviewers' suggestions and feedback in the final camera-ready version of the manuscript.

Reviews: review: title: title: This study aims to connect recent ideas from interpretability literature on language model abstraction with how the brain might be processing language (using fMRI data), tying the notion of dimensionality of activations to explain recent empirical observations that have yet to be explained.
main_review: main_review: **Quality**:
The study appears to be well-executed with rigorous methodology. The authors use state-of-the-art language models, appropriate dimensionality estimation techniques for the activations, and established fMRI encoding methods. The analysis is thorough and considers multiple models, metrics, and controls. 

**Clarity**:
The paper is generally well-written and logically structured. The methods and results are explained clearly, with helpful figures to illustrate key findings. Some technical details are appropriately relegated to appendices, including details related to compute requirements, dimensionality formulas, etc. However, a few sections (e.g., parts of the discussion) could benefit from minor clarification or expansion. For instance, the low correlation in Table for PR-d and the encoding performance is not explained later in the results section.

**Originality**:
The work presents several novel contributions to my knowledge:

* Linking intrinsic dimensionality of language model layers to their ability to predict brain activity (correlate with encoding model performance)
* Providing evidence for a two-phase abstraction process in language models using fMRI data, via the notion of dimensionality
* Demonstrating what this reveals about the abstraction and prediction process in LLMs

These insights offer a new perspective on why intermediate layers of language models are often best at predicting brain activity, along with what the higher-level language areas might be doing in the brain. 

**Significance**:

This research has potentially important implications for both neuroscience and AI, and the particular topic of this workshop:

* It provides new evidence about how the brain processes language, suggesting similarities with the abstraction processes in language models.
* It offers insights into the internal representations of language models that could inform future model development and interpretation.

**Pros**:

1. Novel integration of intrinsic dimensionality analysis with fMRI encoding models
2. Comprehensive analysis across multiple language models and model sizes (from 125m to 13b)
3. Examination of how representational properties emerge during training via layer-wise analysis 
4. Strong empirical support for the main hypotheses using multiple metrics to build robustness, including evidence to suggest the compositional phase contributing to increased encoding performance (while prediction phase deteriorates performance)
5. Many hours of fMRI recordings being studied across subjects, over diverse sets of stories

**Cons**:
1. fMRI data from only 3 subjects, which may limit generalizability of the results
2. Focuses mainly on correlational relationships between dimensionality and encoding model performance; whether this is the mechanism behind alignment is still unclear (Figure 2 particularly appears to suggest some uncertainty on this from later layers)
3. Some aspects of the discussion could be expanded, particularly regarding alternative interpretations of what the dimensionality indicates about the representations in the brain.

In conclusion, this paper presents a high-quality, original study that would be a great fit to the UniReps workshop to scale our understanding and alignment of language processing in both artificial and biological systems. While there are some limitations, the strengths of the work outweigh its weaknesses, and it opens up promising avenues for studying the mechanisms behind the trends in LLM-brain alignment in recent years, in addition to providing support for new hypotheses on how the brain processes language.
title: title: new approaches to understanding the stages of LLM inference
main_review: main_review: Paper explores the relationship between the representations learned by large language models (LLMs) and human brain activity as measured by fMRI by proposing both follow a two phase process of abstraction and prediction. Though the idea that LLMs process in multiple stages is not new (eg: Lad 24 = https://arxiv.org/abs/2406.19384) - this paper puts this work into new perspective in several ways. The first is by looking this phenomenon with two different metrics - intrinsic dimensionality and surprisal; I found both of these interesting ways of framing and measuring the workings of the inner layers. The second is by comparing both to fMRI data to make an argument that this transition point of highest intrinsic dimensionality is the area with the most brain similarity. I also found the study on how these stages develop during the course of training to be strong support for their framework.

I am not well versed in interpreting fMRI data and not convinced that these two phases of abstraction and prediction are universal. But the authors make an interesting case for this and I otherwise found this paper clear and on topic for the workshop as well as densely packed with results in the context of a four page extended abstract.
title: title: A study of why certain layers are more aligned with the brain than others, revealing the distinct roles that certain layers play and the effects that that role has in shaping dimensionality and alignment.
main_review: main_review: Synopsis:

This paper studies the phenomena of brain-LLM similarity, but goes beyond the 'what' (similarity measurement) to an investigation of the 'why'. The authors look into different reasons why middle layers may be better aligned with the brain, finding that earlier layers function as an abstraction process and later layers function as an extraction process for prediction (roughly speaking). The layers that work to develop better abstractions of the input are well aligned. The authors back this up with strong methodology and thorough experiments.


Pros:
- I think that the methodology is very strong. Measuring between different models (OPT and Pythia) is good, and the use of the different checkpoints of Pythia to see how learning effects intrinsic dimensionality and brain-model alignment is very interesting. I think this could nicely lead into some future work if you were to also study grokking, etc.
  
- Tying this into some of the LLM interpretability literature is a good touch. I have seen some studies describing different phases / phase changes in models as they learn, and I was pleased to see that your results provide further evidence of this.
  
- I also think that this paper is quite well written, with very good graphics. The contextualization of the problem is also particularly well done. 
  
- It seems interesting that next-token prediction seems to explain this less than the intrinsic dimensionality. Perhaps a good direction for future work would be to explore how different training objectives alter the intrinsic dimensionality through time. This may require quite a lot more experiments and compute though - I'm not sure if models with provided checkpoints like Pythia are abundant for studying this yet.


Cons:

- Training on some more models would have been nice, especially with different training objectives or input modalities. However, given the compute already expended for these results, this is not a major objection.

Overall:

Overall I'm really quite pleased with your work. I think it's easily strong enough to have gone in the proceedings track instead of the extended abstracts track, but perhaps you want to extend on it and publish a lengthier version elsewhere. I'm quite happy with your decision to dig deeper into why certain layers may be better aligned, and I haven't seen much work (if any) in this direction yet. Keep up the nice work!
title: title: The topic and results are extremely important and interesting for our understanding of the brain and LLMs.
main_review: main_review: It’s a very interesting study that compares the brain and large language models (LLMs) and explores the underlying reasons for their similar functions. Overall, the author presents a well-structured and engaging paper. I enjoyed reading it and didn’t encounter any major issues. However, I have a few minor suggestions that could help further improve the manuscript.

1. In the neural encoding model training section, you briefly mention [4] for the specific methods used in the encoding model. However, including a brief description of these methods in the appendix or main text would be beneficial, as it would help readers better understand your paper.

2. If Figures 1C, E.2, and E.3 use Pycortex for visualization, please cite the relevant paper.

3.  Could you add a table or brief descriptions for each LLM in the appendix if there isn’t enough space in the main text? This would help readers quickly grasp the key information about each model and enhance the clarity of your paper.

4. Please add an "Appendix Section" text label on page 7, as it appears to be missing.