SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Sivaramakrishnan Swaminathan Antoine Dedieu Rajkumar Vasudeva Raju
Murray Shanahan Miguel Lázaro-Gredilla Dileep George
{sivark,aderieu,rajvraju,mshanahan,lazarogredilla,dileepgeorge}@google.com
In-context learning (ICL) is one of the most powerful and most unexpected capabilities to emerge in recent transformer-based large language models (LLMs). Yet the mechanisms that underlie it are poorly understood. In this paper, we demonstrate that comparable ICL capabilities can be acquired by an alternative sequence prediction learning method, namely clone-structured causal graphs (CSCGs). A key property of CSCGs is that, unlike transformer-based LLMs, they are interpretable, which considerably simplifies the task of explaining how ICL works. We show that ICL in CSCG uses a combination of (a) learning template (schema) circuits for pattern completion, (b) retrieving relevant templates in a context-sensitive manner, and (c) rebinding novel tokens to appropriate slots in the templates. We go on to marshal evidence for the hypothesis that similar mechanisms underlie ICL in LLMs. For example, we find that, with CSCGs as with LLMs, different capabilities emerge at different levels of overparameterization, suggesting that overparameterization helps in learning more complex template (schema) circuits. By showing how ICL can be achieved with small models and datasets, we open up a path to novel architectures, and take a vital step towards a more general understanding of the mechanics behind this important capability.
In a pre-trained sequence model, in-context learning (ICL), or few-shot prompting, is the ability to learn a new task from a small set of examples presented within the context (the prompt) at inference time. Surprisingly, large language models (LLMs) trained on sufficient data exhibit ICL, even though they are trained only with the objective of next token prediction [1, 2]. A good deal of the ongoing excitement surrounding LLMs arises from this unexpected capacity, since it dramatically enlarges their set of potential applications. Ongoing attempts to understand this capability take a variety of forms, including higher-level normative accounts using Bayesian inference [3], and mechanistic explanations involving implicit gradient descent [4] or induction heads [5]. Despite this, the mechanisms that underlie ICL in LLMs remain somewhat mysterious.
We take an alternative approach, studying a sequence-learning model called a clone-structured causal graph (CSCG) [6, 7] to reveal the conditions that drive ICL. We show that ICL can be explained as a combination of (a) learning template circuits for pattern completion, (b) retrieving relevant templates in a context-sensitive manner, and (c) rebinding of novel tokens to appropriate slots in templates [8]. Unlike n-gram models, CSCGs allow transitive generalization in the latent space: they assign semantically sensible non-zero probabilities to sequences never seen during training to ensure that the contexts (prompts) used for retrieval are not pure memorizations. In addition, the binding of novel tokens to slots in learned templates allows the same structural knowledge to be applied to entirely novel inputs. We hypothesize how similar mechanisms could exist in transformer-based LLMs. By
elucidating the principles that underpin the mechanics of ICL, we hope to pave the way for the design of novel architectures for abstraction and generalization, while the building blocks we identify guide the search for mechanistically interpretable [9] and editable [10] circuits in transformers [11].
2 Rebinding algorithm for clone-structured causal graphs
2.1 Background on clone-structured causal graphs (CSCGs)
Consider an agent executing a series of discrete actions  $a_1, \ldots, a_{N-1}$  with  $a_n \in \{1, \ldots, N_{\text{actions}}\}$ , e.g. walking in a room. As a result of each action, the agent receives a perceptually aliased observation [12], resulting in the stream of random variables  $X_1, \ldots, X_N$  with observed values  $x_1, \ldots, x_N$ , where each  $x_n \in \{1, \ldots, N_{\text{obs}}\}$ . CSCG [6] is a probabilistic sequence learning model that introduces a latent explanatory variable  $Z_n$  at each timestep  $n$ , with values  $z_n \in \{1, \ldots, N_{\text{latent}}\}$ , to model the action-conditional stream of observations as
$$
P (x _ {1}, \ldots , x _ {N} | a _ {1}, \ldots , a _ {N - 1}) = \sum_ {z _ {1}, \ldots , z _ {n}} P (x _ {1} | z _ {1}) P (z _ {1}) \prod_ {n = 2} ^ {N} P (x _ {n} | z _ {n}) P (z _ {n} | z _ {n - 1}, a _ {n - 1}).
$$
A transition tensor  $T: T_{ijk} = P(Z_n = k|Z_{n-1} = j, a_{n-1} = i) \forall n$  represents the action-conditional dynamics.  $T$  defines a directed multigraph, whose nodes correspond to the values of  $z$ . Conditioned on an action, each entry of  $T$  is the weight of a directed edge between two nodes (from the row index to the column index of that entry). A CSCG can thus recover a graph that represents the latent causal structure [13] of the environment (see Fig. 1D for an example) which can then be used for planning.
An emission matrix  $E: E_{ij} = P(X_n = j|Z_n = i) \forall n$  represents the observation probabilities. CSCGs have a deterministic observation model: for any latent value  $z$ , the same observation  $x$  is always emitted. Multiple values of  $z$  can result in the same observed  $x$ , making the model overcomplete [14]. The restriction to deterministic  $E$  makes CSCGs less general than a hidden Markov model (HMM), but easier to learn [6]. A CSCG can disambiguate multiple aliased percepts (same observation  $x$ ) into distinct causes (different latent values  $z$ ) given a sufficiently long context. If the observations correspond to word tokens, CSCGs can also be used as a language model, with a single action that accesses the next token<sup>1</sup>.
2.2 Rebinding in CSCGs
On encountering a new environment with a similar structure but different observations, an agent can learn that environment faster by reusing the latent graph  $T$  from prior experience and relearning just the emission matrix, through a rebinding process [15]. Rebinding can be interpreted as a soft intervention on the agent's model [16, 17]. See Fig. 1D & F for examples of two rooms that share the same latent structure but different observations. When a new emission matrix binds to an existing schema, it has to respect the clone structure of the original emission matrix (Fig. 1E). The clone structure function  $\mathcal{C}(\cdot) \in 1, \ldots, N_{\mathrm{obs}}$  partitions the latent state in  $N_{\mathrm{obs}}$  slots: two latent states  $z = i$  and  $z = i'$  belong to the same slot iff  $\mathcal{C}(i) = \mathcal{C}(i')$ . An emission matrix respects the clone structure  $\mathcal{C}$  if  $\mathcal{C}(i) = \mathcal{C}(i') \Rightarrow E_{ij} = E_{i'j} \forall i, i', j$ . The 3-tuple  $\{T, \mathcal{C}, E\}$  defines a grounded schema, the tuple  $\{T, \mathcal{C}\}$  defines an ungrounded schema with clone structure and  $T$  alone is a schema [15].
2.2.1 Fast rebinding by attending to surprise
Often, environment changes are localized such that most of the latent structure and observation mapping is preserved while just a few observations need to be rebound: for example, just replacing the carpet in a room while the wall colors remain the same, or getting exposed to a new word in a familiar context. This insight can be utilized to derive an algorithm that focuses the update of the emission matrix only to those observations that were found surprising by the existing model.
Suppose that at test time, a grounded schema  $\{T,\mathcal{C},E^0\}$  is exposed to a sequence with novel observations. Algorithm 1 proposes a fast procedure to update the emission matrix to the new observations by only performing local updates, and to bind the updated emission matrix to the existing schema  $T$ , defining a new grounded schema  $\{T,\mathcal{C},E^{\mathrm{rb}}\}$ . We call this process fast rebinding.
<sup>1</sup>This can be generalized to include actions that skip over one or more next tokens.
Figure 1: A. Inducing room structure (cognitive maps) from sequential sensory observations is challenging due to perceptual aliasing – local observations do not identify locations uniquely. B. Cloned hidden Markov models (HMMs) [7]. Each observation is mapped to multiple clone states in the latent space. C. Graphical model for CSCGs [6], extending cloned HMMs by incorporating actions. CSCGs utilize the latent space to overcome the perceptual aliasing problem. Different clones learn to represent different temporal contexts to recover the latent structure of the room. D. Learned CSCG for the room shown in panel A consists of a latent transition matrix and an emission matrix. We visualize the model in two ways: (i) stacking clone states for respective observations into columns, and (ii) clones as nodes in a transition graph, colored with their respective emissions. E. The emission matrix imposes a slot structure – nodes within the same slot are constrained to bind to the same observation. A new environment with the same latent structure but different observation mapping (Room 2) can be learned quickly by freezing the transition graph and slot structure, and learning a new emission matrix by rebinding slots to a new set of observations. F. CSCG for a different room learned purely through rebinding.
Given a prompt  $(x_{1},\ldots ,x_{N})$  and a surprise threshold, Algorithm 1 proceeds by (a) identifying emission matrix entries that need updating then (b) updating these entries using the Expectation-Maximization (EM) algorithm [18]. The conditional probability  $P(X_{n} = j\mid x_{\backslash n})$  of tokens at timestep  $n$  given all other tokens is used to identify timesteps and latent states that are surprising. Step 3 identifies anchors, i.e., latent states corresponding to observations that are correctly predicted with high confidence: anchors are not rebound. Step 4 identifies candidates for rebinding as latent states (a) not among the anchor states and (b) corresponding to timesteps at which observations are incorrectly predicted with high confidence. Finally, instead of re-learning the whole emission matrix [15, Appendix A.2], Step 5 (detailed in Appendix A) locally updates the emission matrix by only applying EM on the latent states and timesteps identified in Step 4. As a result, only a small subset of rows differ between  $E^0$  and  $E^{\mathrm{rb}}$ . Protected rows correspond to either (a) anchors in the current prompt or (b) slots not relevant to the current prompt but possibly relevant to future observations. Section 6 discusses how a similar mechanism could be implemented in transformers.
Algorithm 1 - Fast rebinding algorithm
Input: Grounded schema  $\{T, \mathcal{C}, E^0\}$ , pseudocount  $\epsilon$ , prompt  $(x_1, \ldots, x_N)$ , surprise probability threshold  $p_{\text{surprise}}$ .
Output: Rebound emission matrix  $E^{\mathrm{rb}}$
1: Define  $\tilde{E}^0\propto E^0 +\epsilon$  , with normalized rows.
2: For timestep  $n$ , use the emission matrix  $\tilde{E}^0$  to compute
$$
P \left(X _ {n} = j \mid x _ {\backslash n}\right) = P \left(X _ {n} = j \mid x _ {1}, \dots , x _ {n - 1}, x _ {n + 1}, \dots , x _ {N}\right), \forall j \leq N _ {\mathrm {o b s}}
$$
3: Identify latent states and timesteps that can act as anchors:
$$
\mathcal {A} = \left\{\left(i, n\right) \mid P \left(X _ {n} = x _ {n} \mid x _ {\backslash n}\right) > p _ {\text {s u r p r i s e}}, \text {a n d} \mathcal {C} (i) = x _ {n} \right\}
$$
4: Identify latent states to be rebound (and their timesteps):
$$
\mathcal {R} = \left\{(i, n) \mid P \left(X _ {n} = j \mid x _ {\backslash n}\right) > p _ {\text {s u r p r i s e}}, j \neq x _ {n}, (\cdot , n) \notin \mathcal {A}, (i, \cdot) \notin \mathcal {A} \text {a n d} \mathcal {C} (i) = j \right\}
$$
5: Fix  $T$ , and use EM to update the emission matrix (initialized with  $E^0$ , and without using any pseudocount) by only using the beliefs for latent states  $i$  and timesteps  $n$  such that  $(i, n) \in \mathcal{R}$ .
Figure 2: A. CSCGs allow both separation of contexts and transitive generalization. The word "bank" is wired to multiple clones corresponding to the different contexts it is used in. If "milk and honey", and "bread and butter" are seen in training, transitive generalization occurs if they get wired through the same "and" clone: "bread and honey" and "milk and butter" appear as valid sequences. B. Probabilistic branching & merging of sequences. C - F. Exemplar CSCG circuits for copying a sequence, parity operation, reversing a list with exactly five elements, reversing lists with a variable number of elements. G. Rebinding to new observations: dashed gray arrows correspond to old emissions while green arrows correspond to newly rebound emissions.
After rebinding, we complete the prompt by performing MAP inference conditioned on the provided prompt in the rebound CSCG. We run the max-product algorithm [19] forward (the backward messages are all uniform) thus generating a series of MAP observations for the tokens following the prompt. We stop once we generate a delimiter token. See Algorithm 2 in Appendix B for details.
3 Outline of the overall argument using CSCG
3.1 Context-dependent latent representations and transitive generalization
The clone structure of CSCGs allows context-based separation and appropriate blending for language modeling. For example, the sense of the word "bank" in "bank robber" is different from the one in "river bank". CSCG learning disambiguates these contexts in the latent space by wiring them to different clones to improve predictive accuracy. In Fig. 2A, the sentences "river bank resort", and "one bank robber" use different clones of "bank". Sequences can have probabilistic branching: "one bank robber" can terminate at "\n", or continue to "eating at river bank resort" or "eating bread and honey", or "eating bread and butter at river bank resort" (Fig. 2B). CSCGs also allow the merging of contexts that result in transitive generalization: even if training data has only the sequences "bread and butter", and "milk and honey", if they go through the same clone state "and", the model will generalize to "bread and honey" and "milk and butter", assigning non-zero probability to those sequences. Due to the combination of context-sensitive separation and transitivity, related topics, concepts, and algorithms get clustered into sub-networks that pass through the same clones. A prompt's context would activate its sub-network, and transitive generalization allows for prompts that are not exact memorizations. As we show in Section 4, the Bayesian inference perspective on ICL [3] corresponds to this context-sensitive and transitively generalizing storage and retrieval alone, and is insufficient to explain the ICL properties we consider in the next sections.
3.2 Learning flexible schemas (template circuits) and rebinding
Just like learning room layouts, CSCG can learn automata circuits [20] for sequence-to-sequence (seq2seq) algorithms. See Fig. 2 for CSCG circuits for computing parity, copying a sequence, and reversing sequences of multiple lengths. The list reversal circuit in Fig. 2E is bound to the specific symbols  $A, B, C, D, E$  used in training. For use as a template, slots in this graph must be able to appropriately bind to contents (arbitrary symbols) that occur in context at test time [8, 21]. The rebinding mechanism (formalized in Algorithm 1) can intuitively be understood as operating based on prediction errors – when the latent context strongly predicts the latent state corresponding to a time instant, but the actual observation is mismatched, rebinding adjusts the emission matrix to wire all the clones of that latent state to the surprising observation. Such a mechanism to mix and gate previous knowledge with new content allows circuits learned during training to become flexible templates
with slots that can dynamically bind to new inputs as required. For example, in the list reversal schema in Fig. 2F, tokens “[”, and “[” are prior contents that detect the beginning and end of the list – these act as anchors for grounding the schema in the observations. Probabilistic branching based on the end of list token “[” allows for length generalization, whereas absorbing arbitrary symbols into the slots corresponding to  $A, B, C, D, E$  allows the algorithm to generalize to new symbols. Fig. 2G illustrates the outcome of this rebinding mechanism where the slots emitting  $A, B, C, D, E$  are respectively rebound to symbols  $K, M, N, P, R$  from the input prompt. Similarly, in the sentence “I wrote in a notebook using a dash”, rebinding can absorb the new token “dax” into the context by binding it to a clone corresponding to “pencil” or “pen”, and use the new word in those contexts.
3.3 Instruction-based or content-based retrieval and completion of tasks
Zero-shot task recognition as content-based retrieval using rebinding: Many striking examples of zero-shot learning involve recognizing tasks from prompts, and repeating them on new inputs. For example, given a prompt "Input: [p, q, r, s] Output: [p, p, q, q, r, r, s, s]; Input: [l, m, n, o] Output: [l, l, m, m, n, n, o, o]" LLMs can infer the task as repeating the elements of the sequence, and apply that to complete the output for a new input prompt even when the tokens "p, q, r, s, l, m, n, o" were not seen during training, in association with this task. The rebinding mechanism offers a natural explanation for this. Given the prompt, expectation maximization (EM) [18] simultaneously evaluates the different rebindings to multiple latent algorithm schemas to infer the best binding, which is then applied to complete the query prompt.
Instruction-based retrieval: When algorithms are trained with prefixed language instructions, CSCGs learn instruction sub-networks that directly point to the circuits that represent the algorithms (see Section 4.2). The algorithm can be retrieved by direct prompting with language instructions that can be significantly different from training instructions due to transitive generalization and rebinding.
3.4 Emergence
We hypothesize and empirically demonstrate in Section 4, that emergence is explainable as the combined effects of the above properties (context-separation, transitive generalization, schema-formation, and rebinding), model capacity, and patterns in the data. Training on a bigger dataset results in the induction of more templates that might not have occurred in the smaller dataset. Learning the schematic circuits for more complex algorithms or more patterns in the data requires greater model capacity because overparameterization helps in the optimization process.
We substantiate the above argument using empirical results on three datasets: (a) the GINC benchmark introduced in [3], (b) a suite of algorithm learning tasks that we introduce in our LIALT datasets, and (c) a zero-shot word usage induction task on a CSCG language model.
4.1 Context-sensitive retrieval on GINC dataset matches Bayesian inference explanation
Dataset: The GINC dataset [3] introduced for studying ICL, is generated from a uniform mixture of five factorial HMMs [22]. Each factorial HMM is referred to as a concept. A document is created by concatenating independent sentence samples from a concept. The in-context test prompts have examples of lengths  $k \in \{3, 5, 8, 10\}$ , varying in number from  $n = 0$  to  $n = 64$ , with 2500 prompts for each setting  $(k, n)$ . Each prompt uniformly selects a concept, samples  $n - 1$  examples  $x_{:k}^{(1)}, \ldots, x_{:k}^{(n - 1)}$  of length  $k$ , and one example  $x_{:k - 1}^{(n)}$  of length  $k - 1$ . The in-context task is to infer the most likely last token of the last example, i.e.,  $\operatorname{argmax}_{x_{k - 1}^{(n)}} p\left(x_{k - 1}^{(n)} | x_{:k}^{(1)}, \ldots, x_{:k}^{(n - 1)}, x_{:k - 1}^{(n)}\right)$ . Since the vocabulary is shared among different latent concepts, observations in GINC are aliased like in natural language, and solving the task requires the model to disambiguate the aliased observations to correctly infer the latent concepts.
Training: We train a single CSCG with 50 clones on the GINC dataset for 100 full-batch EM iterations using a pseudocount [6] of  $\epsilon = 10^{-2}$ . Given a test prompt, CSCG infers the most likely hidden sequence for that prompt, then predicts the next most likely observation.
A Transition graph of the CSCG trained on the GINC dataset
C Latent state inference for different prompt lengths
B In-context accuracy on the GINC dataset
Figure 3: A. Visualizing the transition graph of a CSCG with 50 clones trained on the GINC dataset from [3]. The clones cluster into five groups - one per concept. B.[Left] In-context accuracy averaged over the GINC test dataset (with  $95\%$  confidence intervals (CIs) as in [3]), for the same model For contexts of 8 and 10 tokens, the model predicts the most likely next token at least  $95\%$  of the time-including in the zero-shot regime. [Right] In-context accuracy decreases when we reduce the number of clones to 10-for  $k \in \{8,10\}$  it drops from above  $95\%$  to below  $75\%$ . The numerical values are reported in Appendix C, Table 1. C.Decoded latent state distributions (increasing intensities of black for higher density) for the CSCG with 50 clones, for an  $n = 0$  &  $k = 10$  prompt "o y w r m ro y aj", when truncated to different lengths ( $k = 2,3,5,8,10$ ). Longer prompts improve latent state estimation--resulting in better concept retrieval, and next token prediction.
Results: CSCG learns different latent sub-networks corresponding to the five latent concepts in the GINC dataset (Fig. 3A), and inference on a supplied prompt retrieves the correct latent sub-network (Fig. 3C). Increasing the prompt length improves the localization of the sub-network and the particular states within the sub-network. Figure 3C visualizes the decoded latent state distribution for an example prompt in the zero-shot setting  $(n = 0)$ . The decoding starts out uncertain, and improves as the prompt gets longer. This localization (on the graph) results in effective schema retrieval, and hence accurate prompt completion. Figure 3B[left] reports the in-context accuracy—defined as the average ratio of correct predictions—for each  $(k,n)$  pair of the GINC test set. CSCG in-context accuracy matches the patterns exhibited by LSTMs and transformers in [3], while slightly improving their performance. Fig. 3B also shows that a CSCG with larger capacity, i.e. with 50 clones per token, better separates the latent concepts and significantly outperforms a CSCG with only 10 clones per token. Fig. 9[left] in Appendix C displays the CSCG in-context confidence: for larger contexts, CSCG is better at disambiguating aliasing and the averaged predictions probabilities are higher. Finally, Fig. 9[right] shows that similarly to the transformer and LSTM in [3], CSCG fails at ICL when test prompts are sampled from concepts unseen during training. The GINC results match the context-based retrieval argument in Section 3.1: ICL in this setting is the retrieval of a shared latent concept between the prompt and the model. By using the long-range coherence of concepts in the training documents, the model learns to separate concepts into different latent representations. Despite the train and prompt distribution mismatch [3], CSCG succeeds at prompt completion because the representation allows transitive mixing.
4.2 Learning schemas for seq2seq algorithms and generalization using rebinding