On the Power of Context-Enhanced Learning in LLMs

Summary

• Context-enhanced learning is ordinary gradient training on input–output pairs, except that helpful text is placed in the context and no loss is computed on it. It is a gradient-based analog of in-context learning (ICL), and a version for language models of learning using privileged information.
• On a multi-level translation task of depth 5, a Llama 3.2 3B model that can already translate given the phrasebooks in context learns a fixed translation with about 10 times fewer samples than plain fine-tuning, when phrasebook excerpts are in the context and randomly dropped. At test time it needs no phrasebook. Without ICL capability, correct rules, or dropout, the gain disappears.
• The ICL-capable model reads the five phrasebooks at increasing depth. After context-enhanced learning, each phrasebook is stored in a few layers, and each group of layers ends at (or, for the last phrasebook, one layer before) the layer where the ICL-capable model starts reading that phrasebook.
• In a surrogate model, learning from input–output pairs alone needs \(n^{\Omega(d)}\) samples, while dropping one rule at a time from the context needs a polynomial in \(n\) times \(d\log d\). With a single rule missing, the gradient on that rule's weights points at it; with many missing, it becomes much less reliable.
• The rules seen only in the context are hard to recover from the trained model's output probabilities: 0.00% for the intermediate phrasebooks in the main setting, and at most 3.18% across all tested models even with a filter that knows every output alphabet.

Learning with material in the context

Let \(f_\theta\) be an autoregressive language model and \(g\) a task that maps inputs \(x\) to distributions over outputs. Supervised fine-tuning (SFT) minimizes the autoregressive loss \(\ell_{\text{auto}}\) on pairs \((x,y)\) from a dataset \(D_g\). Context-enhanced learning keeps the same loss tokens but puts curriculum text \(\mathrm{CURR}_g(x,t)\) in front of the input. The curriculum text can be anything helpful, such as explanations, excerpts from a textbook, or worked examples, and it may change with the training step \(t\), which makes it a curriculum. The loss is taken only on \(y\):

There is some evidence that this works. Placing helpful text in the context during fine-tuning, with no loss on it, has been used to internalize symbolic knowledge [1], recurring prompts [2], and prompts that describe how to learn the material [3]. In pretraining, prepending source URLs to documents can improve training efficiency and memorization capacity [4, 5]. A student solving problems with the textbook open, without trying to memorize the textbook, is doing something similar. The classical counterpart is learning using privileged information (LUPI), where extra information is available during training but not at test time [6].

We ask three questions.

• Power. With the loss on the same tokens, can it be much more sample-efficient than plain SFT, and why?
• Capability. Does the model need to be capable of ICL on the task to benefit?
• Privacy. Since no gradient is taken on the curriculum text, can private material be used this way with a lower risk of later leakage?

A task that is hard to learn from examples

To answer these questions precisely, we need a task that is easy to learn with the right material in the context and very hard to learn from input–output pairs. Multi-level translation, \(\mathrm{MLT}(d,n)\), is built for this. It has \(d+1\) alphabets \(A_1,\dots,A_{d+1}\) of \(n\) characters each, and for each level \(i\) a phrasebook \(\pi_i:A_i^2\to A_{i+1}^2\) that maps pairs of characters bijectively to pairs. An input \(s_1\in A_1^L\) of even length \(L\) is translated one level at a time. Each level shifts the string left by one character, wrapping around, and then rewrites every consecutive pair with the phrasebook. In symbols,

where \(T_{\pi_i}\) is the map from \(s_i\) to \(s_{i+1}\) and \(\Pi=\{\pi_i\}\) collects the phrasebooks. The design follows Feistel ciphers [7], although we make no cryptographic claim. Two properties matter. First, \(\mathrm{MLT}_\Pi\) is a bijection on strings, since both the shift and the translation are invertible. Second, the shift mixes neighboring pairs, so each output character depends on \(2d\) consecutive input characters. Without the shift, the whole task would collapse into a single map on pairs. The complete task is written down in \(n^2d\) phrasebook rules, yet we show below that, in the statistical query framework, learning it from input–output pairs alone needs \(n^{\Omega(d)}\) samples.

