Transformer Mechanisms Mimic Frontostriatal Gating Operations When Trained on Human Working Memory Tasks

## Abstract

Abstract Models based on the Transformer neural network architecture have seen success on a wide variety of tasks that appear to require complex “cognitive branching”– or the ability to maintain pursuit of one goal while accomplishing others. In cognitive neuroscience, success on such tasks is thought to rely on sophisticated frontostriatal mechanisms for selective gating , which enable role-addressable updating– and later readout– of information to and from distinct “addresses” of memory, in the form of clusters of neurons. However, Transformer models have no such mechanisms intentionally built-in. It is thus an open question how Transformers solve such tasks, and whether the mechanisms that emerge to help them to do so bear any resemblance to the gating mechanisms in the human brain. In this work, we analyze the mechanisms that emerge within a vanilla attention-only Transformer trained on a simple sequence modeling task inspired by a task explicitly designed to study working memory gating in computational cognitive neuroscience. We find that, as a result of training, the self-attention mechanism within the Transformer specializes in a way that mirrors the input and output gating mechanisms which were explicitly incorporated into earlier, more biologically-inspired architectures. These results suggest opportunities for future research on computational similarities between modern AI architectures and models of the human brain. Keywords: transformers; neural networks; working memory; computational neuroscience; gating; computational cognitive science; mechanistic interpretability

## Introduction

Computational models based on the Transformer architecture Vaswani \BOthers. () have seen success on a wide variety of tasks that appear to require complex “cognitive branching”: the ability to maintain pursuit of one over-arching goal while performing other subtasks along the way. For example, Transformer-based large language models (LLMs) have demonstrated impressive abilities in not just language Brown \BOthers. (), but planning Huang \BOthers. (), navigation Du \BOthers. (), and problem solving Lewkowycz \BOthers. ().

In humans, performance on such tasks is known to depend on a neural mechanism for gating Frank \BBA Badre (); Badre \BBA Frank (); Chatham \BOthers. (); Rac-Lubashevsky \BBA Kessler (); Rac-Lubashevsky \BBA Frank (), which controls whether new information is maintained in working memory or not, the address in memory where it is stored, and the address from which stored information is recalled in response to a task.
Typical Transformer models have no specialized architecture for working memory, in spite of their ability to succeed at tasks which appear to require it. Although some Transformer models have additional built-in structure for memory Dai \BOthers. (); Burtsev \BOthers. (), recurrence Dai \BOthers. (), or hierarchy Y\BHBIS. Wang \BOthers. (), vanilla Transformer models without any such inductive biases remain the dominant architecture for the modern AI systems which are especially visible in their success on complex tasks Brown \BOthers. (); Touvron \BOthers. (). This raises the question: in solving such tasks, does a mechanism for selective input and output gating emerge within the vanilla Transformer?

Transformers are good candidates for learning gating behavior because of inductive biases within the self-attention mechanism, i.e., the Transformer’s defining architectural component. Within self-attention, numerous “attention heads” construct contextual representations for each item in their input sequence through a learned weighted combination of the previous items in the sequence. Attention heads could in principle learn gating behavior by marking sequence elements with a key (input gating) and reading out those values later by querying for those keys (output gating). Moreover, the attention mechanism in Transformers is decomposed in a way that enables it to readily differentiate “reading” and “writing” operations. This behavior is analogous functionally to the neurobiological roles of corticostriatal circuits in humans and other animals, in which isolated clusters of prefrontal neurons represent distinct “addresses” in memory that can be updated or read out from via selective gating actions triggered by basal ganglia and thalamus O’Reilly \BBA Frank (); Frank \BBA Badre (); Calderon \BOthers. (). In computational neuroscience models of this process, the prefrontal clusters (or “stripes”) can also serve as latent abstract “roles” that condition how to interpret content within them, affording functions such as variable binding, indirection, and hierarchical generalization to new situations O’Reilly \BBA Frank (); Collins \BBA Frank (); Kriete \BOthers. (); Bhandari \BBA Badre (). Thus, Transformers may use their attention heads to learn a gating strategy that mimics certain functions of the brain.

In this work, we train vanilla Transformer models with self-attention on a working memory task paradigm that was specifically designed to evaluate models of selective gating and working memory in computational neuroscience O’Reilly \BBA Frank (); Rac-Lubashevsky \BBA Frank (). We use recent techniques from mechanistic interpretability Olah (); Nanda \BBA Bloom () to expose the mechanism that the Transformer uses in order to perform the task.
We find that, as a result of training, the self-attention mechanism specializes in a way that resembles existing models of input-output gating.
Specifically, we find that the trained model uses the key vectors within the attention mechanism to control input gating, i.e., determining which elements in the input to consider vs. ignore, as well as controlling how the information is stored– in other words, assigning it a role such that the model can access it later. The model uses the query vectors to control output gating, i.e., determining which information is accessed in order to complete a task. Our findings highlight the importance of considering the emergent mechanisms that result from training in addition to the innate architectural mechanisms when drawing comparisons between AI systems and human cognitive processes, and opens the door for future analysis and work which can enable more principled studies of the similarities and differences between human vs. machine cognition.

## Background

### Gating Mechanisms

