On the Power of Context-Enhanced Learning in LLMs

Thank you so much for your constructive comments and questions!

Our responses to your questions and concerns are as follows:

---

Applying Context-Enhanced Learning to Real World Data

Context-enhanced learning (CEL) has been shown to outperform standard training or fine-tuning approaches that do not leverage in-context information across a range of real-world tasks, both during pre-training (e.g., Allen-Zhu et al., 2024; Gao et al., 2025) and fine-tuning (e.g., Liao et al., 2024; Zou et al., 2024) (see lines 036–044).

Following your suggestion, we conducted fine-tuning experiments on the GSM8K dataset using two base models—Qwen-7B and LLaMA 3.1 8B—adapting the context-enhanced learning approach from prior fine-tuning works. Specifically, we included helpful in-context examples during training to enable CEL. We experimented with three strategies for selecting in-context learning (ICL) examples:

• Fixed ICL: A fixed set of in-context examples is used for all training instances. We use the ones used by Qwen-2.5 models for their ICL evaluations.

• Random ICL: Five in-context examples are randomly sampled for each training example individually.

• Skill-based ICL: In-context examples are selected based on skill similarity, following the strategy proposed in [1].

In all three cases, we find that CEL outperforms baseline fine-tuning, which only uses question–answer pairs without additional context. Evaluation is performed zero-shot on the trained models. Our fine-tuning approach for CEL follows the Annealing Dropout strategy.

---

Difficulty of Curating In-Context Curriculum

"The method assumes we have additional helpful text for each training example. In practice, obtaining such aligned context might require extra annotation or the output of a stronger model."

Thank you for raising this important point. Our primary motivation was to theoretically demonstrate that models can learn more efficiently when provided with additional supervision. In practice, such supervision could come from various sources—either via outputs from a stronger model or through retrieval of relevant information from a pre-training corpus.

Prior work has offered promising examples of how such context can accelerate learning. For instance, Gao et al. (2025) and Zhu et al. (2024) show that including metadata—such as source information associated with documents—can lead to faster training and improved performance at test time.

While our current work focuses on analyzing these benefits in a controlled setting, we believe that a key avenue for future research is in developing better methods to obtain and integrate high-quality supervision context for training language models more effectively.

---

Alternative Ways for Eliciting In-Context Rules

"When dealing with real language, couldn’t there be clever prompting strategies that an adversary could use to elicit the rules? Without any results real data, the claims about the privacy preserving capabilities of context-enhanced learning feels quite speculative."

We admit that our current evaluation is restricted to verbatim memorization of the phrasebook rules and does not consider adversarial robustness. While our results demonstrate that recovery of phrasebook information is limited in our synthetic setting, we acknowledge that this does not constitute definitive evidence that the information is completely hidden. Additionally, we do not yet provide observations on real-world data. (For the new GSM8K experiments, the context is selected from the same training set, making it difficult to distinguish memorization).

That said, we view this work as the first step toward thinking about privacy preservation in context-enhanced learning settings. Memorization on real world data and in adversarial prompting settings are interesting directions to continue exploring on.

---

Other suggestions on related work and presentation

Thank you for bringing up the related works and the legibility issue of figure 2 (ablation data points are all overlapping around "near random"). We will surely discuss the mentioned related works and address the figure more clearly in future versions of the manuscript.

---

References

[1] Didolkar, Aniket, et al. "Metacognitive capabilities of llms: An exploration in mathematical problem solving." Advances in Neural Information Processing Systems 37 (2024): 19783-19812.

Thank you so much for your constructive comments (especially regarding section 3.3) and questions!

Our responses to your questions and concerns are as follows:

---

Definition of Necessary Rules and Unused Rules (line 165)

For an input sequence $s_1$ of length 20-40 tokens, each translation step will involve 10-20 2-tuple translations specified by the phrasebook of that step. Out of the $n^2$ total rules of that phrasebook, we define the set of necessary rules for that input sequence to be the phrasebook entries involved and the set of unused rules as its complement.

