Dual Process Learning: Controlling Use of In-Context vs. In-Weights Strategies with Weight Forgetting

Authors: Suraj Anand, Michael A. Lepori, Jack Merullo, Ellie Pavlick
Primary area: interpretability and explainable AI
Keywords: Natural Language Processing, In-Context Learning, In-Weights Learning, Active Forgetting

Abstract:
Language models have the ability to perform in-context learning (ICL), allowing them to flexibly adapt their behavior based on context. This contrasts with in-weights learning (IWL), where memorized information is  encoded in model parameters after iterated observations of data. An ideal model should be able to flexibly deploy both of these abilities. Despite their apparent ability to learn in-context, language models are known to struggle when faced with unseen or rarely seen tokens  (Land & Bartolo, 2024). Hence, we study $\textbf{structural in-context learning}$, which we define as the ability of a model to execute in-context learning on arbitrary novel tokens -- so called because the model must generalize on the basis of e.g. sentence structure or task structure, rather than content encoded in token embeddings. We study structural in-context algorithms on both synthetic and naturalistic tasks using toy models, masked language models, and autoregressive language models. We find that structural ICL appears before quickly disappearing early in LM pretraining. While it has been shown that ICL can diminish during training (Singh et al., 2023), we find that prior work does not account for structural ICL. Building on Chen et al. (2024) 's active forgetting method, we introduce pretraining and finetuning methods that can modulate the preference for structural ICL and IWL. Importantly, this allows us to induce a $\textit{dual process strategy}$ where in-context and in-weights solutions coexist within a single model.

Decision: accept (poster)
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730076088446

thread_type:
official_review

<thread_object>
creator: reviewer_zyzg

<content>
summary:
This paper studies the limitations and corresponding solutions for in-context learning with unseen tokens (i.e., structural ICL). Under this scenario, language models (LMs) learn the sentence structure (i.e., simple syntax structure) instead of token semantic meanings. Thus, with carefully designed forgetting strategies, the authors show that they can let LMs perform well under the structural ICL setting.

strengths:
1.  This paper investigates an interesting research question: comparing ICL with parameter-tuning. Some explorations are helpful and provide useful insights. 
2.  Besides, the idea of combining ICL with structures is inspiring. Most previous ICL works run experiments on the sentence classification datasets, which are highly semantic-related. The discussion about sentence structures instead of prompt structures are limited.

weaknesses:
1. Some statements about ICL explanation are not up-to-date. The authors mentioned that typical ICL algorithms do not have stable learning manner and named it as conditional ICL. However, some recent works [1,2] have shown that ICL learns the task composition: the model only learns how to compose the task based on the learned tasks from pretraining, instead of learning from scratch based on the ICL examples. Thus, some claims in paragraph 2 of the introduction should be updated.

2. Experiment settings are a bit problematic. This work is based on previous findings. However, these findings are not general across all settings. For example, Singh et al., 2023 find that ICL slowly dissipates as models are overtrained ONLY for decoder-based LMs without pretraining. In their limitation section, they also mentioned that the conclusion may not hold for pretrained LMs (e.g., Multi-BERT here) under other scenarios. I can’t list all of them, but in general, what I mean is that we should be careful about the assumptions of previous work, and cannot simply generalize and rely on previous conclusions.

    Second, since this paper uses encoder-based LMs, it may not be able to perform ICL as there is little discussion about it. Besides, ICL ability is often used for large models. Implementing experiments on Multi-BERT with probing accuracy instead of generation is not convincing to show the conclusions of this paper is useful in practice.

3. The sentence structure is learned from pretraining instead of predefined. Thus, the model may not treat sentences as human do (i.e., with linguistic structures). Even if for the simple POS triple structure it doesn’t matter, other more general or complex structures may not work.

4. The potential impact of this work may be a bit limited. The motivation of injecting novel tokens under the ICL setting is not clear to me. As aforementioned, ICL is often used for large models, where their tokenizer vocabulary is large enough and it’s not necessary to add extra new tokens (e.g., LLaMA 3-8B). While for small models that may require to extend the tokenizer (frequently), it may not perform ICL well and they are not expected to be used under this scenario. 