Working memory in human brains is known to make use of a gating mechanism, which processes and stores information roughly analogously to gates being opened and closed Rac-Lubashevsky \BBA Kessler (); Rac-Lubashevsky \BBA Frank (); Bhandari \BBA Badre ().
The input gate controls which information is stored or not stored in memory, and if stored, into which “address”. The output gate controls which content within working memory is accessed in order to produce a response or to make subsequent gating operations. Both input and output gating policies are also dependent on the learned task-dependent context (i.e. role) of the information to be stored and accessed in working memory.

In cognitive neuroscience, a variety of tasks are used to study the capacity of working memory. In this work, we focus on a variant of the “reference-back 2” task Rac-Lubashevsky \BBA Frank (), a human paradigm meant to mimic a task designed to showcase the need for selective gating of independent contents of information in frontostriatal neural networks O’Reilly \BBA Frank (). In the reference-back paradigm, symbols such as letters or numbers are viewed one at a time, and the subject must determine whether the current symbol is the same or different as that stored in memory for a given role (letter or number). They also are given a cue to indicate whether to update the current symbol as the “reference” to be compared on subsequent trials of the same role, or if instead they should continue to maintain the previous reference. Thus this task requires selective updating and accessing of information in a role-addressable manner.

### Transformers

Transformers are powerful language models which create contextualized representations of sequences of words, where they learn to to predict the next token one at a time using an “attention” mechanism Bahdanau \BOthers. () to scan the previously seen tokens for relevant information. These models are able to learn and represent any arbitrary sequence modelling task.

For a given prediction, Transformer attention generates three separate vectors at each position in the sequence: a query, key, and value ($q$, $k$, $v$). The query vector scans the previous context (including the current token) for relevant keys, and calculates how much the current prediction should “attend” to those positions. Then, the value vectors at those positions are multiplied by the corresponding weights, summed up, and added to the next representation: for token $i$ at layer $j$, the contextual representation is $\sum_{k}q_{i}^{j}\cdot k_{k}^{j}*v_{k}^{j}$. Thus, the next token prediction includes information from earlier in the sequence by combining the value vectors from previous tokens. In other words, Transformer attention can be viewed as a read/write mechanism: for a given token, the queries and keys dictate which tokens to read from, the values are the content that is read proportional to the attention calculated by the keys and queries, and the summed content is written to a new representation at the given token. As we shall see below, the comparison to role-addressable input and output gating operations is evident, whereby the key vectors form addresses analogous to the PFC “stripes”, with key construction determining input gating, and the query vectors control which of them are accessed, with query construction determining output gating.

#### Limitations as Cognitive Models:

Transformers have properties which make them obviously bad models of human sequence processing. In particular, because Transformers can attend to any part of the sequence when creating a representation, they are not limited by memory representation constraints. Transformers could thus solve tasks that would push the limits of human working memory, but it should be noted they accordingly require large amounts of training data.
The question addressed in this work is orthogonal to these limitations. That is, we focus specifically on if and how Transformers can learn to implement an efficient gating mechanism to solve tasks with human working memory demands.

Figure: Figure 1: Graphical diagram of the path-patching process. Attention heads are represented as circles (layer,head index), and contextual representations of each token (as well as the next token prediction) are represented as rectangles.
Refer to caption: x1.png

### Mechanistic Interpretability

We use a set of recently introduced analysis tools Elhage \BOthers. () which enable us to uncover specific mechanisms defined in terms of model weights within the Transformer.
Specifically, we use path-patching K. Wang \BOthers. (); Goldowsky-Dill \BOthers. (), a generalization of causal mediation analysis Pearl () that allows us to determine which components of a neural model (e.g., attention heads) work together in order to produce observed behavior on a task. The discovered components are referred to as a circuit Räuker \BOthers. ().

Path-patching involves making a incisive edit to the representations of a trained model and observing how the model’s behavior is affected (see Fig. 1). Path-patching typically requires a minimal pair of examples: the “clean” example and the “corrupted” example, in which one token from the clean example is changed, as well as the correct label. Given representations from the model for both the clean and the corrupted examples (the blue and orange components in the figure),
we can chose a specific component anywhere in the model (referred to as the “sender”), and insert the corrupted representation at that component into the clean representation. We then use the model to recalculate the representations up until another component of the model (the “receiver”), thus “patching” the path. In the figure, we send from layer 0, head 0 to layer 1, both heads 0 and 1. All clean representations that are not along this path are not modified and are unaffected by the patch. The model then recomputes all representations after the receiver (the “patched” representations), and arrives at a new prediction. If the model output matches the corrupt prediction rather than the clean one, that prediction is causally dependent on the path from sender to receiver. See K. Wang \BOthers. () and Goldowsky-Dill \BOthers. () for a more comprehensive review of path-patching methods.

Figure: Figure 2: Above: example of textual reference-back task as model input. Below: step-by-step task process; models do not view task-internal grey words. “Update Instruction” executes after “Answer” despite appearing earlier sequentially.
Refer to caption: x2.png

## Task

