Fine-grained Analysis of Brain-LLM Alignment through Input Attribution

TL;DR

This work introduces an end-to-end input attribution framework to dissect why large language models align with human brain activity, and applies it to compare brain alignment (BA) with next-word prediction (NWP) across multiple architectures and brain datasets. By linking LLM embeddings to fMRI responses via a differentiable projection, the authors compute word-level attributions for BA and contrast them with NWP using gradient-based methods, primarily Gradient × Input and validated with Integrated Gradients. They find that BA and NWP rely on largely distinct word sets, with BA drawing more on semantic and discourse-level information and exhibiting broader recency-based integration, while NWP emphasizes syntax and sharp recency/primacy biases, particularly in early layers. Across datasets and models, BA shows higher attribution spread in middle/late layers, suggesting more distributed, high-level linguistic processing underlying brain alignment. The results challenge the notion that BA is fully explained by NWP and provide a generalizable tool to study the cognitive relevance of model predictions across language tasks, with potential to inform brain-tuned model development and cognitive-language research.

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 (27)

• 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 as a function of the top-$t\%$ attribution threshold for the HP dataset. 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 $t\leq10\%$, and then grows linearly as the threshold is increased, but it is always above chance.
• Figure 3: 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 HP 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 ($p<0.001$), assessed via a two-sided paired t-test, with Benjamini-Hochberg correction benjamini1995controlling.
• 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 top $10\%$, $60\%$, and $80\%$ of attribution scores. 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 top-$60\%$ attributed words by distance from the most recent word in the context for two representative models. For each model, we plot the proportion of important words located at each distance bin, comparing BA and NWP. NWP attributions peak at the beginning and end of contexts (primacy and recency), whereas BA peaks are broader and less edge-focused. The dashed lines represent the CoM values, which are shifted toward higher distances for Llama3.2-1B.