[1] Hahn M,; Goyal N. A theory of emergent in-context learning as implicit structure induction[J]. arXiv preprint arXiv:2303.07971, 2023.

[2] Li, J.; Hou, Y.; Sachan, M.; and Cotterell, R. What Do Language Models Learn in Context? The Structured Task Hypothesis. ACL 2024

questions:
See weakness
</content>
</thread_object>

</thread>
<thread>
creation date:
1730546786508

thread_type:
official_review

<thread_object>
creator: reviewer_e33f

<content>
summary:
The paper introduces a distinction between *structural* and *conditional* in-context learning, where structural ICL can generalise to novel tokens because it does not depend on token embeddings. They show that by default, structural ICL emerges and then sharply diminishes during pretraining, and they introduce a novel pretraining method based on periodically reinitialising the embedding matrix, which prevents the loss of structural ICL.

strengths:
- The central thesis is clear
- The contributions are well-grounded in prior literature
- The distinction between structural and conditional ICL is compelling
- The experimental results support the thesis well, and are quite extensive, especially on the dynamics of structural versus conditional ICL

weaknesses:
I feel there are two main weaknesses.

Firstly, it is not clear to me how important structural ICL actually is as an ability for models
  - The connection to glitch tokens is fair, but I don't expect these to be very common, and even when they do occur I'm not sure that preserving structural ICL is the best response - for instance, it wouldn't have fixed all the canonical solidgoldmagikarp problems, because structural ICL only preserves syntactic knowledge of glitch tokens, but fundamentally cannot give semantic knowledge of unknown tokens
  - The application to very rare tokens is likewise fair but by definition not very common, and again it is unclear to me whether in these cases there are any important capabilities which are purely syntactic

Secondly, it is not clear to me how far active/temporary forgetting impairs other abilities in models
- My intuition is that repeatedly reinitialising the embedding matrix is encouraging structural ICL by damaging other kinds of learning
- While the temporary forgetting approach might ameliorate this, it is not clear to me by how much or in what way it does this


Overall, I find much of the paper compelling in its presentation and exploration of structural ICL as a phenomenon, but I would be even more enthusiastic about it if it were clearer to me:
- that structural ICL is an important property (which could be demonstrated by natural examples of its relevance)
- what the side-effects are of the proposed methods for preserving structural ICL

questions:
**Q1**

To what extent does active forgetting impair other model capabilities, if at all? What about temporary forgetting? I'd be interested to hear hypotheses, and particularly any experimental results you have. 

**Q2**

Do you believe the problem of glitch tokens will recur in future, and in such a case, do you think preservation of structural ICL would be a useful strategy?

**Q3**

Can you give an example (or ideally a few examples) of cases where models might have a hard time performing a natural task because they are incapable of performing structural ICL on a rare token? 

---

I have provisionally put down a marginal acceptance, but would be happy to revise my score upwards if these questions are appropriately addressed.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730668568497

thread_type:
official_review

<thread_object>
creator: reviewer_pxza

<content>
summary:
This paper examines "structural" in-context learning (ICL), a specific type of ICL where some tokens in the model’s prompt are randomly initialized. The motivation for studying structural ICL is based on prior observations that ICL behaves unexpectedly when prompts include tokens that are rarely encountered during training, i.e., undertrained tokens. This paper uses synthetic tasks to demonstrate that structural ICL develops early in training but diminishes as training progresses. The authors show that periodically reinitializing token embeddings during training (active forgetting) preserves structural ICL abilities, albeit at the expense of memorizing information in the model’s embeddings. The paper further demonstrates that both structural ICL and memorization are maintained if active forgetting is employed, but only for the first N steps of training. This procedure is referred to as “temporary forgetting.”

strengths:
1. I appreciated that this work distinguishes structural ICL from ICL in general. This seems like a natural and interesting extension to prior work that studies in-weights versus in-context learning.
2. Despite the lack of experiments on downstream tasks, I found the breadth of the experiments satisfactory. I also recognize the creativity employed by the authors given their compute constraints.

