cs.CLOct 14, 2025
SaveFine-grained Analysis of Brain-LLM Alignment through Input Attribution
Organizations: Sapienza University of Rome, Italy · Sony AI, Zurich, Switzerland · Max Planck Institute for Software Systems, Germany
Abstract
Understanding the alignment between large language models (LLMs) and human brain activity can reveal computational principles underlying language processing. We introduce a fine-grained input attribution method to identify the specific words most important for brain-LLM alignment, and leverage it to study a contentious research question about brain-LLM alignment: the relationship between brain alignment (BA) and next-word prediction (NWP). Our findings reveal that BA and NWP rely on largely distinct word subsets: NWP exhibits recency and primacy biases with a focus on syntax, while BA prioritizes semantic and discourse-level information with a more targeted recency effect. This work advances our understanding of how LLMs relate to human language processing and highlights differences in feature reliance between BA and NWP. Beyond this study, our attribution method can be broadly applied to explore the cognitive relevance of model predictions in diverse language processing tasks.
Figures & tables
Figure 1: Overview of the attribution pipeline for BA and NWP. The BA pipeline (pink) extracts token embeddings for the input contexts through a pretrained LLM, aggregates them into TR-level embeddings, that are passed to the brain prediction head. Attribution scores are computed with respect to the mean squared error between predicted and recorded brain responses. The NWP pipeline (green) uses the same input to predict the subsequent word via the model’s language modeling head, with attribution scores computed based on the cross-entropy loss. In both cases, token-level attributions are aggregated to produce word-level importance scores, allowing for direct comparison of the linguistic features driving each task.
Figure 2: Intersection over Union (IoU) between BA and NWP attributions considering top-attributed word sets of increasing size on Harry Potter. For each context, the top- attribution set is defined as the smallest subset of words whose cumulative attribution mass reaches of the total. The plot on the left reports the average IoU across contexts, layers, subjects, and models, with gray shapes indicating per-model IoUs. Gray line: random baseline. The IoU is very low for , and then grows linearly as the threshold is increased, but it is always above chance. The middle plot reports the AUC-IoU, with * indicating significantly higher values than the random baseline. The plot on the right assesses differences in AUC-IoU between models using pairwise FDR-corrected Wilcoxon tests.
Figure 3: Number of top-attributed unique words, averaged across models and contexts, needed to cumulatively reach of the total context attribution, across three representative layers on the Harry Potter dataset. Error bars represent the standard error across contexts. Each plot compares BA and NWP, with the area under the curve (AUC) quantifying the total attribution spread. Asterisks denote significant differences ( ), assessed via a two-sided paired t-test, with Benjamini-Hochberg correction ( Benjamini and Hochberg, 1995 ) .
Figure 4: Distribution of top-attributed words across linguistic feature categories (semantic, discourse, syntactic) for BA-only, NWP-only, and for both. Results are averaged across layers, subjects, and models, with standard errors across models. We show results for the subsets of words that cumulatively account for the top , , and of each task’s total attribution mass. At lower thresholds, NWP prioritizes syntactic features, while BA emphasizes semantic and discourse-level content. As the threshold increases, the overlap between tasks grows, but BA consistently favors meaning-oriented features.
Figure 5: Distribution of the top- most influential words, defined as the smallest subset whose cumulative attribution reaches of the total, by their distance from the most recent word in the context ( last word). Each bin shows, for that distance, the proportion of important words relative to the total number of words occurring at that position. NWP exhibits strong primacy and recency peaks, whereas BA shows a broader, less edge-focused recency profile. Dashed lines mark the CoM, and shaded areas indicate weighted positional standard deviation, which is generally larger for NWP due to its sharp primacy peak.
Appendix figures & tables33 assets
Supplementary material from the paper’s appendix.
Appendix
| Name | # Layers | Architecture | ||
|---|---|---|---|---|
| Falcon3-1B | Transformer | |||
| Gemma-2B | Transformer | |||
| Llama3.2-1B | Transformer | |||
| Mamba-1.4B | SSM | |||
| Zamba2-1.2B | Hybrid | |||
| Gemma-2B-Instruct | Transformer |
Table 1: LLMs used in our study, along with their most relevant characteristics: Falcon3-1B ( Team, 2024 ) , Gemma-2B ( Mesnard et al., 2024 ) , Llama3.2-1B ( Meta AI, 2024 ) , Mamba-1.4B ( Gu and Dao, 2023 ) , and Zamba2-1.2B ( Glorioso et al., 2024 ) .
Figure 6: Extraction and processing of LLM representations. We associate a context of words to each word in the input text, with . We then pass the tokenized contexts to an LLM and extract token embeddings of shape , with being the maximum number of tokens in a context. To get word embeddings of shape , we average the embeddings of the tokens in the last word in each context. Finally, we average the word embeddings of all the words in each TR, obtaining TR embeddings of shape , and concatenate previous TR embeddings to account for the delay in the BOLD response, yielding a final representation of shape .
Figure 7: Illustration of the full brain alignment pipeline.
| Layer depth | |||
|---|---|---|---|
| Model | Early | Middle | Late |
| Falcon3-1B | 5 | 9 | 12 |
| Gemma-2B | 5 | 6 | 12 |
| Llama3.2-1B | 4 | 9 | 12 |
| Mamba-1.4B | 15 | 29 | 32 |
| Zamba2-1.2B | 11 | 17 | 30 |
Table 2: Indexes of the layers used for computing brain alignment attributions. For each model, , with being the number of layers.
Figure 8: Mean percentage increase in cross-entropy loss across masking modalities. Results are averaged over models; error bars show the standard error across models. Masking top-attributed words, either with random tokens or with same-POS replacements, results in large degradation, while masking random words has minimal effect.
Figure 9: Mean absolute drop in Pearson’s across masking modalities, grouped by ROI. Bars show averages across models, while dashed lines represent the original obtained using unmasked inputs. Asterisks denote cases in which top-attributed random or POS-based masking leads to significantly higher drops than the random masking baseline ( ). All perturbations strongly degrade alignment, but top-attributed masking (random or same-POS) consistently matches or exceeds the random-masking baseline, demonstrating that the top-ranked words are indeed functionally informative.
Figure 10: Intersection over Union (IoU) between BA and NWP attributions considering top-attributed word sets of increasing size on Harry Potter, using only 1 or 2 concatenated TRs. For each context, the top- attribution set is defined as the smallest subset of words whose cumulative attribution mass reaches of the total. We report the average IoU across contexts, layers, subjects, and models, with gray shapes indicating per-model IoUs. Gray line: random baseline. The IoU is very low for , and then grows linearly as the threshold is increased, but it is always above chance.
Figure 12: Number of unique important words, averaged across models and contexts, that are needed to cumulatively reach increasing attribution thresholds across three representative layers on the Harry Potter dataset, using 1 or 2 concatenated TRs. Error bars represent the standard error across contexts. Each plot compares BA and NWP, with the area under the curve (AUC) quantifying the total attribution spread. Asterisks denote significant differences ( ), assessed via a two-sided paired t-test, with Benjamini-Hochberg correction ( Benjamini and Hochberg, 1995 ) .
Figure 13: Distribution of top-attributed words across linguistic feature categories (semantic, discourse, syntactic) for brain-alignment-only, next-word-prediction-only, and words prioritized by both, using 1 and 2 concatenated TRs respectively. Results are averaged across layers, subjects, and models, with standard errors across models. We show results for the top , , and attribution thresholds. At low thresholds, NWP prioritizes syntactic features, while brain alignment emphasizes semantic and discourse-level information. As the threshold increases, overlap grows, but brain alignment consistently favors meaning-oriented features.
Figure 15: Distribution of top-attributed words (top 60% attribution) by distance from the most recent word, using 1 and 2 concatenated TRs respectively. For each representative model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. NWP shows a bimodal distribution, with sharp recency and primacy peaks. BA, on the other hand, emphasizes more distributed words, showing a much broader recency peak.
Figure 16: Intersection over Union (IoU) between BA and NWP attributions considering top-attributed word sets of increasing size for the Moth Radio Hour dataset. For each context, the top- attribution (x axis) set is defined as the smallest subset of words whose cumulative attribution mass reaches of the total. We report the average IoU across contexts, layers, subjects, and models, with gray shapes indicating per-model IoUs. The gray line represents a random baseline. The overlap between important words for the two tasks is very low for , and then grows linearly as the threshold is increased.
Figure 17: Number of unique important words, averaged across models and contexts, that are needed to cumulatively reach increasing attribution thresholds across three representative layers on the Moth Radio Hour dataset. Error bars represent the standard error across contexts. Each plot compares brain alignment (BA) and next-word prediction (NWP). The area under the curve (AUC) quantifies the total attribution spread. Asterisks denote significant differences ( ), assessed via a two-sided paired t-test, with Benjamini-Hochberg correction ( Benjamini and Hochberg, 1995 ) .
Figure 18: Distribution of top-attributed words (top 60% attribution) by distance from the most recent word. For each model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. NWP shows a bimodal distribution, with sharp recency and primacy peaks. BA, on the other hand, emphasizes more distributed words, showing a much broader recency peak.
Figure 19: Intersection over Union (IoU) between BA and NWP attributions considering top-attributed word sets of increasing size for the Harry Potter dataset using Gradient × Input and SmoothGrad. For each context, the top- attribution (x axis) set is defined as the smallest subset of words whose cumulative attribution mass reaches of the total. The plots on the left report the average IoU across contexts, layers, subjects, and models, with gray shapes indicating per-model IoUs. Gray line: random baseline. The IoU is very low for , and then grows linearly as the threshold is increased, but it is always above chance. The middle plots report the AUC-IoU, with * indicating significantly higher values than the random baseline. The plots on the right assess differences in AUC-IoU between models using pairwise FDR-corrected Wilcoxon tests.
Figure 21: Number of unique important words, averaged across models and contexts, that are needed to cumulatively reach increasing attribution thresholds across three representative layers on the Harry Potter dataset. Attributions are computed using either Gradient × Input (GXI) or SmoothGrad (SG). Error bars represent the standard error across contexts. Each plot compares BA and NWP, with the area under the curve (AUC) quantifying the total attribution spread. Asterisks denote significant differences ( ), assessed via a two-sided paired t-test, with Benjamini-Hochberg correction ( Benjamini and Hochberg, 1995 ) .
Figure 23: Distribution of top-attributed words across linguistic feature categories (semantic, discourse, syntactic) for brain-alignment-only, next-word-prediction-only, and words prioritized by both. Results are averaged across layers, subjects, and models, with standard errors across models. We show results for the top , , and attribution thresholds. At low thresholds, NWP prioritizes syntactic features, while brain alignment emphasizes semantic and discourse-level information. As the threshold increases, overlap grows, but brain alignment consistently favors meaning-oriented features.
Figure 24: Distribution of top-attributed words (top 10% attribution) by distance from the most recent word using Gradient × Input. For each model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. Both BA and NWP show a strong recency bias, with top-attributed words placed at very low distances from the most recent word.
Figure 25: Distribution of top-attributed words (top 10% attribution) by distance from the most recent word using SmoothGrad. For each model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. Both BA and NWP show a strong recency bias, with top-attributed words placed at very low distances from the most recent word.
Figure 26: Distribution of top-attributed words (top 60% attribution) by distance from the most recent word using Gradient × Input. For each model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. Both BA and NWP show a strong recency bias, with top-attributed words placed at very low distances from the most recent word.
Figure 27: Distribution of top-attributed words (top 60% attribution) by distance from the most recent word using SmoothGrad. For each model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. Both BA and NWP show a strong recency bias, with top-attributed words placed at very low distances from the most recent word.
Figure 28: Distribution of top-attributed words (top 60% attribution) by distance from the most recent word in the context. For each model, we plot the proportion of important words located at each distance bin, comparing brain alignment (BA) and next-word prediction (NWP). NWP shows a strong recency bias, while BA often emphasizes earlier or more distributed words.
Figure 29: Distribution of top-attributed words (top 60% cumulative attribution) by distance from the most recent word in the context, shown separately for each of the 8 subjects for Llama3.2-1B. Each plot compares brain alignment and next-word prediction. The overall shape and center of mass of the attribution distributions are consistent across subjects, indicating low inter-subject variability. This supports the robustness of the observed task-specific attribution differences.
Figure 30: Distribution of top-attributed words (top 60% cumulative attribution) by distance from the most recent word in the context, shown for each of the 8 subjects for Mamba-1.4B. Across all subjects, next-word prediction shows a strong recency and primacy bias, while brain alignment attribution is more concentrated around the recent context and more broadly distributed over recent words. The similarity in curves across subjects confirms that these effects are consistent at the individual level.
Figure 31: Distribution of top-attributed words (top 10% attribution) by distance from the most recent word in the context. For each model, we plot the proportion of important words located at each distance bin, comparing brain alignment (BA) and next-word prediction (NWP). NWP shows a strong recency bias, while BA often emphasizes earlier or more distributed words.
Figure 32: Model-wise brain alignment performance. (a) Noise ceiling computed using cross-subject prediction accuracy ( Schrimpf et al., 2021 ) . (b) Mean normalized correlation between true and predicted brain activity across models. We report the mean across voxels, layers, and subjects, with standard error across subjects. (b) Pairwise significance heatmap showing -values from paired t-tests comparing mean correlation scores between models.
Figure 33: Distribution of top-attributed words by distance from the most recent word in the context for Qwen2-1.5B, shown at thresholds . Each plot compares brain alignment (BA) and NWP. Across thresholds, NWP displays a sharp recency bias and a secondary primacy peak, while BA shows a broader recency distribution. Unlike Llama3.2-1B on the Harry Potter dataset, Qwen2-1.5B does not exhibit oscillatory attribution patterns, instead maintaining smooth and monotonic profiles.
Figure 34: Distribution of top-attributed words (top 10% attribution) by distance from the most recent word in the context, using 80-word contexts. For each model, we plot the proportion of important words located at each distance bin, comparing brain alignment (BA) and next-word prediction (NWP). NWP shows a strong recency bias, while BA often emphasizes earlier or more distributed words.
Figure 35: Distribution of top-attributed words (top 60% attribution) by distance from the most recent word in the context, using 80-word contexts. For each model, we plot the proportion of important words located at each distance bin, comparing brain alignment (BA) and next-word prediction (NWP). NWP shows a strong recency bias, while BA often emphasizes earlier or more distributed words.
Figure 36: IoU of top-attributed words between long (640-word) and short (80-word) contexts as a function of attribution mass threshold . NWP shows high robustness across window sizes, while BA exhibits low overlap, reflecting its more distributed and flexible use of context.
| Task | Model | Time (dd:hh:mm:ss) | Peak Memory (GB) |
|---|---|---|---|
| Brain alignment training | Llama3.2-1B | 00:04:08:16 | 3.16 |
| Gemma-2B | 00:04:47:19 | 3.08 | |
| Gemma-7B | 00:14:46:32 | 3.98 | |
| Gemma-2B-Instruct | 00:04:53:34 | 2.77 | |
| Falcon3-1B | 00:05:42:46 | 3.39 | |
| Zamba2-1.2B | 00:10:24:00 | 3.66 |
Table 3: Per-task compute time and peak memory usage across models on the Harry Potter dataset.
| Attribution Method | Model | Time (hh:mm:ss) | Peak Memory (GB) |
|---|---|---|---|
| Gradient × Input | Llama3.2-1B | 00:22:58 | 2.18 |
| Gemma-2B | 00:08:26 | 2.16 | |
| Gemma- 7B | 01:14:23 | 2.40 | |
| Gemma-2B-Instruct | 00:34:42 | 2.40 | |
| Falcon3-1B | 00:06:07 | 2.17 | |
| Zamba2-1.2B | 00:14:20 | 2.26 |
Table 4: Compute time and peak memory usage for NWP attribution on the Harry Potter dataset.
| Task | Model | Time (dd:hh:mm:ss) | Peak Memory (GB) |
|---|---|---|---|
| Brain alignment training | Llama3.2-1B | 00:08:44:08 | 34.49 |
| Gemma-2B | 00:10:19:46 | 36.05 | |
| Gemma-7B | 02:13:15:03 | 38.53 | |
| Gemma-2B-Instruct | 00:12:37:10 | 29.05 | |
| Falcon3-1B | 01:11:56:51 | 40.98 | |
| Zamba2-1.2B | 02:00:45:34 | 36.24 |
Table 5: Per-task compute time and peak memory usage across models on the Moth Radio Hour dataset.
| Task | Model | Time (dd:hh:mm:ss) | Peak Memory (GB) |
|---|---|---|---|
| Brain alignment training | Llama3.2-1B | 00:05:40:33 | 3.01 |
| Gemma-2B | 00:06:13:33 | 2.99 | |
| Gradient × Input attribution | Llama3.2-1B | 00:09:57:10 | 2.59 |
| Gemma-2B | 00:09:52:24 | 2.57 | |
| NWP attribution | Llama3.2-1B | 00:00:22:58 | 2.18 |
| Gemma-2B | 00:00:08:26 | 2.16 |
Table 6: Per-task compute time and peak memory usage across models for experiments with reduced context length.
More in this topic
LLM Interpretability
Large language model (LLM) interpretability studies how internal representations and computations produce outputs, and how they can be explained or controlled.
See all 778 papers