*We will also use the concept of “necessary rules” and “unused rules” in an additional experiment for introducing proper control on the mechanistic experiment (see below).

---

Proper Control in Section 3.3

To rule out the possibility that the affected depth is only related to the position of perturbation instead of semantic content, we introduce a new set of experiments (see new figure). For each phrasebook $\pi_ i$, we selectively perturb tokens of 10 necessary rules (see the first row) or 10 unused rules (see the second row). Note that $STR(\pi_ i)$ is randomly shuffled, so these necessary rules and unused rules are interleaving in position.

From the figure, perturbing necessary rules result in a qualitatively similar figure as in the current manuscript (with clear difference starting at certain layers), but perturbing unused rules around the same positions yields negligible representation difference across all layers. This suggests that the representation difference is reflecting the model’s processing of the useful phrasebook information with respect to the current translation task. We will replace fig.3 with the new figure in the next version of the manuscript.

---

Real-World Tasks

We have conducted additional experiments using context-enhanced learning to improve math capabilities. Please refer to our response to reviewer PGoK.

---

FLOP Improvement

Empirically, for $n=8, d=5$ (results reported in Figure 2), the SFT baseline with 1000k samples took around 14 hrs while the annealing dropout curriculum with 100k samples took around 2 hrs to complete on the same device.

Theoretically, the full phrasebooks are of length $O(n^2d)$, so including them will lead to a slow down of $O(n^4d^2)$ in FLOPs (quadratic for attention). With large $d$, the $\exp(d)$ sample efficiency improvement will still dominate the polynomial slowdown per sample induced by in-context phrasebooks.

---

Training with Loss on Context

Following your suggestions, we reran the annealing dropout experiments with $n=8, d=5$ while explicitly computing the next-token prediction loss on the context. We report the test accuracy for various number of training samples and compare it with the context-enhanced learning case (in Figure 2).

In general there is no qualitative difference in sample efficiency when we train with loss on context.

To understand what other difference does the additional loss incur, we conduct the following additional experiments:

• We conducted the same verbatim memorization test as in section 4 on these models. The checkpoint trained with just 10k samples already attained a 99.8% recovery success rate, demonstrating very strong verbatim memorization of the phrasebooks despite not giving any additional benefits.
• We reproduced figure 4 for the model trained with loss on context trained with 250k samples and observe similar localized storage similar to the case of having no loss on context.
• We also conducted the No ICL ablation with loss on context, resulting in the near random accuracy for all settings. Even if we are allowed to take loss on the context, the MLT-ICL-capability is still crucial for improved sample efficiency.

The three observations above suggests that verbatim memorizing the phrasebook and being able to translate without phrasebook in context are likely to be detached. Computing loss on the context does not fundamentally change the mechanism of internalization.

---

Bad Performance of No Dropout

We would like to clarify that our central claim (both empirical and theoretical results) all require proper dropout in the helpful context.

In the "No Dropout" scheme, the model would always have very low loss throughout training. This is because we start from an MLT(n,d)-ICL-capable model, which performs the task perfectly when conditioning on full context. Thus there would be little learning signal pushing the model to internalize the phrasebooks if we still always condition on full context. When no context is dropped, we would expect no internalization.

Thank you so much for your comments and insightful questions!

Our responses to your questions are as follows:

---

What accounts for the OOD generalization to empty context?

Short Answer: We believe this OOD generalization is fundamentally a form of compositional generalization, which is achieved by the compositional-friendly sequential translation structure in the MLT(d,n)-ICL-capable model. Following your suggestion, we have tried training an MLT(d,n)-ICL-capable model from scratch, which was proven to be difficult. The current OOD capability is likely a joint result from the base model and the ICL-capablilty training. We expect that as long as the MLT(d,n)-ICL-capable initialization has similar compositional-friendly mechanistic nature, we would see similar OOD behavior.

