Table of Contents
Fetching ...

Disentangling Latent Shifts of In-Context Learning with Weak Supervision

Josip Jukić, Jan Šnajder

TL;DR

This paper tackles instability and inefficiency in in-context learning as demonstrations grow. It reframes ICL as weak supervision and introduces wilda, a teacher–student framework that encodes demonstration-induced latent shifts into lightweight adapters, allowing a query-only student to reproduce teacher-like behavior. The core idea rests on a decomposition that separates zero-shot and ICL contributions, enabling adapter-based knowledge fusion and scalable handling of long contexts via adapter arithmetic. Empirically, wilda improves generalization, stability, and efficiency across ID and near-OOD tasks, often surpassing the teacher through pseudo-label correction and coverage expansion, illustrating a weak-to-strong generalization phenomenon. Collectively, wilda offers a modular, scalable approach to stable task adaptation in LLMs by treating ICL as a source of weak supervision and storing demonstrations as reusable parameter shifts.

Abstract

In-context learning (ICL) enables large language models to perform few-shot learning by conditioning on labeled examples in the prompt. Despite its flexibility, ICL suffers from instability -- especially as prompt length increases with more demonstrations. To address this, we treat ICL as a source of weak supervision and propose a parameter-efficient method that disentangles demonstration-induced latent shifts from those of the query. An ICL-based teacher generates pseudo-labels on unlabeled queries, while a student predicts them using only the query input, updating a lightweight adapter. This captures demonstration effects in a compact, reusable form, enabling efficient inference while remaining composable with new demonstrations. Although trained on noisy teacher outputs, the student often outperforms its teacher through pseudo-label correction and coverage expansion, consistent with the weak-to-strong generalization effect. Empirically, our method improves generalization, stability, and efficiency across both in-domain and out-of-domain tasks, surpassing standard ICL and prior disentanglement methods.

Disentangling Latent Shifts of In-Context Learning with Weak Supervision

TL;DR

This paper tackles instability and inefficiency in in-context learning as demonstrations grow. It reframes ICL as weak supervision and introduces wilda, a teacher–student framework that encodes demonstration-induced latent shifts into lightweight adapters, allowing a query-only student to reproduce teacher-like behavior. The core idea rests on a decomposition that separates zero-shot and ICL contributions, enabling adapter-based knowledge fusion and scalable handling of long contexts via adapter arithmetic. Empirically, wilda improves generalization, stability, and efficiency across ID and near-OOD tasks, often surpassing the teacher through pseudo-label correction and coverage expansion, illustrating a weak-to-strong generalization phenomenon. Collectively, wilda offers a modular, scalable approach to stable task adaptation in LLMs by treating ICL as a source of weak supervision and storing demonstrations as reusable parameter shifts.

Abstract

In-context learning (ICL) enables large language models to perform few-shot learning by conditioning on labeled examples in the prompt. Despite its flexibility, ICL suffers from instability -- especially as prompt length increases with more demonstrations. To address this, we treat ICL as a source of weak supervision and propose a parameter-efficient method that disentangles demonstration-induced latent shifts from those of the query. An ICL-based teacher generates pseudo-labels on unlabeled queries, while a student predicts them using only the query input, updating a lightweight adapter. This captures demonstration effects in a compact, reusable form, enabling efficient inference while remaining composable with new demonstrations. Although trained on noisy teacher outputs, the student often outperforms its teacher through pseudo-label correction and coverage expansion, consistent with the weak-to-strong generalization effect. Empirically, our method improves generalization, stability, and efficiency across both in-domain and out-of-domain tasks, surpassing standard ICL and prior disentanglement methods.
Paper Structure (39 sections, 16 equations, 2 figures, 13 tables)

This paper contains 39 sections, 16 equations, 2 figures, 13 tables.

Figures (2)

  • Figure 1: Illustration of wilda. The teacher processes a concatenation (denoted by $\mathbin{ \text{\textcircled{\scriptsize$\boldsymbol{;}$}} }$) of demonstrations $\mathbf{X}_d$, consisting of $n$ demonstrations $[\mathbf{x}_1, \mathbf{x}_2, \dots, \mathbf{x}_n]$, and the query $\mathbf{x}_q$. The student, using only the query, fine-tunes its adapter weights to produce outputs $\mathbf{y}_s$ aligned with the teacher's pseudo-labels $\mathbf{y}_t$ by minimizing the cross-entropy loss $\ell_\text{CE}$.
  • Figure 2: Empirical analysis of wilda-s on the aggregated GLUE datasets for Llama 3: (a) Histogram of approximated Lipschitz constants across datasets, computed as the Frobenius norm of the input--output Jacobian matrix; (b) Rate of pseudo-label correction over training epochs (shaded areas indicate the standard deviation over $10$ runs); Corrected and corrupted prediction rates for (c) ID examples and (d) OOD examples, based on the normalized Euclidean distance to the nearest correctly pseudo-labeled neighbor. Error bars indicate standard deviation over $10$ runs.