SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Training dataset: To test the ability of CSCG to learn algorithms 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. 4A[top-left]. A demonstration consists of a multi-word language instruction—each algorithm has five different instructions—followed by 10 input-output examples of that algorithm. See Tables 2 & 3 in Appendix D.1 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 “/” delimiter. Demonstrations are separated by a “\n” delimiter. The input lists and matrices values are created by uniformly sampling from a vocabulary of 676 tokens, created by random pairings of uppercase letters. List operations examples vary in length from 3 to 6, and the matrix operations are of sizes  $2 \times 2$  or  $3 \times 3$ . Fig. 4A [bottom-left] shows the training data format.
A Language Instructed Algorithm Learning Tasks
Figure 4: A. [Top-left] List and matrix algorithms used in the LIALT dataset. Format of the training set [bottom-left] and examples of the two LIALT test sets [right]. B. Example of a learned circuit for the "reverse" algorithm, displayed by stacking clones [left] or unrolling them [right].
Figure 5: [Left] In-context accuracy (ICA) with  $95\%$  CIs after a single EM iteration, as a 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-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, overparametrization 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 11 in the Appendix visualizes the same quantities after EM convergence; the similarity demonstrates that the fast rebinding algorithm is not just localized in its updates, but also rapid.
Test dataset: LIALT has two test datasets, respectively containing: (a) instruction-based 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 example-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 to 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, a CSCG allocates its number of clones proportionally to the number of distinct contexts in the training data in which it occurs<sup>2</sup>. 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  $\epsilon = 10^{-6}$ . 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 to the next  $\frac{1}{r}$  delimeter.
Figure 6: A. Transition graph of the CSCG model learned on the LIALT dataset, visualized using the Kamada-Kawai algorithm. B. 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. 4B shows the extracted circuit for the list reversal algorithm. Fig. 5[left] presents the in-context accuracy of CSCGs (using  $\epsilon = 10^{-6}$  and  $p_{\text{surprise}} = 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 a single input-output pair for example-based 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. 6A shows the transition graph using the Kamada-Kawai algorithm [24]. It reveals thirteen loosely connected clusters corresponding to the thirteen algorithms present in the LIALT dataset. Fig. 6B 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 to schemas compatible with the prompt structure—in this case based on brackets & delimiters. However, the structure is insufficient to 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. In 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—and complete the prompt appropriately. Note that a single EM step is sufficient to derive the correct rebinding in these examples. Compare Figs. 5 & 11, and the
tables in Appendix Sec. D.3 for how the in-context 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 on long-term coherence (sufficient for the GINC results in Section 4.1) does not explain the remapping of a latent representation to completely new tokens as required for generalizing on the LIALT algorithms. Without rebinding, even a 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. 6B, 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. 6B, second row). CSCGs are thus able to learn seq2seq algorithms and generalize those algorithms to novel tokens using rebinding.
Emergence: ICL performance of CSCG on the LIALT dataset shows characteristics attributed to emergence. In-context accuracy has a 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 a 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 a 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 1 with  $\epsilon = 10^{-6}$  and  $p_{\text{surprise}} = \frac{1}{16}$  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. 7 shows these results.
Figure 7: Examples of the dax test on a CSCG trained on the PreCo dataset. In each row, the novel word in red (e.g. "terras") is absorbed by binding it to the clones of the corresponding word in blue (e.g. "planets"). The CSCG can then use the new token in similar contexts, as demonstrated by the fill-in-the-blanks probing.
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 a pre-trained model to learn a new task given only a few examples. [29, 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 be 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 a mapping that minimizes a loss on the training data: gradient methods are a popular paradigm [31, 32, 33]. In FSL, a model learns to rapidly adapt to a new task from a limited number of supervised examples [34, 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-learning paradigm aims at learning to adapt to a new task with only a few examples [39, 40, 41] by using multiple learning experiences. In contrast, ICL directly emerges from the pre-trained 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 new tasks at inference time.
How ICL works: [3] explained ICL as implicit Bayesian inference and constructed the GINC dataset (see Section 4.1) for demonstrating ICL. [44] abstracted ICL as an algorithm learning problem and found that a transformer can implicitly infer a hypothesis function. Similarly, [45] showed that a transformer can be trained to perform ICL of unseen linear functions, with performance
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 a 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 meta-gradients. Finally, [5] showed the existence of "induction heads" in transformers, that emerge during training, copy previous patterns, and drive ICL capacities.
What influences ICL: [1, 47] indicated that LLMs' ICL performance "emerges" then keeps improving when the model size increases. [48] proposed a 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 ICL emerge when a model is trained on a combination of multiple corpora, and that low perplexity and ICL performance do not always correlate. [51, 52] 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.
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, 4] do not distinguish the role of sequence positions vis-a-vis contents; we argue that theories might need to emphasize this distinction (see Fig. 8A) to fully understand the inductive biases behind ICL. We conjecture (see Fig. 8B) that layers of the transformer implement multiple mixed templates of positions and content, evaluated at different offsets of a 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.
Figure 8: A. Learned templates in a transformer could involve content, position, or a mix of both. B. Activations in the forward pass of a 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 a factorized latent space and adding skip connections would also allow compositionality while enabling scalability. Further, while we have illustrated here the concept of rebinding to attach new symbols to 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 a 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.
We thank Stephanie Chan, Andrew Lampinen, Anirudh Goyal, Dharshan Kumaran, Neel Nanda and Guangyao Zhou for helpful discussions and comments on the draft.
[49] Stephanie Chan et al. "Data distributional properties drive emergent in-context learning in transformers". In: Advances in Neural Information Processing Systems 35 (2022), pp. 18878-18891.
[50] Seongjin Shin et al. "On the effect of pretraining corpora on in-context learning by a large-scale language model". In: arXiv preprint arXiv: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 et al. "Fantastically ordered prompts and where to find them: Overcoming few-shot prompt order sensitivity". In: arXiv preprint arXiv:2104.08786 (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.06804 (2021).
[55] Andrew K Lampinen et al. "Can language models learn from explanations in context?" In: arXiv preprint arXiv:2204.02329 (2022).
[56] Ryan Schaeffer, Brando Miranda, and Sanmi Koyejo. "Are Emergent Abilities of Large Language Models a Mirage?" In: arXiv preprint arXiv:2304.15004 (2023).
A Locally updating the emission matrix with the transition matrix fixed
We reuse the same notations as in [15], Appendix A.2. The authors describe the EM algorithm for learning the emission matrix of a CSCG with a fixed transition matrix. In particular their M step defines the new emission matrix as:
$$
E (j) = \sum_ {n = 1} ^ {N} 1 _ {X _ {n} = j} \gamma (n) \oslash \sum_ {n = 1} ^ {N} \gamma (n), \forall j
$$
where  $E(j)$  is a column of the emission matrix corresponding to the emission  $j$ ,  $1_{X_n = j}$  is an indicator function,  $\varnothing$  is the element-wise division and  $\gamma (n)$  is derived by the authors from the forward and backward probabilities. The  $i$  entry of the vector  $E(j)$  is then defined as:
$$
E _ {i j} = \frac {\sum_ {n = 1} ^ {N} 1 _ {X _ {n} = j} \gamma_ {i} (n)}{\sum_ {n = 1} ^ {N} \gamma_ {i} (n)}.
$$
In contrast, in Section 2.2, Step 5 of Algorithm 1 only updates the row  $i$  for which we can find a pair  $(i,n) \in \mathcal{R}$ , by only using the beliefs at timestep  $n$ . For this row, the  $j$ th entry becomes:
$$
E _ {i j} = \frac {\sum_ {n : (i , n) \in \mathcal {R}} 1 _ {X _ {n} = j} \gamma_ {i} (n)}{\sum_ {n : (i , n) \in \mathcal {R}} \gamma_ {i} (n)}.
$$
The pseudocount used in Step 1 of Algorithm 1 is an uncertainty parameter that lets the model smooth over incorrect observations. More details of this parameter are available in [6].
B Prompt completion algorithm
Algorithm 2 describes the prompt completion algorithm introduced in Section 2.2. It implicitly considers a single action, which takes the next sequence element.
Algorithm 2 - Prompt completion
Input: Grounded schema  $\{T, \mathcal{C}, E^{\mathrm{rb}}\}$  with rebound CSCG emission matrix  $E^{\mathrm{rb}}$ , delimiter token  $x_{\emptyset}$ , prompt  $x^{(\mathrm{prompt})} = (x_1, \ldots, x_N)$
Output: A completed prompt  $x^{(\text{completed})} = (x_1, \ldots, x_N, x_{N+1}, \ldots, x_{N+P} = x_{\emptyset})$
1: Run max-product for MAP inference and return  $z^{\mathrm{MAP}} = (z_1, \ldots, z_N) = \arg \max_z P(z|x^{(\mathrm{prompt})})$ .
2: Set  $\ell = 0$ . While  $x_{N + \ell} \neq x_{\emptyset}$ , increment  $\ell \gets \ell + 1$  and sample the next most likely observation:  $z_{N + \ell} \in \operatorname{argmax}_j T_{z_{N + \ell - 1}, j}$  and  $x_{N + \ell} \in \operatorname{argmax}_j E_{z_{N + \ell}, j}^{\mathrm{rb}}$ .
C Additional materials for the GINC dataset
First, we present two additional plots for the GINC experiment.
In-context performance on the GINC dataset
Figure 9: [Left] In-context confidence for the CSCG with 50 clones on the GINC test dataset, defined as 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 the transformer and LSTM reported in [3], CSCG fails to extrapolate and has a low in-context accuracy when the test prompts are sampled from five novel concepts, unseen during training.
Second, we present the table of results associated with Fig. 3 for the CSCGs with 10 and 50 clones.
Context length	No. of examples	CSCG with 10 clones	CSCG with 50 clones
3	0	0.509 (0.020)	0.534 (0.020)
1	0.351 (0.019)	0.445 (0.019)
2	0.366 (0.019)	0.453 (0.020)
4	0.356 (0.019)	0.468 (0.020)
8	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)
5	0	0.682 (0.018)	0.927 (0.010)
1	0.640 (0.019)	0.927 (0.012)
2	0.629 (0.019)	0.904 (0.012)
4	0.654 (0.019)	0.883 (0.013)
8	0.627 (0.019)	0.894 (0.012)
16	0.637 (0.019)	0.902 (0.012)
32	0.634 (0.019)	0.901 (0.012)
64	0.637 (0.019)	0.899 (0.012)
8	0	0.696 (0.018)	0.969 (0.007)
1	0.694 (0.018)	0.972 (0.007)
2	0.686 (0.018)	0.972 (0.006)
4	0.681 (0.018)	0.978 (0.006)
8	0.690 (0.018)	0.973 (0.006)
16	0.686 (0.018)	0.975 (0.006)
32	0.676 (0.018)	0.968 (0.006)
64	0.694 (0.018)	0.975 (0.007)
10	0	0.684 (0.018)	0.975 (0.006)
1	0.705 (0.018)	0.977 (0.006)
2	0.674 (0.018)	0.971 (0.006)
4	0.713 (0.018)	0.974 (0.006)
8	0.690 (0.018)	0.977 (0.006)
16	0.689 (0.018)	0.977 (0.006)
32	0.712 (0.018)	0.978 (0.006)
64	0.690 (0.018)	0.978 (0.006)
Table 1: In-context accuracy for a CSCG with 10 clones and a CSCG 50 clones trained on the GINC dataset, averaged (with  $95\%$  confidence 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  $n = 0$  than on few-shot prompts  $n = 1,2,\ldots$ . We hypothesize that the difference between the training and the prompt distributions creates a gap that lowers few-shot in-context accuracy. The performance gap disappears
for larger contexts  $k \in \{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 element at index zero of the list'   
"print the first element from the list'   
"return the leading element from the list'   
"find the head element from the list'   
'retrieve the starting element from the list'   
"print the element at index two of the list'   
"find the third element from the list'   
"locate the third element from the list'   
"output the third item from the list'   
"return the element in third place from the list'   
"duplicate each list item'   
"replicate every element in the list'   
"make a copy of each element in the list'   
"clone each element in the list'   
"create a second instance of every element in the list'   
"print every other member in the list starting with the second member'   
"retrieve alternate items in the list starting with the second item'   
"return every other object in the list starting with the second object'   
"retrieve every other entry in the list starting with the second entry'   
"output odd indexed elements'   
"rotate the list elements one place backward'   
"move the list elements one position to the left'   
"change the items of the list one position backward"   
"displace the elements of the list one index backward"   
"roll the list items one position backward