More details:

Mechanism of the OOD Generalization

We believe two mechanistic properties of the reported model played an important role in this generalization phenomenon (inferred from section 3.3, with additional evidence in response to reviewer vWg8):

1. The sequential structure of translation in the model, in which each step can take information from context or the weights.

2. Phrasebook information at different translation stages are internalized in a localized manner, compensating for the corresponding dropped phrasebook information in the context.

During training, when certain phrasebook entries are dropped from the context while others remain, the dropped information will be learned into the weights of certain transformer layers. When we perform sufficiently many random dropouts on the context, all phrasebook entries will be dropped and internalized at some point despite the context being never empty. Therefore the model will be robust to complete removal of the context since every phrasebook information it needs has been internalized.

Our theoretical model (Surr-MLT, see Def 5.2) is built upon this rationalization. It can achieve full accuracy when conditioning on nothing at test time despite that only one rule is being dropped in each step in training.

“If one were to train an MLT(d, n)-ICL-capable model from scratch instead of fine-tuning a pretrained Llama model, would it still be efficiently fine-tuned to become $MLT_{\Pi^*}$-capable”

Response: We tried training the same 3B model to be MLT(5, 8)-ICL-capable from scratch with various configurations. However even with a million random set of phrasebooks the model is still unable to learn any pattern. Thus unfortunately we cannot empirically answer this question. However if the model trained from scratch also shares similar mechanistic structure as discussed above, then similar sample efficiency may appear.

”Is the generalization driven by the base model’s capabilities (Llama 3.2-3B), or by its MLT(d, n)-ICL capability?”

Response: As discussed above, it is hard to separate the effect from the base model and the MLT(d, n)-ICL capability since the base model affects how the specialized MLT(d, n)-ICL capability (i.e. the sequential translation structure) is formed in the model parameter. Whether the base model is strictly necessary remains unclear.

---

SQ dimension and Sample Complexity of SGD

Short Answer: We stated in the main text that “any algorithm attempting to learn an $MLT_{\Pi^*}$-capable SURR-MLT requires at least $n^d$ sample complexity.” This was an informal statement. A more precise version is: any algorithm that attempts to learn the task by querying distributional properties of $MLT_{\Pi^*}$​​ and receiving noisy estimates in return—i.e., algorithms operating within the statistical query (SQ) framework—requires at least $n^d$ sample complexity.

More details: The formal SQ dimension (as defined in Section F.1 of the appendix) characterizes the computational difficulty for algorithms in the SQ model. In this framework, algorithms interact with an oracle by querying expectations over properties of the target task on data distribution and receive approximate (noisy) answers. A prototypical example is stochastic gradient descent (SGD), where the model estimates gradients based on average loss over a batch. The estimation noise depends on the batch size, along with potential adversarial or quantization noise. For problems with high SQ dimension (e.g., exponential in input size), it can be shown that the number of samples or gradient steps required by such algorithms becomes exponential.

However, these lower bounds do not extend to algorithms that lie outside the SQ framework. For instance, if an algorithm has access to auxiliary information that directly hints at the target function (like phrasebooks in the MLT task), it can bypass the SQ limitations by exploiting this extra structure. In such cases, the learning process can be significantly more efficient due to this inductive bias.

---

Presentation

Thank you for your advice on the presentation! We will add illustrative figures in future versions of the manuscript.

Thank you so much for your comments and insightful questions!

Please find our responses to your questions / concerns below.

---

Generalization in Context-Enhanced Learning

"The MLT task does not shed light into the generalization capability of the technique ... the setup only shows quicker memorization (not generalization capabilities)."

Response: Our main message goes beyond quicker memorization of the phrasebooks. We also highlight that:

1. Even when no next token prediction loss is not computed on the phrasebook tokens, their simple presence in context can still accelerate learning.

2. The model does not need to verbatim memorize the phrasebooks (as shown in Section 4). This is beyond the usual notion of memorization in language model training.

