Disentangling Latent Shifts of In-Context Learning with Weak Supervision

Thank you for your thoughtful comments and for noting both the strengths and areas where clarification is needed. Below, we offer two complementary perspectives in response to your question about unseen domains, and we provide additional context on our theoretical analysis and empirical results.

Perspective 1: Preservation of general capabilities and flexible few‑shot use

Like other parameter‑efficient fine‑tuning methods, WILDA’s adapters add only 0.1–0.3 % of the base model’s parameters. This design means that the underlying language model remains intact: during training we toggle the adapter on and off, and at inference time users can disable it to recover the original zero‑shot behavior.

Prompted by your comment and a similar observation from another reviewer, we conducted additional experiments with Llama 3 (8B) to verify that activating an adapter trained on one domain does not hinder zero‑shot performance on unrelated tasks. Using non‑overlapping GLUE dataset pairs, we observed that performance differences remain within one standard deviation (10 different seeds), confirming that WILDA preserves general zero‑shot capabilities on unrelated queries. The method thus provides a modular, task‑specific memory of demonstration effects that can be activated when relevant and safely bypassed otherwise. The results are as follows:

Moreover, WILDA does not preclude standard ICL. Once the adapter is trained, we can still supply new demonstrations in the prompt and combine them with the adapter’s latent shift. Appendix D.2 (Table 9) reports results where we encode 16 demonstrations in the adapter and then provide another 16 demonstrations in the prompt at inference. This WILDA‑S (16/16) configuration outperforms the standard 32‑shot ICL across all evaluated datasets. In other words, the student can benefit from both the encoded shift and additional demonstrations.

Perspective 2: Domain relevance and the need for task‑specific adapters

Your hypothetical scenario – training an adapter on math demonstrations and then providing physics demonstrations at test time – highlights an important limitation of any ICL method: the usefulness of demonstrations depends on their relevance to the query. In standard ICL, if you mix math and physics demonstrations, the model will mostly attend to the physics examples when answering a physics question, because those tokens more closely match the query. Similarly, a WILDA adapter trained on math will encode how the model’s activations should shift when it sees math‑related queries. Applying that adapter to a physics query does not magically inject physics knowledge; the shift is largely ignored by the model because it does not align with the query’s content. To handle physics problems, one should train a separate physics adapter (on unlabeled physics queries with corresponding demonstrations) and switch it on for physics questions. Importantly, this does not diminish the utility of WILDA for math tasks; you are free to toggle the math adapter on for math queries or turn it off entirely. This also relates to the other reviewer’s question about the impact on the base model’s general capabilities: the model effectively ignores the adapter’s contribution when the query domain does not align with the adapter’s training domain. Thus, the method remains meaningful: it provides a parametric memory of demonstration effects that can be activated when relevant and bypassed when not.

On the weak‑to‑strong analysis

We agree that a deeper theoretical connection between weak-to-strong generalization (W2S) and in-context learning would be valuable. Our paper provides an empirical analysis of W2S phenomena and draws parallels with known theory. Specifically, we approximate the Lipschitz constant of the model by the Frobenius norm of its input–output Jacobian and find that WILDA exhibits a significantly lower Lipschitz constant than standard ICL or PBFT. A low Lipschitz constant implies that the model’s predictions vary smoothly with small changes in input, a key ingredient for correcting noisy pseudo‑labels. We also track the rate of pseudo‑label correction over training epochs and show that it increases steadily, indicating that the student is refining the teacher’s labels rather than merely copying them. Finally, we discuss coverage expansion, where local corrections propagate to nearby inputs in representation space, gradually broadening the region where the model makes confident predictions. These empirical observations align with the W2S framework: local consistency facilitates pseudo‑label correction, which in turn leads to broader generalization. Formalizing this link between W2S and ICL remains an exciting direction for future work.

Thank you for your thoughtful review. We appreciate your recognition of our work’s clarity and empirical effectiveness, and we address your points below.

Computational cost of learning the adapter

We share your view that one of the major appeals of in‑context learning is its flexibility. WILDA was designed to minimize overhead while maintaining this flexibility. The adapter we learn uses LoRA modules that constitute only 0.1–0.3 % of the base model parameters: for Llama 3 (8B), this corresponds to ~21 M parameters, and only ~4.5 M for Phi‑3. Because the base model is frozen, only these additional weights are updated.

