Reconstructing the cascade of language processing in the brain using the internal computations of a transformer-based language model

Piecing together the meaning of a narrative requires understanding not only the individual words but also the intricate relationships between them. How does the brain construct this kind of rich, contextual meaning from natural language? Recently, a new class of artificial neural networks—based on the Transformer architecture—has revolutionized the field of language modeling. Transformers integrate information across words via multiple layers of structured circuit computations, forming increasingly contextualized representations of linguistic content. In this paper, we deconstruct these circuit computations and analyze the associated “transformations” (alongside the more commonly studied “embeddings”) at each layer to provide a fine-grained window onto linguistic computations in the human brain. Using functional MRI data acquired while participants listened to naturalistic spoken stories, we find that these transformations capture a hierarchy of linguistic computations across cortex, with transformations at later layers in the model mapping onto higher-level language areas in the brain. We then decompose these transformations into individual, functionally-specialized “attention heads” and demonstrate that the emergent syntactic computations performed by individual heads correlate with predictions of brain activity in specific cortical regions. These heads fall along gradients corresponding to different layers, contextual distances, and syntactic dependencies in a low-dimensional cortical space. Our findings provide a new basis for using the internal structure of large language models to better capture the cascade of cortical computations that support natural language comprehension. Introduction Language comprehension is a fundamentally constructive process. We resolve local dependencies among words to assemble lower-level linguistic units into higher-level units of meaning (Chomsky, 1965; MacDonald et al., 1994; Partee, 1995; Goldberg, 2006; Berwick et al., 2013; Christiansen & Chater, 2016), ultimately arriving at the kind of narratives we use to understand the world (Bruner, 1985; Graesser et al., 1994). For example, if a speaker refers to “the secret plan,” we implicitly process the relationships between words in this construction to understand that “secret” modifies “plan. ” At a higher level, we use the context of the surrounding narrative to understand the meaning of this phrase—what does the plan entail, who is keeping the secret, and who are they keeping the secret from? This context may comprise hundreds of words 1 unfolding over the course of several minutes. The human brain is thought to implement these processes via a series of computations that transform acoustic speech signals into actionable representations of meaning (Hickok & Poeppel, 2007; Hasson et al., 2015; Ding et al., 2016; Friederici et al., 2017; Martin & Doumas, 2017; Pylkkänen, 2019; Martin, 2020). Traditionally, neuroimaging research has used targeted experimental manipulations to isolate particular linguistic computations—for example, by manipulating the presence/absence or complexity of a given syntactic structure—and mapped these computations onto brain activity in controlled settings (Bookheimer, 2002; Vigneau et al., 2006; Hickok & Poeppel, 2007; Friederici, 2011; Price, 2012). While these findings laid the groundwork for a neurobiology of language, they have limited generalizability outside the laboratory setting, and it has proven difficult to synthesize them into a holistic model that can cope with the full complexity of natural language. This has prompted the field to move toward more naturalistic comprehension paradigms (Hamilton & Huth, 2020; Nastase et al., 2020; Willems et al., 2020). However, these paradigms introduce their own challenges: principally, how to explicitly quantify the linguistic content and computations supporting the richness and expressivity of natural language. In recent years, the field of natural language processing (NLP) has been revolutionized by a new generation of deep neural networks capitalizing on the Transformer architecture (Vaswani et al., 2017; Devlin et al., 2019; Radford et al., 2019). Transformers are deep neural networks that forgo recurrent connections (Elman, 1990; Hochreiter & Schmidhuber, 1997) in favor of layered “attention head” circuits, facilitating self-supervised training on massive real-world text corpora. At a high level, the Transformer architecture represents the meaning of words as numerical vectors in a high-dimensional “embedding” space where closely related words are located nearer to each other. Unlike the previous generation of word embeddings (Landauer & Dumais, 1997; Mikolov et al., 2013; Pennington et al., 2014), which assign each word a single static (i.e. non-contextual) meaning, Transformers process large blocks of text simultaneously to assign each word a context-sensitive meaning. The core circuit motif of the Transformer—the attention head—incorporates a weighted sum of information exposed by other words, where the relative weighting “attends” more strongly to some words than others. The initial embeddings used as input to the Transformer are non-contextual. Within the Transformer, attention heads in each layer operate in parallel to update the contextual embedding, resulting in surprisingly sophisticated representations of linguistic structure (Manning et al., 2020; Linzen & Baroni, 2021). The success of Transformers has inspired a growing body of neuroscientific work using them to model human brain activity during natural language comprehension (Jain & Huth, 2018; Toneva & Wehbe, 2019; Antonello et al., 2021; Caucheteux et al., 2021a, 2021b, 2021c; Heilbron et al., 2021; Lyu et al., 2021; Caucheteux & King, 2022; Goldstein et al., 2022). These efforts have focused exclusively on the “embeddings”—the Transformer’s representation of linguistic content—and have largely overlooked the “transformations”—the computations performed by the attention heads. Recent work in NLP, however, has revealed that particular linguistic operations are approximated by particular transformations (Clark et al., 2019; Tenney et al., 2019): for example, attention head 10 in the eighth layer of BERT is specialized for resolving direct object relations, whereas head 11 in the same layer closely tracks noun modifiers (Clark et al., 2019). Although the individual heads that implement these computations operate independently, their transformations are ultimately “fused” together to form the resulting embedding. Thus, unlike the embeddings, the transformations at a given layer can be disassembled into the specialized computations performed by the constituent heads. Can we leverage