Schema-learning and rebinding as mechanisms of in-context learning and emergence
Authors: Sivaramakrishnan Swaminathan
Antoine Dedieu
Rajkumar Vasudeva Raju
Murray Shanahan
Miguel Lazaro-Gredilla
Dileep George
Venue: NeurIPS 2023 spotlight (2023)
Keywords: mechanistic interpretability
in-context learning
emergence
large language models
TL;DR: Interpretable mechanistic model explaining emergent behavior and in-context learning in large language models.

Abstract: 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 using clone-structured causal graphs (CSCGs). Moreover, a key property of CSCGs is that, unlike transformer-based LLMs, they are {\em interpretable}, which considerably simplifies the task of explaining how ICL works. Specifically, we show that it uses a combination of (a) learning template (schema) circuits for pattern completion, (b) retrieving relevant templates in a context-sensitive manner, and (c) rebinding of novel tokens to appropriate slots in the templates. We go on to marshall 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.

Schema-learning and rebinding as mechanisms of in-context learning and emergence
Sivaramakrishnan Swaminathan
Antoine Dedieu Rajkumar Vasudeva Raju Dileep George
Murray Shanahan Miguel Lázaro-Gredilla
Google DeepMind
{sivark, adedieu, rajvraju, mshanahan, hag.aaaaaaaaaaaaaaaaaaaaaggrreea aeamaaann.aarrgeeelleeeeeeeeeego
Abstract
In-context learning (ICL) is one of the most powerful and most unexpected capabil- ities to emerge in recent transformer based large language models (LLMs). Yet the mechanisms that underlie it poorly understood. In this paper, demonstrate that comparable ICL capabilities can acquired by an alternative sequence pre- diction learning method, namely clone-structured causal graphs (CSCGs). A key property of CSCGs that, unlike transformer- based I.Mss, they are interpretable, which considerably simplifies the task of explaining how ICL works. We show that ICI CSCG uses a combination of (a) learning template (schema) circuits for pattern completion, (b) retrieving relevant templates
sentsxt-sensitive manner, and (c) rebinding novel tokens to appropriate slots in the templates. We go on marshall evidence for the hypothesis that similar mechanisms underlie ICL in LLMs. For example, we find that, with CSCGs as with LLMs, different capabilities emerge different levels of overparameterization, suggesting that overparameter- ization helps in learning more complex template (schema) circuits. By showing how ICL can achieved with small models and dltasstts, we open up a path novel architectures, and take vital step towards a more general understanding of the mechanics behind this important capahility.
Introduction
In a pre-trained sequence model, in-context learning (ICL), or few-shot prompting, is the ability to learn new task from 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 21 A good deal of the ongoing excitement surrounding L.LMs arises from this unexpected capacity, since 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 induction heads (5). Despite this, the mechanisms that underlie ICL in LLMs remain somewhat mysterious.
We take an alternative approach, studying
gruuence-learingg model called clone-structures causal graph (CSCG) 66 [ reveal the conditions that drive ICL. We show that ICL can be explained as combination of (a) learning template circuits for pattern completion, (b) retrieving relevant templates a context-sensitive manner, and (c) rebinding of novel tokens to appropriate slots templates 88. 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-basee LLMs. By
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

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 [T0] circuits in transformers
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 ....... aN-1 with On E E(1,...,,Naciioon J. e.g. walking in room. As a result of each action, the agent receives a perceptually aliased observation [12], resulting in the stream of random variables XI...., XN with observed values T1,.... NN where each 2n, E {l...., Nobss ). CSCG (6) is probabilistic sequence learning model that introduces a latent explanatory variable Zn at each timestep n, with values 2n E {I...., action-conditions stream observations as
( Namem J, to model the
P(Ii,,...,
IN-
Σ P(m))z)P((2) Rm
II P(rnlzn) (,, n=
A transition tensor T: Tyjk- -1 -, i) Vn represents the action-conditionaal dynamics. defines a directed multigraph, whose nodes correspond to the values of z. Conditioned on an action, each entry of I' 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 13 of the environment (see Fig. DD for example) which can then be used for planning.
structure
emission matrix E E,, PX,, zl i) Vn represents the observation probabilities. CSCGs have deterministic observation model: for any latent value 2, the same observation 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 hidden Markov model (HMM), but easier to learn (6]. A CSCG can disambiguate multiple aliased percepts (same observation æ) into distinct causes (different latent values 2) given sufficiently long context. the observations correspond to word tokens, CSCGs can also be used as a language model, with a single action that accesses the next token")
2.2 Rebinding in CSCGs
On encountering new environment with 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 rebinding process (15). Rebinding can be interpreted as a soft intervention on the agent's 's model [16] 17]. See Fig. IDP & 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. [E). The clone structure function C(-) E I..... Nobs partitions the latent state in Noss slots: two latent states 2 and i' belong to the same slot iff C(i) C(i'). An emission matrix respects the clone structure if C(i) = C(i') = Eij = Ezg V i i' j. The 3-tuple (T,C, C, E} defines a grounded schema, the tuple (T, C) defines ungrounded schema with clone structure and T alone is a schema 15
2.2.1 Fast rebinding by attending to surprise
Often, enviromment 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 room while the wall colors remain the same, or getting exposed 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.
new word in a
Suppose that at test time, a grounded schema {T,C, E is exposed to a sequence with novel observations. Algorithm proposes a fast procedure to update the emission matrix 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, C, Erb We call this process fast rebinding.
'This be generalized to include actions that skip over
next tokens.

---

B
C
(x)))) of tokens at Istent shattes
diserations actions
clones
observations
1.0
E
alt
space)
leame From data deservaliot
Room
obsenations
Room
different Blor mail colors
Figure I: 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) 77. Each observation is mapped 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 leamn 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 transition graph, colored with their respective emissions. F. 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 new emission matrix by rebinding slots new set of observations. F. CSCG for different room learned purely through rebinding.
Given a prompt (FI......N and 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, timestep given all other tokens used to identify timesteps and latent states that are surprising. Step 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 [5] Appendix A.2], Step (detailed in Appendix A locally updates the emission matrix by only applying EM on the latent states and timesteps identified in Step 4. As result, only small subset of rows differ between E and Erb, Protected rows correspond to either (a) anchors in the current prompt (b) slots not relevant to the current prompt but possibly relevant to future observations. Section 6 discusses how a similar mechanism could be implemented transformers.
Algorithm 1
Fast rebinding algorithm
Input: Grounded schema {T", C, E*), pseudocount e, prompt (II,...,,N) surprise probability threshold Psurprise-
Output: Rebound emission matrix E
1: Define E" x E + with normalized rows.
2: For timestep n, the emission matrix E to compute
=() P(X, = j æi,..., -, Ix 2n ,,y,, Vj s Nas
3: Identify latent states and timesteps that can act as anchors:
A = {(i.m) P(X,, = 2n 2..)) > Paurpise: and C(t) = a..
4: Identify latent states to be rebound (and their timesteps):
R = ((i,n) P(X, = j (..)) > Pourpise: j # zn. (sm) $ A, (i.-) & A and C()) =3
5: Fix T, and use EM to update the emission matrix (initialized with E", and without using any
pseudocount) by only using the beliefs for latent states and timesteps n such that (i, n) E R.
3

---

Ieporation ConTRCC
generalcation
Recbly branchin/mmeeg sequercec dooing
mar
C copy sequence
C
parity
reverse Sequence
revers lengths
reverse with ropid rebinding
Figure 2: 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 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 rehound 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 generate delimiter token. See Algorithm 2 in Appendix B for details.
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. DA 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 "cating at river bank resort" "cating bread and honey", or "caaiig 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 "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 ceeuuence----sqqeeee (seq2seq) algorithms. See Fig. 2 CSCG circuits for computing parity. copying sequence, and reversing sequences of multiple lengths. The list reversal circuit in Fig. 2 bound the specific symbols A, B. C, D. E used in training. For use as template, slots in this graph must be able to appropriately bind to contents (arbitrary symbols) that occur in context at test time (8] (T]. The rebinding mechanism (formalized in Algorithm [T] 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 the surprising observation. Such 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 imputs as required. For example, in the list reversal schema in 2F tokens "T", and "I" 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 "I" 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 dax", rebinding can absorb the new token "dax" into the context by binding it to clone corresponding to "pencil" "pen", and use the new word in those contexts.
3.3 Instruction-based or content-based 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. r, r, S, s}; Input: [l. m, n, o] Output: [I, I, m, m, n, n, o, ol" LLMs can infer the task repeating the elements of the sequence, and apply that to complete the output for a new input prompt even when the tokens "p. q. I, s, I, m, n, o" were not seen during training, in association with this task. The rebinding mechanism offers natural explanation for this. Given the prompt, expectation maximization (EM) (18) simultaneously evaluates the different rehindings multiple latent algorithm schemas to infer the best binding, which is then applied to complete the query prompt.
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.
Instruction- based retrieval:
3.4 Emergence
We hypothesize and empirically demonstrate in Section 4 that emergence is explainable as the combined effects of the above properties (context-separatioon transitive generalization, schema- formation, and rebinding), model capacity, and patterns in the data. Training on 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 more patterns in the data requires greater model capacity because overparameterization helps in the optimization process.
Results
We substantiate the above argument using empirical results on three datasets: (a) the GINC benchmark introduced in (3) (b) suite of algorithm learning tasks that we introduce in our LIALT datasets, and (c) 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 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 & E (3,5,8, 10). varying in number from n = 0 to n = 64, with 2500 prompts for each setting (ks, n). Each prompt uniformly selects concept, samples n 1 examples
(1) (n-1) of length k. and one example " of length k - 1. The in-context task is to infer 2 (n-1)
T. ................................ L.
the most likely last token of the last example, i.e., argmax = P (=
x
x-1)
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 single CSCG with 50 clones on the GINC dataset for 100 full-batch EM iterations using a pseudocount (6) of e = 10 2. Given a test prompt, CSCG infers the most likely hidden sequence for that prompt, then predicts the next most likely observation.
5

---

A Transition graph of the CSCG Irained on the GINC datasel
B
In-context accuracy on the GINC dataset
CSCG with 50 clones
CSCG with 10 clones
contoot: accuracy
0
No. of examples
60
examples
60
Latent state inference for different prompt lengths
ooy
y winroye a]
o yyvr y
oyvenre
oyvEnroys
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 33. for the same model For contexts of 8 and 10 tokens, the model predicts the most likely next token least 95% of the time-including the zero-shot regime. [Right] In-context accuracy decreases when we reduce the number of clones to for & E {8, 10) drops from above 95% to below 75%. The numerical values are reported in Appendix C Table | C. Decoded latent state distributions (increasing intensities of black for higher density) for the CSCG with 50 clones, for an 0 & k 10 prompt "o y prompts improve latent state estimation resulting in better concept retrieval, and next token prediction.
y aj". when truncated to different lengths (k- = 2, 3, 5, 8, 10). Longer ceemmatioo-rsultiingg
Results: CSCG learns different latent sub-networks corresponding to the five latent concepts in the GINC dataset Fig. AA, and inference on a supplied prompt retrieves the correct latent sub- network (Fig. BC). Increasing the prompt length improves the localization of the sub-network and the particular states within the sub-network. Figure BC 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 3 reports the in in-context accuracy- decfrncy-effined as the average ratio of correct predictions-fre for each (k:, n) pair of the GINC test set. CSCG in-context accuracy matches the patterns exhibited by LSTTSs and transformers in B while slightly improving their performance. Fig. 33 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. [eeft] in (erenddi([ddislaayy the CSCG in-context confidence: for larger contexts, CSCG is better at disambiguating aliasing and the averaged predictions probabilities are higher. Finally, Fig. 9 righht] 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 shared latent concept between the prompt and the model. By using the rangr 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
To test the ability of CSCG to learn algorithmss that generalize to novel inputs not seen during training, we construct the Language Instructed Algorithm Learning Tasks (LIALT) dataset. The LIALT training set contains demonstrations of 13 list and matrix algorithms displayed in Fig. -eto--left] demonstration consists of a multi-word language instruction-een each algorithm has five different instructions- followed by 10 input-output examples of that algorithm. See Tables 2 in Appendix for the complete list of instructions used. For each instruction, the dataset contains 20 demonstrations. Within a demonstration, the language instruction and the examples are separated by a "Y" delimiter. Demonstrations are separated by a "\n" delimiter. The input lists and matrices values are created by uniformly sampling from vocabulary of 676 tokens, created by random pairings of uppercase letters. List operations examples vary length from 3 to 6, and the matrix operations are of sizes 2 2 or 3 x 3. Fig. 4A [bottom-left] shows the training data format.
Training dataset:

---

Language Instructed Algorithm Learning Tasks
Algorithms
Test set 1: instruction based retrieval language instruction novel input completion
repeat
index
print altemate even/ /odd
reverse
roll columns 1 step transpose diagonal
prompt
shift forward/ haarkaadd return rthh element
reverse the list [XY KL MR] [RR KL XY]
Training set format examples
Test set 2: example based retrieval in algou(in,) ing completion
10 language description / algoy(in.) /.../ ingalgo/(inn]
five variations
proppt
reverse the let / PZ LM RT [RT LM P2 / [Q
flip the list QM AY JQ HH
AY
[abldm] x/ X
Example learned circuit
the
Chralled viewr
reverse the list.
Figure A. [Top-left] List and matrix algorithms used in the LIALT dataset. Format of the training [bottom-left] and examples of the two LIALT test sets [right]. B. Example a learned circuit for the "reverse" algorithm, displayed by stacking clones [left] or unrolling them [right].
In-context accuracy by task
Instruction
based prompts
In-coneet
oouraay
0.1
16
Overallocation ratio
0-1
0-3
100
30
Example
based prompts
Incontes
scaresy
10t ctem
2nd stem.
105
3rd stemm
160
100
115
alt. 600
100
CrC with Fe
Overallocation ratio
drc
Figure S: [Left] In-context accuracy (ICA) with 95% CIs after single EM iteration, function of the overallocation ratio for a CSCG trained on LIALT and averaged [top] on the instruction-based LIALT test set [bottom] on the example-base based LIALT test set. ICA increases with model capacity. [Right] ICA with standard errors per task on the two LIALT test sets: for each task, overparamerrization improves performance, Invisible bars indicate zero accuracy for the respective combinations of model and task. All the numerical values are in Appendix [D.3 Figure in in the Appendix visualizes the same quantities after EM convergence; the similarity demonstrates that the fast rehinding algorithm is not just localized in its updates, but also rapid.
Test dataset: LIALT has two test datasets, respectively containing: (a) instruction basedd retrieval prompts, and (b) example-based retrieval prompts. An instruction- based retrieval test prompt consists of a natural language instruction followed by a single input. An srample-based retrieval test prompt consists of a first input-output example of an algorithm, without any natural instruction, followed by a second input. All the lists and matrices in the two test datasets contain novel tokens. For both types of prompts, the in-context task is predict the algorithm's output when applied to the (last) input. Note that for an example-based prompt, CSCG has to infer the algorithm used from the first example. Each test set contains 100 prompts, constructed by uniformly sampling instructions, and list or matrix tokens. Fig. 4A [right] shows the formats of these two test sets.
Training: For each token, CSCG allocates its number of clones proportionally to the number of distinct contexts in the training data in which it occurs"] We parameterize CSCG capacity via this proportionality factor the "overallocation ratio". We train CSCGs for an increasing sequence of overallocation ratios on the training data with 500 EM iterations and a pseudocount of = 10 After running EM, we run 10 iterations of Viterbi training 23
As the same token might occur in different contexts in the training data, knowing the context allows predicting the sequence of following tokens, up the next "Y" delimiter.
7

---

odd
matrix operations list operations
repeat twice
elem at idx
disgonal
al. even
revense
transpose
3id elem.
1st elem.
cire shift
rall colurnss
cire shift for
2nd elem.
P.
2.))
|rtgosk] [krtgoo [ [xjd]
[Etgosk] [kaogtt]]koogss//EEE
after 1 EM iter. before rebinding
before rebirdding
... [xxa][[FFL
.../[xa]](xx]
... [_./[x]]][j3]
after EM Itr.
Figure 6: Transition graph of the CSCG model learned the LIALT dataset, visualized using the Kamada- Kawai algorithm. Visualizing the inferred probability of the observation at timestep n, conditioned on observations at all other timesteps, before rebinding. This drives the identification of anchors and slots selected for rebinding. C. Decoded latent state distributions (and predicted prompt completions) for the two different example-based LIALT prompts specified in subfig. B: (top) before rebinding, and (bottom) after one iteration of EM. Fig. 12 in Appendix D.3.2 extends the same visualization to EM convergence. The left prompt corresponds to the operation of circularly shifting the list forward, and the right prompt corresponds to reversing the list.
Results: CSCGs with sufficient model capacity successfully learn the algorithms in the training set, and rebinding generalizes those algorithms to novel tokens. Fig. 4 shows the extracted circuit for the list reversal algorithm. Fig. 55 left] presents the in-context accuracy of CSCGs (using e 10-6 and Paurprise = 0.1) on the two LIALT test sets: the best performing CSCG (a) successfully rebinds the learned schemas to the test prompts' novel tokens and (b) correctly infers the algorithm from single input-output pair for example-basse prompts. Fig. 5 also shows that model size drives ICL performance [left] even when breaking down the performance by tasks [right].
The learned CSCG (initialized with an overallocation ratio of 3) is visualized in Fig. 10 in the Appendix, using stacked clones. Fig. 66A1 shows the transition graph using the Kamada-K Kawai algorithm [24). reveals thirteen loosely connected clusters corresponding to the thirteen algorithms present the LIALT dataset. Fig. [3 illustrates the rebinding process, with the decoded distributions over latent states of the learned CSCG model, for two different example-based prompts. Even before any rebinding, the identification of anchors and slots already restricts the decoding schemas compatible with the prompt structure-in this case based on brackets & delimiters. However, the structure insufficient disambiguate completely between the compatible schemas (list operations corresponding to reversal, circular forward shift, and circular backward shift), and both the chosen prompts result in the same latent state distribution. Hence, the decoded distribution after the first E-step localizes to the three compatible schemas. In the M-step that follows, the slots in all three schemas will be rebound for this prompt. At the end of the first EM iteration, the new bindings for slots in the correct schema will be highly certain given the consistent evidence, while inconsistent evidence will lead to uncertain bindings for the other slots. the E-step of the second iteration, the respective levels of certainty in the bindings then help promote the correct algorithm schema to become the most likely decoding- complete the prompt appropriately. Note that single EM step is sufficient to derive the correct rebinding in these examples. Compare Figs. 5 & II and the
8

---

tables in Appendix Sec. D.3 for how the in in-conextt completion performance after the first EM step in the rebinding process is very similar to that at the end of EM convergence.
The LIALT results substantiate the arguments we made in Sections 3.2 and 3.3 Bayesian inference of the latent context based long-term coherence (sufficient for the GINC results in Section (.1)) does not explain the remapping of latent representation to completely new tokens as required for generalizing on the LIALT algorithms. Without rebinding, even prompt containing a full-length example of an algorithm but with novel tokens does not retrieve the correct algorithm schema or produce the correct completion based on inference over the latent states alone (Fig- cB, first row). By contrast, simultaneously inferring the rebindings and the latent states results in accurate retrieval of the algorithm schema and the correct prompt completion (Fig. (B, second row). CSCGs are thus able to learn seq2seq algorithums and generalize those algorithms to novel tokens using rebinding.
Emergence: ICL performance of CSCG the LIALT dataset shows characteristics attributed to emergence. In-context accuracy has clear dependency on the level of overparameterization of CSCG, offering evidence in support of our hypothesis in Section 3.4
4.3 Dax test
In language, the "dax" test [25] is used to demonstrate the capability of model to absorb the usage of an entirely new word from a single presentation. To test for this capability, we train a CSCG on the PreCo dataset [26). which is large-scale English dataset for coreference resolution. We then test the model on five word-replaced query prompts, where certain words in the prompts do not appear in the training set. We use Algorithm I with 10-6 and Psurprise 1 to rebind the emission matrix on each of these prompts, each time probing the model for completing a sentence by filling in the blanks (uncertain inputs) using MAP inference. Fig. shows these results.
unseen word
Replacement query prompts
alter rebinding
Fill in the blanks probing
warning
hesting
global beating may big problem.
>>
victors new computer ...
contest victors announced
planets
>> terras
temas yet the of AG intellience.
seven
also go around
vilfiiil >> AG
dangerous this intelligece progress then
bikes
>> cycles
most people ride cycles to school
more people ride their sydes around the world
Figure 7: Examples the dax test on a CSCG trained on the PreCo dataset. In each row, the novel word red (c.g. "terras") is absorbed by binding it to the clones of the corresponding word in blue (c.g. "planets") "). The CSCG can then use the new token in similar contexts, as demonstrated by the
probing.
Related work
In-context learning: Similar to how humans learn by analogy 27 and how synaptic plasticity allows the brain to rapidly adapt to a new task (28) ICL capabilities [1] allows pre-trained model
learn a new task given only few examples. 229 30 showed how demonstrations that explicitly guide the reasoning process improve the ICL performance of transformers on new complex tasks. We clarify below some concepts that should not confused with ICL, and then discuss some works that aim at understanding ICL. and the factors that influence it.
Supervised learning (SL) and few-shot learning (FSL): SL approaches learn mapping that minimizes a loss on the training data: gradient methods are popular paradigm [31] 32 33 In FSL,
model learns rapidly adapt to new task from a limited number of supervised examples 33 35
36). and performs this same task at inference. In contrast, ICL tasks are only revealed at inference. 37
38
showed that finetuning transformers on ICL instructions improves their ICL performance.
Meta-learning: The meta-lnarning paradigm aims at learning to adapt to a new task with only a few examples [39 40 by using multiple learning experiences. In contrast, ICL directly emerges from the traitrd model. [42 43 proposed a meta-learning framework for ICL where the model is fine-tuned: it learns to leverage few-shot examples and to adapt to tasks inference time.
3 explained ICL as implicit Bayesian inference and constructed the GINC dataset (see Section 4.1 for demonstrating ICL. (44) abstracted ICL as algorithm learning problem and found that a transformer can implicitly infer hypothesis function. Similarly, [45) showed that transformer can be trained to perform ICL of unseen linear functions, with performance
How ICL works:

---

comparable to the optimal least squares estimator. [46] showed that, in the linear case, transformers implicitly implement gradient descent and train an implicit linear model on the ICL examples. 4 proposed dual between transformer attention and gradient methods and suggested pre-trained models as meta-optimizers. They presented ICL. as implicit finetuning, where the forward pass on the demonstrative examples produces reta-gradients. Finally, (5) showed the existence of "induction heads" in transformers, that emerge during training, copy previous patterns, and drive ICL capacities.
47 indicated that LLMs' ICL performance "emerges" then keeps improving when the model size increases. (48) proposed substitute for the positional encoding, and demonstrated how transformers can learn schemas for algorithmic tasks and generalize to test sequences longer than any seen during training. Some works have highlighted the role of the training data in ICL. [49] showed that ICL emerge when the training data has a large number of rare classes and when examples appear in clusters. while [50] demonstrated that ICI. emerge when model is trained on combination of multiple corpora, and that low perplexity and ICL. performance do
What influences ICL:
not always correlate. [5 [22 found that ICL is highly unstable and is influenced by the prompting template, the selection of in-context examples, and the order of the examples. (53) showed that the
ICL performance is driven by the exposure to the label space, the input distribution, and the overall format of the sequence. Similarly, (54) found that selecting ICL examples with closer embeddings to ICL test sample improves ICL performance, and (55) showed that adding explanations in-context improves performance. Finally, (56) recently claimed that the sharp emergence of ICL in larger models might be an artifact of the metrics, not a fundamental property of the model.
Discussion
With ICL deconstructed into schema learning, schema retrieval, and slot rebinding, an interesting avenue for future work would be to probe various sequence models for how robustly each of these components are manifested or even construct models around these principles. Here we consider how this framework might map to transformers, where the phenomenon of ICL was originally observed. Unlike CSCGs, transformers buffer the inputs and represent location as positional encoding, allowing attention to gate by the structure of the prompt, along with the contents. Prior explanations [3] do not distinguish the role of sequence positions vis-a-vis contents; we argue that theories might need to emphasize this distinction (see Fig. (8)) to fully understand the inductive biases behind ICL. We conjecture (see Fig. 3) that layers of the transformer implement multiple mixed templates of positions and content, evaluated at different offsets of prompt. The template assembly that can auto-regressively match the prompt wins out the competition to gate the content. The rebinding mechanism requires only a few iterations of sparse updates to the emission matrix, and can be temporally "unrolled" into a forward pass, allowing ICL behavior with fixed weights since the slotting process lives in the space of activations.
A
content: based predictor
position based prodictor
content & position based predictor
B
selection among templates
Tomp ate evalusted at diferent offeets to gatire
position
Figure 8: A. Learned templates in transformer could involve content, position, mix of both. B. Activations in the forward pass of transformer could be selected among pre-learned templates that mix content and position to achieve ICL without weight changes.
Coming back to CSCGs, implementations can scale to larger models and datasets by exploiting sparsity and parallelizing computations in the EM steps. Allowing factorized latent space and adding skip connections would also allow compositionality while enabling scalability. Further, while we have illustrated here the concept of rebinding attach new symbols existing slots, rebinding "through time" can also target connections between clones, enabling compositional behavior in- context. We leave these explorations for future research. Our goal here has been to elucidate general framework for ICL behavior, leveraging the interpretability of CSCGs. We hope this demystifies the ICL behavior observed in LLMs by analogy, showcases avenues for further research on ICL capabilities, and provides broad impetus for interpretable methods.
10

---

Acknowledgements
We thank Stephanic Chan, Andrew Lampinen, Anirudh Goyal, Dharshan Kumaran, Neel Nanda and Guangyao Zhou for helpful discussions and comments on the draft.
References
[1] Tom Brown et al. "Language models are few-shot learners". In: Advances in neural information processing systems 33 (2020), pp. 1877-1901.
[2] Taylor Webb, Keith Holyoak, and Hongjing Lu. "Emergent Analogical Reasoning in Large Language Models". In: arXiv preprint arXiv:22 2t..22//2.9///1 (2022).
[3] Sang Michael Xie et al. "An explanation of in-context learning implicit bayesian inference".
In: arXiv preprint arXiv:2 100:2111.00000 (2021).
[4] Damai Dai et al. "Why Can GPT Learn In-Context? Language Models Secretly Perform Gradient Descent Meta Optimizers". In: arXiv preprint arXiv:2212.10559 (2022).
[5] Catherine Olsson et al. "In-context learning and induction heads". In: arXiv preprint urXiv:2209. 1895 (2022).
[6] Dileep George al. "Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps". In: Nature communications 12.1 (2021), p. 2392.
[7] Antoine Dedieu et al. "Learning higher-order sequential structure with cloned HMMs". In: arXiv preprint arXiv: 0055000507 (2019).
[8] Murray Shanahan and Melanie Mitchell. "Abstraction for Deep Reinforcement Learning".
In: Proceedings of the Thirty-First International Joint Conference Artificial Intelligence,
IJCAI-22. 2022, pp.
[9] Neel Nanda et al. "Progress measures for grokking via mechanistic interpretability.. In: arXiv preprint arXiv:2301.00 05217 (2023).
[10] Kevin Meng al. "Locating and Editing Factual Associations GPT". In: Advances in Neural Information Processing Systems 36 (2022).
[11] Ashish Vaswani al. "Attention is all you need". In: Advances in neural information process- ing systems 30 (2017).
[12] Lonnie Chrisman. "Reinforcement learning with perceptual aliasing: The perceptual distinc- tions approach". In: AAAI. Vol. 1992. Citeseer. 1992, pp. 183-188.
[13] Judea Pearl. Causality. Cambridge university press, 2009.
[14] Vatsal Sharan et al. "Learning overcomplete hmms". In: Advances in Neural Information Processing Systems 30 (2017).
[15] Swaroop Guntupalli al. "Graph schemas as abstractions for transfer learning, inference, and planning". In: arXiv preprini arXiv:2302. 07350 (2023).
[16] Jonas Peters, Dominik Janzing, and Bernhard Schölkopf. Elements of causal inference: foun- dations and learning algorithms. The MIT Press, 2017.
[17] Daniel Eaton and Kevin Murphy. "Exact Bayesian structure learning from uncertain interven-
tions". In: Artificial intelligence and statistics. PMLR. 2007, pp. 107-114.
[18] Arthur P Dempster, Nan M Laird, and Donald Rubin. "Maximum likelihood from in- complete data via the EM algorithm"". In: Journal of the royal statistical society: series B (methodological) 39.1 (1977). pp. 1-22.
[19] Judea Pearl. Probabilistic reasoning in intelligent systems: networks of plausible inference. Morgan kaufmann, 1988.
[20] Bingbin Liu al. "Transformers learn shortcuts to automata". In: arXiv preprint arXiv:2210. 10749 (2022).
[21] Oren Kolodny, Arnon Lotem, and Shimon Edelman. "Learning generative probabilistic grammar of experience: A process-level model of language acquisition". In: Cognitive Science
39.2 (2015), pp. 227-267.
[22] Zoubin Ghahramani and Michael Jordan. "Factorial hidden Markov models". In: Advances in Neural Information Processing Systems 8 (1995).
[23] Frederick Jelinek. "Continuous speech recognition by statistical methods". In: Proceedings of
the IEEE 64.4 (1976), pp. 532-556.

---

[24]
Tomihisa Kamada, Satoru Kawai, et al. "An algorithm for drawing general undirected graphs". In: Information processing letters 31.1 (1989), pp. 7-15.
[25] Haley A Vlach and Catherine DeBrock. "Remember dax? Relations between children's cross-situational word learning, memory, and language abilities". In: Journal of memory and language 93 (2017), pp. 217-230.
[26] Hong Chen et al. "PreCo: A large-scale dataset in preschool vocabulary for coreference resolution". In: arXiv preprint 1xX1::811000900 (2018).
[27] Patrick Winston. "Learning and reasoning by analogy". In: Communications of the ACM
23.12 (1980), pp. 689-703.
[28] Katie C Bittner et al. "Behavioral time scale synaptic plasticity underlies CA1 place fields". In: Science 357.6355 (2017), pp. 1033-1036.
[29] Jason Wei et al. "Chain of thought prompting elicits reasoning in large language models". In: arXiv preprint arXiv:2201. 1903 (2022).
[30] Denny Zhou et al. "Least-to-most prompting enables complex reasoning in large language models". In: arXiv preprint arXiv:2205. 10625 (2022).
[31] Léon Bottou, Frank E Curtis, and Jorge Nocedal. "Optimization methods for large-scale machine learming". In: SIAM review 60.2 (2018), pp. 223-311.
[32] Sebastian Ruder. "An overview of gradient descent optimization algorithms". In: arXiv preprint arXiv:1609, 04747 (2016).
[33] Diederik P Kingma and Jimmy Ba. "Adam: A method for stochastic optimization", In: arXiv preprint XrXiv: /41.66991066800 (2014).
[34] Yaqing Wang al. "Generalizing from few examples: A survey on few-shot learning". In:
ACM computing surveys (csur) 53.3 (2020), pp. 1-34.
[35] Li Fei-Fei, Robert Fergus, and Pietro Perona. "One-shot learning of object categories", In: IEEE transactions on pattern analysis and machine intelligence 28.4 (2006), pp. 594-611.
[36] Brenden M Lake, Ruslan Salakhutdinov, and Joshua B Tenenhaum. "Human-level concept learning through probabilistic program induction". In: Science 350.6266 (2015), pp. 1332- 1338.
[37] Jason Wei et al. "Finetuned language models are zero-shot learners". In: arXiv preprint arXiv:2109.00 01652 (2021).
[38] Victor Sanh et al. "Multitask prompted training enables zero-shot task generalization". In: arXiv preprint 2
(2021).
[39] Devang K Naik and Richard J Mammone. "Meta-neural networks that learn by learning". In:
[P (Procediings 1992] IJCNN International Joint Conference on Neural Networks. Vol. 1. JEEE. 1992, pp. 437-442
[40] Sachin Ravi and Hugo Larochelle. "Optimization as a model for few-shot learning". In: International conference on learning representations. 2017.
[41] Sepp Hochreiter, A Steven Younger, and Peter R Conwell. "Learning to learn using gradient descent". In: Artificial Neural Networks
2001: International Conference Vienna, Austria, August 21-25, 2001 Proceedings 11. Springer. 2001, pp. 87-94.
[42] Sewon Min al. "Metaicl: Learning to learn in context".. In: arXiv preprint arXiv:21 15943 (2021).
[43] Yanda Chen et al. "Meta-leaming via language model in-context tuning". In: arXiv preprint 10.0014 (2021).
[44] Yingcong Li et al. "Transformers as Algorithms: Generalization and Stability in In-context
Learning.. In: ()
[45] Shivam Garg et al. "What can transformers learn in-context? a case study of simple function
classes". In: Advances in Neural Information Processing Systems 35 (2022), pp. 50583-30598
[46] Ekin Akyürek et al. "What learning algorithm is in-context learning? investigations with linear models". In: arXiv preprint arXiv:2211.1555 155661 (2022).
[47] Jason Wei et al. "Emergent abilities of large language models". In: arXiv preprint
arXiv:2206.0708 07682 (2022).
[48] Yuxuan Li and James McClelland, "Systematic Generalization and Emergent Structures in Transformers Trained on Structured Tasks". In: arXiv preprint arXiv:2210.0040 (2022).
12

---

[49] Stephanie Chan et al. "Data distributional properties drive emergent in in-context learning in
transformers". In: Advances in Neural Information Processing Systems 35 (2022), pp. 18878- 18891.
[50] Seongjin Shin al. "On the effect of pretraining corpora in-context learning by large-scale language model". In: arXiv preprint WwXi::2204.13509 (2022).
[51] Zihao Zhao et al. "Calibrate before use: Improving few-shot performance of language models".
In: International Conference on Machine Learning. PMLR. 2021, pp. 12697-12706.
[52] Yao Lu al. "Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity". In: arXiv preprint arXiv:2104.0888 088786 (2021).
[53] Sewon Min et al. "Rethinking the Role of Demonstrations: What Makes In-Context Learning
Work?" In: arXiv preprint arXiv:2202. 12837 (2022).
[54] Jiachang Liu et al. "What Makes Good In-Context Examples for GPT-3?" In: arXiv preprint
arXiv:2101.cm 00800000444 (2021).
[55] Andrew Lampinen et al. "Can language models learn from explanations context?" In: arXiv preprint arXiv:2204.02.22 (2022).
[56] Rylan Schaeffer, Brando Miranda, and Sanmi Koyejo. "Are Emergent Abilities of Large Language Models Mirage?" In: arXiv preprint arXiv:2304. 15004 (2023).
13

---

A Locally updating the emission matrix with the transition matrix fixed
We reuse the same notations in [13], Appendix A.2. The authors describe the EM algorithm for learning the emission matrix of CSCG with a fixed transition matrix. In particular their M step defines the new emission matrix as:
N Σ n=l
E
1x.=
7()) 2
Σ (n). Vj
where FG)) is column of the emission matrix corresponding to the emission j. xx,=; is an indicator function, @ the element-wise division and ((n) is derived by the authors from the forward and backward probabilities. The i entry of the vector EGG) is then defined as:
ΣN m-= 1x.=ı R((n)
Eg
Ex (()))
In contrast, in Section 2.2 Step 5 of Algorithm only updates the for which find a pair (i,n) E R. by only using the beliefs at timestep n. For this row, the jth entry becomes:
En (i,))R -x,=1
Eij
En: (L,))R (n2)
The pseudocount used Step of Algorithm iis an uncertainty parameter that lets the model smooth incorrect observations. More details of this parameter available in 66)
B Prompt completion algorithm
describes the prompt completion algorithm introduced Section 2.2 implicitly considers single action, which takes the next sequence element.
Algorithm 2
Algorithm 2 Prompt completion
Input: Grounded schema {T, C, Etb] with rebound CSCG emission matrix ETb, delimiter token
Tp. prompt (prompt) = (#1,..., IN)
Output: A completed prompt x (completed) (.1...., NN ..N+1...,,NN.. = ag))
1: Run max-product for MAP inference and return .MAP argmax, P(= (poop))))))
(Zl,..., =
2: Set l. = While IN+4 # zgg increment & <- E and sample the most likely observation: NN E argmax, T T and IN+E E argmax, Eib "N
14

---

C Additional materials for the GINC dataset
First, we present two additional plots for the GINC experiment.
In-context performance on the GINC dataset
In-context confidence seen concepts
In-context accuracy on unseen concepts
In-context confidence
In-context accuracy
0.30
0.4
0.00
0.15
0.04
0.00
k=8
k-5
0.3
0.02
k=10
0.00
Num examples
Num examples
Figure 9: [Left] In-context confidence for the CSCG with 50 clones on the GINC test dataset. defined the averaged probability of the predictions. For larger values of k. CSCG correctly infers the context of the aliased observations and is more confident in its predictions. [Right] Similar to transformer and LSTM reported in B CSCG fails to extrapolate and has low in-context accuracy when the test prompts are sampled from five novel concepts, unseon during training.
Second, we present the table of results associated with Fig. for the CSCGs with 10 and 50 clones.
Context length
No. of examples
CSCG with 10 clones
CSCG with clones
0.509 (0.020)
0.534 (0.020)
0.351 (0.019)
0.445 (0.019)
0.366 (0.019)
0.453 (0.020)
4
0.356 (0.019)
0.468 (0.020)
0.360 (0.019)
0.454 (0.020)
16
0.354 (0.019)
0.460 (0.020)
32
0.354 (0.019)
0.441 (0.0219)
64
0.369 (0.019)
0.468 (0.020)
(0.018)
0.927 (0.010)
0.640 (0.019)
0.927 (0.012)
2
0.629 (0.019)
0.904 (0.012)
0.654 (0.019)
0.883 (0.013)
8
0.627 (0.019)
0.891 (0.012)
16
0.637 (0.019)
0.902 (0.012)
32
0.631 (0.019)
0.901 (0.012)
64
0.637 (0.019)
0.893 (0.012)
0
0.696 (0.018)
0.969 (0.007)
0.694 (0.018)
0.972 (0.007)
0.686 (0.018)
0.972 006)
4
0.681 (0.018)
0.978 (0.0 (0.006)
0.690 (0.018)
0.973 (066)
16
0.686 (0.018)
0.975 (0.006)
32
0.676 (0.018)
0.968 (0. (06006)
(0.018)
0.975 (0.007)
(0.018)
0.975 (0.006)
1
0.705 (0.018)
0.977 (0 (0.0066
0.674 (0.018)
0.971 (0. (0.006)
(0.018)
0.974 (0.006)
8
0.690 (0.018)
0.977 (0.006)
16
0.689 (0.018)
0.977 (0.00 (0.006)
32
0.712 (0.018)
0.978 (006)
64
0.690 (0.018)
0.978 (0.006)
Table I: In-context accuracy for a CSCG with 10 clones and a CSCG 50 clones trained on the GINC dataset, averaged (with 95% contidence intervals) on each each pair (k, n) of context length and number of examples n of the GINC test set.
CSCG performs better on zero-shot prompts than on few-shot prompts: We observe that, for short contexts, CSCG in-context accuracy is higher on zero-shot prompts = 0 than few-shot prompts n = 1,2,.... We hypothesize that the difference between the training and the prompt distributions creates gap that lowers few-shot in-context accuracy. The performance gap disappears
15

---

for larger contexts k E {8, 10] as they "overpower" the train-test distribution divergence. [3] made a similar observation for transformers. However, their performance gap was also observable for larger contexts.
D
Additional materials for the LIALT dataset
D.1 Natural language instructions
Tables 2 and 3 present the natural language instructions respectively used for the nine list algorithms and four matrix algorithms of the LIALT dataset. Language instructions are grouped in clusters of five: all five instructions within one cluster describe to the same algorithm. As described in the main text, each demonstration of the LIALT training and first test set uniformly selects one instruction.
"find the "print the first elenent from the Hiae" "Yeturr the leading elanent from the Hatt "find the head element from the list" 'retrieve the starting cleaent from
index
Hist"
alement at index of list" "find the second elenent from the liat" "retrieve the sacond alemen fron the list" "oocae the recund iten from the list" "recurn elexent in decond place from list"
list"
"print the elanant at indar teo of the Hat" "find the third elezent from the liet" "locate the
"raverse the liat" "nirror the list" "Mipp list" "flip the order of the list"
elesent fron the 1ist"
"catpat the third isea from the list" "return the elemet in third place from the 1ist"
"reverse the order of the itens in the 1ist"
"duplicate each
itea"
elenenta one place forvard" "roll the list elaxents one position to the right" "suitch the iters of the list one poaltion forward" forward"
"rotate
"replicate overy element is the list"
"make a capy of each slenent in the Hist"
"clone
list" "creata second inatanca of avary alanent the list"
"advance
elerents "nove 1iat alerenta one poaition forward"
starting with "retrieve alternate itens in the list starting with second sten"
"print every
zeaber in
"print every other nember the first nerber"
list starting
the second narbar"!
"find alternate elenents in the liet beginning with the first clement"
"geturn every other object in the list starting with the second object"
"prist every secoad iten the list, starting with the first elanent"
"retrieve every other entry in the list starting with second entry"
"outpos second elenent in the list, starting firat alerent"
"cubpat odd indexed elepents"
"output even indezed elemente"
"rotate the list elenents place backsard" the
elenenta position the left" "change the iteas of the list one position backvard" "displace the elenents of the 11st cns index bacivard"
"roll the list
position backvard"
Table 2: Natural language instructions for the list algorithms used in the LIALT datasct
"return the satrix diagonal"
"return the natrix transpose" "ratrisve tranapose the natrix'
"callact the diagonal valusa the natrix"
"retrieve the diagonal elexente of the natrix'
"get the transposed matrix" "coopptt the transposed forn of the natrig"
"return the di diagonal entries of the natrix" "fetch the diagonal
of the satrix"
"derive the transpose natrix'
"roll the columne of the natrix to the right"
"find the natrix elexent in the recond row and zeccad colum""
"rotata the ratrix calurna the right"
'find the value in the second and second calume of the natrix"
"nove the natrix columne to the right"
"fetch the natrix elenent located in row 2 and colunn 2
"shift the column "spin the natrix colurns to the right"
natrix to righe"
"print the value 2 in natrix" "retrieve the zatrix clement at 2 2"
Table 3: Natural language instructions for the matrix algorithms used in the LIALT datasct
D.2 Learned CSCG model
Our next Figure 10 displays the transition graph of the CSCG model trained on the LIALT dataset with an overallocation ratio of 3, with stacked clones for each symbol.
16

---

Figure 10: CSCG model learned on the LIALT dataset, visualized with stacked clones.
17

---

D.3
Results on the LIALT dataset
D.3.1
After single EM iteration
Presented below are the tables of results associated with Fig. 5 Table contains the in-context accuracies averaged the entire test set, Table [ contains the in-context accuracies per task instructions-baxee prompts, and Table conttiins the in-context accuracies per task on example-base prompts.
Overallocatio ratio
Instruction based prompts
Example based prompis
0.1
0.00 (0,00)
0.00 (0.00)
0.3
0.20 (0.08)
0.09 (0.06)
0.54 (0.10)
(0.10)
3.0
0.89 (0.06)
0.91 (0.06)
Table 4: Average in-context accuracy of each CSCG model-with 95% confidence intervals-as as a function of CSCG overallocation on both (a) the cirruction-based LILAT test set and (b) the example- based based LIALT test set.
Overallocation
Tisk
0.1
0.3
1.0
3.0
list 1st elem.
0.00 (0.00)
0.00 (0.00)
0.89 (0.10)
1,00 (0.00)
list 2nd clem.
0.00 (0.00)
0.60 (0.15)
0.70
1.00 (0.00)
3rd elem.
(0.00)
0.60 (0.15)
(0.15)
1.00 (0.00)
Tist reverse
0.00 (0.00)
0.00 (0.00)
0..0
1.00 (0.00)
repeat
(0.00)
list alt. even
0.00 (0.00)
0.00 (0.00)
(0.20))
1.00 (0.00)
odd
0.00
0.00
0,86
list cire. shift fw:
0.00 (0.00)
0.00 (0.00)
0.00 (0.00)
(0.17)
Tist circ. shifl bw. matrix diagonal
0.00 (0.00)
0.00 (0.00)
(.14))
1,00 (0.00)
0.00 (0.00)
0.50
0,50
1,00 (0.00)
matrix transpose matrix roll columns
0.00 (0.00)
0.00 (0.00)
0.50 (0.20)
1,00 (0.00)
(0.00)
0.00 (0.00)
(0.00)
0.17 (0.15)
matris elem, at ids.
0.00 (0.00)
0.50 0.50
1.00 (0.00)
1.00 (0.00)
Table 5: Average in-context accuracy by task-with standard as a function of CSCG overallncation on instruction-based prompts.
Overallocation ratio
Task
0.1
0.3
1.0
3.0
elem.
0.00 (0.00)
1.00 (0.00)
list 2nd clem.
0.00 (0.00)
0.12 (0.12)
(0.00
1.00 (0.00)
3rd elem.
0.00
0.38 (0.17)
(0.00)
1.00 (0.00)
list reverse
0.00 (0.00)
0.00 (0.00)
(0.12)
0.88 (0.12)
Tist repeat twice
0.00
0.00 (0.00)
(..18)
0.62 (0.17)
all. even
0.00 (0.00)
0.00 (0.00)
0.00 (.00
1.00 (0.00)
list alt. odd
0.00 (0.00)
0.00 (0.00)
0.38 (0.17)
1.00 (0.00)
(0.00)
0.00 (0.00)
(0.00)
0.88 (0.12)
list circ. shifi bw.
0.00 (0.00)
0.00 (0.00)
0.29 (0.17)
1.00 (0.00)
matrix diagonal matrix transpose matrix roll columns
(0.19)
0.67
1,00
0.00 (0.00)
(0.00)
0.43
1.00 (0.00)
0.00 (0
0.00 (0.00)
(.000
0.50 (0.18)
matrix elem, at ids.
0.00 (0.00)
0.38 (0.17)
1.00 (0.00)
1.00 (0.00)
Table 6: Average in-context accuracy by task- standard errors a function of CSCG overallocation example- based prompts.
D.3.2 After EM convergence
present analogues of Figs. 5 but after the EM algorithm in the rebinding process has converged. We that the results are mostly identical. Table contains the in-context accuracies averaged on the entire test set, Table contains the in-context accuracies per task on instructions-basee prompts, and Table contains the in-context accuracies per task on example- saseeeebbased based prompts.
Figs. 12
18

---

In-context accuracy by task
Instruction
based
Inccomtn
500/OCy
01
0.3
16
3.3
Overallocation ratio
0-1
100
30
Example
based
covotted
occueoyy
Overallocation
033 1.0
15t 1srelem
2nd elons
107
ah: ovan
drc shit
mi
made
not Hitr
mare
Figure 11: [Left] In-tontext accuracy (with 95% Cls) after EM convergence, as function of the overallocation ratio for CSCG trained LIALT and averaged [top] the instruction-basee LIALT test [bottom] example-bane LIALT test set. In-context accuracy increases for CSCGs with larger capacities. [Right] In-context accuracy (with standard errors) per task on the two LIALT test sets: for each task, overparametrization improves performance. All the numerical values are in Appendix D.3 Invisible bars indicate zero accuracy for the respective combination of model and task.
Overallocation ratio
Instruction-based prompts
Example-based prompts
0.1
0.00 (0.00)
0.00 (0.00)
0.3
0.16 (0.07)
0.11 (0.06)
1.0
0.54
0.49 (0. 10)
3.0
0.89 (0.06)
0.93 (0.05)
Table 7: Average in-context accuracy of each CSCG model-with 95% confidence intervals-as
as a function of CSCG overallocation on both (a) the instruction-based LII.TT test set and (h) the example-hased LIALT test set.
Overallocation
Task
0.1
0.3
1.0
3.0
Bist 1st elem.
0.00 (0.00)
0.00 (0.00)
0.89 (0.10)
1.00 (0.00)
list 2nd elem.
0.00 (0.00)
0.60 (0.15)
0.70
1,00 (0.00)
3rd elom.
(0.00)
0.60 (0.15)
(0.15)
1,00 (0.00)
list reverse
0.00 (0.00)
0.00 (0.00)
0.70 (0.14)
1.00 (0.00)
repeat twice
0.00 (0.00)
(0.25)
(0.22)
list all. even
0.00 (0.00)
0.00 (0.00)
0,50 (0.20)
1.00 (0.00)
0.00 (0.00)
0.86
list cire. shift fw.
0.00 (0.00)
0.00 (0.00)
0.00 (0.00)
0.56 (0.17)
cire. shift bw. matrix diagonal matrix transpose
0.00
0.00
(..13)
1.00
0.00 (0.00)
0.00 (0.00)
(.50
1.00 (0.00)
0.00 (0.00)
0.00 (0.00)
(..20)
1,00 (0.00)
matrix roll columns matrix elem, at ids.
(0.00)
0.00
(0.00)
0.17 (0.15)
0.00 (0.00)
0.50 (0.18)
1.00 (0.00)
1.00 (0.00)
Table 8: Average in co-tettex accuracy by task with standard errors as a function of CSCG overallocation on instruction-bast prompts.
Overallocation
Task
0.1
0.3
1.0
3.0
list 1st elem.
0.00 (0.00)
0.00 (0.00)
1.00 (0.00)
1.00 (0.00)
Tist 2nd elem.
0.00 (0.00)
0.12 (0.12)
1.00 (0.00)
1.00 (0.00)
list 3rd elem.
0.00 (0.00)
0.38 (0.17)
1.00 (0.00)
1.00 (0.00)
list reverse
0.00 (0.00)
0.00 (0.00)
0.00
0.88 (0.12)
repeat
0.00 (0.00)
list alt. even
0.00 (0.00)
0.00 (0.00)
(0.00)
1.00 (0.00)
odd
0.00
0.00
(.50))
1.00 (0.00)
list cire. shift fw:
0.00 (0.00)
0.00 (0.00)
0.00 (0.00)
0.88 (0.12)
cire. shifl bw. matrix diagonal matrix transpose
0.00 (0.00)
0.00
0.29
1,00
0.00 (0.00)
0.67 (0.19)
0.67 (0.19)
1.00 (0.00)
0.00 (0.00)
0.00 (0.00)
0.43 (0.19)
1.00 (0.00)
matrix roll columns matris elem, at ids.
(0.00)
0.00 (0.00)
(0.00)
(0. 18)
0.00 (0.00)
0.38 (0.17)
1.00 (0.00)
1.00 (0.00)
Table 9: Average in in-context accuracy by task-with with standard emorr as a function of CSCG overallocation on example-hased prompts.
19

---

p. 5. 2..)
[rtgoak] kkrgoos x [xja] 1
[
after EM convg. after EM iter. before rabinding
before rehinding
...[xjd]((xp]
.,./[xj][[jj]
after IEM iter.
/ [x][xj]][aa]]
after EM comug
Figure 12: Extension of Fig. to EM convergence. A. Visualizing the inferred probability of the observation at timestep n, conditioned observations all other timesteps, before rehinding. This drives the identification of anchors slots selected for rebinding. Decoded latent state distributions completions - for the two different example-based LIALT prompts specified subfig. A: (top) before rebinding, (middle) after one iteration of EM, and (bottom) after EM convergence. The left prompt corresponds to the operation of circularly shifting the list forward, and the right prompt corresponds reversing the list.
predicted prompt
D.4 Example failures
Finally, present few examples which illustrate the failure modes of approach. These primarily a consequence of imperfections in the learned CSCG model. Each example is presented in the format (prompt, ground truth correct output, actual model response).
1. For these failures, the instruction circuit has been wired the wrong algorithm circuit (possibly driven by the ambiguity of the forward slash delimiter separating the instruction from the example), resulting in the retrieval of the wrong schema.
output odd indexed elements / [UVBQKI)
[
V Q 1
flip the list [SEJ]
[ E
S S E J JJ
the list / [ B]] [BTR /
[RR B B ] /
mirror list B A OTJ [TOAB]] [BBAAOOTT]]
2. For these failures, the schema has been learned incorrecty..
switch items of the list one position forward [LNGXXT]
[TL N
[T N G X M T
shift the columns the matrix to the right [CCDY] [VF]]
[ F v
[ get
/ [ ] [
/ [BAEFUL]
[ B
W
W L L ] /
[ A F
L L /
C X J
C P
T [VEJPEN]
[FP 1 /
[FPW] [FPR]
20


Reviews:
[{"review_id": "vSLqANy8bF", "reviewer": "Reviewer_PPgs", "summary": "The authors propose the clone-structured causal graph (CSCG) with rebinding as a model of in-context learning in language models. The authors conduct several experiments showing that CSCGs can learn latent graphs corresponding to meaningful concepts seen in the training data, while also generalizing to instantiations of those concepts containing novel tokens seen at test time.", "questions": "How does the size of the latent graph for e.g. list reversal scale with list size?\n\nWhat is the computational cost of fitting a CSCG with EM? PreCo is many orders of magnitude smaller than typical pre-training corpora. Is CSCG plausible even for very large corpora?\n\nAre the results in Figure 6 cherry picked? How often is rebinding + MAP inference successful?\n\nIs there a typo in the emission matrix/clone structure definition? The number of columns of the emission matrix is the size of the observation space when it is defined, but it is the size of the latent space when the clone structure condition is defined later in 2.1.", "limitations": "The authors should add some discussion about the limitations of CSCGs for modeling the behaviors of real-world LLMs, particularly whatever computational challenges may exist in scaling CSCG-like models to billions or trillions of tokens.", "rating": 7, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 3, "strengths": "The presentation of the CSCG as a model of in-context learning is new to me and very interesting, addressing a hot topic in LLMs that is relevant to a large portion of the community. The experiments show that CSCGs can be learned in a way that leads to meaningful concepts, as well as the fact that overparameterization is useful for CSCGs as with typical LLMs. Finally, the experiments show that CSCGs can perform fast binding of novel tokens to previously-seen structures in a \"dax test\", which is an important capability the authors argue is not adequately explained by existing frameworks.", "weaknesses": "Key claims about interpretability are not substantiated. In particular, the abstract and intro clearly emphasize the importance of interpretability in models of language/ICL, but the experiments do not meaningfully explore how CSCGs are more interpretable than existing LLMs.\n\nSome experiments are somewhat superficial or small scale, particularly 4.3, where only a few qualitative examples are shared, without any quantitative analysis showing that the results are representative. The largest datasets are orders of magnitude smaller than the datasets used by real LLMs. This shortcoming isn't absolutely critical, but of course the results would be much more convincing if ICL phenomena could be shown at something closer to real-world scale.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are glad that you consider novel and interesting our application of CSCGs towards a framework for understanding ICL. We will elaborate on the challenges of scaling CSCGs to large datasets in the discussion section, and some potential directions for progress.\n\nTo recap how we leverage CSCG interpretability in the paper: in Fig. 3A, we extract the CSCG transition graph and observe five clusters – corresponding to each of the concepts in the GINC dataset. Similarly, in Fig. 4B, we extract interpretable circuits from a CSCG trained on the LIALT dataset: these circuits explicitly represent sequences of different lengths (relating also to the question of the latent graph size). \n\nWe also present two additional figures in the one page PDF supplement to the global response (which we will include in the final paper), illustrating:\n* How increasing prompt length contributes to schema retrieval\n* The step-by-step functioning of the retrieval and rebinding process, on LIALT prompts\n\nFinally, please note that some previous work on CSCGs [1] has also exploited their interpretable structure to match with cognitive maps in the hippocampus.\n\nCumulatively, we hope this buttresses our claims regarding the interpretability of CSCGs.\n\n\n> How does the size of the latent graph for e.g. list reversal scale with list size?\n\nThe size of the latent graph for the reversal algorithm, as measured by the number of nodes, scales (a) as O(N) when the list size is fixed to N (b) in the worst case, as O(N^2) when the list sizes can vary between 1 and N. We illustrate this scaling (a) conceptually in Figs. 2E and 2F and (b) empirically by extracting the latent graph from a trained CSCG in Fig. 4B: the unrolled view illustrates how lists of different sizes are represented in the latent graph.\n\n\n> What is the computational cost of fitting a CSCG with EM? [...] Is CSCG plausible even for very large corpora?\n\nFor a CSCG with M clones per symbol trained on a sequence of length N, each EM iteration has a computational cost in O(N M^2). Note that the size of the latent space is H = M*E, where E is the total number of distinct symbols. It is possible to exploit the sparsity structure of the problem (as well as other elements, such as factorizing the transition matrix, parallelizing the computation of the EM steps, etc)  to scale the training of CSCGs to large language datasets.\n\n\n> Are the results in Figure 6 cherry picked? How often is rebinding + MAP inference successful?\n\nWe selected the examples in Figure 6 so that, when filling in the words, the context makes it clear what the missing word is. The same is necessary for humans; it is hard to rebind correctly with an ambiguous sentence such as “I went to the dax this morning”. When there is no ambiguity about the missing word, we found MAP inference + rebinding to be successful.\n\n\n> Is there a typo in the emission matrix/clone structure definition? The number of columns of the emission matrix is the size of the observation space when it is defined, but it is the size of the latent space when the clone structure condition is defined later in 2.1.\n\nThank you for pointing out a typo on line 64: we have inverted the rows and columns of the emission matrix. We will fix it.\n"}, {"review_id": "KAMM5yE41k", "reviewer": "Reviewer_ZCs7", "summary": "The paper aims to replicate the in-context learning (ICL) phenomenon in large language models (LLM) with clone-structured casual graphs (CSCG), which is roughly trying to learn hidden states in PODMP such that the transition matrix is invariant new environments, and only the emission matrix needs re-learning.\n\nCSCG has the properties of context-separation, transitive generalization, schema-formation, and refining, which can be used to explain ICL behaviors. Results on three benchmarks show CSCG can match some LLM ICL behaviors.", "questions": "- How is learning (of transition matrix) done? The paper seems to only talk about re-binding for learning the emission matrix...\n\n- Why transition T is assumed to be fixed, and what kinds of train-test split do we assume this? (LLM might be hard to assume this?)\n\n- Is the latent symbol space for z huge? Is \"clone\" just defined as #latent_symbols/#observed_symbols? Also, why not several latent tokens between observation tokens? Might make symbol space exponentially more efficient and more like language compositionality?\n\n- Figure 3, is accuracy kind of low? does it really match Transformer models?\n\n- Is this whole CSCG thing kinda like a latent automation? Do you think LLM is possible to recover such symbolic structures?", "limitations": "Did not see any limitation section in the paper, but there should be some (see weaknesses).", "rating": 6, "confidence": 1, "soundness": 2, "presentation": 2, "contribution": 3, "strengths": "- ICL is important and emerging, and its understanding is timely. The paper contributes to a novel view into ICL with some interesting constructions and experiments.\n\n- The context-dependent latent representation and transitive generalization make sense.\n\n- Experiments seem interesting and help replicate some ICL behaviors on toy datasets, e.g. accuracy of ICL depends on overparametrization.", "weaknesses": "- CSCG is unfamiliar to most readers and there could be some confusions that hinder understanding the paper (see questions). Importantly, **I do not understand how is the transition learned in the first space**, even before any \"rebinding\". And what is even the training data for each experiment??\n\n- The setup of CSCG is still far from LLM ICL, and there should be acknowledgements of such limitations if the goal is to explain LLM ICL via CSCG. For example, learning and architecture, training task, test task...\n\nI confess I do not understand CSCG fully. It might be a super smart and cool idea with a lot of potential, and I'm willing to raise my score these are addressed.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Indeed, it is the combination of context-dependent latent representation and transitive generalization capability that drives the power of CSCGs. Since the thrust of this paper is on in-context learning behavior, page constraints limit us from elaborating on CSCG details. We are happy to add to the appendix a section elaborating on these details, if that might be helpful. Towards the end of this response, we try to clarify the aspects you have asked about. \n\nOur emphasis in this paper has been on leveraging the interpretability of the CSCG model for elucidating in-context learning behavior. While we do not directly interpret ICL behavior in LLM models, we reason by analogy, and hope that this paper will spur more work in that direction. A mechanistic understanding of ICL in transformers is still a work in progress (with several exciting works that we have referenced). We hope also that the interpretability of CSCGs will motivate such an emphasis in LLMs (and sequence models more generally) as they proliferate in application. We will emphasize these aspects in the discussion section. \n\n\n> Also, why not several latent tokens between observation tokens? Might make symbol space exponentially more efficient and more like language compositionality?\n\nYou seem to be suggesting having a CSCG with factorized latent space: this is a relevant modification to the model. We agree that such a structure can allow better compositionality while enabling scalability. We hope to explore such modifications in our future research.\n\n\n> Figure 3, is accuracy kind of low? Does it really match Transformer models?\n\nOn the GINC dataset, CSCG achieves an in-context accuracy above 90% for context-length of 5 and above 95% for context lengths of 8 or 10. \nThese results are higher than the ones reported for Transformers in [2]-which introduced the GINC dataset–and comparable to the best LSTMs reported by the authors. We acknowledge that further tuning might boost transformers’ performance; we are only conveying that these high accuracies illustrate that CSCGs are good at context-specific retrieval (which is essentially what the GINC dataset is probing).\n\n\n> Is this whole CSCG thing kinda like a latent automation? Do you think LLM is possible to recover such symbolic structures?\n\nYou are right, one could think of the action-conditioned transition matrix as a stochastic generalization of a finite deterministic automaton in the latent space. \nSome recent work [4] has demonstrated how LLMs might be able to learn such automata.  The advantage of CSCGs is how easily interpretable the learned schemas are, as graphs. Refer the second figure in our supplemental PDF to the global response, for an example.\n\n\nCSCG details:\n* Train & test data: For the GINC dataset, we use the same training and in-context test set as [2], which is available on GitHub [3]. For the LIALT dataset, we describe the structure of the train and test set (with examples) in fig. 4A, and lines 174-184 and lines 185-193 of the text. One can think of our training sets as analogous to the pretraining datasets used in LLMs, and of our in-context test set as the prompts used to test LLMs for in-context learning.\n* Learning of the transition matrix: Learning in CSCGs uses the Expectation-Maximization algorithm to maximize the log-likelihood of the CSCG on a sequence of symbols. It is therefore similar to HMMs training, with additional computational benefits due to the clone structure. Note that during CSCG training, the emission matrix (which specifies the latent spaces associated with an observation) is fixed, and only the transition matrix is learned. The original CSCG paper [1] has more details.\n* Freezing of the transition matrix: The transition matrix T is only learned during the “training” phase, which we describe in line 155 (for the GINC dataset) and in line 194 (for the LIALT dataset). We demonstrate that ICL behavior in CSCGs,  which emulate  “algorithms” on novel tokens, can be implemented by (a) freezing the learned transition matrix (the schemas, corresponding to algorithms) at test time, and (b) only modifying the emission matrix through the rebinding algorithm.\n* Latent space: “Clones” correspond to different latent states which emit the same observation. For a given dataset, the latent space that needs to be modeled is at least as big as the number N of distinct contexts necessary to correctly predict the next token.  In line 196, we introduce an “overallocation ratio” to parameterize the CSCG capacity proportionally to N. In practice, our latent space (including the overallocation) contains a few thousand states.\n\n\nReferences:\n\n[1] Dileep George et al. “Clone-structured graph representations enable flexible learning and vicarious evaluation of cognitive maps”. In: Nature communications 12.1 (2021), p. 2392.\n\n[2] Sang Michael Xie et al. “An explanation of in-context learning as implicit bayesian inference”. In: arXiv preprint arXiv:2111.02080 (2021)\n\n[3] https://github.com/p-lambda/incontext-learning/blob/main/data/GINC_trans0.1_start10.0_nsymbols50_nvalues10_nslots10_vic0.9_nhmms10\n\n[4] Bingbin Liu et al. “Transformers learn shortcuts to automata”. In: arXiv preprint arXiv:2210.10749 (2022)\n"}, {"review_id": "cRojmev9uk", "reviewer": "Reviewer_kc5K", "summary": "This paper implements Cloned Structured Causal Graphs (CSCG), a model that was previously used in Neuroscience to explain Hippocampal cognitive maps, on language tasks that require in-context learning (ICL). They show the success of CSCG's across a variety of language benchmarks and come up with a theory on the mechanisms behind ICL and validate this theory using the results from ICL in CSCG's. The properties they claim are necessary for ICL include context-based separation, context-based merging for transitive generalization, and the presence of general abstract schema that can implement certain abstract content-independent operations that can also be rebound as necessary. Targeted experiments are conducted to illustrate each of these properties with a handful of relevant benchmarks, both old and new - GINC (old), LIALT (new), and PreCo (old). ", "questions": "What was the process of finding and identifying the schemas shown in Figure 2? Is there a principled way to detect abstract generalizable schemas learned during training? ", "limitations": "I don't think the authors have included explicit sections discussing the limitations of using CSCG's in language modeling or on the broader societal impacts of this work. Would appreciate including this in a rebuttal, possibly with the extra page that is given to authors during the rebuttal phase. ", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 4, "strengths": "* I have not seen CSCG's being used for the language domain, so I think that is a novel contribution inofitself. \n\n* The properties that explain the mechanisms of ICL are very interesting and, although some I think can only be uniquely revealed through the architecture of CSCG's, seem like architecture-agnostic generalizable principles that add to the interpretability literature on ICL. \n\n* Experiments are targeted, well-organized, and thorough. The contribution of the LIALT dataset also can be a valuable test-bed for ICL abilities. ", "weaknesses": "The paper makes some (though I think reasonable) conjectures on how these properties are implemented in transformers. However, there isn't any empirical evidence to verify these claims. I don't think this is a strong weakness, though, because the empirical evidence from CSCG is a valid contribution by itself and this can be the topic of future work. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We are heartened by your emphasis on the strengths of the CSCG approach and the novelty of its application in this setting. We are similarly excited to leverage the interpretability of CSCGs to elucidate by analogy a general framework for in-context learning behavior. You might also find interesting the additional figures in our one-page PDF supplement to the global response, illustrating the process of schema retrieval and rebinding. \n\n> What was the process of finding and identifying the schemas shown in Figure 2? \n\nFig. 2 aims at giving a high-level intuition of CSCG circuits to the reader: we selected simple sequence-to-sequence algorithms and manually designed the schemas. In Figs. 4 and 9, we show that we are able to extract similar schemas from a trained CSCG.\n\n> Is there a principled way to detect abstract generalizable schemas learned during training?\n\nEach schema corresponds to a cluster of latent states with a specific connectivity pattern. Performing community detection on the learned CSCG graph is one simple approach to discovering schemas.\n\n>  I don't think the authors have included explicit sections discussing the limitations of using CSCG's in language modeling or on the broader societal impacts of this work. Would appreciate including this in a rebuttal, possibly with the extra page that is given to authors during the rebuttal phase.\n\nWe appreciate that you consider this of significant social impact to warrant discussion. In addition to elucidating the mechanics of ICL, we hope also that it serves as an exemplar in motivating the pursuit of interpretable models, especially as LLMs proliferate in application. Such interpretability will plausibly help tame biases and enforce guardrails for safety. We shall emphasize this, and the challenges of CSCGs for language modeling, in the discussion section. "}, {"review_id": "aPc1zINY8f", "reviewer": "Reviewer_2Tdj", "summary": "The authors attempt to analyze the reason for the success of In-context learning (ICL) for few-shot learning regimes. They apply ICL to clone-structured causal graphs (CSCGs) that can be used to interpret how ICL works in LLM. The CSCG is constructed as causal graph model where for the task of next token sequence prediction. The authors use CSCGs as interpertable language models to which they apply ICL to analyze key-properties of the learning mechanism. The authors go on to show how CSCG can perform similarly to ICL tasks designed for Transformer models and then show similar properties of CSCG to Transformer models (such as over-parameterization improves performance). The insights can help design new models. ", "questions": "Please see weaknesses. ", "limitations": "None", "rating": 7, "confidence": 3, "soundness": 4, "presentation": 3, "contribution": 3, "strengths": "1. The authors use CSCG an interpretable model to study a learning method (ICL) that is of high importance to the research community. \n2. The experiments are thorough, and the authors show interesting capabilities of CSCG that they compare to a Transformer model.\n", "weaknesses": "1. It is a difficult paper to evaluate as the authors evaluate CSCG ability on standard ICL benchmarks such as 4.3 Dax test and the use of CSCG as a proxy to study ICL, e.g. 4.2 \"Emergence\". Expanding on my point, the authors choose to evaluate CSCG as a direct replacement to LLM (which does not appear to be either their premise or motivation of the work). \n2. Related to 1.; the results and figures, lack conclusions. For example, I would expect to know how / what the experiment explains for how ICL works. An example where this is done successfully is Figure 5. `In-context accuracy (with standard errors) per task on the two LIALT test sets: for each task, overparametrization improves performance.`, but is not present in other figures or sections. \n3. The jump on conclusions from a Transformer model to a CSCG is quite large. As such it appears the authors are under-delivering on their original premise, while the paper is of great interest without making the reference. ", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate the review for distilling the essential concepts we have tried to convey in the paper. Our goal here is indeed to elucidate a general framework for in-context learning behavior, leveraging the interpretability of CSCGs. A mechanistic understanding of ICL in transformers is still a work in progress (with several exciting works that we have referenced); we hope that this paper will spur further work in that direction. We shall emphasize this, and the current limitations of CSCGs as a direct alternative to LLMs, in the discussion section. \n\n> I would expect to know how / what the experiment explains for how ICL works. An example where this is done successfully is Figure 5, but is not present in other figures or sections.\n\nThanks for the comment regarding clarity; we will improve the figure captions to make this more explicit. \n\n* In the GINC experiment, Section 4.1, we show that ICL is driven by (a) learning ​​a transition graph where the hidden variables are clustered by concepts (see Fig 3A) and (b) successful retrieval of the shared latent concept between the prompt and the model. In the first figure in our one-page PDF supplement to the global response, we illustrate the schema retrieval process, and the role played by increasing prompt length.\n\n* In the LIALT experiment, Section 4.2, we show that ICL requires — in addition to (a) and (b) — also (c) rebinding of novel tokens not seen during training to appropriate slots. Both Figs. 3B and 5[left] additionally show that model size drives retrieval and ICL capabilities. The second figure in our one-page supplement to the global response illustrates the schema retrieval and rebinding process, again leveraging the interpretability of CSCGs.\n"}]

Author rebuttal:
We would like to thank the reviewers for their time and thoughtful comments. As the reviews have identified, we have used CSCGs as a sequence model and leveraged their interpretability to deconstruct in-context learning (ICL) behavior into a combination of schema-learning (at training time) and schema-retrieval + rebinding (at prompting time). We hope our work demystifies by analogy the surprising ICL behavior observed in LLMs, showcases avenues for further research on ICL capabilities, and provides impetus for interpretable methods.

In our one page supplement, we present two additional figures that illustrate the mechanics of the schema retrieval and rebinding process, and also highlight the interpretability of CSCGs. We shall make use of the additional page for the camera-ready version of the paper to incorporate these.


Meta-review: This work shows how in-context learning can emerge in clone-structured causal graphs (CSCGs).  Reviewers observed that the conjectured connections with transformers (e.g. Discussion, Section 6) were not empirically supported, and that the claims of greater interpretability were not supported.  Despite these issues, reviewers rated the paper highly, and universally favored acceptance.  

I recommend acceptance, and I also believe that Section 6 adds to the paper, as it highlights an important advantage this work has over several previously proffered explanations of in-context learning.  However, this section can be improved in clarity, and I believe the authors should clarify how CSCGs relate to induction heads, which they note are another explanation for in-context learning, but do not discuss in any depth.  Furthermore, the claims around interpretability should be tempered, and practical limitations of the work should be noted, in line with the comments of Reviewer PPgs.
Decision: Accept (spotlight)