A phrasebook is written as text of the form a b -> C D; e d -> B A; …, with rules in any order, and \(\mathrm{STR}(\Pi)\) is the concatenation of the \(d\) texts. The curriculum text for an input \(s_1\) is an excerpt of \(\mathrm{STR}(\Pi)\).

At test time the phrasebook is absent, so the model cannot be taught an explicit chain of thought that copies rules from the context. Yet the task has \(d\) sequential steps, and the model needs room to compute them. We therefore teach it to emit a fixed number of <THINK> tokens before the answer, a silent chain of thought in the spirit of pause tokens [8]. A model is \(\mathrm{MLT}(d,n)\)-ICL-capable if, for any set of phrasebooks \(\Pi\), given \([\mathrm{STR}(\Pi),s_1]\) it outputs \(\mathrm{MLT}_\Pi(s_1)\) after the <THINK> tokens. Such a model can read phrasebooks. The question is whether that ability helps it learn one particular set of phrasebooks \(\Pi^*\) that it will not see at test time.

Ten times fewer samples in Llama 3.2 3B

We start from Llama 3.2-3B-Instruct [9] with \(d=5\) and \(n=8\) or \(10\). Training has two stages.

Stage 1: an ICL-capable model. We fine-tune on translations with random phrasebooks, all different from \(\Pi^*\), with one example per set so that no phrasebook is memorized: \(3\times10^5\) sets for \(n=8\) and \(10^6\) for \(n=10\), with inputs of length 20 to 40. We follow a standard chain-of-thought internalization pipeline [10, 11, 12]. For the first 10% of training the targets contain every intermediate string \(s_2,\dots,s_d\). Between 10% and 60% these are replaced by <THINK> tokens, first to last, at a rate that grows linearly to 100%. We use AdamW with weight decay \(10^{-4}\), a cosine schedule with peak learning rate \(10^{-4}\) and 6% warmup, and batch size 64.

Stage 2: learning \(\Pi^*\). We fix \(\Pi^*\) and build datasets of \(10^4\) to \(10^6\) random inputs of length 20 to 40, each paired with <THINK> tokens and the answer \(\mathrm{MLT}_{\Pi^*}(s_1)\). Each dataset is used for one epoch, with the same optimizer and no weight decay. The curricula differ only in the curriculum text:

• No Context (plain SFT): empty.
• Fixed Dropout: the rules of \(\Pi^*\) that \(s_1\) uses, with 20% of them dropped at random.
• Annealing Dropout: the rules that \(s_1\) uses plus 25% unused rules, with a dropout rate that rises linearly from 0% to 100% over the first 60% of training and then stays at 100%.
• In both dropout curricula, each training step drops rules either from all five phrasebooks or from one phrasebook chosen at random, keeping the others intact.
• Ablations: No Dropout (all used rules, always), Wrong Context (Annealing Dropout with incorrect rules), and No ICL (Annealing Dropout from the base model, which is not ICL-capable).

We measure next-token accuracy on the answer tokens of held-out inputs, with no curriculum text at all. Figure 2 shows the result.

• Annealing Dropout learns fastest, followed by Fixed Dropout. Both need about 10 times fewer samples than plain SFT. For \(n=10\), plain SFT stays at chance up to \(10^6\) samples.
• The three ablations stay near chance, apart from a single point of Wrong Context for \(n=8\). Correct context, dropout, and ICL capability are each necessary. Without dropout, the model can keep reading the rules from the context and has no reason to store them.
• With only 20% dropout during training, the model is perfect at test time with 100% dropout. It therefore combines rules that never appeared together in the context of a single training example: it learns the rules one by one and recombines them at test time.
• The model stores a rule only if it is needed. When we hold out half of one phrasebook's rules from every training translation but still show them in the curriculum text, the trained model is at chance on test inputs that need those rules (Appendix C.1 of the paper, for each of the five levels). A rule is internalized only when dropping it from the context raises the training loss.

Where the phrasebooks are stored

Reading. We first ask how the ICL-capable model uses phrasebooks in its context. For \(\mathrm{MLT}(5,8)\) we feed the phrasebooks, the input, the <THINK> tokens, and the answer, then replace the outputs of 10 rules of \(\pi_i\) that \(s_1\) uses, and record the \(\ell_2\) change of each layer's hidden states at the <THINK> and answer tokens. The first layer with a clear change is where the model starts reading \(\pi_i\). Later phrasebooks are read by later layers. Perturbing rules that \(s_1\) does not use changes almost nothing, so the depth reflects the content of the rules, not their position in the prompt.