weaknesses:
1. No experiments on downstream tasks. While the motivation to address the issue of undertrained tokens is interesting, it is difficult to assess how training with temporary forgetting will affect language models in real-world applications. Given the author’s limited computational resources, they might consider a variant of temporary forgetting that can be employed as a fine-tuning step, e.g., fine-tune with active forgetting for $k$ steps before resetting the token embeddings to their original states, leaving only the model parameters changed.
2. Limited technical novelty. Unless I misunderstood the paper, temporary forgetting is simply active forgetting applied for $K$ steps. This is only a minor weakness in my view and is excusable if the authors could find a way to perform experiments on downstream tasks.
3. The presentation can be improved. For example, there are instances where the figure legends contain terms that are not defined. In Figure 6 (right), “Vanilla Forgetting” is mentioned. Figure 11 uses the term “Stop Forgetting,” which I believe is intended to mean “Temporary Forgetting.”

questions:
1. What are the limitations of temporary forgetting? Is it strictly beneficial, or did the authors find that some model abilities deteriorated? Given that evaluation wasn’t performed on downstream tasks, what synthetic experiments can be conducted to evaluate the trade-offs of temporary forgetting?
2. Suggestion: Since no evaluation was performed on downstream tasks, this paper could be improved by expanding on why studying undertrained tokens is important (beyond what is already discussed in the introduction). I am familiar with this area, so the introduction was adequate for me, but adding such a section might make this paper more broadly accessible.
</content>
</thread_object>

</thread>
<thread>
creation date:
1730743988699

thread_type:
official_review

<thread_object>
creator: reviewer_szdj

<content>
summary:
This paper studies the emergence and disappearance of in-context vs in-weights learning in masked language models; in particular, a kind of ICL termed “structural ICL”, which refers to the ability to infer (and then use) information about the structural role (e.g., POS) of words from the context. In experiments, this ability appears and then disappears over the course of MLM training. A recently proposed “active forgetting” approach helps keep the ability, leading to models that use both in-context and in-weights learning.

strengths:
- The paper studies a question of substantial interest, nicely continuing prior work on different kinds of learning
- considers both real-world LM and toy models
- Provides novel insights by distinguishing "structural" from "conditional" ICL
- Provides method for mitigating the loss of structural ICL in a toy setup

weaknesses:
- Structural ICL is operationalized using performance on random token embeddings. However, as the paper also demonstrates (Figure 18), such embeddings are drastically OOD for a trained model. Wouldn't a fairer test use novel embeddings that match the distributional properties of the embeddings of the trained model?
- Results on Structural ICL in MultiBERT (Section 3) are based on POS = noun vs adjective as the only property. It remains unclear how robust findings are to other properties, such as other pairs of POS, or other properties, such as grammatical number.
- Results on Structural ICL in MultiBERT (Section 3) are based only on probing accuracy, not behavior of the model. This seems suboptimal given that in-context learning as an emergent property of LLMs is generally thought of as an ability appearing in next-token prediction (e.g., complete a prompt with an appropriate label), One could test for ICL in properties such as POS behaviorally by creating contrastive examples where the key word has different POS and comparing model probabilities.
- The term “Dual Process” here is used to refer to models using both in-weights and in-context learning, and the paper refers to Kahneman 2011; Miller 2000. The link to dual process theory appears tenuous: dual process theory refers to thse use of implicit, unconscious vs explicit, conscious processes; the link to in-weights vs in-context inference is at best hand-waved in line 90-91, but no explicit justification for the link is given. The choice of the term and the link to dual process theory thus appears a bit of a stretch to apply the in this context.
- It remains unclear if the proposed strategy for keeping structural ICL abilities would transfer to real-world LMs. While training an LLM is understandably out of scope, even a modest LM could provide substantial insight here.

questions:
See weaknesses.

Minor note: Line 317-8: it seems the sentence is missing a verb
</content>
</thread_object>

</thread>
<thread>
creation date:
1731996954788

thread_type:
official_comment

<thread_object>
creator: authors

<content>
title:
Main Response

comment:
# Main Response:

We thank the reviewers for their detailed feedback and questions. Their comments and suggestions have helped us to make significant improvements to the work! Below are responses to common questions.

## Motivating Examples/Importance