We create a modified text-based version of the reference-back 2 task Rac-Lubashevsky \BBA Frank () designed to tax selective WM gating O’Reilly \BBA Frank ().The textual reference-back task requires making same/different judgments between incoming symbols assigned to a particular “register” in memory, with respect to those seen previously and linked to those same registers. Like the original tasks, the textual reference-back task is sequential, and requires the maintenance and independent updating of two memory registers, each containing one of $S$ arbitrary symbols at a time. The contents of the registers are updated over the course of the task. At the beginning of each sequence, each register is initialized individually to one $s\in S$ (the pool of symbols is shared between registers, which was shown to more substantively tax gating mechanisms in O’Reilly \BBA Frank ()). Each sequence is composed of $L$ tuples, each containing register address $\text{Reg}_{i}$, symbol $\text{Sym}_{i}$, same/different label $\text{Ans}_{i}$, and update instruction $\text{Ins}_{i}$. For a tuple $i\in L$, the answer $\text{Ans}_{i}$ is a binary value that is either same if symbol $\text{Sym}_{i}$ is currently stored in the register with address $\text{Reg}_{i}$, or different otherwise. The update instruction $\text{Ins}_{i}$ also takes one of two values, evenly distributed: if ignore, then there is no effect further on in the sequence.
If store, then from that point on in the sequence, $\text{Sym}_{i}$ is stored in the register with address $\text{Reg}_{i}$ until otherwise updated. An example is shown in Fig. 2.

We implement each reference-back task example in our data as a single sequence, and measure models’ ability to predict same versus different for each $\text{Ans}_{i}$. Each sequence has 10 same/different answers, and we generate 100,000 train, 1,000 dev, and 1,000 held-out test sequences.

The class balance of same to different answer labels in the train/test datasets is roughly 1:2, making a “maximum class” heuristic solution 0.66 accuracy, 0.33 precision, and 0.5 recall. We test several other heuristics, the strongest of which is predicting same if another tuple including Store and the target register and target symbol exists in the sequence, which scores 0.80 accuracy, 0.82 precision, and 0.85 recall.

## Model

We use vanilla Transformers in order to facilitate interpretability, as done in prior work that analyzes emergent mechanisms Elhage \BOthers. (). Our models contain two decoder-only layers, each with only two heads (four in total), and no multilayer perceptrons or layer normalization, followed by a linear “unembed” layer to project the output of the last decoder into the space of the entire vocabulary at each timestep (^1^11In practice, only ‘same’ and ‘different’ are ever predicted.). Our network uses absolute positional embeddings Vaswani \BOthers. (). The vocabulary contains all possible tokens, represented individually with embedding size $E$. Models are trained to predict the next token with the language modelling objective, meaning if the model is predicting $\text{Ans}_{c}$, it will have access to all ($\text{Ins}_{i}$, $\text{Reg}_{i}$, $\text{Sym}_{i}$, $\text{Ans}_{i}$) tuples where $i<c$, as well as $\text{Ins}_{c}$, $\text{Reg}_{c}$, and $\text{Sym}_{c}$. However, the models only receive loss at positions where a same/different token must be predicted. Furthermore, each layer gets a causal attention mask– when constructing each token representation, it cannot look ahead at tokens further down the sequence.

The models are trained over 60 epochs of the 100k training data points, learning from 6 million examples in total. Models are evaluated on their accuracy (whether the correct $\text{Ans}_{i}$ is predicted for each tuple $i$), measured in precision and recall, as well as the same versus different token logit difference.

## Experiments

First, we select and analyze a single Transformer model which succeeds on the task, and upon investigation of its weights find that it learns a mechanism for input/output gating.
Second, when we conduct a search over more trained models, we find that model performance correlates with markers of learning a gating policy, analogous to findings in the frontostriatal neural networks Frank \BBA Badre ().

Figure: Figure 3: Model behavior when predicting same/different (token 15) is shown. We measure attention visualized as a shade of purple, with deeper shade corresponding to higher attention to that token. We create “corrupted” minimal pairs in which changing a token (light blue) either changes the correct label at index 15 (examples b, c, e) or does not (d, f). We make small path-patching edits with the minimal pair to targeted network components (layer 1 keys for b, c, d, f; queries for e,f). In other words, we replace specific components (denoted with red text) with their corresponding representation from the “corrupted” sequence, but hold all other representations constant, and run the model and get a new same/different prediction. In all test examples, making the small patch successfully results in the model’s prediction changing to align with the “corrupted” example.
Refer to caption: x3.png

We first establish that a Transformer model is able to succeed on the reference-back task.
We perform a small hyperparameter search and select a model that reaches 100% accuracy on the held-out test data for further analysis. We determine the circuit that the model uses to solve the textual reference-back task through an array of path-patching experiments with a simplistic minimal pairs paradigm.
Our “sender” within path-patching is always both attention heads at layer 0, and our “receiver” is always both attention heads at layer 1.

At layer 0, the model learns to condense the task-critical information from each tuple into one embedding, at the position for $\text{Sym}_{i}$(^2^22Redundantly, the model does the same at the position for $\text{Ans}_{i}$. Through additional experimentation, we determine that this is a quirk of Transformer learning, and does not impact our analysis.). At this layer, the model pays 85.8% of total attention to the task-critical information to that tuple, and just 14.2% of attention to other tuples.

