Schema-learning and rebinding as mechanisms of in-context learning and emergence

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.

Strengths

1. The authors use CSCG an interpretable model to study a learning method (ICL) that is of high importance to the research community. 2. The experiments are thorough, and the authors show interesting capabilities of CSCG that they compare to a Transformer model.

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). 2. 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. 3. 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.

Questions

Please see weaknesses.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

4 excellent

Presentation

3 good

Contribution

3 good

Limitations

None

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).

Strengths

* I have not seen CSCG's being used for the language domain, so I think that is a novel contribution inofitself. * 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. * 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.

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?

Rating

8: Strong Accept: Technically strong paper, with novel ideas, excellent impact on at least one area, or high-to-excellent impact on multiple areas, with excellent evaluation, resources, and reproducibility, and no unaddressed ethical considerations.

Confidence

4: You are confident in your assessment, but not absolutely certain. It is unlikely, but not impossible, that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work.

Soundness

4 excellent

Presentation

4 excellent

Contribution

4 excellent

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.

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. CSCG 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.

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. - The context-dependent latent representation and transitive generalization make sense. - 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?? - 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... I 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.

Questions

- How is learning (of transition matrix) done? The paper seems to only talk about re-binding for learning the emission matrix... - 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?) - 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? - Figure 3, is accuracy kind of low? does it really match Transformer models? - Is this whole CSCG thing kinda like a latent automation? Do you think LLM is possible to recover such symbolic structures?

Rating

6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.

Confidence

1: Your assessment is an educated guess. The submission is not in your area or the submission was difficult to understand. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

2 fair

Contribution

3 good

Limitations

Did not see any limitation section in the paper, but there should be some (see weaknesses).

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.

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. Some 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.

Questions

How does the size of the latent graph for e.g. list reversal scale with list size? What 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? Are the results in Figure 6 cherry picked? How often is rebinding + MAP inference successful? 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.

Rating

7: Accept: Technically solid paper, with high impact on at least one sub-area, or moderate-to-high impact on more than one areas, with good-to-excellent evaluation, resources, reproducibility, and no unaddressed ethical considerations.

Confidence

3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.

Soundness

2 fair

Presentation

3 good

Contribution

3 good

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.

Author-Reviewer Discussion phase (Aug 10-16)

Today begins the Author-Reviewer Discussion phase, which lasts 1 week (**Aug 10-16**). I ask the reviewers to please **carefully read all other reviews and the author responses and (if appropriate) respond to author responses promptly.** If you've read the author response, please take the time to leave a comment, even if you have nothing to add. I also encourage both authors and reviewers to monitor OpenReview for further comments in order to enable as much back-and-forth as possible during this short period.

Thanks

I've increased my score from 5 to 6, but I hope authors work on clarity during revision to make the interesting topic more accessible and understandable!