Storage. We then train this model with Annealing Dropout on 100,000 samples and ask where \(\Pi^*\) ends up. For every range of layers \(L_{\text{start}}\) to \(L_{\text{end}}\) we build a stitched model, the ICL-capable model with those layers replaced by the trained model's layers [13], and evaluate it with \(\mathrm{STR}(\pi_i^*)\) removed from the context and the other phrasebooks present. A range that restores accuracy stores \(\pi_i^*\).

\(\mathrm{MLT}(5,8)\); Llama 3.2 3B has 28 layers. The layer where the ICL-capable model starts reading each phrasebook is annotated in the paper's Figure 3, and the layers that store it after context-enhanced learning (Annealing Dropout, 100,000 samples) in its Figure 4.

Each phrasebook is recovered by a small group of layers, the groups do not overlap, and each group ends at the layer where the ICL-capable model begins reading that phrasebook (25 against 26 for \(\pi_5\)). Our reading is that the model does not store the phrasebooks as one block. It re-learns each translation step locally, to make up for that step's rules when they are dropped from the context, and it does so in a group of layers that ends where the ICL-capable model starts reading that step's phrasebook. We conjecture that context-enhanced learning helps by localizing learning in parameter space in this way. Across dataset sizes, later phrasebooks are internalized with fewer samples than earlier ones. With Fixed Dropout the locations are similar, so they are set mainly by the ICL-capable initialization rather than by the dropout schedule.

An exponential gap in a surrogate model

Gradient training of a 28-layer transformer is out of reach for current theory, so we build a surrogate that has the structure we just observed: one layer per translation step, with each layer reading its phrasebook from the context. The surrogate, SURR-MLT, uses the alphabet \(\{1,\dots,n\}\) at every level and encodes a pair \((a,b)\) as the one-hot vector \(v(a,b)=e_a\otimes e_b\in\mathbb R^{n^2}\). A string becomes a matrix \(V_1\in\mathbb R^{n^2\times L/2}\) with one column per pair, and each layer computes

Shift turns two neighboring pairs \((a,b),(c,d)\) into the straddling pair \((b,c)\), with the last column paired with the first, and HardMax makes each column one-hot at its largest entry. \(C_i\) is the phrasebook given in the context, the matrix whose column for input pair \(p\) is one-hot at \(\pi_i(p)\), with a zero column for every dropped rule. \(W_i\) is trainable. With \(W=0\) and the phrasebooks in context the model solves any task, so it is ICL-capable. With no context and \(\operatorname{HardMax}(W_i)\) equal to the matrix of \(\pi_i^*\), it solves \(\mathrm{MLT}_{\Pi^*}\) on its own. A transformer with \(2d\) attention layers, \(2d\) MLP layers, and embedding dimension \(2n^2+2d+4\) can simulate SURR-MLT exactly, so the results below apply to that transformer as well.

The statistical query framework [14, 15] covers algorithms that only use expectations over samples, which includes minibatch gradient methods, and the step from dimension to SGD follows the argument for sparse parities [16]. For \(n=2\), two random sets of phrasebooks give translations whose outputs are uncorrelated with probability at least \(1-\tfrac13(\tfrac79)^{d-1}\). The Lovász local lemma then yields exponentially many pairwise uncorrelated tasks, and larger \(n\) reduces to this case through binary encodings.

The gradient result is proved only for \(d=2\), but it holds in simulation at larger depth. The paper trains a softmax relaxation of SURR-MLT (column softmax at temperature \(1/25\), cross-entropy, learning rate 100, weights clipped to \([0,1]\)) with a curriculum that cycles through every rule of every level and drops one at each step. On \(\mathrm{MLT}(10,10)\), every layer learns its full phrasebook, whether each step updates only the layer whose rule is missing or all layers. The same holds for \(\mathrm{MLT}(5,10)\) and \(\mathrm{MLT}(20,10)\), and, with Adam and rules dropped from randomly chosen layers, for the hand-built transformer on \(\mathrm{MLT}(5,10)\). Figure 3 runs this procedure in the browser at a smaller size, next to the same model trained without context.