To quantify compute, we measure training cost relative to the cost of processing a 16‑shot prompt on the same backbone. For Llama 3, a 16‑shot prompt takes roughly 120x longer than a zero‑shot query because attention cost grows quadratically with prompt length. Fine‑tuning the adapter for ten epochs on 100 unlabeled queries requires about 2100 forward passes of a 16‑shot prompt. In other words, the training cost is equivalent to answering around two thousand 16‑shot queries. Once the adapter is trained, we no longer feed demonstrations at inference, so each subsequent query runs at zero‑shot speed. For workloads serving more than a few thousand queries, the one‑off cost is amortised, and the total compute is lower than repeatedly processing long prompts.

Moreover, long prompts often exceed the model’s context window or lead to unstable performance; to address this, WILDA supports partitioning the demonstration set across multiple adapters. Each adapter is trained on a manageable subset of demonstrations, and these adapters can later be fused by parameter addition. This strategy not only mitigates context‑length limitations and reduces memory overhead but also preserves stability: instead of processing a single long, order‑sensitive prompt, we store several compact “latent shift” adapters and combine them as needed, fitting within the available context window and ensuring consistent performance.

Task specificity and generalization

An adapter is by design intended to specialize the model toward a given task, and in WILDA, this is achieved by encoding the latent shift induced by task‑specific demonstrations. Our experiments suggest that this specialization is both effective and reasonably robust: across seven GLUE tasks and two MMLU subsets, WILDA‑S improves over standard ICL by 2.6–11.9 points, and on out‑of‑domain task pairs (e.g., training on QNLI and testing on RTE), it maintains higher accuracy than baseline methods. This indicates that the adapter likely captures a task‑level representation rather than memorizing individual examples. Additionally, our “adapter arithmetic” experiments show that multiple adapters trained on different demonstration subsets can be fused, allowing the model to combine information from larger or more varied prompts without exceeding the context window. We acknowledge that transfer across entirely different tasks remains limited, and exploring multi‑task adapters or hierarchical parameterizations is an interesting direction for future work.

Impact on the base model’s general capabilities LoRA adapters are implemented as additive residuals and do not modify the base model’s weights. During training, we toggle the adapter on (student) and off (teacher), and at inference, it can be enabled or disabled without altering the core model. Importantly, additional experiments (see table below) we conducted in response to your remark show no regression in zero‑shot performance, even when an adapter trained on one domain remains active during evaluation on a completely different domain (a scenario conceivable in practice, where the model may be presented with ad hoc queries) – accuracy differences remain within normal variability.

To quantify this effect, we evaluated Llama‑3 (8B) on unrelated, non‑overlapping dataset pairs from GLUE, where both 16‑shot ICL and WILDA‑S are exposed to task‑specific demonstrations from a different domain. Results show that performance differences remain well within one standard deviation (averaged over 10 seeds), demonstrating that WILDA‑S does not harm general zero‑shot capabilities. We will include these results in the revised manuscript.

Relation to “Implicit In‑context Learning”

The recently published Implicit In‑Context Learning (I2CL) paper (Li et al., 2025) proposes condensing all demonstration examples into a single context vector. Specifically, each demonstration is passed through the base model to produce a demonstration vector; these vectors are then aggregated in a permutation‑invariant manner (e.g., via averaging and a small MLP) to form the context vector. At inference, this context vector is linearly combined with the query’s activations and injected into each residual block of the model. The combination weights (“calibration coefficients”) are learned for each layer using a small calibration set. This design eliminates the need to prepend demonstration tokens, while also reducing memory usage and speeding up inference. However, a closer look at their results reveals that I2CL regularly underperforms compared to standard ICL, indicating a trade‑off between efficiency and accuracy.

WILDA‑S, in contrast, uses a teacher–student framework: a teacher model generates pseudo‑labels from few‑shot prompts, and a student fine‑tunes a small adapter to mimic and improve upon those outputs. This mechanism captures the non-linear latent shift induced by demonstrations, extending beyond the attention layers to include feed-forward networks and activation functions, and enables weak-to-strong generalization, where the student corrects noisy pseudo-labels and expands coverage. Empirically, WILDA‑S consistently surpasses standard ICL and enables handling long context by splitting the demonstrations into multiple subsets and then merging them through adapter arithmetic. Furthermore, the student can still benefit from new demonstrations supplied in the prompt: when we encode 16 demonstrations into the adapter and provide 16 more in the prompt, WILDA‑S (16/16) outperforms a 32‑shot ICL baseline across all datasets (Table 9 in Appendix D.2).

Thank you for your constructive feedback and for noting the strengths of our approach. We appreciate your recognition of WILDA’s ability to mitigate order sensitivity and outperform baseline methods. Below, we respond to your concerns and clarify the motivations behind our design choices.

Computational overhead and the “investment” perspective