At layer 1, the attention heads learn to attend to the $\text{Sym}_{i}$ key vector representing the tuple where information was last stored in the target register. The heads pay 70.2% of total attention to this tuple (the “stored” tuple), and only 29.8% of attention to all other tokens. This behavior is tied to the target register matching the register in the stored tuple, which is analogous to gating of the relevant role-addressable PFC stripe. We focus our analysis on the Layer 1 representations which exhibit this learned gating policy, shown in Fig. 3.

### Input Gating through Key Vectors

Input gates in working memory control what incoming information is remembered and “role-addressable”– i.e., stored in memory in such a way that it is able to be freely accessed later when it is needed for task completion. In a Transformer, the key vectors serve the role of addresses (analogous to PFC “stripes”), which are retrieved based on their match to a query vector from the current or later timesteps during self-attention. Thus the input gating in Transformers is controlled through key construction; the composition of the output of the Layer 0 attention controls which content is stored in the key vector at Layer 1 for later use. At a later timestep, a query vector will address the information in the key vector.

In our model analysis, we find that key composition at the $\text{Sym}_{i}$ position (positions 2, 6, 10, and 14 in Fig. 3) roughly represents each tuple. A query’s ability to address this position depends on whether the represented tuple contains a Store or an Ignore. Key vectors representing an Ignore tuple receive very little attention (0.4% of layer 1 attention averaged over test set), whereas those representing a Store tuple receive the bulk of the attention (86.8%). We determine this effect causally with path-patching (Fig. 1).
First, we create clean sequences sampled from our test set, and then corrupt these sequences by switching a Store within tuple $i$ to an Ignore. We then path-patch only the key vectors of $i$. We expect, if the key controls input gating, that patching these key vectors should “block” attention to all of tuple $i$. An example attention pattern is in Fig. 3, examples a and b. We find that the model’s attention shifts away from the tuple accordingly in 100% of patched instances.
The presence of an Ignore or a Store within a tuple controls whether the key construction acts as an open input gate or a closed input gate.

Key construction also depends on the role of the represented content; within our task, that means whether $\text{Reg}_{i}$ is Reg0 or Reg1. When making a same/different prediction, key vectors representing a tuple that matches the target register receive most of the model’s attention (92.5% of total attention), while those that do not match are not attended to (3.3% of total attention). Similarly to input gating, we use path-patching to determine that key construction encodes roles. This time, given a target tuple $i$ with a target register, we corrupt the register of the stored tuple, changing it from Reg0 to Reg1 or vice versa; to predict the answer for $i$, the model must attend to an earlier tuple with the target register. An example is Fig. 3, row c. The model’s attention shifts away accordingly across every example in the test set. Note that the stored tuple must be modified; if the same corruption is made earlier (as in row d), attention does not shift. This behavior shows that the gating within self-attention is role-addressable; the registers within the task function as roles, and are embedded within the key vectors as part of the representation.

### Output Gating through Query Vectors

Within working memory, output gates control which addressable information is accessed in order to complete a task. Given that key vectors serve the role of addresses, query vectors in turn control which key vectors are accessed, through the final Q*K dot product in attention. Query construction thus performs the role of output gating within Transformers.

The query composition controls which addressable $\text{Sym}_{i}$ representations are attended to based on the identity of the target register; changing Reg1 to Reg0 controls which role-addressable content is accessed. We determine this through a final set of path-patching experiments, an example of which is shown in Fig. 3, row e. Rather than editing the register in the stored tuple as done in row c, we corrupt the target register itself; this means that the query must now find representations corresponding to the other tuple. In row e, the model finds $Sym_{5}$, and predicts different; and in row f, we patch in $Sym_{4}$ at index 2, and the model predicts same.
Upon inspection, the query is constructed to correspond with key vectors that represent tuples which also contain the target register. When we edit the target register in minimal pair experiments, we observe that the attention shifts from the original stored tuple (74.1% of attention) to the stored tuple that matches the edited register, and successfully makes an updated same/different comparison to the symbol in the edited register in every instance.

Editing aspects of the target tuple other than target register has minimal effect on the query construction behavior. No edits to the query cause the model to attend to a Ignore tuple, further evidencing of output gating behavior– only content that has been made “addressable” can be accessed for a response. Furthermore, we find that the target instruction and symbol do not factor into the query composition– changing them through path-patching to the query does not affect attention. This is notable because the model could employ other strategies for determining which tuples are eligible to be the stored tuple; e.g. attending to all symbols to match if any of them are the same as the target symbol.

### Successful Task Performance is Related to Discovering Gating Policies

To further identify how readily Transformer models learn a gating policy, and how useful such a policy is to succeed on the task, we train new models with the same hyperparameters across many different random seeds, and measure their performance on the target task as well as on markers of the gating policy.
We train 20 new models on the same textual reference-back data, each with a different random initialization, and measure both training loss and test set accuracy. 5 of the models succeed 100% of the time, and the other 15 models succeed between 94%-99.99% of the time, with a mean of 97.72% and a standard deviation of 2.03. The models are trained on the same amount of data (6 million examples).

We observed in the prior sections that the trained Transformer model uses its key composition to control input gating and its query composition to control output gating and role addressability. To identify whether the new models learn to gate similarly, we evaluate the key and query composition of all 20 new models by making minimal pair path-patching edits for every test example, where the answer changes from same to different or vice versa, similarly to Figure 3.