Below, we provide a few instances where models have difficulty because they cannot perform structural ICL on a rare token. These are mainly symbolic reasoning tasks, which are important for widespread use cases, including those enumerated below.
1. **Abstract reasoning tasks:**
Dasgupta et al., 2022 evaluate large language models on abstract reasoning tasks and find many imperfections, which according to them “is frequently highlighted as a crucial missing component of current AI.” For instance, they test Natural Language Inference (NLI) on nonsense words (e.g. If vuffs are bigger than feps, then feps are smaller than vuffs) and find that models cannot accurately predict statement validity above chance (Dasgupta et al., 2022 - Figure 3). While nonsense words are not natural, we would often want to introduce words unfamiliar to the model such as names, places, events, etc and would like to preserve logical reasoning. Moreover, it is highly likely that additional words become a part of the common lexicon as language is continually evolving while the training data has a static cutoff.

2. **Low-resource languages:** 
Large language models perform poorly on in-context learning tasks involving low-resource languages, such as Swahili. Various works explore X-ICL (cross-lingual in-context learning) methods, which incorporate source language exemplars into context (Winata et al., 2021b; Shi et al., 2023). However, these still fail to compete with a simple translate-test, which translates the query and then conducts zero-shot inference in a high-resource language (Cahyawijaya et al., 2024). These languages have much difficulty as tokens are quite rare in pretraining. Structural ICL would enable much higher ICL performance. 

3. **Code Generation:** 
Models are often used for code generation, where function and variable names are often rare or unique (but properties should be derivable based on the context). Wen et al., 2024 analyzed 14 LLMs and found ~22% of code generation errors attributable to NameErrors. Structural ICL would ensure robustness to naming in code generation, which would likely mitigate many of these errors.

We have modified the manuscript to briefly mention these examples in lines 516-517 and will expand on these examples in the camera-ready version.


## Reproducing Structural ICL Transience on Generation Task

Multiple reviewers had concerns about the generalizability of our MultiBERT experiments, asking us to examine behavior. We designed another task to do just this: 

Our task requires abstract reasoning on untrained tokens in a decoder-only transformer. We give the model the syllogism template 

<X> is <Y>. 

<Y> is <Z>. 

Therefore, <X> is __

The correct answer is <Z>. We examine accuracy, which we find by looking at the probabilities of predicting <Z> versus <Y>. We test baseline performance over training steps where <X>, <Y>, <Z> are chosen from the set of tokens representing A-Z, and we test unseen token performance by replacing <X> with an unseen token. We test on checkpoints released from Pythia-1.4B (Biderman et al., 2023), a decoder-only language model, and reproduce the trend of Structural ICL emerging and then quickly disappearing (as seen by unseen token performance). Additional details about the setup and results are located in Appendix A.3 of the revised manuscript.


## Probabilistic Temporary Forgetting during Fine Tuning of GPT-2

Most reviewers deemed that it would be very valuable to test how temporary forgetting might work in a natural language model setting. We strongly agree, and thus designed an experiment on GPT-2 large, a decoder-only language model using the same syllogism task. 

In our experiment, we finetune GPT-2 large (Radford et al., 2019) on wikitext sentences for 2000 steps. To accommodate the fine-tuning setting, we use a probabilistic variant of temporary forgetting: every step, we replace tokens in the batch with p=0.10 with randomly initialized embeddings. After the step, we set the embedding matrix back to its original values, hence maintaining the character of temporary forgetting.

We find that syllogism accuracy with unseen tokens jumps from 0.02 to 0.927 while the baseline syllogism accuracy remains the same. This jump does not occur with vanilla fine-tuning.

We believe that this experiment shows that temporary forgetting is useful in real language modeling settings on downstream tasks. Additional details are included in Appendix A.3.3, and we plan on polishing these experiments for the camera-ready version. Moreover, we plan to put our fine-tuning experiment in the main body of the paper. Please let us know if you have any questions!

Once again, we sincerely appreciate your feedback!
</content>
</thread_object>

</thread>


Metareview:
metareview:
The authors continue the thread of existing work on ICL and token-sensitivity to define "structural in-context learning". They show that this has a different pattern of emergence than general ICL and suggest a pre-training method to improve it. Issues raised about applicability to different models and real-world relevance are discussed.

The paper is relevant and the experiments solid if not revolutionary, recommend accept as poster.