Training In-Context and In-Weights Mixtures Via Contrastive Context Sampling

TL;DR

This paper explores training strategies for large language models to balance in-context learning and in-weights learning by using contrastive context sampling, enhancing model flexibility and stability.

Contribution

It introduces a simple contrastive context method that enforces similarity-based contrasts to promote stable ICL-IWL mixtures during training.

Findings

Contrastive context sampling improves ICL-IWL balance.

Models trained with contrastive context maintain stable in-context and in-weights learning.

Empirical results show enhanced performance across multiple tasks and models.

Abstract

We investigate training strategies that co-develop in-context learning (ICL) and in-weights learning (IWL), and the ability to switch between them based on context relevance. Although current LLMs exhibit both modes, standard task-specific fine-tuning often erodes ICL, motivating IC-Train - fine-tuning with in-context examples. Prior work has shown that emergence of ICL after IC-Train depends on factors such as task diversity and training duration. In this paper we show that the similarity structure between target inputs and context examples also plays an important role. Random context leads to loss of ICL and IWL dominance, while only similar examples in context causes ICL to degenerate to copying labels without regard to relevance. To address this, we propose a simple Contrastive-Context which enforces two types of contrasts: (1) mix of similar and random examples within a context to evolve a correct form of ICL, and (2) varying grades of similarity across contexts to evolve ICL-IWL mixtures. We present insights on the importance of such contrast with theoretical analysis of a minimal model. We validate with extensive empirical evaluation on four LLMs and several tasks. Diagnostic probes confirm that contrasted contexts yield stable ICL-IWL mixtures, avoiding collapse into pure ICL, IWL, or copying.