To evaluate the ability to open and close input gates, we corrupt the stored tuple’s Store to Ignore, and path-patch only to the stored tuple’s key vectors. To evaluate the role addressability of the content during output gating, we corrupt the target register (changing Reg0 to Reg1 or vice versa), and path-patch only to the query vector for making the same/different judgment. This is a more challenging task than patching to all of the keys or queries respectively; a model will only succeed on these subtasks if it implements a gating policy with the same markers that the model analyzed in earlier sections does. A model makes the patched prediction successfully if its prediction matches the corrupted sequence and not the original sequence– i.e. the targeted path-patch was sufficient to change its same/different prediction.

Figure: Figure 4: Model performance over training on patching subtasks. Each row corresponds to an individual model’s training loss (solid line) and subtask accuracy (dashed line) over time; blue and orange lines respectively correspond to models which reach 100% test accuracy and to those that do not.
Refer to caption: extracted/5405402/latex_cogsci_fullpaper/imgs/expt2_horizontal.png

We visualize the 5 runs that reach 100% accuracy on the test data in Figure 4, as well as 5 randomly selected runs that do not reach 100% accuracy on the test data, and plot their training loss versus their accuracy (between 0 and 100) on the two patched subtasks over the course of training.

Two trends become apparent from the data: first, models which score a perfect test accuracy appear to succeed at the subtasks more readily than models which do not. Of the former, 3 of 5 models reach 100% accuracy on both subtasks readily, plateauing less than halfway through training duration. However, examples from the latter category of models do not reach such immediate success at the patched subtasks (including the 10 not pictured in this graph); in fact, many categorically fail, scoring as low as 49% accuracy. These results do not indicate that this class of models’ representations are useless for the task– they all score between 94% and 99.99%, well above heuristic performance. Failing to succeed at the targeted patching subtasks reveals that these models may implement some other strategy that is not gating, and may be brittle or heuristic in some way. We take these results as evidence that learning a robust policy for gating correlates with model performance at the textual reference-back task.

The second emergent trend is that many models across both classes have a sharp decline in training loss, which correlates with a similarly steep increase in accuracy on both subtasks. Sudden jumps in performance is a noted phenomenon that has been observed in other Transformer models in cases of e.g., grokking Power \BOthers. (). We interpret this phase transition as suddenly learning a gating mechanism. Models that do not exhibit phase transitions to the same degree take longer to fit the task, and do not reach high subtask accuracy. The cause of phase transitions and what is learned during this process is left for future work.

## Summary and Discussion

In this work, we investigate Transformer models for emergence of a learned gating mechanism; a network component performing role-addressable gating, similar to that in working memory of humans. We observe that the model readily learns a gating policy, and upon training more models find that task performance is correlated with gating ability. Our results show how competence at cognitive branching tasks can emerge in Transformers, and suggest that integrating Transformer components may improve existing computational neuroscience models of working memory.

The Transformer models which were trained on the textual reference-back task are capable of making use of key composition for input gating and query composition for output gating. We find that making precise corruptions to specific architectural elements of the network causes the model’s prediction to change from same to different or vice versa, indicating that those components are causally responsible for the gating mechanism. The architectural biases of attention within the vanilla Transformer model lend themselves well to representing role-addressable content, as the learnable nature of keys, queries, and values allows the model to learn to create internal representations in a manner which allows it to represent roles and addresses, mimicking the variable binding and input / output gating mechanisms in biological neural networks O’Reilly \BBA Frank (); Frank \BBA Badre (); Collins \BBA Frank (); Kriete \BOthers. ().

In this work, we focused on characterizing the mechanism that the Transformer model learns as a result of training on cognitive branching tasks, and did not evaluate the robustness of the mechanism. The textual reference-back task is similar in nature to the FFLM task Liu \BOthers. (), comprised of 1 register, 2 symbols (0 or 1), and long sequences. found that similar Transformer models were able to succeed on FFLM data, but struggled to generalize outside of their training distribution. We leave experiments characterizing generalization of Transformer mechanism behavior on data with the reference-back paradigm– e.g. more than 2 registers, distinct sets of symbols, novel symbols introduced at test time– for future work.

When we trained more models on the task, we found that the models which perform best on the task correlate with the markers of gating we observed in our circuit analysis, and that the learning trajectory shows a steep decrease in training loss and a steep rise in patched subtask accuracy simultaneously, suggesting that the model has learned a component of the gating policy at that time. Both findings are analogous to those of \citeAfrank2012mechanisms, in which they find that networks which learned a hierarchical gating policy performed better at a hierarchical learning task, and humans that learn this policy also show a sharp decrease in loss when they discover it.

Ultimately, finding connections between emergent behavior of Transformer models and human working memory serves to benefit both computational cognitive neuroscience and artificial intelligence. Although Transformer models themselves are limited in their biological plausibility, in this setting they learned behavior mimicking the functionality of working memory, and their application within computational models of the brain should be further explored. From the perspective of artificial intelligence, understanding the strengths and limitations of Transformer models on cognitive branching tasks may inform model analysis across the many diverse settings in which these models are applied.

## Peer Reviews