In this simulation, at the default settings and with context, the curves for all layers climb to 1 within the cycles shown. Without context, which the paper does not simulate, the model gets the same inputs and targets and the same learning rate, yet its weights match the phrasebooks no better than chance, \(1/n^2\) per rule.

Why one missing rule at a time

The proof of Theorem 2 rests on one fact. When a single column of \(C_1\) is zeroed out, the gradient with respect to the matching column of \(W_1\) points toward the missing rule. This relies on the other rules being present. Intuitively, the network then computes everything except one step, so the error at the output can be traced back to it. When many rules are missing, errors from different steps mix.

To measure this, we define the gradient prediction accuracy. Drop each column of \(C_i\) independently with probability \(p_i\), and for a dropped column \(j\) of \(C_1\) check whether the largest entry of the negative gradient on \(W_1^{(j)}\) is the correct rule:

We use the first layer because it is the deepest, where noise from dropped rules downstream should intuitively accumulate most, and we evaluate at initialization (\(W=0\)), where the model reads everything from the context. The loss can be the expectation over inputs or the average over a random batch.

• Accuracy falls as the dropping rate rises. With a few rules missing, the gradient almost always names the right rule, and with most rules missing it is close to chance.
• Larger batches help, but they do not fully make up for a high dropping rate.
• Dropping rules from several phrasebooks at once lowers accuracy quickly, even at moderate rates.

This suggests why the dropout curricula work in the Llama experiments, although the paper measures gradient quality only in the surrogate model and leaves the LLM case open. In the surrogate, correct context with few rules missing gives an informative gradient for exactly the missing rules. Plain SFT is the extreme case in which every rule of every phrasebook is missing. The paper defines the accuracy with the argmax, as above, while the released code checks whether the gradient entry of the correct rule is negative, for which chance is 1/2 rather than \(1/n^2\). The floor of about 0.5 in the paper's Figure 6 matches the second criterion. Figure 4 shows both, and they agree on every trend. Also as in the released code, Figure 4 drops a fixed number of columns at each rate rather than each column independently, and in its second view it drops the same columns from every affected phrasebook.

What the trained model reveals about its context

The phrasebooks were in the context throughout training but never received any loss. We test whether the trained model reproduces them. We feed it the phrasebook text and, for each rule a b -> C D, check whether it completes a b -> with C D, by greedy decoding (both tokens are the top prediction) and by sampling at temperature 1 (the product of the two probabilities). An adversary can also filter the output: forbid the <THINK> token, or, knowing the intermediate alphabets, keep only characters of the rule's output alphabet. Each rate is averaged over 20 orderings of the phrasebook text and over every rule. A random guess of the pair succeeds with probability \(1/n^2\), 1.56% for \(n=8\) and 1% for \(n=10\).