Within MLT there is also OOD generalization behavior: the model can perform the task with empty context at inference despite only 20% of the rules are dropped during training. This suggests that it has learned to generalize by composing rules at inference time in an OOD manner (see lines 248–249).

Please also check our response to reviewer evJV on understanding this OOD phenomenon and response to reviewer PGoK demonstrating effectiveness of context-enhanced learning on real-world data.

---

What is the Motivation of Having no Context at Test Time

Response: Our motivation is twofold: philosophical and theoretical explanation of existing empirical works leveraging context-enhancement.

First, we aim to examine whether models can internalize and benefit from additional information presented during training, even when that information is unavailable at test time. This setting mirrors the human paradigm of open-book learning followed by closed-book testing, and our framework seeks to assess whether models can exhibit similar learning behaviors.

Second, we aim to offer a theoretical perspective on the empirical success of context-enhancement in LLM training. Prior works (e.g. Allen-Zhu et al. (2024) and Gao et al. (2025)) demonstrate that pre-training with auxiliary information improves both training efficiency and final performance — even when such information is absent during inference. Moreover PromptIntern (Zou et al. 2024) and SKIntern (Liao et al. 2024) suggest that context-enhancement can enable more efficient inference by reducing the prompt length. Our work provides a simple theoretical model to help explain these empirical findings and unveil the theoretical potential on sample efficiency.

---

How Phrasebook Information are Provided to Surrogate Model

"In the surrogate model the phrasebooks are combined with the model weights, a more natural representation would be augmentation of these at the input space."

Response: The surrogate model is motivated by our mechanistic experiments in section 3.

First, the ICL-capable model utilizes the in-context phrasebooks in a sequential manner: phrasebooks of later steps are involved in the translation process in later layers (fig 3). During context-enhanced learning, the phrasebook information of a certain translation step is locally internalized to a small set of layers coupled to the layer processing the in-context phrasebooks for that step (fig 4). Thus we parameterize each translation step as a coupling of an in-context representation $C_i$ and an in-weight representation $W_i$, which is the minimal model capturing the mechanistic behavior of the model before and after the training process.

"The transformer model that mimics this surrogate learning also does not include the phrasebooks at the input place, so the claim that standard transformer model can replicate the learning seems misleading."

Response: We would like to highlight that the transformer model we constructed in theorem 5.3. is taking all context information from the input sequence. (see the segment 1 of input in Algorithm 6 and 7, P70-71). Our construction of the transformer is based on an exact reparameterization of the surrogate model:$\def\R{\mathbb{R}} \def\bm#1{\begin{bmatrix} #1 \end{bmatrix}} \def\p#1{\left(#1\right)} \def\x{\times}$

Let $e_i\in\R^d$ be the basis vector and let $0_{m\x n}$ be an $m\x n$ all-zero matrix. Let the context-augmented input be $X_1 = [C_1,\dots, C_d, V_1]\in\R^{n^2\x (dn^2+L)},$ then the reparameterized surrogate model is

$\$

X_{i+1} &= X_i\bm{I_{dn^2} & 0\\0 & 0_{L\x L}} + \p{\p{X_i\bm{e_i\otimes I_{n^2}\\0_{L\x n^2}} + W_i}Shift\p{X_i\bm{0_{n^2d\times L}\\I_L}}}\bm{0_{dn^2\times dn^2} & 0\\0 & I_{L}}\\ &=[C_1,\dots, C_d, 0_{n^2\x L}] + [0,\dots, 0, \p{C_i + W_i}Shift\p{V_i}]\\ &=[C_1,\dots, C_d, V_{i+1}].
$$

This model outputs $[C_1,\dots, C_d, V_{d+1}]$ with input $[C_1,\dots, C_d, V_1]$, where all phrasebooks are provided in-context. Please refer to algorithms 6 and 7 on how we modeled the computation above using self-attention, MLPs, and residual connections in a standard transformer.