confidence: 4
contribution: 2
initial_rating: 5
presentation: 3
questions: 1.	Given that the layer 2 head primarily computes functions linearly (Elhage et al 2021), has the author considered potential failure modes? For example, layer 2 head attends more to the Store tuples than Ignore tuples (based on layer 1 head's output?), and attend more to the most recent Store tuples than more distant Store tuples (based on position embedding?). Consider the scenario of a Store tuple followed by many Ignore tuples—could a recent Ignore tuple receive higher attention than distant Store tuples if position embedding outweighs the Store feature (due to linearity in key-query operation)? Such a failure mode would be significant, as RNNs, unlike Transformers, are less susceptible to this (assumed learned attractor dynamics within the hidden state representation).

2.	Could the authors consider using mechanistic interpretability tools, such as QK and OV circuit analysis, or representational geometry mapping of key, query, and value spaces? These analyses help further clarify how the mechanisms are implemented within each head and layer.

3. For other questions please see weaknesses.
soundness: 3
strengths: Originality: The study’s focus on emergent mechanisms in Transformer models trained on working memory tasks is novel. The observation that Transformers struggle with the three-register task is surprising. This research adds a unique perspective to both AI and cognitive science by applying Transformers to a working memory framework typically explored in neuroscience.

Quality: The study employs a simplified, attention-only Transformer, enhancing interpretability and enabling the identification of specific mechanisms within the model. The identification of Transformer mechanisms sheds light on how these models might solve working memory tasks and how specific task demands can trigger such emergent functionality. Control tasks are designed to isolate the conditions necessary for mechanistic emergence, strengthening the quality of the findings.

Clarity: The analysis and results are clearly presented, making it accessible for broad readers.

Significance: Understanding how Transformers solve tasks requiring working memory could substantially impact Transformers interpretability and development. By drawing parallels between Transformer mechanisms and human frontostriatal gating, this paper contributes valuable insights for interdisciplinary research. The study’s findings bridge AI with cognitive neuroscience, offering a model that informs both fields and fosters further exploration into the intersection of biological and artificial intelligent systems.
summary: This paper explores how Transformers, when trained on human working memory tasks, develop mechanisms that relate to frontostriatal gating operations observed in human cognition. The study focuses on understanding whether Transformers can solve working memory tasks using self-attention to mimic input and output gating, which are key to human working memory function. By training a small, attention-only Transformer model on tasks requiring selective memory gating, the authors find that certain task conditions lead to the emergence of gating-like behaviors in the model. The results suggest that these emergent mechanisms enhance the model’s generalization and task performance, potentially bridging cognitive neuroscience and artificial intelligence.
weaknesses: •	While the use of a small Transformer model with few heads allows for easier interpretability, the study does not fully analyze or demonstrate the distinct functions of each attention head in each layer. For instance, why exactly are two heads required per layer, and what distinct functions do each head serve? More detailed analyses of each head’s role (e.g., see below) would enhance the interpretability of the model’s mechanisms, and thus the impact of the work.

•	A comparative diagram for understanding “working memory” mechanisms between Transformers and RNNs—often used in biological working memory research—would provide useful context for readers.

•	In Figure 2, while path patching examples are shown, it is implied rather than explicated stated which layer and head are depicted, making it somewhat unclear which specific head is involved.

•	For brevity, I use x^0_i, x^1_i, x^2_i to specify the residual stream at tuple i before the first layer's output, after the first layer's output, and after the second layer's output, respectively. The layer 1 head appears to bind the components (Ins_i, Reg_i, Sym_i) into a tuple and write it to the residual stream x^1_i of Sym_i. However, it is not fully explained if the layer 1 head writes the Store or Ignore Ins_i to different directions within the residual stream. If so, it may function as part of the "input gating" mechanism (suggesting that the Store direction in layer 1 head’s output, but not Ignore direction, is similar to a "working memory" space). Further, do the two heads function similarly or differently? 

•	The functions of the layer 2 head, especially how it compares Reg_i and Reg_j (based on Ins_j=Store), are not fully elucidated. A detailed explanation of how the layer 2 head compares the current Sym_i and a prior Sym_j to output same/different answers would clarify its role in “output gating”. Do the two heads function similarly or differently? The study could use mechanistic interpretability tools, such as analyzing QK and OV circuits, or examining the geometric alignment of vectors in key, query, and value subspaces to address these gaps.

•	While transformers can be trained on working memory tasks, they inherently differ from biological systems as they have access to all previous residual stream positions. The most suitable counterpart of “working memory” might be x^1_i, which layer 1 head writes into and layer 2 head reads from. However, x^1_i is different for each position i, different from a true hidden-state-like "working memory". This divergence raises significant concerns about labeling the mechanisms as "input gating" and "output gating" or referring to Transformer mechanisms as "working memory". Using different terms across the manuscript, or discussing the limitations of these analogies would make the paper’s framing more precise.

•	The introduction lacks a review of previous work on comparing Transformers to brain/neuroscience (to name a few, https://arxiv.org/abs/2112.04035, https://arxiv.org/abs/2405.14992), which could more accurately depict the gap and contextualize the paper’s contributions.

---

confidence: 3
contribution: 2
initial_rating: 3
presentation: 2
questions: Major question:
- Do the authors see their key contribution in understanding transformers of making a strong point that transformers work like the brain? If it is the latter, I would expect a more detailed discussion of neuroscientific theories, at least a comparison of Transformers with models which are more typically considered working memory models with hidden states, and ideally some direct comparison with data from cognitive neuroscience. If it is the former, then I suggest the neuroscience link should be heavily deemphasized to not be misleading about the similarity to the brain. At the same time, I think in that case a reviewer who is well-versed in the MechInterp field should be included in the decision around acceptance. At the very least, the neuroscience content should largely be replaced by an actual overview of related MechInterp papers.

Minor questions:

- Can you add more infos on the actual setup of the data, for example how many timesteps does one trial have? I do not find that information.
- Is there any control to make sure the trials in the validation set are sufficiently different from the training set? Given you use a finite set of symbols with no noise, I wonder whether there are trials identical across datasets? The probability of this is hard to judge given the information about the length of trials seems to be missing. I am sorry if it is there and I just cannot locate it.
soundness: 1
strengths: I personally enjoyed the authors investigations on identifying how task-dependent gate mechanisms develop through training and how it can be identified through behavioural investigations using thoughtfully constructed task variants. I do not necessarily find their findings surprising, but I acknowledge that having evidence for expected mechanisms is worthwhile scientific work and hence would not let that influence my rating.
summary: In their paper “Transformer mechanisms mimic frontostriatal gating operations when trained on human working memory tasks” the authors train simple transformer-based networks a task used in cognitive neuroscience to study working memory. By using behavioural analysis in combination with some derivative tasks and the path-patching technique from MechInterp they show that transformers can learn to solve the task with Gating operations. They draw conceptual conclusions to human neuroscience where Frontostriatal circuits are believed to implement gating operations as part of working memory as well.
weaknesses: Unfortunately, I do see quite major weaknesses in this paper in its current state. In the following I group my concerns into a block of neuroscience-related concerns and Mechanistic Interpretability concerns.

*Neuroscience-related concerns*

The authors spend quite a bit of their manuscript making conceptual links to the brain’s gating mechanism in working memory but I struggle to see the relevance of their investigations to the brain’s working memory system, for the following reasons:

- Using transformers for working memory investigations: The authors declare themselves that a weakness of their investigations is that transformers have access to the entire input sequence, unlike the brain. In fact, the very idea of working memory is that the immediate past needs to continually compressed into a latent state of activations or rapid connectivity changes (e.g. see Stokes 2015 TICS). I wonder why the authors would not opt to use Mamba-like State Space Models which would seem to be much closer to brain like processes while also allowing for input dependent processing, and potentially slots through the distinct hidden states in Mamba. I see that models need to be abstractions and that one can in principle study working memory with transformers, but I am not aware of any prior investigations showing such links and hence the study here would need to provide such links themselves, which brings me to my next point:
- Comparison to the brain is purely conceptual: Given that the authors focus on the brain so much in both their title, abstract, and introduction I would have expected to see some actual data comparisons to cognitive neuroscience data, but the authors do not seem to provide any. The only data-like comparison they have is that the model struggles more to learn the task with more working memory items which supposedly is similar to humans but that simply seems like a general task difficulty effect and does not link to gating specifically. 
- Poorly referenced neuroscience work: The authors heavily rely on the idea of ‘Stripes’ in frontal cortex for working memory for their model to be relevant, as it is about recalling content from distinct memory slots. Recent neuroscientific investigations call the idea of discrete slots implemented through distinct groups of cells into question and instead think of working memory as being implemented by dynamic and distributed population codes (Meyers 2018 JNeurophys; Miller 2013 Dialogues in ClinNeuro). In case that authors make to want a serious link to neuroscience I think they should discuss how their model could be reconciled with the dominant idea of population codes. Also, the 3-4 chunks given as working memory capacity is on the lower end on the scale though I acknowledge that there is some discussion around it. Of course, the classical number to use is 7 +/- 2 (as discussed in the reference the authors give).

*MechInterp-related concerns*

So, the above points are to be taken seriously to make a believable link to neuroscience. Of course, the neuroscience-link is only conceptual, and the main work of the paper actually is in the MechInterp world showing how gating mechanisms develop in trained transformers. As said in the strengths section, that seems to be an interesting analysis though I am not really well qualified to judge how new that finding is. It seems intuitive to me that this should happen and, as the authors mention themselves, models like LSTM were actually constructed specifically with such ideas in mind. If the authors see their main contribution in studying how a gating mechanism develops in transformers, then I would suggest to strongly deemphasize the neuroscience narrative and instead contextualise the research more in the context of existing MechInterp findings. For example, work like [1] from earlier in the year looks at reasoning mechanisms which at least partially seem to rely on mechanisms similar to the ones proposed in this paper and I assume there is additional other work which I am not aware of.

[1] Brinkmann et al, 2024 https://arxiv.org/abs/2402.11917

---

confidence: 4
contribution: 1
initial_rating: 5
presentation: 1
questions: My questions revolve around addressing the 3 primary weakness I mention above.

1. Besides neuroscience, how does this current work relate to past and ongoing work in mechanistic interpretability? E.g., Meng et al., 2023a, 2023b; Olsson et al., 2022; Li et al., 2023. More specifically, how does the path patching approach used here compare to methods and insights in those papers?

2. Is it possible to include additional figures/visualizations to improve understanding of the paper? More specifically, is it possible to create a clearer visual explanation of the reference-back-2 task or a diagram showing how the query and key vectors function as gating mechanisms? Since the paper aims to interpret mechanisms of transformers, the contribution of the paper hinges on its clarity and presentation.

3. What is the relation of the final result (Fig 5) to curriculum learning?

Minor question: What do the authors hypothesize would happen for multiheaded transformers? What does each attention head do?
soundness: 3
strengths: Altogether, the results form an intuitive and helpful story to improve the understanding of the transformer’s self-attention mechanism. The combination of the simplicity of the task, and the path-patching approach on the attention-only transformer are instructive tools for understanding.
summary: This manuscript aims to understand the conditions by which interpretable gating mechanisms emerge in self-attention using attention-only transformers. The authors train simple, attention-only transformers on a simple reference-back-2 task. Using path-patching approach to interpret transformer mechanisms, the authors characterize what the query and key vectors do/implement during the reference-back-2 task and related control tasks. Overall, they find that the query vector accesses information from specific tokens within the context window (output gating), whereas the key vector provides the address to store an item to gate information (input gating).
weaknesses: Much of the paper focuses on trying to understand the mechanisms of self-attention. Besides a general concern that I find the results somewhat unsurprising, I have three other weaknesses I’d like to highlight.

1. Novelty. There are a number of papers in mechanistic interpretability that aim to study how transformer mechanisms (including self-attention) influence output/performance via path-patching, e.g., Meng et al., 2023a, 2023b; Olsson et al., 2022, Li et al., 2023 etc. However, most the background in the paper references neuroscience theories (which is fine, but perhaps misleading as to existing work in mechanistic interpretability). A more thorough contextualization of mechanistic interpretability papers would be beneficial.

2. Though the paper is geared towards understanding self-attention mechanism through experimentation, I found the manuscript to be difficult to understand. For example, I found the presentation of the task (both text description and figure 1) to be quite confusing. Only when I reviewed the text and figure from another paper (Rac-Lubashevsky & Frank, 2021) did I understand the task. It shouldn’t be necessary to review a prior paper to understand the experiment. I found figure 2 to be also quite dense and confusing. (I have previously found the figure presentation in Meng et al., 2022, Fig 1 to be quite instructive.) Finally, many of the analyses (e.g., Queries gate outputs; keys gate inputs) are only reported in text. Would it be possible to include a visual figure/understanding that is intuitive as to what the query and key vectors are actually computing? 

3. The pretraining result is generally unsurprising (fig 5), and is consistent with prior results in curriculum learning, which is not mentioned. 
This is perhaps more of a comment than an explicit suggestion, and something for the authors to consider: In some ways, I think contextualizing the work using the background of ‘frontostriatal gating’ mechanisms may be confusing to some readers. I understand that trying to demonstrate this biological/neural mechanism in transformers may have been the motivation to the authors, yet to a reader who is unfamiliar with neuroscience, this may not add much. The explanation of gating mechanisms via path-patching should in theory computationally sufficient, and doesn’t require reference to neuroscience (which can often be obscure).

Meng, Kevin, David Bau, Alex Andonian, and Yonatan Belinkov. “Locating and Editing Factual Associations in GPT.” arXiv, January 13, 2023. http://arxiv.org/abs/2202.05262.

Meng, Kevin, Arnab Sen Sharma, Alex Andonian, Yonatan Belinkov, and David Bau. “Mass-Editing Memory in a Transformer.” arXiv, August 1, 2023. http://arxiv.org/abs/2210.07229.

Olsson, Catherine, Nelson Elhage, Neel Nanda, Nicholas Joseph, Nova DasSarma, Tom Henighan, Ben Mann, et al. “In-Context Learning and Induction Heads.” arXiv, September 23, 2022. https://doi.org/10.48550/arXiv.2209.11895.

Li, Kenneth, Aspen K. Hopkins, David Bau, Fernanda Viégas, Hanspeter Pfister, and Martin Wattenberg. “Emergent World Representations: Exploring a Sequence Model Trained on a Synthetic Task.” arXiv, February 27, 2023. https://doi.org/10.48550/arXiv.2210.13382.

Meta-review: This work shows that Transformers, when trained on tasks requiring working memory gating, develop input and output gating mechanisms that resemble those in human frontostriatal systems.

The paper is generally well-written and aims to bridge the gap between gating mechanisms in AI Transformers and frontostriatal systems in biological brains. The approach provides valuable insights into the interpretability of the self-attention mechanism in Transformers.

However, all three reviewers rated the paper below the acceptance threshold due to several unresolved issues, even after the rebuttal:

Lack of comparisons with other AI architectures, such as Memba and RNNs, and insufficient justification for selecting Transformers as the focus.
Insufficient references and comparisons to existing works linking Transformers to neuroscience.
Absence of behavioral or neuronal-level comparisons with neuroscience data from biological brains.

The authors are encouraged to address these concerns and include the suggestions in future submissions.

Decision: Reject