For \(\mathrm{MLT}(5,10)\) trained with Annealing Dropout on 100,000 samples (the paper's Table 2), the rules of the intermediate phrasebooks \(\pi_1\) to \(\pi_4\) are recovered 0.00% of the time without a filter, by greedy decoding and by sampling. The last phrasebook is near chance, at 0.20% and 0.89%. The strongest filter brings the intermediate phrasebooks to 1.66% with greedy decoding, above chance but far from recovering the phrasebook. Across all ten models in the paper's Tables 4 and 5, every one of which translates almost perfectly without context, intermediate recovery without a filter is at most 0.26%, and with any filter at most 3.18%.

A model can thus rely on material it saw only in its context while revealing little of it through its output probabilities. This could let privileged data be used in training with a lower risk of leakage. The flip side is that text which should not be used for training could be used in a way that such queries might not detect. Both points are hypothetical: we did not try adversarial prompting, and the setting is synthetic.

Related work

Privileged information. Learning using privileged information was introduced for kernel SVMs [6], analyzed for classification [17], and used to explain knowledge distillation [18]. For autoregressive language models two new questions arise: whether a loss on the extra text is needed, and how the extra text changes training. Context-enhanced learning places no loss on the extra text at all.

In-context and in-weights learning. Properties of the training data affect whether transformers learn in context or in their weights [19, 20, 21]. Here strong ICL is used to speed up in-weights learning of a specific task. The silent computation builds on chain-of-thought internalization [10, 11, 12] and pause tokens [8].

Hardness and localization. The lower bound uses statistical query arguments developed for parity-like problems [14, 15, 16], and the layer analysis uses stitching, as in skill localization after fine-tuning [13]. Unlike masked language modeling and infilling, no loss is ever placed on the text that is removed from the context.

Limitations

The task is synthetic, chosen to be absent from pretraining and too hard for plain SFT, and whether the gains carry over to tasks such as math or coding is open. The convergence analysis uses a surrogate model, and the gradient result is proved only for \(d=2\). The detectability test uses likelihood-based queries, not adversarial prompting. Implications for privacy and copyright [22, 23] are hypothetical in this toy setting.

References

1. Liao et al. SKIntern: Internalizing Symbolic Knowledge for Distilling Better CoT Capabilities into Small Language Models. arXiv preprint arXiv:2409.13183, 2024.
2. Zou et al. PromptIntern: Saving Inference Costs by Internalizing Recurrent Prompt during Large Language Model Fine-tuning. arXiv preprint arXiv:2407.02211, 2024.
3. Choi et al. Teaching LLMs How to Learn with Contextual Fine-Tuning. arXiv preprint arXiv:2503.09032, 2025.
4. Allen-Zhu and Li. Physics of Language Models: Part 3.3, Knowledge Capacity Scaling Laws. arXiv preprint arXiv:2404.05405, 2024.
5. Gao et al. Metadata Conditioning Accelerates Language Model Pre-training. arXiv preprint arXiv:2501.01956, 2025.
6. Vapnik and Vashist. A New Learning Paradigm: Learning Using Privileged Information. Neural Networks, 22(5–6), 2009.
7. Knudsen. Practically Secure Feistel Ciphers. Fast Software Encryption, 1993.
8. Goyal et al. Think Before You Speak: Training Language Models With Pause Tokens. arXiv preprint arXiv:2310.02226, 2023.
9. Dubey et al. The Llama 3 Herd of Models. arXiv preprint arXiv:2407.21783, 2024.
10. Deng et al. From Explicit CoT to Implicit CoT: Learning to Internalize CoT Step by Step. arXiv preprint arXiv:2405.14838, 2024.
11. Pfau et al. Let's Think Dot by Dot: Hidden Computation in Transformer Language Models. arXiv preprint arXiv:2404.15758, 2024.
12. Hao et al. Training Large Language Models to Reason in a Continuous Latent Space. arXiv preprint arXiv:2412.06769, 2024.
13. Panigrahi et al. Task-Specific Skill Localization in Fine-tuned Language Models. International Conference on Machine Learning, 2023.
14. Kearns. Efficient Noise-Tolerant Learning from Statistical Queries. Journal of the ACM, 45(6), 1998.
15. Blum et al. Weakly Learning DNF and Characterizing Statistical Query Learning Using Fourier Analysis. ACM Symposium on Theory of Computing, 1994.
16. Edelman et al. Pareto Frontiers in Neural Feature Learning: Data, Compute, Width, and Luck. arXiv preprint arXiv:2309.03800, 2023.
17. Pechyony and Vapnik. On the Theory of Learning with Privileged Information. Advances in Neural Information Processing Systems, 2010.
18. Lopez-Paz et al. Unifying Distillation and Privileged Information. arXiv preprint arXiv:1511.03643, 2015.
19. Chan et al. Data Distributional Properties Drive Emergent In-Context Learning in Transformers. Advances in Neural Information Processing Systems, 2022.
20. Singh et al. The Transient Nature of Emergent In-Context Learning in Transformers. Advances in Neural Information Processing Systems, 2023.
21. Reddy. The Mechanistic Basis of Data Dependence and Abrupt Learning in an In-Context Classification Task. International Conference on Learning Representations, 2024.
22. Carlini et al. Extracting Training Data from Large Language Models. USENIX Security Symposium, 2021.
23. Karamolegkou et al. Copyright Violations and Large Language Models. arXiv preprint arXiv:2310.13771, 2023.