We acknowledge that WILDA introduces a lightweight fine‑tuning step and understand the reviewer’s concern about compute cost. However, we believe the initial compute cost should be viewed as an investment rather than a recurring cost. The adapter is extremely small; LoRA modules constitute only 0.1–0.3% of the base model parameters, which makes the process computationally lightweight. Once trained, the adapter can be toggled on or off at inference. As our limitations section notes, processing a 16-shot prompt on Llama‑3 is about 120× slower than zero-shot because the self-attention cost grows quadratically with prompt length. Fine-tuning the adapter with 100 unlabeled queries costs the equivalent of ~2,100 such 16-shot inference runs. Thus, in scenarios where one expects to answer more than a few thousand queries, the one‑time training cost gets amortised, and the overall wall‑clock time is reduced because the model no longer needs to re‑process demonstrations for every query.

In addition to amortising the one‑time training cost, partitioning demonstrations across multiple adapters also mitigates context‑window limitations. As Table 4 and the accompanying discussion explain, each adapter is trained on a manageable subset of demonstrations; merging them via adapter arithmetic allows us to incorporate many examples without exceeding the model’s context length. This strategy not only reduces memory overhead but also avoids the quadratic scaling of self‑attention for long prompts. By distributing the prompt across several small adapters, we fit within a single GPU and lower inference cost. At inference, we load only the necessary adapters and avoid the quadratic cost of long prompts, which can be beneficial even for edge devices. We acknowledge that WILDA is less suitable when one has abundant labeled demonstrations and must process new demonstration sets for every query, where we need to fine-tune an adapter for each new prompt and may negate the amortization advantage. In many real‑world applications, however, labeled demonstrations are scarce, and the same few examples are reused across many queries.

Moreover, long demonstration sequences themselves can lead to unstable predictions. Standard ICL is highly sensitive to the order and selection of demonstrations, and performance often degrades as prompts grow longer. Our experiments show that WILDA’s parametric approach, encoding the latent shift into a fixed adapter, significantly improves stability: across multiple random seeds, WILDA‑S exhibits markedly lower variance than standard ICL. This stability stems from learning a reusable representation of the task rather than relying on the precise ordering of a long prompt. Thus, WILDA not only alleviates context‑length constraints but also delivers more consistent behavior.

These additional points strengthen the view of WILDA as an upfront investment: by training a small adapter once and partitioning demonstrations when necessary, we both reduce inference cost and improve stability, making the method practical even when context windows are limited or prompts are long.

Presentation of results and “average scores”

We report results as means over multiple runs (with different random seeds), along with standard deviations, for each dataset. For example, Table 1 includes averages over ten runs. If the reviewer is referring to averaging scores across different tasks into a single number, we deliberately avoid this because it may obscure task‑specific behavior and can be misleading when tasks have different scales or importance. We are therefore unsure what kind of average the reviewer would find helpful. We welcome further clarification and will be happy to include additional aggregates in the final version if that aligns with community standards.

Evaluation and harder datasets

We chose GLUE and selected MMLU subsets because they are well‑established NLP benchmarks that allow fair comparison with prior work. In this work, we focused on models under 10 billion parameters due to computational constraints; such models are typically better aligned with GLUE/MMLU‑style tasks than with extremely challenging reasoning benchmarks. Importantly, WILDA is task‑agnostic, and there is nothing in the approach that fundamentally limits its application to harder datasets.

Nonetheless, prompted by your remark, we evaluated WILDA on the ARC-Challenge dataset. Using Llama‑3 (8B) across all methods, we observed that WILDA-S achieves the highest accuracy with improved stability, outperforming the baselines and related approaches (see table below). This trend is consistent with our results on GLUE and MMLU. We will include these findings and the corresponding table in the revised manuscript.

Just a quick follow-up to kindly check if you’ve had a chance to review our rebuttal. We’d be happy to provide any additional clarification if needed and hope our response addressed your concerns.

We thank the reviewer for the feedback and appreciate the opportunity to clarify our contributions further and address the concerns raised.

WILDA and distillation

While distillation is a useful tool and WILDA’s teacher–student setup resembles it, our goal is fundamentally different. Rather than compressing a large model into a smaller one, WILDA focuses on capturing the latent shift induced by demonstrations, which can mitigate the problems with long context down the line. In our setup, the teacher processes the full demonstration–query prompt and outputs a probability distribution, while the student, equipped with a small LoRA adapter (0.1–0.3% of parameters), sees only the query and learns to approximate the teacher’s distribution. By keeping teacher and student architectures identical, we isolate the effect of the adapter and show that gains come from encoding latent shifts, not differences in model capacity. This approach is orthogonal to model compression: in principle, one could use a smaller student model while still leveraging WILDA to capture demonstration‑induced shifts. Importantly, matching the teacher’s outputs is a necessary step for the student to internalize these shifts, but WILDA goes beyond simple imitation. In fact, the student often surpasses its teacher, with WILDA‑S improving standard ICL by 2.6–11.9 points on GLUE and MMLU for Llama‑3. This behavior reflects weak‑to‑strong (W2S) generalization: during training, WILDA corrects noisy pseudo‑labels, expands coverage, and achieves lower Lipschitz constants (Section 4).

To further demonstrate that the adapter encodes demonstration content rather than merely copying the teacher’s outputs, we conducted an experiment (Appendix D.3) treating the adapter‑induced latent shift as a “task vector.” Retrieving tokens via nearest‑neighbor search over the vocabulary produces text that closely matches the original demonstrations. This aligns with the intuition that, because the student only observes the query while the teacher has access to both demonstrations and the query, it cannot reach teacher‑level performance unless it internally represents the missing demonstration information. Our empirical results confirm this: zero-shot performance is consistently worse than few-shot teacher predictions, and the student closes this gap, often surpassing the teacher, by encoding and leveraging these demonstration-induced shifts. This corroborates the disentanglement we describe: the adapter captures the demonstration signal, while the query is directly provided as input.

Zero-shot baseline

We appreciate the suggestion to include a zero-shot teacher baseline. In fact, our experiments already consider this setting: pattern-based fine-tuning (PBFT) fine-tunes an adapter using zero-shot prompts and template verbalizers, which is effectively equivalent to distilling from a zero-shot teacher with gold labels. As shown in Table 1, WILDA‑S consistently outperforms PBFT across all GLUE and MMLU tasks, with substantial gains for both Llama‑3 and Phi‑3. This indicates that incorporating few‑shot demonstrations through WILDA provides meaningful advantages over zero‑shot distillation.

Student with new demonstrations within the prompt

WILDA enables the student model to incorporate additional demonstrations at inference time. New demonstrations can be supplied in the prompt and are naturally combined with the adapter’s latent shift. Empirically, prompting the WILDA student with extra demonstrations (few‑shot WILDA) consistently improves accuracy. For instance, a WILDA adapter trained with 16 demonstrations and prompted with 16 additional demonstrations outperforms standard ICL using 32 demonstrations. For clarity, we reproduce the table from the paper below. The notation $(n/d)$ denotes the configuration, where $n$ is the number of shots in the prompt and $d$ is the number of encoded demonstrations in the adapter.

Adapter merging (Table 4)

In WILDA, each adapter is trained on a distinct subset of demonstrations and thus captures a specific latent shift. Table 4 does not simply test whether individual adapters are strong, as already shown in Table 1, but instead examines how effectively these independently learned shifts can be combined. When we fuse 2, 4, or 8 adapters by summing their parameters, WILDA’s performance increases markedly. If each adapter merely replicated the same information, we believe that summing their weights would not lead to such notable gains. The observed improvements suggest that each adapter captures complementary information from its subset of demonstrations, effectively expanding the total amount of usable demonstration signal beyond the constraints of the context window. The fusion aggregates these distinct latent shifts, resulting in stronger performance. MMLU datasets

We chose these two MMLU subsets to complement the GLUE tasks used in our main experiments. While GLUE focuses on natural language understanding tasks, we included “elementary math” (MATH) and “miscellaneous” (MISC) to evaluate cross-domain generalization in multi-choice question answering. The MATH subset targets math-oriented reasoning skills, whereas the MISC subset spans a diverse range of topics, providing broader coverage beyond standard GLUE benchmarks.

Experiment suggestion – 8-shot with noisy demonstrations

We thank the reviewer for suggesting that we include an experiment with noisy demonstrations. Our original experiments already show that WILDA is robust to random ordering and selection of demonstrations. Across multiple seeds, it exhibits significantly lower variance than standard ICL, suggesting it learns to focus on informative signals. To further test whether WILDA can denoise irrelevant demonstrations and upweight useful ones, we conducted an additional experiment following the reviewer’s proposal. Using Llama‑3 (8B), we compared a 1‑shot prompt with a single high-quality (gold) example and an 8‑shot prompt containing the same gold example with seven instances from an unrelated domain, evaluating both standard ICL and WILDA‑S on GLUE tasks (10 runs). Results show that WILDA‑S generally matches or slightly outperforms the clean 1‑shot setup and remains more stable in the presence of noise, while vanilla ICL shows a modest regression when noisy demonstrations are added. These findings suggest that WILDA captures latent shifts from informative demonstrations and shows the ability to emphasize relevant signals, providing additional evidence of weak‑to‑strong generalization. The results are summarized in the following table: