On the Power of Context-Enhanced Learning in LLMs
Authors: Xingyu Zhu
Abhishek Panigrahi
Sanjeev Arora
Venue: ICML 2025 spotlightposter (2025)
Keywords: In-Context Learning
Sample Efficiency
Knowledge Internalization
Privileged Information
Multi-step reasoning
Primary area: deep_learning->theory
TL;DR: We formalize Context-Enhanced Learning, a paradigm where LLMs leverage context without direct gradient updates, achieving exponential sample efficiency gains over standard autoregressive training when the model has ICL capabilities

Abstract: We formalize a new concept for LLMs, **context-enhanced learning**. It involves standard gradient-based learning on text except that the context is enhanced with additional data on which no auto-regressive gradients are computed. This setting is a gradient-based analog of usual in-context learning (ICL) and appears in some recent works.

Using a multi-step reasoning task, we prove in a simplified setting that context-enhanced learning can be **exponentially more sample-efficient** than standard learning when the model is capable of ICL. At a mechanistic level, we find that the benefit of context-enhancement arises from a more accurate gradient learning signal.
We also experimentally demonstrate that it appears hard to detect or recover learning materials that were used in the context during training. This may have implications for data security as well as copyright.

On the Power of Context-Enhanced Learning in LLMs
Xingyu Zhu Abhishek Panigrahi! Sanjeev Arora
Abstract
placement of additional helpful material in the context, but without actual auto-regressive gradient updates on this material. We will call this form of learning context- enhanced learning. Because the material used for context- enhancement can evolve over the course of training, this approach naturally aligns with the idea of a curriculum.
We formalize a new concept for L.Mss, context- enhanced learning. involves standard gradient- based learning on text except that the context enhanced with additional data on which no auto- regressive gradients are computed. This setting is gradient-based analog of usual in-context learn- ing (ICL.) and appears in some recent works.
Context-enhanced learning intuitively mirrors how humans learn: when solving problems, they refer to texthooks or demonstrations for guidance, yet they do not seek to mem- orize these resources per An analogous concept Learn- ing using Privileged Information (or LUPI), has been well- studied in the context of kernel SVMs (Vapnik & Vashist, 2009) and classification models. Our work adapts this con- cept for LLMs, and surfaces the following questions:
Using multi-step reasoning task, we prove in simplified setting that context-enhanced Iearn- ing can be exponentially more sample-efficient than standard learning when the model is capa- ble of ICI., At mechanistic level, we find that the benefit of context-enhancement arises from a more accurate gradient learning signal. We also experimenaally demonstrate that it appears hard to detect recover learning materials that were used in the context during training. This may have im- plications for data security as well as copyright.
Q1: Even though autoregressive loss is computed the same set of tokens, can context-enhanced learning be significantly more powerful than usual auto-regressive learning that has no additional in-context materials? If so, can we theoretically characterize and understand the mechanism behind such improvement?
1. Introduction
Do models need a certain capability level benefit from context-enhanced learning? This is natural question, since leveraging in-context information (e.g.. ICI.) likely requires a minimum capability level or model size (Brown et al., 2020; Wei et al., 2022).
Q2:
Pre-trained LLMs (Brown et al., 2020; Touvron et al., 2023; Team al., 2023) show strong capability to learn new ma- terial at inference time, for instance via in-context-leaminng (ICL). There is also emerging evidence that gradient-based learning on piece of text (say, math Q&A) can be hanced if additional helpful text is placed in the context, even though no auto-regressive loss is computed on this helpful text (Liao et al., 2024; Zou et al., 2024; Choi et al., 2025). Such strategies have also been shown to benefit pre-training as prepending source URLs documents can enhance the model's training efficiency and memorization capacity (Allen-Zhu & Li, 2024; Gao al., 2025).
Q3: context-enhanced learning a viable way to use privileged/private information during learning? Providing such privileged information in the context could conceivably enhance the model's learning, but since auto-regressive gradient updates happen on the privileged/private infor- mation, there might be a lower risk of leakage of such information via API calls.
Section 2.1 formally defines context- enhanced learning. To allow rigorous understanding of the power of context-enhanced learning, Section 2.2 introduces multi-step reasoning task called Multi-layer translation.
Paper Overview:
In this paper we seek to formally study this phenomenon, whereby L.Ms' gradient-based learning is enhanced via
'Equal contribution with theoretical framework and setting up the problem. XZ undertook most of the experiments. 'Princeton Language and Intelligence, Princeton University. Correspondence to: < <xxingyu.zhu, ap34) } @ (erincetoon.du>>.
This is a synthetic setting involving d 1 languages L1, L La+1 over finite alphabets. For each i, there is a simple phrasebook that describes how to translate from L; to Li+1, and the mapping from L1 to La+1 is a sequential application of the set of phrasebooks.
Proceedings of the nnnd International Conference Machine Learning. Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).
The goal is to learn how translate text from L2 into Lat1

---

On the Power of Enhance- Learning in LLMs
without explicitly writing down intermediate steps. The learner is provided excerpts from these phrasehooks as help- ful information in the context during training, but allowed no auto-regressive gradient updates on these tokens.
Definition 2.1 (g-capable model, informal). A language model fo is g-capable for a language task if fo is close to y, as measured by suitable metric on Xg-
If we train with auto-regressive loss on translation output conditioning on the phrasebooks' excerpts and the input, a model with certain ICL capacity level may quickly learn the translation task by leveraging the in-context phrasebooks. However, this learning could be brittle, that the model be- comes reliant on having the phrasebooks' excerpts in con- text. This reliance can be weaned off by use of probabilistic dropout on phrasebooks tokens in context. Intuitively, this curriculum forces the model to not only read phrasebooks' excerpts, but also gradually internalize the phrasebooks contents. Over time, the model's s ability to translate from L to La+1 will become robust to the dropout of phrasebooks' excerpts, and eventually, their complete removal.
Vanilla supervised fine-tuning (SFT) aims to create g g-capable model by minimizing auto- regre-regeessvee loss euuo on a supervised dataset Dg = {(ri, yt) )i==- where the label y for each 2; E X, is sampled from g(z;).
Context-enhancee learning involves augmenting the supervi- sion with additional curriculum-tes that depends the task 9. imput æ, and training step t. We denote curriculum-ter as CURR, (z, (). which could be anything (helpful explana- tions, excerpts from textbooks, worked-out examples, etc.),
Algorithm 1 Context-Enhanced Learning In contrast to standard SFT, it relies on curriculum-tes in context on which no auto-regressive loss is computed.
Experiments show that this training strategy indeed works when the learner is a pre-trained LLM that is capable of ICL (but fails when LLM is incapable of ICL). Even when train- ing with 20% dropout rate, the model can perfectly translate strings from L1 La+1 without any phrasebooks' excerpts at test time. The rest of the paper is structured as follows:
Input: Supervised dataset Dg, curriculum-text CURRg initialization 6. total steps T
for 1 to T do
Sample (I, Dg Compute loss 1 <- Eato (Jo((CURR,(1r, (), v)).
Update parameters 0 with gradient Vol end for Return
Section 3 details our experiments and the findings sketched above. Experiments show that an ICL-capable model follows an intuitive sequential processing of the phrasebooks provided in-context, whereby transformer layers approach stages of translation in an intuitive way; e.g., La + T4 is done after L.2 La (Section 3.3).
On sample (ar, y) drawn from the supervised dataset, we use auto-regressive loss for model's prediction on y conditioned on [CURR, f), x] to train our models. Note that no loss is computed for curriculum textttetttttttttttteee text tokens. We denote this loss as fasto (fo([CURR, (2, 1), x. v)).
Section 4, shows that after context-enhanced learning, the output probabilities of the model reveal little about the phrasebooks rules that were seen during training.
In Section 5, we propose a theoretical framework using a surrogate/simpliiied model that represents an ideal LLM for the translation task (Section 5.1). This framework shows an exponential gap in sample complexity depending on whether the model is trained with or without in-context information on phrasebooks (Sections 5.2 and 5.3). Experiments reveal that the mechanism behind the increased sample efficiency of context-enhancee learning is an improved gradient signal, measured by gradient prediction accuracy (Section 5.4).
2.2. Multi-level Translation (MLT)
To study the power of context-enhanced learning, we intro- duce a multi-step translation task that is easy to learn with a straightforward curriculum in the context, but very difficult to learn with just input-output examples.
The task is inspired by encryption methods such as the Feistel cipher (Knudsen, 1993). The multi-level translation (MLT) task involves a bijective mapping from strings to strings that is a composition of 2d simpler bijections, cach involving simple shift by 1, or transforming bigrams (2- tuples of characters) via bijection. The depth-d translation can be described by (dd) bits. But we will show that learn- ing the task only from input-output pairs would require A(d) sample complexity in the SQ-learning framework (Kearns, 1998) (see Theorem 5.4).
2. Setup
2.1. Context-Enhancee Learning
Let X be the space of all possible text strings and let y be the space of all possible distributions over texts. Let g be a language task mapping inputs x E Xg C X to a distribution Y E y: Let So X + y' be general auto- regressive language model. We characterize fo's capability on task g as follows:
Note that there is no proof that current cryptographic tasks are fundamentally difficult for computers. Here the task makes sense as an example of multi-step reasoning and we are interested in sample-complexity lower hounds.
2

---

On the Power of Context-Enhancee Learning in LLMs
We denote II {mT)] the set of phrasebooks of all levels, and denote MLTT 81 (L)T((11) T Txp (s1) the mapping from input to output s4+1 using II. We use MLT(d, n) to refer to the family of translation tasks involving d step and n characters in each alphabet.
Table Important notations for defining MLT
Depth of translation task
Number of characters in each alphahet An alphahet set
A phrasebook between 2-tuples in two alphabets A set of phrascbooks { defining a translation task Translation task with set of phrasebooks IT Family of translation tasks of depth and n characters
11
MLTn
We note that (d, n) has two key properties: 1. Once the phrasebooks are fixed, the translation task defines a bijection between input and output strings since Circular shift and Translate are invertible (see Lemma E.1). 2. Each character in the output string depends on 2d characters from the input text string (see caption of Figure 1), making learning from input-output pairs very difficult (Theorem 5.4).
(L)(d,))
Concretely, let .......,A4.. be d + alphabets all of the same size with characters. For every consecutive pair of alphabets A, and A++1, we fix a phrasebook Ty : A A? 1 bijective mapping from 2-tuples in A; to 2-tuples in Ai+1. Each phrasebook T. can be represented by binary stochastic matrix Matrin Matrix((T:) with rules represented as one-hot columns (see Definition G.4).
For each phrasebook Tis we compose its textual repre- sentation STR (T::) be of the form which lists (insensitive of
a ->
D; e -> B A; ordering) phrasebook rules between 2-tuples in the previous alphabet and the next alphabet. Moreover, we denote the concatenation [STR will be used to define curriculum-text.
$1
A
D
B
H
STR (Tæ)] as STR (II), and
Circular Shift
B
Phrasebook x
BC - eb DE hd
2.3. Needed: Curriculum without Explicit CoT
FG - at HA - gi
$2
g
To teach the model a particular translation task MLT from input-output pairs of the form (s1, MLT we can train it with relevant sections of the phrasebooks STR (II) in context as curriculum, but at test time it would not have access to the phrasebook so it is important not to teach it explicit tfthought (CoT) containing in-context information. (Another consideration is data privacy, with the phrasebook being considered privileged information.) However, a dual use of CoT is to provide the model extra compute at inference time (Goyal et al., 2023), which is needed here since the translation task has d stages. To facilitate such silent computation output fixed number of <THINK> tokens, sometimes refered to silent CoT or internalized CoT.
Circular Shift
3/
b
h
f
Phrasebook x bh - yB da on
Eg - (u
53 Y B
m:
e
1
Figure 1. Illustration of an M2T 8) instance with input
quence of 8 tokens. Input sequence & went through 2 translation input characters,
steps 83. Each output character depends for example, character is derived from c, e in 82, which, in turn, are computed from 2, B, in 81.
teach the model to
The input of the translation process is an even-length se- quence, referred to as $1 E Af where L is the sequence length. The translation process modifies 81 recursively. For every i E [d], 8; E A will be transformed to 8;+1 E A/ using phrasebook Ti through the following sub-processes:
2.4. ICL-capablity for AL,(d,,m) n))
To learn from books, one needs to know how to read. The analogous notion under study here is whether context- enhanced learning requires capability to sort-of "understand" the in-context material (Q2). In the context of MLT, we formalize such capability as being able to achieve low loss on the translation task when provided with the relevant phrasebook sections in context while allowing silent CoT.
Circular shift: The characters in 8; E Af are shifted by character leftward (and wrapped around to the end if necessary) to give sequence á; E At. Formally, for each j E L] we have št.j = $1,
Using the phrasebook ,, A 7 A 1 we translate 2-tuples (bigrams) of consecutive characters in sequence 5; create 8++1. That is, for every odd j E [1, L]. 15411j.8441..
Translate:
Definition 2.2 (MLT n) m)JIIL-ccpaability informal), language model So is MLT W)CLC--cppaaee if for any set of phrasebooks II in MLT (d, n), fo([STR [I]) is close to $2+1 MLT (si) dis- regarding the <THINK> tokens, when measured by a discrepancy metric over all valid input strings 81.
(5. 11.)+1).
We denote the mapping from 8; to 8;41 8 = Tz, (8:). The d-step translation is defined as the composition sd+1 = To.oT ox. 81 to sd+1 through d translation steps. Please see Figure for a visual illustration with d = 2, n
(8)) which converts the input sequence
3

---

On the Power of Context-Enhancd Learning in LLMs
3. Experiments and Observations
Wrong Context (ablation): Equivalent Annealing
Dropout but the rules in the curriculum are incorrect.
In this section, we fix a set of phrasebooks II" and study context-enhanced learning on MLT
3.2. Experiment Results
We first introduce the preparation of an MLT n)-ICL- capable model. We then introduce a context-enhanced learn- ing curriculum involving random dropping of phrasebook rules in context. We then present empirical evidence for sig- nificant sample efficiency of context-enhanced learning. We conclude the section with mechanistic insights into context- enhanced learning concerning internal representations and evolution of parameters.
To check the sample efficiency benefit of context-enhanced learning (Q1), we construct supervised datasets Drr- with 10 to 10 unique samples and train the models for one epoch on each. 2 We report the next-token prediction accu- racy on the final answer tokens (ignoring thought tokens) for held-out samples when conditioning on no curriculum- (100% dropped-out) and compare against the supervised dataset size. To check the necessity of proper ICL capability (Q2), ablate with Annealing Dropout but starting from mon-MLT( ICC-Capbble 3B base model.
We use the Llama 3.2-3B instruction-tumel model (Dubey et al., 2024) as the base model and fix d = 5 with n = 8 or 10. Detailed configurations are available in Appendix B.5.
No Context (baseline)
No Dropout (ablation) Wrong Context (ablation) No (ablation)
3.1. Experimental Setup
Fixed Dropout
Annealing Dropout
(i) Preparing an (LT n) n)-ICI--Capable Model:
d=5, n=8
d=5, 5, n=10
The Llama 3.2B model is (LT -)-ICL-capable as it has not seen the task during training. To make it ICL-ca for our purpose, we use SFT on other random translation tasks with random phrasebooks II,,..., MI.. following common CoT internalization pipeline (Deng et al.. 2024; Pfau et al., 2024; Hao et al.. 2024). We use training example per set of phrasebooks to prevent memorization of specific phrasebooks. At the end of training, given input [ITR( ] for any 8i and II the model can generate <THINK>, MLTT( correctly. Details on the first stage of training are described in Appendix B.4.
(without conteee)
1.00
Test Accuracy
0.75
0.50
0.25
0.00
Ra Ram0am
10a
10a
10
Number of Samples [Pr1l
Number of Samples [Prr
Figure 2. Context--nhancce learning of MLTII- when 8 and
T. = 10. Annealing Dropout learns the fastest followed by Fixed Dropout, requiring 10x less samples compared to vanilla SFT. Ablations show the necessity of correct context, proper dropout, and sufficient ICL capability as a good initialization. Note that there are multiple ablations overlapping around random guesses.
(ii) Setting up context-enhanedt learning for IT We use the MLT -a--Caaeee model above ini- tialization and train for MLTn-.. Supervised dataset DrT- is curated with input-label pairs of the form (s1, [<THINT,,,> ... MLT (s)))), where $1 is a random string sampled from A,, with length between 20 and 40.
Figure 2 demonstrates the significant sample efficiency of context-enhanced learning. Moreover, models trained with subsets of phrasebooks and just 20% dropout give perfect heldout test-time accuracy with 100% dropout rate. Thus they are able to effectively use phrasebook rules from sub- sets of the phrasebook that did not co-oceur in the same train- ing sample. Clearly, the model has learned the phrasebook atomically, and can combine the rules needed test time.
We define curriculum-tex CURRI- ($1,t) using excerpts from phrasebooks (TT( (selected based on $1) with random dropout of rules (parameterized by training step f). We explore the following curriculum and study their impact:
No Context (vanilla SFT); Empty curriculum-text.
In an ablation (Appendix C), we show that context-enhanced learning only internalizes the rules whose dropout from curriculum-text leads to an increase in loss on training data.
Fixed Dropout: A simple strategy independent of step I; given 81, only curate rules in II" used in the translation of 81, then randomly drop 20% of the curated rules.
The experiment results can be summarized as follows:
Annealing Dropout: A better strategy: for 81, select the necessary rules from II plus 25% unused rules. Apply random dropout on these rules, increasing linearly from 0% to 100% over the first 60% of training, then maintain 100%.
(i) Context-enhanced learning from an ICL-capable model greatly improves training sample efficiency. (ii) The phrasebook rules are internalized atomically. and only when missing them incurs an increased loss.
No Dropout (ablation): Given $1, always provide all rules in II" used in the translation of 81 in curriculum-text.
"We fix I epoch for fair comparison on sample efficiency.
4

---

On the Power of Context-Enhancee Learning in LLMs
Perturbing
Perturoing
Perturbing its
Perturbing n
Perturbing fts
26
100
80
Layer Index
75
75
Layer 22
60
20
16
Layer 19
11
40
50
50
40
Layer
20
25
25
20
Thought Troem
Though: Triers <te...
Ans
Thought Tecams <cTe- Ans
Thought hroers <T>... Ans
Thought Tolers <Te
Am
50 100 150 200
250
50 100 150 200 250
50 100 150 200 250
100 150 200 250
100 150 200 250 Output Tokens
Output Tokens
Outout Toloons
Output Takens
Output Tokkns
Figure 3. Evidence for sequential processing in (L, 8)-ICI--capaeee model. Fach entry is the 12 norm hetween the latent repre- sentation pre and post-perturbing कi (after certain layer at certain token position). Perturbing later phrascbooks in the context changes representations in the later layers, suggesting that later translation steps happens in later layers, To rule out the possibility that the affected depth only dependent on position of the perturbation instead the semantic content, we conduct the same set of experiment except that we only perturb rules that are not used when translation 81, which yields negligible representation difference (see comparison in Figure 8),
Drapping STR(m))
Dropping STR(m;) )
Dropping STR(ng')
Dropping STR(m))
Dropping STR(1)))
1.0
20
20
25
22-
25
22
Layer10
Lay 19
22
Layer20 -aayi22
22
Layer25 -aayr255
Exalusdion Accuacc
19
Layer14 Layer17
19
16
16
16
16
16
13
10
Layerd Layer10
10
10
10
10
7
4-
4
4-
4
0.2
10
13 25
10 13 16
19 25
10 13 16
19 25
28
10 16
19
10 13
19 25
0.0
Lend
Land
Land
Lend
Land
Figure 4. Starting from the MLT 8)-ICL-capable model fo evaluated in Figure 3, we fix a set of phrasebooks IT", train with Annealing Dropout curriculum on 100k samples, and obtain MLT -capable model fo-. We construct "stitched models" by selectively substituting layers from layer Lumn to layer Lend in fo by fo- and evaluate on ML.T with certain phrasebooks dropped from context. Bright colors elose to the diagonal reflects that knowledge required compensate the dropped phrasebook in context can be localized in small subset of layers in fa-. It is particularly worth noting that for any level i, the end of the group of layers responsible for storing sf matches the start of the layers responsible for reading (see Figure before context-enhanced learning.
3.3. Mechanistic Insights: Layer by Layer Mapping of the Translation Task
STR(T;) with STR STR(ft). The first layer showing a significant difference in its output representation is identified as the layer where the model begins processing the phrasebook.
To understand the behavior of context-enhanced learning, we probe into the hidden representations and weights of models before and after context-enhanced learning.
In Figure 3, we show that an ICL-capable model processes phrasebooks in curriculum-texx sequentially, with earlier phrasebooks read by earlier layers. Formal details are in Ap- pendix B.2 and additional experiments in Appendix C.2.
Sequential Processing in ICL-capable model
First, we look into how phrasebooks in curriculum-tex are used by (LT n)-ICcccae model during the trans- lation task via the following controlled experiment. Fix a random set of phrasebooks II = {TII...,, T and an input sa+
Localized Storage after Context-Enhances Learning
Here, we verify whether a similar sequential pattern for in- ternalizing phrasebooks is used in a MLT MLTn--apaable model. We denote the cap-ca model as fo and its post context- enhanced learning counterpart as fo- As fa- 's capability no longer depends on textual representation of the phrasebooks, our analysis focuses on the model's 's parameters.
sequence $1, we pass [ITR $1, <THINK>, into the ICL-capable model, and record the model's s hidden representations (output of each transformer block) for the 8.41-
tokens <THINK>,
Next, apply independent perturbations to each phrase- book n. in the curriculum-text and measure the change in the model's representations. That is, for each level < i < d, we randomly sample 10 rules in n, used in the translation of 8i and replace their image by other randomly chosen tuples. We denote the perturbed phrasebook (*;) and compute the norm difference between the model's representations at tokens <THINK>, .... 8d+1 before and after replacing
We assess the importance of each layer in model fa- for each phrasebook by constructing "stitched" models and measur- ing their output behavior. For every 1 < L S Lend S 28 (total layers in Llama3.2 2 3B), we replace layers L to Lend of the ICL model fa with corresponding layers from for to create stitched model. This strategy has been used in prior works on localizing information after SFT (Gong et al., 2022; Panigrahi et al., 2023; Wei et al., 2024).
5

---

On the Power of Context-Enhanced Learning in LLMs
For each level 1 S S d. we evaluate the stitched mod-
Table 2. Recovery Success Rate (Rounded 2 decimals)
els on II. using the following in-context informa-
tion: (STR(TI), (TTR()))) STR(T(41),--- (TR(=))) Note here the textual representation of कद has been dropped. If a stitched model shows high accuracy on inputs that use :; for translation but contain only the above in-context in- formation, then the layers selected from sa- to create the stitched model can be deemed responsible for storing infor- mation on x in their parameter space.
Greedy Decoding
Sampling
Token Filter
x5
None
0.00%
0.20%
0.00%
0.89%
No <THINK>
0.00%
0.20%
0.00%
0.90%
Only from A,+1
1.66%
0.20%
1.28%
0.94%
Figure 4 demonstrates that information of all phrasebooks can be localized to few mutually disjoint layers in fo.. Moreover, the end of the group of layers where x, is local- ized in fo- marks the start of the layers that begin processing level i phrasebook the ICI--capable model So (e.g. com- pare the role of layer 17 in Figures and 4). Formal details and additional experiments are in Appendices B.3 and C.2.
to a strong adversary which knows the alphabet set of the intermediate phrasebooks, but not the rules.
As shown in Table 2, recovering rules from intermediate phrasebooks (T.....,, R44 is nearly impossible without token filtering. For the final phrasebook (T5)) results remain near-random, where random guess probability is 1% for 2-tuple when = 10. Even with filtering (ii, recovery success rates remain only slightly above random. Setup details and more results are available in Appendix C.3.
This suggests that instead of storing phrasebooks as a single chunk in its parameters, the model re-learns each transla- tion step locally to compensate for missing information when rules are dropped. Thus, we conjecture that context- enhanced learning leverages curriculum-text to improve training by localizing learning in the parameter space. We build a surrogate model in Section 5 to show how such localized learning can prove beneficial for faster training.
5. Mathematical Analysis
Having empirically demonstrated the sample efficiency benefits of context-enhanced learning, now formalize them through a mathematical lens. We first define the sample complexity of an algorithm for learning the task.
4. Curriculum-text: Detectable from Queries?
Definition 5.1 (informal). Sample complexity for an algo- rithm to learn MLT, is defined as the minimum total length of all (possibly repeated) input sequences required by the algorithm to return an MLT ML.T capaable model fo-
Context-enhanced learning for MLT uses the phrasebooks in its curriculum-text. Can rules of the phrasebooks be re- covered post-hoc querying the model (Q3) using likelihood- based methods for detecting training data?
Analysis of gradient-based learning on a multilayer trans- former (let alone a 28-layer model like Llama 3,2-3B) is an open mathematical question. We use our mechanistic find- ings (i.e.. translation layers map onto transformer layers) to propose surrogate model to think about how transformers learn MLT(d, n). In this surrogate model we demonstrate that learning task MLT via vanilla SFT will require n R(d) samples. Then, we prove that, with context-enhanced learning, the surrogate model can learn MLT with a sam- ple complexity of ((poly(n) (nd log d).
We take an MLTn--capable model (with d d 5, n = 10) trained with context-enhanced learning from Section 3.2. Let STR (II*) denote the concatenation of all phrasebooks. For each rule in STR (IT*), which are of the form "a b -> D"3, we measure the probability of the model generating the ground truth tokens "C D" after observing "a b ->" We compute (1) the fraction of cases where the model's top-1 prediction matches the ground truth (greedy decoding) and (2) the probability of generating ground truth tokens via random sampling (with temperature set as 1). These are standard tests on textual description, and we would expect high probabilities for the correct tokens.
5.1. Surrogate Model (SURR-MLT)
Formalization of the surrogate model, in short SURR-MLT, relies on observations from Section 3.3, which reveal that an ICL-capable model (Definition 2.2) performs the trans- lation task stsp-by-step, with earlier phrasebook processed by lower layers. This aligns with how the model stores internalized knowledge in layers after context-enhanced learning. Our SURR-MLT represents an idealized and sim- plified transformer that has already been "pre-conditioned" to solve MLT in this sequential fashion. Using SURR-MLT we will show the benefits of context-enhanced learning.
We also test two adversarial strategies with additional token filtering* in generation: (i) setting probability of <THINK> tokens to zero (ii) when querying for a rule in T; with output alphabet A141, setting probability of all tokens outside A to zero". Note that this corresponds
"a, by C, D are generic tokens for presentation *See detailed characterization in Appendix C.3.
'We did not explore the full spectrum of attacks (e.g. adversar- ial prompt engincering) and leave that interesting future work.
6

---

On the Power of Context-Enhanced Learning in LLMs
Without loss of generality we assume the alphabet sets as A1,...,,Adt Ad+1 == A (1,2,...,,1) nn. SURR-MLT will represent length-L sequence 8, (86,1.....1... as an embedding matrix Vi E R=nxx//2 that uses {v (,,2,, ....), , (s,,)=1,,,LL) as columns. Here, for any 2-tuple (a, b). v(a, b) E IR"" represents a one-hot vector with 1 at dimension an + b. SURR-MLT operates on em- bedding matrices, transforming V + V/2 Each layer i will be primarily defined by two matrices, C. W; E I="xx" C. presents a (possibly partial or completely dropped) phrasebook that is provided contontex, and W: is trainable parameter storing phrasebook information during context-enhanced learning".
Theorem 5.3 (cf Lemma H.4). There exists transformer
that can simulate SURR-MLT with 2d self-attention and 2d MLP layers with embedding dimension 2n + 2d + 4.
5.2. Sample Complexity for Vanilla SFT
For the surrogate model, vanilla SFT corresponds to always setting in-context representations {Ci) to Os when training for MLT While we use the surrogate model for consistency in our discussion, the argument generalizes to any model learning MLTHI with vanilla SFT.
-.... - Va+1-
Our analysis is built on the Statistical Query (SQ) framework (Kearns, 1998), which measures the difficulty of learning tasks using algorithms that rely on expectation estimates of specific functions to approximate the true solution. Gradient- based methods, such as Stochastic Gradient Descent (SGD), fall under this framework as they compute gradients by es- timating expectations of loss functions and their derivatives.
Definition 5.2. SURR-MLT with trainable parameters (Wi) and in-context representation {C: is rep- resented by its operation on an input $1 as
V4+1 SURR-MLT (Wi)kz) ((C)) Vi), where
The complexity of learning a task is quantified by the SQ di- mension, which measures the number of candidate functions that are pairwise uncorrelated under the input distribution and difficult to distinguish with limited samples. A higher SQ dimension implies a richer hypothesis class, that requires more samples to identify the correct function. We show in the following theorem that the SQ dimension of MLT n) grows exponentially with task parameters.
V++1 = HardMax(C; + Wi) Shift(Vi), for 2 1,
and V/2 is the embedding matrix for the input string 81-
Here Shift represents Circular shift operation and is de- fined as Hadamard product on the embedding matrices (details in Definition G.2). Hdrd represents hard-max function converting C; + W; to a binary column stochastic matrix. In the following discussion, we show 2 examples where the surrogate model can perfectly represent an ICL. capable model and - model.
Theorem 5.4. SQ dimension of MLT nz) under uniform input distribution is at least n R(d)
Case 1 (Representing MLT n)) -nJ-IC--appalle:: The trainable matrices are all 0's as the model hasn't undergone context-enhanced learning. To produce the output for task MLTH. SURR-MLT takes phrasebooks into in-context representations by setting each C, as a stochastic matrix Matrix(m;). which represents rules of Ty as one-hot columns (please see Definition G.4 and Lemma G.5).
Informally this implies that any algorithm that tries to learn -capaaaa SURR-MLT with train- able parameters (W.) and in-context information {C always fixed at Os, by minimizing loss across samples will require at least n ((d) sample complexity.
A corollary is that for vanilla SFT with SGD, sample complexity to learn MITHI- can be least R(d) This is adapted from Edelman al. (2023), who analyse for sparse parity that has similar SQ dimension (Corollary F.6).
Case 2 (Representing -apablap For model that learning on translation
has performed task MLT no in-context information will be provided to the surrogate model and so {c,) will be all Os. A capabeeeaalll should contain the phrasebooks as [Matrix(m;)) || its trainable parameters {W.)
Informal proof for SQ dimension: Our proof extensively analyses the case where number of characters is 2 (lem. F. F.10), on which we will build proof for general n (lem. F.20). The proof for n = 2 leverages uncorrelations between two randomly selected translation tasks; i.c. we show that for two random set of phrasebooks II II the translation tasks MLTH1 MLT will have 0 output correlation with probability at least 1 - 2-9(d) w.r.t. random choice of II II (Lemma F.13). We then show that we can pick exponentially many such random set of phrasebooks for which the translation tasks will be pairwise uncorrelated. This translates to a high SQ dimension.
SURR-MLT Ideal Transformer for MLT: The fol- lowing theorem constructs transformer that can simulate SURR-MLTT Thus, while our discussions and proofs focus on the surrogate model for simplicity, they remain fully applicable to the transformer architecture.
'For simplicity of presentation C, 's are provided directly to the corresponding layers. We note that SURR-MLT can be exactly re- parameterized that C,'s 's are provided in context (i.e. concatenated with Vi). Please see Definition G.8 for more discussions.
7

---

On the Power of Context-Enhanced Learning in LLMs
5.3. Sample Complexity of Context-Enhanced Learning
Layerwise GD
Full Param GD
1.0
1.0
W:
Here, we show that context-enhanced learning substan- tially improves the sample complexity of learning in SURR-MLT. An (LT n) -I----ccaaablc model gets set of phrasebooks II by setting {Matrix(r;)) as in-context representations {C:) When a curriculum is followed such that a translation rule is dropped from a phrasebook, say n.. SURR-MLE will set the corresponding column its in-context representation C; as O's. Denote the zero-ed out column by c for some generic column index j. we note that the loss will be low if and only if the corresponding column in the learnable parameters wili) exactly matches Matrix ())
0.8
0.8
W.
0.6
0.0
Wr
W-
0.-
0.2
0.0
0.0
Wx
7500
25:00
7500 10000
Steps
Steps
Figure S. Percentage of columns in HardMax (W.)) matching Matri x(m²) when optimizing SURR-MIIT for (10,10,, 10). Learnble parameters of all layers quickly converge to the correct values, regardless of whether layer-wise updates are imposed.
5.4. Insight from SURR-MLT: Gradient Quality
A heuristic search algorithm, that searches among n² pos- sibilities for the dropped rule and stores its one-hot repre- sentation in the corresponding column in W, can be used to minimize the loss. By sequentially dropping rules, fol- lowed by search and store process, the algorithm achieves polynomial sample complexity.
In this section, we show the major difference between context-cnhanced learning and vanilla SFT to be the amount of predictive information in gradients for the trainable parameters. Our theoretical analysis in Theorem 5.6 on SURR-MIL primarily shows that when a single translation rule in a phrasebook पद dropped, equivalently a column in the in-context representation, say c(O) is zero-ed out, the gradient of the corresponding column in trainable parameter w aligns strongly with the one-hot vector representa- tion of the dropped rule, Matrix(m;)) ) However, this strong alignment heavily relies on the presence of other rules io-toxtext When more rules are dropped, the gradi- ent signal gets increasingly "inaccurate". We quantitatively characterize such degradation as follows:
Theorem 5.5 (Informal; ef Corollary G.19). For any task ML there is a heuristic search algorithm paired with a curriculum of iteratively dropping rules from phrasebooks, that can learn a MLTTI--capai -capable SURR-MIL with sample complexity On ad log d) with high probability.
The enumerative step in the heuristic search algorithm re- quires O(n²) steps on average, as the algorithm needs to search over (n²²) possibilities when a rule is dropped from a phrasebook. Instead, we show that gradient descent re- quires only a few steps per dropped rule. Dropping a rule sets the respective column in the in-context representation to O, causing the gradient for the corresponding column in the trainable parameters to strongly align with the missing column. We formally present results for d = 2; due to expo- nentially growing number of terms to analyse with higher d, we keep the result for general d as a conjecture.
Consider an ICL-capable SURR-MLT, we focus on col- umn (with index j. denoted as superscript) w in the first layer with ground truth Matrix(m;) We define gradient prediction accuracy (see formal definition in Def- inition B. for w(()) as the probability that the argmax entry of the negative stochastic batch gradient on w(6) matches the argmax entry of Matrix(Ti) Intuitively, higher gradient prediction accuracy will make learning Mat Matrix (ix((")) easier. We track the accuracy metric when only c(u) is zero-ed out as well as when taking more ag- gressive context dropout schemes.
Theorem 5.6 (Informal; Corollary G.25). When d = 2, there is a gradient descent based algorithm, paired with curriculum of iteratively dropping random rule from phrasebooks, that can return MLT -capbble SURR-MLT with sample complexity Onn²) with high probability.
Batch size
Rule Dropping Rate
Prediction Accuracy
10a
10a
10a
10a
0.1
0.2
0.3
0.4
0.5
While theoretical analysis of GD dynamics beyond d = 2 is challenging, empirically we show that when training with SGD, the trainable parameters for much deeper SURR-MIL MLT can quickly learn the set of phrasebooks in II. In Figure 5, we show context-enhanced learning results for SURR- MLT a randomly selected set of phrasebooks II" in MLT(10, 10). Regardless of whether we learn each layer separately with layer-wise SGD or all layers simultancously with SGD in the surrogate model, the trainable parameters learn the phrasebooks in II very quickly. More discussions are deferred to Appendix G.5.
1.0
Gradient
0.8
0.6
0.4
0.00
0.25
0.50
0.75
1.00
6
Rule Dropping Rate
Max. Phrasebook Index
Figure 6. Gradient prediction accuracy when (left) varying drop- ping rates for rules in n and batch sizes of gradient computation, and (right) varying dropping rates and max phrasebook index (k) for rules dropped from Ti,... Tk. Higher dropping rates always lead lower gradient prediction accuracy.
8

---

On the Power of Context-Enhancd Learning in LLMs
In Figure 6, we report the gradient prediction accuracy with dropout schemes involving (1) multiple rules dropped out from the first phrasebook and (2) rules dropped from mul- tiple phrasebooks. We can see that higher dropping rates significantly degrade the gradient prediction accuracy, high- lighting the necessity of proper in-context information in order to obtain the optimization benefit. More details on the metric and experiments are deferred to Appendix B.I. We note that finding above is based observations the simplified surrogate model. A quantitative characterization of the optimization benefit for context-enhanced learning in the LLM regime is left as an interesting future work.
edge memorization) during pre-training of language models (Chan et al., 2022; Reddy, 2024; Singh et al., 2023; 2024; Nguyen & Reddy, 2024). These works show that data dis- tribution properties affect the behavior of the model during training. Our work can be thought of as contemporary to the works above, where we show that strong ICL capabili- ties can be utilized for improving knowledge on a task by context-enhanced learning.
Benefits of in-context learning: In-context learning has been primarily studied in the context of few-shot prompting of large language models. Better supervision with in-context supervision can help in improved performance (e.g. some representative works (Arora et al., 2022: Si et al., 2022: Wu et al., 2022; Lu et al., 2021; Su et al., 2022)), OOD generalization and factuality (reduced hallucination) (Yang et al., 2023; Dhuliawala et al., 2023; Chen et al., 2023; Didolkar et al., 2024), and more structured latent represen- tations (Park et al., 2024) for large language models. On the other hand, we show that improved supervision with in-context supervision can also help a model learn faster in SFT, while seemingly not leaking the in-context information in its output probabilities.
6. Related Works
In this section, we discuss related works to this paper. More related works regarding (1) compositional and OOD gener- alization and (2) mechanistic understanding of transformers are discussed more extensively in Appendix D.
Learning using Privileged Information (LUPI)
LUPI was formally introduced by Vapnik & Vashist (2009) in kernel SVMs, and it is heavily related to concepts of Learning with Side Information (Kuusela & Ocone, 2002; Jonschkowski et al., 2015; Zhang et al., 2018). Primarily, both concepts refer to training a model with additional infor- mation that could help training but may not be available at test time. This framework has been extended theoretically for classification tasks (e.g. Pechyony& & Vapnik (2010); Momeni et al. (2018)) and has been used to explain benefits of knowledge distillation (Vapnik et al., 2015; Lopez-Paz et al., 2015). To name a few applications, this concept has been heavily studied for improving boosting for classifica- tion tasks (Chen al., 2012), visual and video encoders (Hoffman et al., 2016; Cheng al., 2020; Xu al., 2017), human preference predictions (Farias & Li, 2019), multi- agent games (Sessa et al., 2020), speech recognition (Syn- naeve et al., 2014), and medical recognition (Ceccarelli & Maratea, 2008; Sabeti et al., 2020).
7. Discussion, Limitations, and Future work
Some experimental works have implicitly used the notion of context-enhanced learning but the current paper formalized this notion for ruto-regressive models and showed, using MLT, that this form of learning can be exponentially more sample-efficient than standard SFT. At the end of training it is hard to recover the in-context information seen during training from the model's output probabilities. We note that this finding appears to have implications about copyright law (e.g.. whether or not LLM training amounts to "transforma- tive use" of text (Carlini et al., 2021; 2022; Karamolegkou et al., 2023)) whose further study left for future work.
Our experiments focus on synthetic MLT task for a few reasons: (1) to ensure that the task is absent from LLM pre-training, which allows precise quantification of benefits of context-enhanced learning, including not revealing the curriculum text at inference time. (2) the task is too difficult (at least for Llama 3.2 3B model) to learn via vanilla SFT, but is learnable via context-enhanced learning. Extending these findings real-world complicated tasks (e.g., in math and coding) is left for future work.
While LUPI has rich application in classification, extending LUPI to LLMs introduces unique challenges due to their ruto-regressive training. Such concept raises questions on whether applying auto-regressive loss on the additional information is necessary to get the benefits of additional supervision information in context, and how the additional information changes the training behavior of LLMs. Hence, our work is nontrivial generalization of this framework to LLMs that connects to their in-context learning strengths.
Our convergence analysis for context-enhanced learning re- lies on surrogate model, and extending it to an actual trans- former remains an open challenge for theory of deep learn- ing. Extending formalization of context-enhancee learning to explore LLM training in multi-agent settings-where where mod- els collaborate and learn from each other to discover novel condepts-would be an exciting avenue for future research.
In-context learning and memorization
Recent works have studied emergence ICL and competi- tion with in- inweights learning (IWL) (equivalent to knowl-
9

---

On the Power of Ennhace-Enhaneed Learning in LLMs
Acknowledgment
language models. Advances in Neural Information Pro-
cessing Systems, 35:38546-88556, 2022.
We thank Yun Cheng. Simon Park, Tianyu Gao, Yihe Dong, Zixuan Wang. Haoyu Zhao, and Bingbin Liu for discus- sions, suggestions, and proof-reading at various stages of the paper. AP, SA acknowledge funding from NSF, PLI, DARPA, ONR, and OpenAI. XZ is additionally supported by Gordon Y.S. Wu Fellowship in Engineering.
Arora, S., Narayan, A.,
Chen, M. F., Orr, L., Guha, N., Bha- tia, K., Chami, I., Sala, F., and Ré, C. Ask me anything: A simple strategy for prompting language models. arXiv preprint arXiv:22 10.02441, 2022.
Bavarian, M., Jun, H., Tezak, N., Schulman, J., McLeavey, C., Tworek, J., and Chen, M. Efficient training of language models to fill the middle. arXiv preprint arXiv:2207 2022.
Impact statement
We formulate a basic notion "context-enhanced learning". and study how it is different from usual learning. One consequence of our study is that enhancing the context with good quality data can enhance the quality of the learner.
Bhattamishra, S., Ahuja, K., and Goyal, N. On the ability and limitations of transformers to recognize formal lan- guages. In Proceedings of the 2020 Conference on Empir- ical Methods in Natural Language Processing (EMNLP), pp. 7096-7116, 2020.
This enhancement could be used with privileged information and it appears that the training can be done in such a way that the model does not leak this privileged information. This can be seen as enhancing privacy, since there is lower chance of leaking privileged training data.
Blum, A., Furst, M., Jackson, J., Kearns, M., Mansour, Y., and Rudich, S. Weakly learning dnf and characterizing statistical query learning using fourier analysis. In Pro- ceedings of the twenty-sixth annual ACM symposium Theory of computing, pp. 253-262, 1994.
The flip side is that it suggests albeit in very toy and synthetic setting, with full study left for future work - that model training could use off-limits data (albeit with no gradient updates on it) and this use might not be detectable from querying the trained model. But this is hypothetical at this point since the paper concerns very toy setting.
Brown, T. B., Mann, B., Ryder, N., Subbiah, M., Kaplan, J., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. arXiv preprint arXiv:2005.. 14165, 2020.
References
Carlini, N., Tramer, F., Wallace, E., Jagielski, M., Herbert-
Voss, A., Lee, K., Roberts, A., Brown, T., Song, D., Erlingsson, U., al. Extracting training data from large language models. In 30th USENIX Security Symposium (USENIX Security 2/), pp. 2633-2650, 2021.
Akyurek, E., Schuurmans, D., Andreas, J.. Ma. T., and Zhou, D. What learning algorithm is in-context learn- ing? investigations with linear models. In The Eleventh International Conference on Learning Representations, / wwwpoppeeww.eeeeeeeeeeeee id=0g0X4HByN4I.
2023. URL
Carlini, N., Ippolito, D., Jagielski, M., Lee, K., Tramer, F., and Zhang. C. Quantifying memorization across neu- ral language models. arXiv preprint arXiv:2202.077 07646. 2022.
Akyürek, E., Wang, B., Kim, Y., and Andreas, J. In-context language learning: Architectures and algorithms. arXiv preprint ar arXiv:2401.12973, 2024.
Ceccarelli, M. and Maratea, A. Improving fuzzy clustering of biological data by metric learning with side informa- tion. International Journal of Approximate Reasoning, 47(1):45-57, 2008.
Allen-Zhu, Z. and Li, Y. Physics of language models: Part I, Context- free grammar. arXiv preprint arXiv:2305.13673, 2023a.
Chan, S., Santoro, A., Lampinen, A., Wang, J., Singh, A., Richemond, P., McClelland, J., and Hill, Data distri- butional properties drive emergent in-context learning in transformers. Advances in neural information processing systems, 15:18878-11 18891, 2022.
Allen-Zhi Zhu, Z. and Li, Y. Physics of language mod- els: Part 3.2, Knowledge manipulation. arXiv preprint arXiv:2309. 2tti:2309..4402, 2023b.
Allen-Zhu. Z. and Li, Y. Physics of language models: Part 3.3, knowledge capacity scaling laws. arXiv preprint arXiv:2404.05405, 2024.
Chen, J., Liu, X., and Lyu, S. Boosting with side infor- mation. In Asian Conference on Computer Vision, pp. 563-577. Springer, 2012.
Anil, C., Wu, Y., Andreassen, A., Lewkowycz, A., Misra, V., Ramasesh, V., Slone, A., Gur-Ari, G., Dyer, E., and Neyshabur, B. Exploring length generalization large
Chen, J., Pan, X., Yu, D., Song. K., Wang, X., Yu, D., and Chen, J. Skills-in-contex prompting: Unlocking
10

---

On the Power of Context-Enhanced Learning in LLMs
compositionality in large language models. arXiv preprint arXiv:2308. 00:04, 2023.
Association for Computational Linguisties: ACL 2022,
pp. 1479-1472, 2022.
Cheng, I.., Zhou, X., Zhao, I.,, Li, D., Shang, H., Zheng, Y., Pan, P., and Xu, Y. Weakly supervised learning with side information for noisy labeled images. In Computer Vision-ECCV 2020: 16th European Conference, Glasgow, UK, August 23-28, 2020, Proceedings, Part XXX 16, pp- 306-321. Springer, 2020.
Goyal, S., Ji, Z., Rawat, A. S., Menon, A. K., Kumar, S., and Nagarajan, V. Think before you speak: Train- ing language models with pause tokens. arXiv preprint arXiv:2310.0222 2023.
Hao, S.. Sukhbaatar, S., Su, D., Li, X., Hu, Z., Weston, J., and Tian, Y. Training large language models to reason in a continuous latent space. arXiv preprint arXi::2412066763 2024.
Choi, Y., Asif, M. A., Han, Z., Willes, J., and Krishnan, G. Teaching Ilms how to learn with contextual fine-tuning. arXiv preprint arXiv:2503.09032 2025.
Hendrycks, D. and Gimpel, K. Gaussian error linear units (gelus). arXiv preprint arXiv:1606.0844 2016.
Deng, Y., Choi, Y., and Shieber, S. From explicit cot to implicit cot: Learning to internalize cot step by step. arXiv preprint arXiv:2405. 14838, 2024.
Hoffman, J., Gupta, S., and Darrell, T. Learning with side in- formation through modality hallucination. In Proceedings of the IEEE conference on computer vision and pattern recognition, pp. 826 834, 2016.
Dhuliawala, S., Komeili, M., Xu, J., Raileanu, R., Li, X., Celikyilmaz, A., and Weston, J. Chain-of-verification reduces hallucination in large language models. arXiv preprint arXiv:2309. 1495, 2023.
Jonschkowski, R., Höfer, S., and Brock, O. Patterns for learning with side information. arXiv preprint arXiv:1511.06422 2015.
Didolkar, A., Goyal, A., Ke, N. R., Guo, S., Valko, M., Lillicrap, T., Rezende, D., Bengio, Y., Mozer, M., and Arora, S. Metacognitive capabilities of llms: An explo- ration in mathematical problem solving. arXiv preprint arXiv:2405. 12205, 2024.
Karamolegkou, A., Li, Zhou, L., and Søgaard, A. Copy- right violations and large language models. arXiv preprint arXiv:2310.1377/ 2023.
Kearns, M. Efficient noise-tolerant learning from statistical queries. Journal of the ACM (JACM), 45(6))983-1006 1998.
Donahue, C., Lee, M., and Liang, P. Enabling language mod- cls to fill in the blanks. arXiv preprint 2rX5v:2005.05339, 2020.
Kenton, D. M.-W. C. and Toutanova, L. K. Bert: Pre- training of deep bidirectional transformers for language understanding. In Proceedings of naacL-HLT, volume 1. Minneapolis, Minnesota, 2019.
Dubey, A., Jauhri, A., Pandey, A., Kadian, A., Al-Dahle, A., Letman, A., Mathur, A., Schelten, A., Yang, A., Fan, A. et al. The lama herd of models. arXiv preprint arXiv:2407.21788. 2024.
Knudsen, L. R. Practically secure feistel ciphers. In In- ternational Workshop on Fast Software Encryption, pp. 211-221. Springer, 1993.
Edelman, B. L.,, Goel, S., Kakade, S., Malach, E., and Zhang, C. Pareto frontiers neural feature learning: Data, com- pute, width, and luck. arXiv preprint arXiv:2309.03880) 2023.
Kuusela, P. and Ocone, D. Learning with side information: i.rt 02 2002.
Eldan, R. and Li, Y. TinyStories: How small can lan- guage models be and still speak coherent English? arXiv preprint arXiv:2305.0773
Li, D., You, J., Funakoshi, K., and Okumura, M. A-tip: attribute-aware text infilling via pre-trained language model. Proceedings of the 29th International Con- ference on Computational Linguistics, pp. 5857-5869, 2022.
2023.
Farias, F.F. and Li, A. A. Learning preferences with side information. Management Science, 65(7):3131-3149, 2019.
Li, Y., Li, Y., and Risteski, A. How do transformers learn topic structure: Towards a mechanistic understanding. In International Conference on Machine Learning, pp.
Gao, T. Wettig, A., He, I.., Dong, Y., Malladi, S., and Chen, D. Metadata conditioning accelerates language model pre-training. arXiv preprint arXiv:2501.0195: 2025.
19689-19729, PMLR, 2023.
Liao, H., He, S., Hao, Y., Li, X., Zhang, Y., Liu, K., and Zhao, J. Skintern: Internalizing symbolic knowledge for distilling better cot capabilities into small language models. arXiv preprint arXiv:2409. /3/83, 2024.
Gong, Z., He, D., Shen, Y., Liu, T.-Y., Chen, W., Zhao, D., Wen, J.-R., and Yan, R. Finding the dominant winning ticket in pre-trained language models. In Findings of the

---

On the Power of Context-Enhanced Learning in LLMs
Liu, Y., Ott, M., Goyal, N., Du, J., Joshi, M., Chen, D., Levy, O., Lewis, M., Zer, L., and Stoyanov, V. Roberta: A robustly optimized bert pretraining approach. arXiv preprint arXiv: petp:::0777999999 2019.
Press, O., Zhang, M., Min, S., Schmidt, L.,, Smith, N. A., and Lewis, M. Measuring and narrowing the com- positionality gap in language models. arXiv preprint arXiv:2210.03350, 2022.
Lopez-Paz, D., Bottou, L.., Schölkopf, B., and Vapnik, V. Unifying distillation and privileged information. arXiv preprint arXiv: wr.Xiv:511.03060 1.06643, 2015.
Raffel, C., Shazeer, N., Roberts, A., Lee, K., Narang, S., Matena, M., Zhou, Y., Li, W., and Liu, P. Exploring the limits of transfer learning with unified text-to-text transformer. Journal of machine learning research, 21 (140):1-67, 2020.
Loshchilov, and Hutter, F. Sgdr: Stochastic gra- dient descent with warm restarts. arXiv preprint arXiv:1608.03988) 2016.
Ramesh, R., Lubana, E. S., Khona, M., Dick, R. P., and Tanaka, H. Compositional capabilities of autoregressive transformers: A study on synthetic, interpretable tasks. arXiv preprint arXiv:231 1,1997,, 2023.
Loshchilov, I. and Hutter, F. Decoupled weight decay reg- ularization. In International Conference on Learning Representations, 2019. https:/ ppppeeeeeeeeeeeeeeoeeeeee www.ooooummonueiiiiiiiooo
Reddy, G.
The mechanistic basis of data dependence and abrupt learning in in-context classification task. In The Twelfth International Conference on Learning Represen- tations, 2024.
Lu, Y., Bartolo, M., Moore, A., Riedel, S., and Stenetorp, P. Fantastically ordered prompts and where find them: Overcoming few-shot prompt order sensitivity. arXiv preprint MiXv:221040087866 2021.
Sabeti, E., Drews, J., Reamaroon, N., Warner, E., Sjod- ing, M. W., Gryak, J., and Najarian, K. Learning using partially available privileged information and label uncer- tainty: application in detection of acute respiratory dis- tress syndrome. IEEE journal of biomedical and health informatics, 25(3):784-796 2020.
Momeni, A., Tatwawadi, K., and Stanford, U. Understand- ing lupi (learning using privileged information). Tono- sphere, 201(7);6, 2018.
Motwani, R. Randomized Algorithms. Cambridge Univer- sity Press, 1995.
Sessa, P. G., Bogunovic, I., Krause, A., and Kamgarpour, M. Contextual games: Multi-agent learning with side information. Advances in Neural Information Processing Systems, 33:21912 229222, 2020.
Nanda, N., Chan, L., Lieberum, T., Smith, J., and Steinhardt, J. Progress measures for grokking via mechanistic inter- pretability. In The Eleventh International Conference Learning Representations, 2023.
Si, C., Gan, Z.., Yang. Z., Wang, S., Wang, J., Boyd-Graber, J., and Wang, L. Prompting gpt-3 to be reliable. arXiv preprint arXiv:22 10.09/50 2022.
Nguyen, A. and Reddy, G. Differential learning kinet- ics govern the transition from memorization to gen- eralization during in-context learning. arXiv preprint arXiv:2412.00104, 2024.
Singh, A., Chan, S., Moskovitz, T., Grant, E., Saxe, A., and Hill, The transient nature of emergent in-context learning in transformers. Advances in Neural Information Processing Systems, 36:27801-27819, 2023.
Panigrahi, A., Saunshi, N., Zhao, H., and Arora, S. Task- specific skill localization in fine-tuned language models. In International Conference on Machine Learning, pp. 27011-27033. PMLR, 2023.
Singh, A. K., Moskovitz, T., Hill, F., Chan, S. C., and Saxe, A. M. What needs to go right for an induction head? A mechanistic study of in-context learning circuits and their formation. arXiv preprint Xiv::2404.07122 2024.
Park, C. F., Lee, A., Lubana, E. S., Yang. Y., Okawa, M., Nishi, K., Wattenberg, M., and Tanaka, H. Ielr: In-context learning of representations. arXiv preprint arXiv:2501.00070 2024.
Spencer, Asymptotic lower bounds for ramsey functions. Discrete Mathematics, 20:69-76. 1977.
Pechyony, D. and Vapnik. V. On the theory of learnining with privileged information. Advances in neural informa- tion processing systems, 23, 2010.
Su, D., Sukhbaatar, S., Rabbat, M., Tian, Y.. and Zheng, Q. Dualformer: Controllable fast and slow thinking by learning with randomized reasoning traces. arXiv preprint arXiv:2410.0991 2024.
Pfau, J., Merrill, W., and Bowman, S. R. Let's s think dot by dot: Hidden computation in transformer language models. arXiv preprint orXiv:2404 2400.25444.5... 2024.
Su, H., Kasai, J., Wu, C. H., Shi, W., Wang, T., Xin, J., Zhang, R., Ostendorf, M., Zettlemoyer, I.., Smith, A.,
12

---

On the Power of Context-Enhanced Learning in LLMs
et al. Selective annotation makes language models bet- ter few-shot learners. arXiv preprint arXiv:2209.01975, 2022.
Yang. S., Gribovskaya, E., Kassner, N., Geva, M., and Riedel, S. Do large language models latently perform multi-hop reasoning? arXiv preprint 142.177402.16837, 2024.
Synnaeve, G., Schatz, T., and Dupoux, E. Phoneties em- bedding learning with side information. In 2014 IEEE Spoken Language Technology Workshop (SLT), pp. 106- 111, 2014. doi: 10. 00.1099/ST22211.777755.
Yao, S., Peng, B., Papadimitriou, C., and Narasimhan, K. Self-attention networks can process bounded hierarchical languages. In Proceedings of the 59th Annual Meeting of the Association for Computational Linguistics and the 1ith International Joint Conference on Natural Language Processing (Volume 1: Long Papers), pp. 3770-3785, 2021.
Team, G.. Amil, R.,. Borgeaud. S., Alayrac, J.-B.. Yu, J.. Sori- cut, R., Schalkwyk, J., Dai, A. M., Hauth, A., Millican, K., et al. Gemini: a family of highly capable multimodal 2023.
models. arXiv preprint arXiv:2312.
Yu, D., Kaur, S., Gupta, A., Brown-Cohen, J., Goyal, A., and Arora, S. Skill-mix: A flexible and expand- able family of evaluations for ai models. arXiv preprint arXiv:2310.1753 2023.
Touvron, H., Lavril, T., Izacard, G., Martinet, X., Lachaux, M.-A., Lacroix, T., Rozière, B., Goyal, N., Hambro, E., Azhar, F., et al. Llama: Open and efficient foundation lan- guage models. arXiv preprint 2.2.:::22302.39977 2023.
Yu, P., Xu, J., Weston, J., and Kulikov, I. Distilling system 2 into system 1. arXiv preprint arXiv:2407.06023. 2024.
Vapnik, V. and Vashist, A. new learning paradigm: Learn- ing using privileged information. Neural networks, 22 (1-6):5445557 2009.
Zhang, and Sennrich, R. Root mean square layer nor- malization. Advances in Neural Information Processing Systems, 32, 2019.
Vapnik, V., Izmailov, R., et al. Learning using privileged information: similarity control and knowledge transfer. J. Mach. Learn. Res., 16(1) (6(1):2023-2044 2015.
Zhang, R., Nie, F., Guo, M., Wei, X., and Li, X. Joint
learning of fuzzy k-means and nonnegative spectral clus- tering with side information. IEEE transactions Image Processing, 28(5): 2018. 2018.
Wang, B., Yue, X., Su, Y., and Sun, H. Grokked transform- are implicit reasoners: A mechanistic journey to the edge of generalization. arXiv preprint arXiv:2405.15071, 2024.
Zhao, H., Panigrahi, A., Ge, R., and Arora, S. Do trans- formers parse while predicting the masked word? In Proceedings of the 2023 Conference on Empirical Meth- ods in Natural Language Processing, pp. 16513-16542, 2023.
Wei, B., Huang, K., Huang, Y., Xie, T., Qi, X., Xia, M., Mittal, P., Wang, M., and Henderson, P. Assessing the brittleness of safety alignment via pruning and low-rank modifications. arXiv preprint 102.5 2024.
Zhao, H., Kaur, S., Yu, D., Goyal, A., and Arora, S. Can models learn skill composition from examples? In The Thirty-eighth Annual Conference on Neural Information Processing Systems, 2024.
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., al. Emergent abilities of large language models. arXiv preprint 22X.::22006.078222 2022.
Zhong, Z., Liu, Z., Tegmark, M., and Andreas, J. The clock and the pizza: Two stories in mechanistic explana- tion of neural networks. Advances in neural information processing systems,
Wu, Z., Wang. Y., Ye, J., and Kong, L. Self-adaptive in- context learning: An information compression perspec- tive for in-context example selection and ordering. arXiv preprint ar arXiv:2212. 10375, 2022.
2023.
Zhou, H., Bradley, A., Littwin, E., Razin, N., Saremi, O.. Susskind, J., Bengio, S., and Nakkiran, P. What algo- rithms can Transformers leam? A study in length gener- alization. arXiv preprint arXiv:23 16028, 2023.
Xu, H., Gao, Y., Yu, F., and Darrell, T. End-to-end learning of driving models from large-scale video datasets. In Proceedings of the IEEE conference on computer vision and patlern recognition, pp. 2174-2182, 2017.
Zou, J.. Zhou, M., Li, T., Han, S., and Zhang, D. Prompt- intern: Saving inference costs by internalizing recurrent prompt during large language model fine-tuning. arXiv preprint arXiv: 01 2024.
Yang. L.. Zhang, S., Yu, Z., Bao, G., Wang. Y., Wang, J., Xu, R., Ye, W., Xie, X., Chen, W., et al. Supervised knowledge makes large language models better in-context learners. arXiv preprint arXiv:2312. 15918, 2023.
13

---

On the Power of Fontext-Enhaneel Learning in LLMs
Appendix
A. Overview of the Appendix and Common Notations
Here, we outline the structure for the appendix for easier readability. Due to space constraints, we had to defer a lot of details from the main paper. Appendix B.1 contains details on experiments on gradient prediction accuracy from Section 5.4. Appendix B.4 shows the details on CoT internalization pipeline that we followed to get an (L.T n)-ICL-capaable model. Appendix B.5 gives more details on the context-enhanced learning experiments conducted in Section 3.
Appendix C shows mechanistic experiments, designed in Section 3.3, for additional experimental settings. Appendix C.3 present additional details and results for information recovery through querying experiments conducted in Section 4. Appendix D presents additional relevant related works, including discussion on OOD generalization of language models and mechanistic understanding of large transformer architectures. Appendix discusses few properties on the MLT. Appendix presents the formal theorems on the computational hardness of learning the translation task and their proof', which were informally outlined Section 5.2 in the main paper. Appendix G then presents the formal statements and proofs for context-enhanced learning in the surrogate model, which were informally outlined in Section 5.1 in the main paper. Finally, Appendix H presents the theoretical construction of an ideal transformer that simulate the surrogate model. We present extensive details on prompts, and examples of training sequences, that we used for context-enhanced learning in Appendix I.
Table 3. Important notations
Scope of Notation
Symbol
Description
A
A set
A
A matrix
General Notations
A
The j-th column of matrix A
ej
One-hot embedding vector with 1 at dimension k:
X
All possible text strings (input space for causal LM)
y
All possible distribution over texts (output space for causal LM)
fo
Causal LM parameterized by 0
Language Modeling
Language task mapping imputs E Xg CX to distribution Y e y Auto-regressive cross entropy loss
Cauto CURR, C t)
In-context curriculum for learning task g with input æ at step t
d
Depth of translation task
72
Number of characters in each alphabet
A
An alphabet set
8
sequence in alphabet
phrasebook between 2-tuples in two alphabets, e.g. (m : Af - AB).
MLT Translation Task
B
Space of all possible phrasebook on A - A3
II
of phrasebooks { defining translation task
MIT
Translation task with set of phrasebooks II
MLT n)
Family of translation tasks of depth d and characters
STR STR(T) (T)
Descriptive text for phrasebook π (see Section 2.2)
V
Embedding matrix representing sequence 8, (Definition G.I)
C
In-context information matrix for level Trainable parameter matrix for level Effective translation matrix for level (Definition G.6)
W;
Surrogate Model
P.
HardMax
Column-wise hard-max function
Matrix(T))
Matrix representation of a phrasebook æ (Definition G.4)
SURR-MLT (W:)
Surrogate model parameterized by W,'s (Definition G.7)
14

---

On the Power of Context-Enhanced Learning in LLMs
B. Deferred definitions, and experimental details from the main paper
B.1. Gradient Prediction Accuracy
Our theoretical analysis in Theorem 5.6 was built on the fact that when a translation rule in a phrasebook is dropped by zeroing out a column in an in in-context representation C₁ the gradient for the corresponding column W; points the direction of the dropped rule.
On the other hand, show that when multiple rules simultaneously dropped from the phrasebooks, the gradients for the trainable parameters become increasingly noisy. To quantify this degradation, we compute gradient for each column of the trainable parameters for an ICL-capable model, when the corresponding rule is dropped from phrasebooks by zeroing out the relevant column in the in-context representations. We then compute whether the computed gradient points to the right rule. By progressively increasing the number of simultaneously dropped rules, we measure the resulting degradation in the accuracy of the gradient's 's predictions.
More formally, denote RANDOM- DROP an operation that takes in column dropping rates per layer p1.". Pa, set of phrasebooks II", and returns in-context representations {C+] such that c( = Matrix (T;) (1) (jth columns of C₁ and Ma Matrix (x.) are equal) with probability Pt and 0 otherwise. Then,
Definition B.1. For set of column dropping rates per layer P.. Pd, set of phrasebooks II' and MLT n)- ICL-capable model, predictive accuracy of gradients is defined as
E (C = RandOM DROP(P1. -Pa.II") HardMax (-v wVr JE(I,?
= (Matrix (T()) C) (-vwao
f= E= SURR-MLT (Wi)2, ((C)) Vi) ML. (+1))
where f, adapted from Section 2.1, computes cross-entropy loss on the predicted output embeddings of surrogate model using true output string and denotes the indicator function.
The above definition computes gradients on the expected loss of the model. We primarily focus on predicting the trainable parameters of the first layer, i.e. Wi. as that is the deepest layer in the surrogate model and intuitively should suffer the most with noise accumulation from dropped rules. On the other hand, we can further adapt the definition to compute the accuracy for batched gradients, where the gradients are computed using average loss on randomly sampled batch of input sequences.
In Figure 6, we report the predictive accuracy of gradient for an MLT dd n) -CCL-clbee model, and its behavior with varying batch size and the column dropping rates. We report for two cases, where column dropping rates is non-zero only for the first phrasebook, and one where increase the number of phrasebooks for which rules are independently and uniformly dropped. In both cases,
Increased column dropping rates leads to noisier gradients and reduced prediction accuracy.
Larger batch sizes improve gradient accuracy but cannot fully compensate for high dropout rates.
Dropping rules from multiple phrasebooks significantly degrades gradient prediction accuracy.
15

---

On the Power of Ennance-Enhaneed Learning in LLMs
B.2. Hidden Representations of Model
A transformer fe with embedding dimension p and K layers takes any input sequence æ, say of length L, converts to an embedding matrix H IR lxx,, and modifies the embeddings using succession of K transformer layers; which will denote by 1 _. .... SB _(k)) We refer to the hidden representations for the input æ, with embedding matrix Hi, as the output of the model after every layer. We will denote them as H++1 TR. Lxp for the output of layer f(d) That is,
H
(
(1) (H)), for all i 122 1. 1 5
22-norm in change in hidden representation with perturbation in in-context information For an MLT n)- ICL-capable model, we supply in-context information for the textual description of a set of phrasebooks II as (TR(II) = [STR(T)).... (TR((2))- Our inputs to the transformer for an input string 8) will be of the form [STR (II). 81, <THINK> ".. where su+1 = (LTn By the definition of hidden representations, H2,- HK+1 will denote the output of the transformer layers for this input string. However, we will be only in- terested in the hidden representations for the tokens involved in the tokens for <THINK>, .... 8441; and we will refer to the corresponding subsets of H2. .K++1 that represent these specific tokens as V2 VK+1-
Now, suppose we randomly take phrasebook T; in II and change to a random phrasebook ä, The corresponding textual description that will augment the context for an input string will then be [STR(a)), ... (1) STR(),))R()..... STR((()))ST(TTTT STR(((+1), ..... STR((T)) If V2,..., VK+1 now denote the hidden representations that represent the tokens for <THINK> 8d+1, then the Ez-norm in the change of the hidden representation after layer j (for any S j S K) with the perturbation in II will be given by | -V3
B.3. Definition of "stitched" model
We reuse notations from Appendix B.2. Suppose we have an MLT ni-ICL-capablee model fo and an MLT mI- capable model fa.. Their corresponding transformer layers are denoted by f _ ... ((()) (K) and (()) c(2)
f0- Each
k
model takes in an input sequence and processes them with their transformer layers.
Formally, we will write for the ICL-capable model. It takes input sequence æ, and converts to an embedding matrix, say H1, and the output after the Klayers are given by:
H 16()
H
Process of "stitching": The process of stitching takes in two parameters Lytat and Lend and replaces all layers from Lstart to Lend in fo with the corresponding layers in fo- to give a "stitched" model, say _o The output of the "stitched" model f0.g beat on an input sequence æ will be given by
HK+1 1 .... f (L
(Leaa)
(Lam)
off offleen-1 o 1!() 0 (H1).
(1)
Layers are replaced by layers from for
B.4. Pipeline on CoT Internalization
We randomly sample M sets of phrasebooks Wi...., TM not equal to II", For each set of phrasebooks n,, we randomly sample a single input sequence 81 and compute all the intermediate translation steps , We first train the model to do robust explicit CoT by auto-regressive training on sequences [STR(T;), $1, $2,..., Sd, s1+1] with loss computed over ......, 8441. Then we follow common CoT internalization strategies (Deng et al., 2024; Hao et al., 2024; Yu et al., 2024; Su et al., 2024) and gradually replace the intermediate sequences by <THINK> tokens in training. After all intermediate sequences have been replaced, the model has low loss $2+1 with input [ITR(II,). 81, <THINK>,-........ .... <TH <THINK>, s4+1), satisfying Definition 2.2. Since we only sample sequence per set of phrasebooks, there is little memorization on particular phrasebooks.
Details on training hyperparameters: We use M = 3 x 10 random sets of phrasebooks with length between 20 and 40, for getting MLT (5, ICCL-caaeel model, and M = 10 random sets of phrasebooks with length between 20 and 40, for getting (L, 10)-CCL-eae model. We use cosine learning rate schedule (Loshchilov & Hutter, 2016), with peak
16

---

On the Power of Context-Enhancd Learning in LLMs
learning rate 10-4 and 6% warmup phase, where learning rate is linearly increased from 0 the peak. We AdamW optimizer (Loshchilov & Hutter, 2019) with weight decay fixed 10-1 We batch size of 64 for training.
CoT internalization curriculum: For the first 10% fraction of training, we train the model with explicit CoT tokens that contain the intermediate steps in translation. Then between 10% to 60% fractions of training. CoT tokens are gradually replaced by <THINK> tokens, with the rate of replacement increasing linearly from 0% to 100%. We follow deterministic first-to-last order for replacing CoT tokens; earlier CoT tokens are replaced first with <THINK> tokens. After that, the model trained with the <THINK> CoT tokens till the end of training.
B.5. Experiment configuration for context-enhanced learning
For context-enhanced learning, we create supervised datasets Dr- of different sizes; each containing between 10 to 10th samples. When performing Annealing Dropout Fixed Dropout, each step of training, we randomly perform dropout the rules of all phrasebooks apply dropout the rules of randomly sampled phrasebook to define curriculum-text That is, if 7i, .......> # represent the phrasebooks, then at each step of training, we either randomly drop rules uniformly from all of ii ........ n or just drop from one of the phrasebooks randomly selected from Ti "i..... T5. while keeping the rules of all other phrasebooks intact, to create curiculum-tee -extt
Hyperparameters: Training hyperparameters are set equal to the optimization hyperparameters used in preparation of ICL-capable training phase (Appendix B.4), except we set weight decay to in experiments. report the performance of the trained model after single epoch of training on each Dn- and plot against the size of the dataset in Figure
17

---

On the Power of Context-Enhanced Learning in LLMs
C. Additional Experiments Results
C.1. Selective Internalization of Context
In this experiment, we test whether the model can internalize rules that don't incur an increase loss when dropped during training. To do so, we ablate on Annealing Dropout. We select a phrasebook and create splits of rules in the phrasebook; one set of rules will be utilized by training samples for performing the translation task (which we call the training split), while other set of rules will appear in curriculum-text during training but never utilized for the translation task (which we call the heldout split).
At test time, we measure the performance of the trained model on 2 sets of evaluation examples, one that only use rules from the training split for their translation (equal to training distribution), and other that uses rules only from the heldout split for their translation (different from training distribution). The model being measured without any phrasebooks information at evaluation. We conduct the above experiment for each phrasebook, i.e. we create 5 sets of experiments where we only create heldout split for one specific level of phrasebook. In Figure 7, we show that the model fails perform any translation that the rules from the held-out split all settings. This shows that the model only internalizes those rules that important for the translation task for the training samples, and which incurs increase in training loss when dropped.
Only
training examples
Without any rules training examples
Holcing out m
Holding 1/2 of m
Holding out m
Holding out 1/2 m
Holding out 1/2
context)
Test Accuracy
0.75
0.50
Stets
0.00
10"
10
10f
10*
100
103
10
16f
10f
10*
10a
10f
10f
10f
10f
Number
Number Samples
Number
Number
Number Samples (Py!)
Figure 7. Ablations on Annealing Dropout for (5, 10) assess whether models internalize rules not used in training samples. (Left to right) < < 5: Five independent experiments where held-out split is created for phrasebook nī and training excludes samples that use the translation rules in the heldout split of nद. Evaluation is conducted on two sets: one where samples do not use held-out rules from n. and one where only held-out rules from n; are used. The model's random performance on the latter indicates that it internalizes only rules used during training. particularly those whose removal increases loss.
C.2. Mechanistic Insights
Here we provide additional figures corresponding to Figure 3 and Figure 4, but in more settings (n Annealing Dropout vs Fixed Dropout).
10 VS n2 = 8,
Parturbing (used)
Porturbing (used)
Perturbing (used)
Perturbing (used)
Porturbing
26
Layw
21
16
11:
50
300
150
200
250
50
150
15e
200
250
50
160
150 200 250
250 Perturbing Ru (unused)
so
100
156 z00
50
100 15e 200 250
Perturbing in (unused)
Perturbing ra (unused)
Perturbing mn (unused)
Perturbing m (unusca)
26-
Laya
56
Laver
Thought Tahlns <TN..
Are
Thought Toaura
Ara
Thought Mite <b...
Ara
Thought Toens <Ts..
Ara
Thought Reara <T
Ara
100
200
250
250
100
Taken
output
Output lokers
Ourput Tokins
Output
do
75
20
10
1g
Figure 8. Comparison between perturbing used rules (top row, identical Figure 3) and unused rules (bottom row) in context. When perturbing rules used in the translation at later phrasebook, representation changes at later layers. However only perturbing unused rules leads to negligible representation changes. This experiment rules out the possibility that the affected depth only dependent on position of the perturbation instead of the semantic content.
18

---

On the Power of Context-Enhanced Learning in LLMs
Perturbing Dictionary
Perturaing Dictionary R:
mote
80
20
15
15
100
10
20
Perturbing Dictionary
Perturbing Dictionary
100
150
80
1009
20
15
60
15
10-
10
50
20
Parturbing Dictionary
Perturbing Dictionary
100
150
80
Layer motes
100
15
60
15
10.
10
20
Perturoing Dictionary
Perturbling ng Dictionary
80
25-
oyer mats
60
40
10-
10-
20
0
Perturbing Dictionary
Perturbing Dictionary
25
Laver reex
20
15
Layer
10.
10
10-
5
0
The toeesants Genues
The Pecoosatoon mas Smh a
c9 S
03 098
slestles seeee
<T> Siaure
soppii
<tb
Cutput Token Index
Output Token Index
(a) MLT (5, 8CCL--alale Model
(b) MLT(5. 10)-ICL-capable Model
Figure 9. Evidence for sequential processing in MLT-ICL capable models (n 8, = 10). The left figure is identical to Figure 3, except substitute the entire phrasehook n; with another random phrasebook the level f.. We observe the same hehavior for ML.((5, 10)-ICl-cepable model: perturbing later phraschooks in the context changes output representations in the later layers.
19

---

On the Power of Context- Enhaexee- Learning in LLMs
Dropping R
Dropping STR(n)))
Dropping STR(m)) )
Dropping STRin;?
Dropping STR(n;)
20
28
28
28
28
25
16
16
36
16
16
13
02
1922
1013161922200
1013161992 16119922 25
101316 29 25
10 16 19
Land
Land
Land
Land
Lund
(a) Dn- with 10000 samples
Dropping ST
Dropping STR(
Dropping STR(m;)
Dropping STRin;
Dropping STRin;)
28
28
28
28
Brokerton Accuraoy
16
16-
36
16
16-
13
13
13
13
13
10
13 19:61922222 19:222522
10 16193919222
101111613
10 25
10 16 22
Lend
Lend
Lend
Land
Lend
(b) Dr with 25000 samples
Dropping STR(n""
Dropping STR(n)))
Dropping STR(n;)
Dropping STRIn;
Drapping STR
28
25
25
25
25
25
0.8
22
Eyslunton Accuu
10
13
10
13
10
13
10
13
13
10
0.4
101316 19
22
Lend'
Lend
Lend
Lend
Lend
(c) Dr- with 50000 samples
Dropping STR(m""
Dropping STR(n;))
Dropping STR(m;)
Dropping STR(m))
Dropping STR(n))
25
25
25
25
25
0.8
22
22
22
19
13
13
13
13
13
10
10
10
10
0.4
25111619222522
101111
Land
Lend
Lond
Lend
Lond
(d) Dr- with 100000 samples
Dropping STR(m()
Dropping STR(m))
Dropping STR(m;)
Dropping STR(m))
Dropping STR(n)')
25
25
25
0.8
Eelootton Accrroy
22
22
22
22
19
30
33
13
10
30
10
10
0.8
0.0
Land
Lend
Lend
Lend
Lend
(e) Dr with 250000 samples
Figure Repeated experiments from Figure models trained with context-enhanced learning at varying supervised dataset sizes (Plots for row (d) are identical the plots in Figure 4). We observe that phrasebooks are progressively internalized with the number of training samples available during context-enhanced learning; later phrascbooks are internalized with fewer samples than the carlier ones. We observe similar localization patterns across layers from different phrasebooks across the models, however, we also observe that the localization patterns get increasingly sparser as training continues.
20

---

On the Power of Context- Enhanxee- Learning in LLMs
Dropping R
Dropping STR(n)))
Dropping STR(m)) )
Dropping STRin;?
Dropping STR(n;)
20
28
28
28
28
25
16
16
36
16
16
13
02
10113661922
1913661922220
11113151972 25.88
101316 10 25-2
10 19
Land
Land
Lend
Land
Lund
(a) Dn- with 25000 samples
Dropping ST
Dropping STR(
Dropping STR(m;)
Dropping STRin;)
Dropping STRin;)
28
28
28
28
Bokkattnnn Accurooy
16
16-
36
16
16-
13
13
13
13
13
10
13 19 25
10 1619
10161111
1013 16 22
10 16 22
Lend
Lend
Lend
Land
Lend
(b) Dr with 50000 samples
Dropping STR(n""
Dropping STR(n)))
Dropping STR(n;)
Dropping STRIn;)
Drapping STR
28
25
25
25
25
25
0.8
22
Eyslunton Acc.A
10
13
10
13
10
13
10
13
13
10
0.4
101366192225
101316
22
Lend'
Lend
Lend
Lend
Lend
(c) Dr with 100000 samples
Dropping STR(m"")
Dropping STR(n;')
Dropping STR(m;)
Dropping STR(m,)
Dropping STR(n))
25
25
25
25
0.8
22
22
22
22
19
13
10
13
10
13
10
13
10
0.4
1011100922222
101115100
Land
Lend
Lond
Lend
Lond
(d) Dr- with 250000 samples
Dropping STR(m()
Dropping STR(m))
Dropping STR(m;)
Dropping STR(m))
Dropping STR(n)))
25
25
25
0.8
Eelootton Accrroy
22
22
22
22
19
30
33
13
10
30
10
10
0.8
0.0
Land
Lend
Lend
Lend
Lend
(e) Dr with 500000 samples
Figure Repeated experiments from Figure models trained with Fixed Dropout instead Annealing Dropout. Observations remain similar. This suggests that the position of phrasebook internalization is primarily decided by the MLT(d, -IIIL-caaaa initialization, and less dependent on the specific dropout curriculum we use for context- learning.
21

---

On the Power of Context-Enhanced Learning in LLMs
C.3. (Non) Verbatim Memorization of Phrasebook Rules
In this subsection, we provide more detailed explanation for the evaluation on the feasibility of recovering phrasebook rules from models trained with context-enhanced learning with phrasebook excerpts in context.
model fo- trained with context-enhanced learning where phrasebook rules from (TR( are provided within the context during training, we test if the model retains explicit memory of the textual format of rules, Namely, for each of the phrasebook rule of the form "a b -> C D in STR( whether it can complete the ground truth output tokens "C D" providing "a b ->" and additional context of the phrascbook STR STR(II") before "a b -> D". We conduct the test by computing the forward pass through fo- with STR( as the input (see exact imput format in Figure 17).
Given a
Now formally define the query strategies and metries used for creating Table 2. Suppose there are h unique tokens and STR " contains L tokens, let S R denote the corresponding output logit score matrix when we pass STR(II') into fo-. For a translation rule with ground truth output tokens of index a,, Q2 E A C [h], let ( and (*++1) E R denote the logit vector corresponding to predicting these two entries.
we are doing greedy decoding, then we will recover the correct phrasebook rule if and only if we greedily select both a1 from ((k) and 02 from S(k+1), so the probability correctly recovering (a1, a2) conditioned on
P (rreedy-RRecove((1..00 02)] = 1
arg max (*k)) = an iE(A)
arg max S ((++1) = 02 xE(M]
Meanwhile, if we apply random sampling with softmax temperature of 1, the probability of correctly recovering (41, a2) conditioned on S is just then
ss())
)
exp (shht!) ( (((tt)))
exp
P (Sampping-Recove((11,..3
Σ s!(tt)
EE( exp
EiE(N) exp
Recall from Section that we have also introduced two stronger adversaries with additional token filtering when doing decoding: (I) setting probability of <THINK> tokens to zero (2) when querying for a rule n, with output alphabet A;+1+ setting probability of all tokens outside A4+1 to zero.
Denoting the token index of <THINK> as OęTнINES, then filter 1 corresponds to
PP Sreeedyy-Fiter1-Reeoveeeeeo a2)] -
arg max REIN (acrmsss }
S(k)
aj
1
arg max selhi\farrmmn (aezuno
1(k+1)
.12
S
(5(())
exp
exp
P Fileriing-illeer--eeeeeeeeeee az)l
(s((+11)
exp (s(++1))
Σ HEEhhMMaamnnnn exp
Similarly, for filter 2, the probabilities are
IP [[rreddyFiiler2-Rcoovrreeeeel]] =
arg max S(K) a1
11
arg max S =02 rEA: 1
iEAr11
(s(ti)
((()) exp
exp
=
ΣlEdi exp (s(ht1))
ElEdi exp (((k-1)
Note that the second filter is very strong adversarial assumption which assumes the user already has side information on the set of tokens contained in alphabet A; and A,+1- To compute the final statistics, we sample 20 permutations of STR for forward passes and compute the mean of the above statistics over all atomic phrasebook rules appearing in the context. That is 1280 entries for each phrasebook in the case of and 2000 entries for each phrasebook in the case of n 10.
22

---

On the Power of Context-Enhanced Learning in LLMs
Table Recovery Success Rate For 8, d = Cases (Rounded to decimals) Random Guess Baseline: 1.56%
Greedy Decoding
Sampling (7
1)
Curriculum
# Training Samples
Query Method
4
7T5
W1 R4
75
Base
0.00%
0.00%
0.00%
0.01%
Annealing Dropout
50000
Rule out <THINK>
0.06%
1.95%
0.00%
0.54%
Only keeping A,
3.18%
1.95%
1.82%
1.49%
Base
0.00%
0.00%
0.00%
0.64%
Annealing Dropout
100000
Rule out <THINK>
0.00%
0.08%
0.00%
1.25%
Only keeping A,
1.37%
0.08%
1.69%
1.80%
Base
0.00%
0.23%
0.00%
1.25%
Annealing Dropout
250000
Rule out <THINK>
0.00%
0.23%
0.00%
1.28%
Only keeping A, 2.95%
0.23%
2.53%
1.38%
Base
0.00%
0.00%
0.00%
0.00%
Fixed Dropout
50000
Rule out <THINK>
0.08%
0.78%
0.00%
0.09%
Only keeping A; 0.82%
0.86%
1.48%
1.43%
Base
0.00%
0.00%
0.00%
0.21%
Fixed Dropout
100000
Rule out <THINK>
0.43%
0.00%
0.03%
0.80%
Only keeping A, 2.36%
0.00%
1.95%
1.32%
Base
0.00%
0.47%
0.02%
0.51%
Fixed Dropout
250000
Rule out <THINK>
0.68%
1.09%
0.11%
0.73%
Only keeping A, 2.23%
1.09%
2.19%
1.10%
Table 5. Recovery Success Rate For n= 10, d = Cases (Rounded to 2 decimals) Random Guess Baseline: 1%
Greedy Decoding
Sampling (T'
1)
Curriculum
Training Samples
Query Method
T T
T W4
π2
Base
0.00%
0.20%
0.00%
0.89%
Annealing Dropout
100000
Rule out <THINK>
0.00%
0.20%
0.00%
0.90%
Only keeping A,
1.66%
0.20%
1.28%
0.94%
Base
0.00%
1.80%
0.00%
1.12%
Annealing Dropout
250000
Rule out <THINK>
0.00%
1.80%
0.00%
1.13%
Only keeping A;
3.05%
1.80%
1.72%
1.23%
Base
0.00%
1.60%
0.00%
1.07%
Fixed Dropout
250000
Rule out <THINK>
0.05%
1.60%
0.01%
1.07%
Only keeping A,
2.46%
2.15%
1.99%
1.39%
Base
0.26%
2.05%
0.05%
1.13%
Fixed Dropout
500000
Rule out <THINK>
0.33%
2.05%
0.07%
1.13%
Only keeping A;
2.11%
2.05%
1.91%
1.34%
Here we report the query success rate for a variaty of models trained with context-enhancel learning using different curriculum on different datasets sizes. models reaches nearly perfect test accuracy when no context is provided (see Figure 2), i.c., the context phrasebook rules played significant role in the learning of the models.
For all runs, phrasebook rules for hidden steps ( T4) even when we provide the correct partial phrasebooks in context (see columns corresponding to Query Method "Base"). Ruling out <THINK> token when sampling also did not significantly increase the recovery rate. We note that the phrasebook knowledge are not memorized in an completely undetectable manner, as the strongest token filtering gives non random probability of outputting the correct target 2-tuple. However the probability is still very low (< 3%), which can be considered as negligible recover the full correct phrasebooks STR(II*).
that it is nearly impossible (with recovery probability 0.1% in most cases) recover the correct
23

---

On the Power of Context-Enhanced Learning in LLMs
D. Additional Related Works
Differences with Masked Language Modeling (MLM) and Language Infilling
MLM models like BERT, RoBERTa, and T5 (Kenton Toutanova, 2019; Liu et al., 2019; Raffel ct al., 2020) train models by either masking or removing tokens from a sequence and compute loss on the model's prediction on the missing tokens. This concept has been adapted for training auto-regressive models via language infilling task (Bavarian et al., 2022; Li et al., 2022; Donahue et al., 2020; Li et al., 2022). The primary difference from these works is that context-enhanced learning does not take loss on the context tokens when they are removed from curriculum-text.
Compositional and OOD generalization
Measuring generalization for a transformer beyond training distribution has been study of interest in many prior works. OOD generalization is measured by training transformer simpler examples and measuring its performance on harder ones. Prominent studies include length generalization, informally defined as the ability of the model to reason longer than what has been trained on (Zhou et al., 2023; Anil et al., 2022), and compositional generalization on concepts, defined as the ability of the model to reason on composition of the concepts that has seen during training (Press et al., 2022; Allen-Zhu & Li, 2023b; Ramesh et al., 2023; Yu et al., 2023; Zhao et al., 2024; Wang et al., 2024; Yang et al., 2024). Our experiments on Fixed Dropout in Section 3, where we train with 20% dropout on the phrasebooks information but measure performance with 100% dropout at test time, measures compositional OOD generalization behavior of the language model. The results show that the model internalizes the rules from the phrasebooks in an atomic way, and re-compose them together as necessary at test time.
Mechanistic behavior of transformers with synthetic datasets:
Our work builds on a growing body of research exploring the behavior of transformers trained on synthetic datasets. Prior studies have examined tasks such as modular addition (Nanda et al., 2023; Zhong et al., 2023), context-fre grammars (Zhao et al., 2023; Allen-Zhu & Li, 2023a), regular and n-gram languages (Bhattamishra et al., 2020; Yao et al., 2021; Akyürek et al., 2024; Li et al., 2023), and synthetic article-style datasets (Allen-Zhu & Li, 20236; 2024; Eldan & Li, 2023). While our work is structurally similar to these studies, it investigates mechanistic study on context-enhanced learning that has not been explored in previous works.
E. Properties of MLT
MLT n)) is defined by the phrasebooks T1,". Ta at each of its translation layers. We use D as all possible set of bijective maps that can be used to define the phrasebooks. We will use variable æ to refer to an arbitrary bijective map from the set _ For simplicity of proof, we will refer to any alphabet set A of size n as {0, 1, n 1}.
Here, we formally mention some of the properties of MLT(d, n).
Lemma E.1 (Invertibility of sequence translation). For any level E fixing {m,, m441,.... ma) gives a bijection between
8, and 8841-
Proof. All of the four operations involved in the mapping from 8; to 8:41 are invertible.
Structure of the mappings: The mappings n, are selected as random bijective maps between 2-tuples of characters in A? and 2-tuples of characters in A A combinatorial argument can then give the number of such possible phrasebooks to be ?!
Lemma E.2 (Number of mappings in each level). The number of possible bijective maps between 2-tuples of characters from alpbhabet sets A;, Aitt, each being of size n, is n21, i.e. 18TT = 2
Proof. Each alphabet set contains n² possible 2- 2-tuples of characters. If we fix an order in which the 2-tuples appear in A241 then the number of bijective maps between A? and A2 can be reduced to the number of possible ordering of the 2-tuples in A3. The number of possible orderings is ?!
24

---

On the Power of Context-Enhanced Learning in LLMs
Importance of Circular shift: The composition of the d random bijective maps can be demonstrated result in another random bijective map. Consequently, without the Circular shift, each character in the output sequence $d+1 depends only two characters from the input sequence 81 via shared random map across all the 2-tuples.
Incorporating Circular shift on the other hand enables each character in the output sequence to depend on 2d characters from the input sequence. This because the character positions are shifted to the right at each step, causing the input 2-tuples to the bijective map to also shift to the right at each stage. As we will elaborate later, incorporating Circular shift increases the required number of training samples to learn the set of phrasebooks from input and label pairs to n ( whereas without Circular shift, this requirement is only Onn)).
Representing the mappings 2-tuples in-context: Each map be defined using Onn²) characters, it can be simply defined by the n² relations each connecting 2 unique random 2-tuples from their corresponding alphahet sets. Thus, defining d maps in-context will require Onn²d) characters in curriculum-text In contrast, describing a completely random bijective mapping that maps d-tuples of characters in input sequence to character in output sequence will require (()())) bits. Thus, MLT with d translation steps involving random mappings on 2-tuples and Circular shift helps define a mapping where each character in the output sequence can depend on d character in the input sequence, and the set of phrasebooks can be described using (ndd) characters in curriculum-text.
F. Lower Bound: Hardness of Learning MLT(d, n) without Context
F1. Brief introduction to SQ framework
Statistical query (SQ) bounds : The statistical query (SQ) framework measures the computational hardness of learning a task in the presence of noise. It measures the hardness of learning a task by the number of statistical queries needed by a learning algorithm to learn the true function. Statistical queries are defined by some Q of labeled instances and tolerance parameter the algorithm receives response from the oracle within error of the true value. The statistical dimension, SQ-dim, is measured in terms of the number of functions in the hypothesis class that the learning algorithm needs distinguish and the number of queries necessary to do the same. Correlation between two functions is used to define the statistical query dimension.
property [0, over (I, y) D where the data distribution. For query,
Definition F.1. Correlation of two functions f1, f2 on a domain X' with respect to a distribution D is given by
Correlation(f1..a. D) = Pr (f = IED
= fa(x)] - Pr (f)) fa(x))
=ED
For functions : {0, 1), the above definition is also equivalent to Correlation D) = 28 (D)(((() f2(1)]
Remark F.2. Two functions f1, f2: : X + 10, 13 are said to be uncorrelated w.r.t. D if
Pr (r] [2(x)] Pr (1(x) + fa(x)) =ED
ZED
1
(alternately) E.ew ( f2(=)] 2'
On the other hand, if
(racolfix)) fe(x)] = (or 0) (alterntely) Enepl @ f2(x) = 0 (or 1),
then (frrelation(11.2..1. - 1.
We take the following formal definitions of SQ-dim and its relation to computational hardness in the SQ framework from (Blum al., 1994).
25

---

On the Power of Context-Enhanced Learning in LLMs
Definition F.3 (Definition 2 in Blum al. (1994)). For function class F of boolean functions over {0. 1 and D distribution over (0, 1)", SQ-dim(J, D), the statistical query dimension of F with respect to D, is defined be the largest f with the property that for all i + we have:
natural number y such that F contains μ functions fi,
Correlation (f (f., = Pr fs = 1 Pr + fil DED
= p
RED
Theorem F.4 (Theorem 12 in Blum et al. (1994)). Let F be class of functions {0, 1)th and D distribution such that SQ (D-dim(F.21) D) μ 2 16. Then if all queries made with a tolerance of atleast 1.73. at least 1t/3/ /2 queries are required to learn F with error less than 1/2 1/1²² in the statistical query model.
E.2. Lower bound lemma
Here, we mention the 2 main theorems that study the SQ dimension of the MLT task at hand. The first theorem shows the SQ dimension bound when number characters n = 2, which is then adapted to get the SQ dimension bound for general n. Proofs of both the theorems are given in Appendix F.4 and Appendix F.5 respectively.
Theorem F.10 (SQ dimension for n = 2). The family of translation task MLT(d, 2) on input distribution u({0, 1324) has statistical query dimension SQ-dim( D0-dim(MLT(dd 2)) atleast 2
Theorem F.20 (SQ dimension for general n). For the translation task MLT n) that has depth and characters per level, the statistical query dimension SQ-dim (MLTM n)) is atleast ((d))
Notations: We will require the following notations for proving the above theorems. First, we will denote a 2-tuple that contains arbitrary characters and as (a, b). For bijective map π, π ((a, b)); will represent ith character in the output of π on any 2-tuple input (a, b) and any i E [1,2), 23. Similarly, for any set of phrasebooks II, we will use MLT (81); to denote the ith character in the output of MLTH on any L length input sequence s and any i E [1,
Recall that for any input sequence $1, the output of the ith level of MLT task will be denoted as sj. Furthermore, to denote the jth character (arbitrary) in the sequence 8i, we will use Sij-
F.2.1. BOUNDS FOR SGD
The following corollary measures the sample complexity needed to learn MIT by SGD. It has been adapted from proposition 3 in Edelman et al. (2023), who study the sample complexity for learning d-sparse parity task on n length sequences, whose SQ dimension is (*) = =) We simply state the corollary without specifying the proof.
Loss function and SGD updates: Consider training of a model fo with parameters that trained with mean squared error, i.c. la ([CURR,, t), æ, w)), y) = [ t), 22, y]) w/|2 The empirical loss on batch of batch size B from the supervised dataset DT-- will be denoted by (sJoo MLTTT- E( (fo([CURR, (), ar, v)), y) SGD updates are of the
MLTTT)) ). [[ (z, r, w]) y|l²; the population loss will be denoted by Luttonne (L.TTTT form:
P..1 0. - (.-(nVVaa (VLs, (fo, MLT + R(O,) + Ge).
for some sample Sr, step size VJ, regularizer RZ-),, and adversarial noise &. E r]". For simplicity, we assume the gradient VoLs, (0t) is bounded on all parameters in parameter space of the model.
Fake trajectory: Suppose denote the constant function that maps all inputs to 0. Then, with the contemporary losses Ls(Je, 0) and ( 0) that computes difference from this constant function, consider the following trajectory 0,,-.. 0z, ... starting from the same initiation 0o:
det 0. m/ (V.L.. (fa, 0) R(O)))
Assumption F.S. For all I, suppose
VoLs, (So, 0)0| S T/2.
Corollary F.6 (Sample complexity bounds for SGD). Fix an initialization Oo such that Soa is statistically independent (correlation 0) from all possible MLT Under this assumption. if LTB S n (dd) /r. then there exists least one task
The results holds for any loss that satisfies
OCanly 5y
DEmuply .0) By
26

---

On the Power of Conhaxe-Enhaneed Learning in LLMs
MLTH. MLT(d. n)) for which the functions obtained after the first T SGD updates, Jor Jo. remain statistically independent (correlation from MLTTI. despite training it.
LTB the product of input length, total training steps, and batch size, which represents sample complexity used by the SGD algorithm.
F.3. Useful lemmas for n 2
Here, we mention some useful lemmas for charaterizing the bijective maps on 10 - {0, that we will regularly use for the proof of Theorem F.10. The first lemma will show that each bijective map can be represented by three operations on imput characters, copy, xor, and not operations. The second lemma measures correlation on the outputs of two randomly sampled bijective maps. The proofs are given in Appendix F.6.
We define the following necessary operations to describe the random bijective maps on {0, + {0,
1. copy: Given tuple (11, 22), and position argument E [1, the operation returns value of x; as output. We will use copy, and copy2 to indicate the copy operation on positions and 2 respectively.
2. not: Given variable zi, this operation returns flipped value of x;. That is, if x; 1, then returns and vice-versa.
3. xor: Given a tuple (2), this operation returns 21 @ T2.
Lemma F.21 (Formulation of bijective maps for 2). Any bijective map {0, 1)2 -> {0, 1) can be expressed using copy, not, and xor operations. Furthermore, from 24 possible maps for n,
1. There are maps 1 2..... for which characters in the output tuple can be defined by copy and xor operations on the characters in the input tuple.
2. Mirror maps: For each map x E A2. As). there exist mirror maps ((1) (22), ((3) whose output on each imput tpple can be defined by selective not operations on either or both characters of the output tuple of n. We call {T, T(1): T(2): ( as mirror map set of π, in short, Mirrorset(
Remark F.7. Mirror map set definition is general and isn't restricted to maps E {Ai, Ag, Ao). Because mirror maps are defined in terms of not operation, a mirror map set Mirrorset(T) for E {A, Ag.... ... As} is also equivalent to Mirrorset( (() for i E 21,
Remark F.8. Lemma F.21 can be re-stated as follows: the 24 possible bijective maps {0, - {0, 1)2 can be grouped into 6 family of maps, each containing maps and represented by unique map from {A1. A2, As). Each family is defined by mirror map set of their corresponding representative map.
Lemma F.22 (Correlation of bijective maps for 2). For two randomly selected maps T" ,",T :001132 + {0, the following hold true.
1. Correlation at atleast one output character pair: Fix an i,j E {0, 1}. With probability 1 w.r.t the random selection of the maps, the ith character in the output of ntth has perfect correlation to jth character in the output of 3 Lee.
(areelation(**)))) 8 (),,(((0,1))))
/ ( 1
Pr [x)
Pr 3
= +x)2)))) =1
x_U((0.13))
2. Correlation at both output character pairs: With probability' 3. w.r.t. the random selection of the maps, both the characters the outputs of n., n have perfect correlation, i.e. either one of the cases hold true
Correlation (T (-), & ()),2(((0))))) ())))))))))))))) = 1.
Correlation (-),, 8 2100, 133) 1.
27

---

On the Power of Context-Enhanced Learning in LLMs
Correlation ( (1) (1, = 1.
Correlation ( 7 ()1, 1J9) =
Other cases are not possible, i.e. for any E {1, 2}. both () ()) can't 't hold true.
())
= and
Remark F.9. Implication Lemma F.22 follows: With probability at most 3.. output of two random maps can stay correlated at either output character pair or both pairs of output characters. This would suggest that the probability of the output of MLT under two random set of phrascbooks staying correlated should decay exponentially with the depth of the task.
F.4. Proof for Statistical dimension lower bound for
We repeat the lemma of interest for presentation.
Theorem F.10 (SQ dimension for n = 2). The family of translation task MLT (, 2) on input distributtion U({0, 1)2d) has statistical query dimension SQ-dim(MLTT( 2)) atleast 20(d)
Proof. We narrow our argument to the first character output of the task. The proof goes through major steps.
1. First, we show that two random instances of MLT d. 2). defined by set of phrasebooks { ( क and " will be uncorrelated with probability at least (1/3) w.r.t. the selection of the random set of phrasebooks. The lemma is formally given Lemma F.13.
2. Then, we can apply Lovász local lemma (Theorem F.11) to show that we can create 20(d) instances of MLT 2) that will have zero pairwise correlation of their output (Corollary F.12).
By the definition of SQ-dim from Definition E.3, the above observations suggest that SQ-din(MIT((.. 2)) 2 20(4).
Theorem F11 (Lovász local lemma, theorem 1.5 in Spencer (1977)). Let A,.... Ar be events in some probability space with Pr(A,) < p. 1 < i < k such that each event is dependent on atmost ky other events. If ep(ko + 1) < 1, then Pr(-A1 A A2 A A3 ... AA O.
Corollary F.12 (Number of uncorrelated instances of MLT(d, 2)). There exists set of instances in MLT(d, 2) that are pairwise uncorrelated to each other on input distribution ((00 2d)
II represent & randomly sampled set of phrasebooks. For each < k, we will denote event Ai as the event that the outputs of MLTH, and MLTn, are uncorrelated on input distribution 2({0, 1]2).. This happens with probability p = - (7/3)(7/9)d3 from Lemma F.13. Because each event can atmost depend on atmost k(k + 1)/2 k22 events (total number of events), by Theorem F,
Proof. Suppose II,
ep(k2 + +1) 11
then there exists a list of such set of phrasebooks which are pairwise uncorrelated w.rtt. the outputs of their corresponding translation tasks input distribution ((() 2d)). Solving the above for k. we can set k ((ep)-1 1)2/2 20(() I
1.4.1. AUXILIARY LEMMAS
Here, we will prove the following primary lemma, that is used to prove Theorem F.10.
28

---

On the Power of Eentext-Enhhneed Learning in LLMs
Lemma F.13. With probability at least - (g)=² d- w.r. random map selection, the following holds true for 2
sets of random phrasebooks II { ... ; and II { ]
Correlation MLT ML (1=(()))--1111---- O1, MLT 11 (= ((0,0)))))) =0 0.
Proof. We first dive into the dependencies between characters in input and output sequences in the translation process. In Lemma F.15, we show that at any level 1 S é < of MLTTI with phrasebooks II {me]²_1- the following relation holds true on an input $1 {0, 1
$1412) = (5,111,,
We will use superscripts and 3 to differentiate the intermediate outputs of MLT when the phrasebooks are set as IT {मf,... a and II [mf. = respectively. We will use induction strategy find the correlation of 8-1,1 and & 1+1.1 To do so, we will require the following variables:
1. Plafi For one pair of i,j E (2,3). this represents the probability at a level 2 < E < 1+ d w.r.t. the randomness of define Pe,1; that
"",,",,""" and {mī, m that sY, and si, are perfectly correlated. Additionally,
represents the probability at level 1 < l < + d w.r.t. the randomness of (хf ". x and (œf ".,...,,, t-1 "" that
88,1 and A E,1 are perfectly correlated.
2. Pebomh This represents the probability at level 2 < l S 1 - d w.r.t. the randomness of { ; T and that either 87 S are correlated and 8 8 are correlated, or 87 6 are correlated and 87,3. 22 are perfectly correlated.
There are three relations that we need to keep in mind, before we proceed with the induction proof.
1. First, the correlations are circular nature, i.e., for j E (2,3) and any integer k. Pt.s = Pe:::++2yy2ed (Lemma F.17). We will use this relation to connect Pr.1:1 with Pe,3,3, i.c.
P4111
Pr333-
(I)
2. Second, Pt,ooth < max; Pl.s Intuitively, this is because getting correlations at both positions must be at most as probable as getting correlations at one of the positions.
both S max Pe,egj
(.j
3. Third, cross-correlation probabilities given by P22233 and P2,32 must be S (P + Pt,3.) (Lemma F.19).
| PEAi S (press + PESs), for any i,j E (2,3), j.
Base condition: We will use the values of the variables translation is same both MLTHI, M. Then, for input $1.
our base condition. The input for first layer of
(s = x ((s,2,, (s 822) -f((81,2,811)) 813))
We can then use Lemma F.22 to show that
1
P2
for all i,j {1,2),
3
bo.
3"
29

---

On the Power of Context-Enhanced Learning in LLMs
Connecting the variables at l and e We 244.222 to P2,23 P0,32 P233: and beth This will undergo case by case analysis.
8 1. First, under phrasebooks -f.-. - and Ti." TT
(either) 8² and 22 are perfectly correlated, and o, and o are perfectly correlated (or) 87 and 2 are perfectly correlated, and 87 and s are perfectly correlated,
(4)
we can use Lemma F.14 to show that with probability 1/ w.r.t. the selection of मद and T.. 8 and 8 1+122 will be correlated. Conditions necessary for Lemma F.14, i.e. uniform distribution of the characters in the input sequence, are shown to hold true in Lemma F.16. The probability w.r.. the selection of phrasebooks , .... - and π. such that Equation (4) holds true is given by the variable Pt,both-
2.
8 On the other hand, if there exists pair i,j {2,3), such that under phrasebooks -f,'" Th and 11 $7, and St, are perfectly correlated, then with probability å w.r.t the selection of Tे and T 86-12 and 1,1 will be correlated. We again refer to Lemma F.14 for this statement. Probability that this condition happens under phrasebooks "i."" ... ; and || ... - given by the variable pr.s.
2-1
3. If none of the above situation occurs, then for all pairs i, (2,3), 8i; and 8 are uncorrelated, and so, by Lemma F.22, correlation between 24,, and 11-122 will stay 0 for any choice of σ. and n..
Thus, combining the cases, we must have
P2+1,2;2 S
3 pt,bo
Σ
(5)
(.E3/2,3)
3 botb
+
max Peçt: (2,3)
5 max (2,3)
Peiu + max Praj
max
(6)
9 (2.3)
9 (.jE(2.))
Reasoning for each step is as follows:
1. Equation (5) follows from conditional probability computations using the 3 cases that we discussed before.
2. Equation (6) uses Equation (2) to connect both to Peçid
We can give the same inequality for P1+1,1:1- As P141111 = Pe-1.333 from Equation (1) and Pr,2;3 and P2,3;2 must be atmost the average of Pr.2.2 and Pr,3;3 (Equation (3)), we can write
7 max Pr.n (JE()33)
max PE+1AS (2,3))
S
9
This implies that the probability of correlation at any output character under the two random set of phrasebooks decays at the rate of 7. g- Solving the recurrence will give:
7 d-1
7
d-1
1
Pa+1is (2.3)
S
P2,i.j
9
3
As 1 1.3:3 should be equal to Pd+ 1,111, this implies that with probability at least 1 P.411333 = =1 1-d--1 must hold true:
d1) 3. the following
Correlation PorrelationMMITT--//e (MTT
MIT (.T 1j2d) 0.
30

---

On the Power of Enhanet-Enhaneed Learning in LLMs
Lemma F.14. Suppose f {0,1]* 13*h + (0,1) denote functions for some arbitrary k, satisfying conditions for uniformity in output distribution, i.e.
(== = 1/2, [0-))(001))))0))) = =1/2
(=-3m1 @ f(x)2 1/2, (011411))))) g(x)2 1/2
for all j E (1, Then, the following holds true:
1. Both output character pairs have correlations between j and ga If
(either) ((tplatoonnnff......c)))00000000 (or) Correlation((()1.0)0)))22220
1 and Correlationl (f (())) 1}))
= and (oorelation(((())2000)))))1 1) = =,
then for two randomly picked phrasebooks x" ,T
(a) Correlation of an output character pair of TT" (f()) and T (g())) Fix an i.j {1,2]. With probability 1/3
W.I.t the random selections of n",,,"
Correlation( (x" n (g ((10))))) =1
(b) Correlation of both character pairs in output of n" (f()) and x (g())) With probability 1/3 w.r.t the random selections of Th T², one of the following conditions hold true.
(either) Correlation T" (()) ())) and Correlation x ((()), 13
(or) Correlation T (f()) 1) and Correlation x ((0)).
2. Only pair of characters have correlations between f and 9: If there exists only one pair i,j E (0, 1) such that
Correlationí f()), ((,
1,
and for all other i j', Correlation (f g 1}=)) = 0, then
(a) Fix any j / (1,2). With probability w.r.t. the random selections of T. 7
Correlation( (f))))))(()))))
If the above condition holds true, for all other 7,J pairs,
Correlation(?? (J(-), ( (g()))
0.
3. No correlations between f and for all pairs i,j E (1,2), Correlation (rteeaatioo(()...))))),))(0 = 0, then for all pairs i'," J E {1,2].
=
Proof. We will prove each case separately.
1. Both output character pairs have correlations between and gi We will consider the case when Correlation((()) (() ((0,, 1}²))= 1 and (orrelation((())...)))22. 1110113)) 1, proof for the other case
is similar. Then, there are cases possible.
Subcase f(x); [()), for 10
Subcase 2: f(x)1 not(g(=))..(()> = g(x)2 for fraall e {0,
Subcase 3: f(x) - for E 1)
=
Subcase 4: f(x) not(g(=))), f(x)2 f()) not(g(z)2) for all (0, 10
31

---

On the Power of Context-Enhanced Learning in LLMs
Because f()) and g(-) are inputs the phrasebooks x and x respectively and f(+) and g(-) are related by selective not operations in their outputs in each of the possible subeases, the output of 7 (g(-)) be replaced by # (J()) for mirror map #8 of T². We showcase this formally for one subcase, say subcase 2; argument for other subcases are similar.
Suppose Subcase 2 true: Then, for any two phrasebooks x² and n suppose # denotes a mirror map of T such that
(x)) - not( (x))), ()))) x (2)2, for x E {0,
Such a map exists by Lemma F.21. Then, for a pair i, {1,2),
Correlation (
(7)
Correlation (orcelation hcapa//aooffffaaooo/o//oa///atod
(8)
Correlation (aaaaatooono(((())))aa)aaaaa((() (100),/000))))).F170000000044444
Correlation n (+),, (1))0))))))))))))))((((()))))))
(10)
Reasoning for each step as follows.
Because the outputs of f and are 2-tuple on any input, Equation (7) simply replaces f(x) (similarly g(z)) as (f(x), f(x)2) for any input x.
Equation (8) uses the relation between the outputs of f and g when subcase 2 is true.
Equation (9) simply replaces T² with # due to their relation as mirror phrasebooks.
Finally, because of the assumption on the output of f. i.c. (o-=2((..))))((();; 1/2, 1/2.)) @ f(z)2 1/2, the distribution of outputs of is identical to (({001)²)
Hence, Correlation Tt (T (f())),, x (g())) ((()) 1)²)) boils down to Correlation (F) then use Lemma F.22 show the probabilities of each of the desired conditions.
We can
2. Only pair of characters in output have correlations between and gi For typographical simplicity, simpliicc consider the case when
Corccatton((())))))))
1,
and for all pairs i', with atleast one of them being not equal to 1, Correlation([())e (0),),,(( 0. The argument when the condition holds for other i.j pairs can be similarly handled. Then, there are 2 cases possible.
Subcase 1: f(x) g(x)1, for all a E 1)
Subcase 2: f(x) = not(g()))), for all x E 10
while in each of these pairs, (f(-)s, g())2,, (f(-),, (()))((())))))))))))))))) and (f(-)1.g())) the output bits of f and g completely independent of each other. We only consider subcase 1; subcase can be handled using mirror phrasebooks similar to the proof for case 1 (in particular with Equation (9)).
For simplicity, we will argue for i', j 1, 1; the argument about other i',j' pairs are similar. The proof will follow by two steps.
Step (a): We first argue about the probability with which (oreeation((=(((()))))))))))))) to 1,
is equal
Step (b): We then argue that if the condition (a) holds true, then Corellaioo((:*"(((()))))))) must be 0 for any other pair i' j' where atleast one of them not equal to
Step (a): For simplicity, we will assume that N" T E {Ai A2, handled similarly as Lemma F.23.
Ag} defined Table 7; other cases can be
32

---

On the Power of Ennanee-Enhaneed Learning in LLMs
= Correlation ( (00hah/0000000..00//00//0h0///ha
(1I)
Correlation
( (Wh..hh0000000000000000000000000
(12)
Pr 1=)
77" [
(13)
Pr m-M(11111
Pr 2-
Pr (0.11)
T (<((,)))
7 ((z.
Pr v.y ((
x [>)((,,)) ( ((ax)/))aa]
(14)
The reasoning for each step is follows:
Because the outputs of f and g are a 2-tuple on any input, Equation (11) simply replaces fz) (similarly g(x)) as
(f(x),; (x))2 for any input x.
Equation (12) uses the relation between the outputs of f and g when subcase is true.
Equation (9) simply writes the definition of correlation.
Finally, because of the assumption on the output of f, i.e. Eg- 1/2, E. (((.s)))))( @ f(x)2 1/2, the distribution of outputs of can be shown to be identical to ({)
A in Table 7, the first character in the outputs of x" and & can be expressed using 3 operators on the input characters, which are copyı, copy2, and xor. These operators are independently selected for T. and T as they are randomly picked from these 6 possibilities. Formally, for any tuple of variables (xr y) and 3").
From the definitions of n", , E
12,
copy1 (z. v) = x, copy>(x, = y. xor(x, v) = x + y, copy (z,8') = x, copy2(2,y') =y,, xor(ar
Byy
However, because v, y, yf independent variables, only operations copy (xr, y) and copy (x, y') for (2,, M/ v') u(fo, 1)a) will be correlated. This can be verified by writing the definition of correlation and using copy. for the first character output of a and 8:
Pr (=201(0.13)
Pr ((0.))((
Pr (T = = ((x, way -(.
TT (TI, = = (
Pr (=
Pr (1.1))((0011))))
(c)py1(x,)) copy copy) (2, v. v.
Pr
(2, + co =1
as the first term 1 and the second term is O. However, if you pick any other pair of operations, the correlation will be 0. We demonstrate by using copy, and xor for xº and T respectively.
Pr m-1(((011))
Pr Jnn 19.1))
[ ((x,))) = =
Pr M.y/ (0,
((x,)))
Pr xm11((0.1))
Pr Un ((1.)(())1))3)
(copy: (I, xor(2,8')
Pr N/w 200 2001133
[copy (z, xor(z,v))
Pr [.,1{(
Pr B.. (110)))
Pr [x= 2o v.y [1._)1)(
x x x
= =0,
as both terms are equal to 5 in the final step. By a counting argument, one can reason that the probability of (fr)) T (g(()))())))))))))) (((0, 1}* being 1 is equal to the probability of copy, being selected to define the first characters of both and R which will be equal J.
Step (b): Now, say we have selected pair x" and T such that (orcelation( (o()), 11 = 1.
From the proof of step (a), this is only possible when copy: was selected to define the first characters in the outputs of
33

---

On the Power of Context-Enhanced Learning in LLMs
T and 79. Any other operation pairs for defining the first characters in the outputs of क and x would have meant correlation to be O.
We can use this same argument to show that for any other pair ,', where atleast one of them is not equal to 1. Correlation (x" (f())vv, x" < (g (( 1)()) must be This is because once copy1 has been used define the operation for the first character, can't 't be used define the operation for the second character (please refer at the truth then show that correlation between output characters (J()),, (g))))) will be O, as copy: can't 't be used to define at least one of these characters.
tables for 1A 1,.... A in Table 7). Following similar argument,
3. The third case is when none of the above conditions hold true. That is, for any pair i, E {1,2},
Corrclation(f((): 9( Corc:laioo((()),.)))),,,,,1,,111
O.
In such case, following similar arguments case and 2, one can show that for any i.j E {1, for any choice of n and 9
(Fo
Pr a'
Pr
(Fx( = x ((x')
Pr 10
Pr
(t'((,,)))+ + T" ((x' v/))g]]
= 0,
due to independence inputs to n. and 8
Lemma F.15. At any step of MLT 2) with phrasebooks {Te) the following relation holds true for the intermediate d+1
outputs {s, {s:) input 81 E 10.
(1+111 ( 81,3)).
Proof. MLT 2) has 2 primary steps: Circular shift and Translate. By Circular shift, first, we first get sequence ū;, where for any j E [L] we have ر.ا 81,(+1)%L. After Translate step,
(.112)))112)) (,,(5,,1, s12)) = (5,(,,,,
Lemma F.16. At any step i of (L, 2) with phrasebooks {me)d= the following conditions hold true for the intermediate output 8i.
1 for all L
E, (s,,2((0,))11811 2' for
E ( 1 8i.j' = 2. for all j j
The above conditions are equivalent to showing that sú,j behaves like uniformly random boolean variable, independent of any other character 8i,' for all coordinates j' + j.
Proof. The proof will follow by induction on the output the translation task each step. We will show the result for coordinate in s;; similar argument holds for other coordinates j.
Base condition: At layer = 1, the s1 represents the input sequence from 4((0,1))). By definition of uniform distribution, the conditions hold true for the input.
34

---

On the Power of Context-Enhanced Learning in LLMs
Operator
Expected value
Expected + value with operator
copy:
copy2
xor
copy: copy2 xor
Es, 8i-18 1,2 = 1/2
Ex, 1,2 8i-1,3
1/2
Es 8i- 1,3 1/2
Ex 8,3 11/2
E. 1,2 @ 8, 1,3 1/2
Es 8i 1/2
8 8,- + 8,3 = 1/2
g, 1,313 3 1/2
Es 81-1.2 11/2
Table Est,1 (similarly E8,,2) and 18,,,1 @ 812 under different m; phrasebooks, defined by copy and xor operations on 8,-1,2 and 81-133
Induction step: Argument for general i> 1: Suppose the conditions are true for all layers 1 I é i. Then for layer i, we will provide an argument for the condition hold true for j = 1 and j 2, arguments for other js will extend similarly. By Lemma F.15,
(81,1, 84,2) mi (8: 1,3)).
From Lemma F.21, we have each output character can be defined in terms of copy, xor, and not operations on input characters. Then, the relations for 84,1, 81,2 can be computed follows:
If map x" is selected from Mirrorset(rt T) for some {AI..... b then can show that the conditions hold true for = T if the conditions hold true for и. This follows because the output of d² follows from the output of π by selective not operations to the output of π. As not operation won't change the expected values of a variable which behaves like a random boolean variable, the argument follows.
Now, we show that for मं = π for some E {AI,.... AoJ. the conditions hold true. For these phrasebooks, the output characters are defined by copy and xor operations. We use the definitions of these phrasebooks from Table 7 and show the expected values 81, in terms of expected values of 8;-1,1 and 8;-1,2, and the expected values 81, @ 8 in terms of expected values of 8;-1,1 and 8,-1,2 in Table 6.
Both of the above arguments then can be combined show that
Eg )
for j E [1,2]
E Eg ( 1 8122
We can similarly extend the argument for expectation of each individual character other positions 2, i.e. (x Sig for all other possible js. We can also similarly extend the argument for joint expectation of two consecutive characters 81 8,,k++2 for any general k.
The remaining argument will be to show that characters that don't form consecutive 2-tuple are going to be independent of each other as well. Consider any two 2-tuples ($1,2841, 81,24+2) and (81,2/411 $,,2442,, with k * k'. By adapting Lemma F.15, one can show that
15.2k41, = ny T, 81,)))3))
(8.2k44, 42) = n, =={((1,1,,4444) $1,2443))
The primary thing to note here is that both the tuples depend on two distinct tuples in layer - 1. By induction assumption, characters across these two 2-tuples are independent of each other. As म, is bijective map, this will also suggest that characters across the 2-tuples in the resulting output must also be independent of each other. This will give the final argument.
Lemma F.17. Under the assumption that sequence lengths L are even, for any Siiid d 1, where ST S., denote the output after 1st translation step for a random sequence 81 ~ ((O 1 at any position S S under the two set of phrasebooks II := { {- 1-1
===1 and II: : {mT ) the following holds true for all positions j:
Correlation( sf.,, S. u((0,,11 1) Correlation (, (.1+21y%1+ s 2. 1)")), for all integer
35

---

On the Power of Context-Enhanced Learning in LLMs
Proof. Fix integer k. By Definition F.I, the necessary condition would be show
[ 10 e
22EE (0.1) 1,2+2R))6 + 8 1,0++2892
We will look at the behavior of the function h(81) = + s Denote C2 as a circular operation that takes sequence 81 and returns shifted sequence O S1, i.c. if S = Cx(s1) then for all S1.. $1.(+2%)% Note that, we can also define an inverse function C and 81 = Cx (Cx().
From the definition of MLTn for any phrasebook II, we have for any input 81:
8; T- o T (81),
where for any level i, xx denotes the translation step at level that includes Circular shift and Translate using n In Lemma F.18, we show that for any translation task MLTI and any input $1.
Ta,
Tap (81) = C (Tm. ..... o Tap o Co (81)).
On input 88 Ca(s)), respectively, then the above statement says that
if of and o of represent the output of translations using II; == {mp)²=] and II (m (m/)-)
of - (0,)),
=C= ("of]
(or) Cs (s Ssf
Cs
(of) g
Thus, for any position j. we will have Cs (8f") input S8 follows:
We can then compute the value of on 8 (and similarly for O s.
h( $1) SSf S
= C (sf), & (
- @ (2)-2))%
The last step follows from using the definition of Cg. On the other hand.
Ee - 81 E,, (((
as the uniform distribution can be shown to not change under circular function Cs. This will imply:
3 ((. (0+2))) (.(j42))%L E (0.10.0 @ 8/ 1 [.
1 - 2E 9 -
2E, ~ (0.1 (1+28)) + 8 -10--2))%
Correlation E (10,1) Correlation (, (0 U(00 13))
Lemma F.18. For any translation task MLT at any level S d and any input 81
Tx. .... Ta (81) = C, (Tx, ..... o Tx o C+ (81))
Proof. Denote Cs circular operation that takes a sequence and returns a shifted sequence Os i.e. if = C+(81) then for all j. $1.j = 81,0+28)%.. Note that, we can also define an inverse function C,¹ and 81 CT (Ca(s)))
Recall that for any level i, Tn, denotes the translation step at level i that includes Circular shift, and Translate with n,. The argument will again follow by an induction step.
36

---

On the Power of Contexe-Enhanced Learning in LLMs
Base condition: i= 1: We need to show that
Ta, (81)= C Cx (Tr, o Cp (81)).
To do so, we will look at the behavior of the first characters. Argument for others can extended. If If Ca (81).
$2 Tx (81), 02 Tx s1)1). then by Lemma F.15,
(52.1, 82,2) = тi ((5,,2 s,,, O 82,2) ( 81,2 81,3))
But by definition Ck.
8,,,, 81,3) 21,
On the other hand, Lemma F.15 can be adapted to give
82, (2,1-22
Thus, we can show by combining the above 3 steps that
( $2,11 92222
1 (1 s113))
(22.(142)) S2,22)))))
We can extend the above argument to show that for any position j.
82 (2.(++)))
which by definition of Ck. implies
32 = Cx(22)
or 82 = CT (082).
which can be further simplified (using the notations:
Osi Ca (81), 82 T (81), 08z (s)))
82 Cx (2x2)
C (Ta, (os1)
C (I', (C& (81))) = C (Tr, oc (81)).
General argument for i: Suppose the induction condition holds true for all layers < i. Then,
TH T ($1) T (Cr (T T o C& (81))).
We can then follow the same argument the base condition, and show that
Tx. (CT' (Trm- o
Tu, o C& (s1)) = CT (Tr, ... Tx, Cs (81)).
Lemma F.19. For any Sii d, 8 2. denote the oulput after 1st translation step for a random sequence 81
u(10, 1 at any position 1 S j S L under the two sets of random phrasebooks II == { [ and II. the following holds true for all positions j:
{ 1
Correlation (Bi.) ( u((0.134) =1]
Pr IT
Correlation (Si ((1 =1l)
Correlation (st,,
=1] Pr ?
Pr I.
2
37

---

On the Power of Context-Enhanced Learning in LLMs
Proof. We will prove the required result with counting argument. We will create family of phrasebooks: let FAMILY(j) and Fj + 1) denote sets of phrasehooks, such that for every phrasebook I in FAMILY()), there exists least one phrasebook I in FAMILY(J 1) such that
Correlation
S ( (((0.114)) 11
(Equiv. to saying translations for II; and II have correlations at position j and j + 1)
where " Li+1 1 are outputs on a random sequence 81 ~ u(10, 132) corresponding to using II; and II, respectively.
We then apply a grouping algorithm GROUP to group correlated phrascbooks together in each family. That is, in FAMILY()). we create groups of phrasebooks {S1,S2,"* "., such that for any two phrasebooks II and II that belong to group S,
Correlation
و ((), 1,
(Equiv. to saying translations for II; and II have correlations at position j)
where { i.j+1 are outputs on a random sequence 81 ~ -M((0,1)) 132) corresponding to using HI and IT respectively. We call the resulting output of this operation as GROUP(FAMILY Similarly, we compute GROUP(FAMILYú + 1)).
We can observe the following two characteristics of GROUP(FAMILY( and GROUP(FAMILY() 1)):
1. For every set S GROUP(FAMILY()) there will exist one set S2 E GROUP(FAMILY(( + 1)), such that for all phrasebooks II S and II E S2, correlation will be 1 for output at positions j and j - 1.
2. For every set S1 GROUP(FAMILY()) there can't 't exist two sets S2, S E GROUP(FAMILY() + 1)). such that the phrasebooks 8] are correlated to phrasebooks from both S2, S for output at positions j and j + respectively. Otherwise, we could have merged S2 and S under the GROUP operation.
Let CORRELATION-MAP denote the map between GROUP(FAMILY(i) and GROUP(FAMILY() + +1), which connects sets Si GROUP(FAMILY(j) to set S2 E GROUP(FAMILY() + 1)) such that any two phrasebooks in S₁ and S2 have correlations for output positions j and j + respectively.
The result will then follow from counting argument. The number of possible pairs (can be identical phrasebooks) that can give correlations for output at position are given by: (GROU((RMMX()))) |111. Similarly, the number of possible pairs that can give correlations for output at position j + I are given by: Es, Essccroou((AMILYG-U SS 2 On the other hand, the number of possible pairs that can give correlations for output at position and j respectively are given |S (12) Applying the AM-GM inequality, can show that the average of the number of pairs for which correlation is 1 for output characters at either position or j + 1 is higher than the number of pairs for which correlation is 1 for output at positions j and j 1.
by: Es, (GRoUP(): :0S==CCOORLLLOOO-AAESSS
F.5. Proof for Statistical query lower bound for general n
We present the main theorem statement again for readability.
Theorem F.20 (SQ dimension for general n). For translation task MLT(d, n) that has depth and n characters per level, the statistical query dimension (22-ii(MMLTV (d,n)) is atleast n ((d))
Proof. We adapt the SQ-dimension proof for MLT to show the SQ-dimension proof for MLT n). We will design family of set of phrasebooks II { {0, 1......... n 1)2 {0, 1,.... on top of translation task in MI.T (log, 2).
_] -=) where phrasebooks in II are built
For a character a E {0, 1... ... n - 1}. suppose (IT(a) E {0, 1) indicates its binary representation, and NUMERIC represents the map from binary representation to its corresponding numeric representation. Then, we design each phrasebook æ using a random translation task v E MI.T(log, n, 2). For any tuple (a, b) {0, 1,.,,,..,, ,,, 1}2, output of π given as
38

---

On the Power of Sontext-Enhanced Learning in LLMs
[(a b) = (01,00), where
(ā, @ b23 1)
))
01 = NUMERIC 1
02 = NUMERIC
6)
ā = BIT(a)
BIT(b)
Primarily, the phrasebooks are defined as follows:
1. On 2-tuple of characters (a, b), we first compute their binary representations (BIT(a), BIT(6)). We then compute two intermediate outputs, one where xor operation is applied on (a), (bT at each bit, and another where BIT(b) is simply copied. This operation is equivalent to applying a deterministic map on 2-tuples of binary characters (identical to Ag from Table 7), where 2-tuples are created by pairing binary bits in binary representation of and b.
2. We then apply random MLT task of depth log, n on each of the intermediate outputs. This applies a random bijective map on the sequence of bits in the intermediate outputs, using a translation task in MLT n, 2) (Lemma E.
3. The final tuple of characters is returned by applying a NUMERIC operation on the binary representations.
Thus, we have narrowed our focus on special group of tasks from MLT n) that applies MLT n, 2) on the binary representations of the characters at each level. By Theorem F. at any level S S d, we can create 20 n) phrasebooks, the output of which are pairwise uncorrelated. Now, we can compose these uncorrelated phrasebooks to give multiple set of phrasebooks that are uncorrelated. That is, following similar proof Theorem F.10, we show that can create (20(log, 2)) RQ() = (d) set of phrasebooks, using the above restriction, that are pairwise uncorrelated any bit in the binary representation of the output of their corresponding translation tasks. This will translate to the output of the translation task in numeric form as well, as the mapping between binary representation and numeric form of a digit is bijective.
F.6. Proofs of Useful lemmas
Here we give the proofs for the useful lemmas necessary to prove Theorem F.10. We repeat the lemma statements for easier readability.
Lemma F21 (Formulation bijective maps for n = 2). Any bijective map x {0, 1}2 + 1)2 can be expressed using copy, not, and xor operations. Furthermore, from 24 possible maps for n,
1. There are 6 maps A A..... for which characters in the output tuple can be defined by copy and xor operations on the characters in the input tuple.
2. Mirror maps: For each map x E (A1, Ag. As), there exist mirror maps T(1), (T2)) T(3) whose output on each input tuple can be defined by selective not operations on either or both characters of the output tuple of n. We call (n, T(i) T(2): ( as mirror map set of n, in short, Mirrorset(
Proof. There are 4 possible tuples (1, (1, By Lemma E.2, the number of possible bijective maps (phrasebooks) that connect 2-tuples are 4! = 24. will show that the output tuple each map can be represented by operations.
Operation not creates mirror maps: For each map n. there exists 3 alternative maps T(1); ((2) (() such that if T(a, b); represents the ith character in the output tuple for input tuple (a, b):
39

---

On the Power of Context-Enhanced Learning in LLMs
Map name (T)
Output for corresponding input tuple π (a, b)
General formulation on output for map
(0,0)
1)
(1,0)
(1,
I(a,b) b)
π (a, b)2
A
(0,0)
(0,
(1,0)
(1,1) copy, (opy)(a;1) b)
ca,y2(a, b)
2
(0,0)
(0,1) (0,1)
(1,0)
copy (a,
xor(a,
A
(0,0)
(1,0) (0,1)
(1,1) copy
copy, (a,
A
(0,0)
(1,1) (0,1)
(1,0) copy2(a,
xor(a. b)
A5
(0,0)
(1.0) (1,1)
(0,1)
xor(a, b)
copy (a, b)
A
(0,0) (1,1)
(1,0)
(0,1)
xor 6),66)
copy (a,
Table 7. The table captures the definition 6 bijective maps (phrasebooks) on 2-tuples {0, 1}2 + {0, 1) whose output characters can be defined in terms of copy and xor operations on the input characters. Any other bijective map be shown to belong Mirrorset of one of these maps.
T(1) selectively applies the not operation the first character in the output tuple of on any input tuple, i.e.
(a)(a,, 6) = (a)) ((a,b) 6))
for all tuples (a, b) E {0, 133
((2) selectively applies the not operation to the second character in the output tuple of on any input tuple, i.e.
(a2) (m(a, b)1, (a(()()))))
for all tuples (a, b) E 133.
m(3) selectively applies the not operation to both characters in the output tuple of i on any input tuple, i.e.
(a) b) = (not( 6)1), not((((,,))))) 6)))
for all tuples (a, E (0, 1)2
Thus, for each map there exist other alternative maps that simply modify the output of map x with the not operation.
After removing the mirror maps: We now show that there 6 possible maps that apply either characters to get the output characters. We name them A, A Ag. We give the output of each map on the 4 tuples in Table 7 and show that the each character in the output tuple can be represented using copy and xor operations.
copy the input
Lemma F.22 (Correlation of bijective maps for 2). For two randomly selected maps T", , :0, + (0, 10, the following hold true.
1. Correlation at atleast one output character pair: Fix an i, E {0, 1}. With probability 1/2 w.r.t. the random selection of the maps, the ith character the output of x2" has perfect correlation to jth character in the output of 3 i.e.
Correlation( (+), (),,(((0,11)))) (5,
Pr (x); (<),))
Pr (x); S
=1.
2-21((0113))
2~1((001)2))
2. Correlation at both output character pairs: With probability w.r.t the random selection of the maps, both the characters in the outputs of n² R have perfect correlation, i.e. either one of the cases hold true
(orrelation (+),, T (1),2(((011)))) = 1.
Correlation (-)2, &
1.
40

---

On the Power of Ennnaxe-Enhaneed Learning in LLMs
Correlation (T (+), ()2 = 1.
Correlation (T (). =
Other cases are not possible, i.e. for any E (1. both Correlation( (-)s,, 1)s.T" (1,, = and Correlation(T (+) ()),, 132)) = can't hold true.
Proof. We prove the lemma case I, cases and 3 can be similarly proved. From Lemma E.2, x and x can be randomly selected from set of 24 possible candidates. On the other hand, Lemma F.21 shows that there are maps {A..... A whose output can be defined in terms of copy and operations of characters in the input tuple. For each in this set, there mirror maps T()), T(2), T(3) whose output are defined by selective not operations the output of æ,, and the set of 4 maps is represented by Mirrorset(r)..
The proof will follow from 2 steps: first, we argue about correlations when n" and T are selected from { {Ai," As). and then we argue about the general case when x² and T7 are selected from the general set of bijective maps.
In Lemma F.24, we show that for two maps that are randomly selected from { A the correlation 1 with probability 1/3.
The remaining possibility when x² and n belong to Mirrorset(a.) and Mirrorset(Tx) for some Ty, Ty E ..,,,,, As). We show Lemma F.23, correlation R² and n will be equal to correlation of иx; Ty-
Thus, we can combine all the observations to show that for two random maps T", T that belong to Mirrorset(m.) and Mirrorset(TT)) (Ty) for some Ty. Tyy E ....... As).
Correlation (() π ()), 1)²)
1, w.p. 1/3 w.r.t. randomness in T. My
= + 1)9)
0, otherwise.
Lemma F.23. The following holds true for any maps Ta and To with no x , E (Ai,... ,.)) and for all za E As)
Mirrorset( ") 7 Mirrrrrst(
Correlation x
1)2)) Correlation (-) (0,,
Proof. We prove as follows:
=)) 1)3)
Pr T (x)) + (x)) (0.1(
Pr x-10((0113)))
[=((x) T (x)1]
Pr
(z)1 # T (z)1 -
(15)
2 Pr Pr (x)) # (z))
1 if condition "c1" is true
(16)
PreNu TT (x) म (x) if condition "c2" true
Pr (x)) + (x)1 2_.11 (0.1)
Pr [# # (2)1
Correlation # (01, の (4) 4 ()01))))
(17)
where the second and the penultimate steps follow from the law of total probability. Here, condition "c1" holds when either case is true,
* (x)) = not " (z)1), #9(x) not ( (x)1). for all x E (0,1)²
70(z) = x (x)1, # & (x)1, for all (0,

---

On the Power of Context-Enhanced Learning in LLMs
and condition "c2" holds when either case is true,
7 = #²(2) not(TT())) (x)1). for all E {0,
* = not n (x)1, #0(x) = (x),, for all x E (0,1)²
One of condition "c1" condition "c2" is true because #* and x² (similarly, H and T²) are mirror maps.
Lemma F.24. The following holds true for two randomly selected maps Th and R with Tº, T E (Ai.... As)
With probability 1/3 w.r.. random selection,
Correlation( (), >0 (=1). 1).
With probability 1/3 w.r.t. random selection, s,,ecto,,
Correlation(a"()))
> 0 (= 1).
Proof. We prove for case 1, proof for case 2 is analogous.
Among {A,.... A we can create three family of maps: Peopy, : [ A2). Peopy: { A.). Fxor ('s,, Ao} that are identical in operation (copy), copy2, xor respectively) at the first character in output tuple. This would imply, if the T² and 8 both belong to one these families, Correlation(="1)) (-) x ((10, will 1.
On the other hand, one can show that for any operation f1. E {copy1, copy2, xor) with /1 + J2 will have Correlation((1,,2
f2. ((, 1}3)) 0. That would then suggest that if x" and T belong to different families among Feepy, F Fxor then Correlation(7*((). T ()), 132)) will be 0.
By simple counting argument, with probability 3.2 two maps T and T randomly selected from {Ai, A.s will have non-zero correlation.
42

---

On the Power of Context-Enhanced Learning in LLMs
G. Upper Bound: Context-Enhance Learning of MLT(dd,), n) with Simple Surrogate Model
G.1. Setup of Surrogate Model with In-context Capability
Given d + alphabets Al...., Ad+1 of size n and bijective phrasebooks π; A - A4+1 The input the translation process is an even-length sequence the first alphabet, which we denote as $1 Af where the sequence length. The translation process modifies the input string recursively from 8, 8,+1 through the following 2 sub-processes:
1. Circular shift:
The characters in 8; A shifted by character leftward (and wrapped around the end if necessary) give sequence &, A/. Formally, for each E [1, we have $t.j *1.1+11%%.-
2. Translate: Using the phrasebook n; : A A2+1 translate 2-tuples (bigrams) of consecutive characters in sequence 8; create 81+1- That is, for every odd j [1, DJ. (81+1,j, 8i+1.++) (8 $i.+t).
Now let us revisit the surrogate model introduced in Section 5.1. Without loss of generality let Ai A =..... Aa+1
A {1, ...... n (ca) = 1. For any 2-tuple (a, b) E A², let its vector representation be a -dimensional vector v(a, b) - ea 8 es. Note that va,, (e, b) is also one-hot vector where
For any single character A, let its vector representation be one-hot vector R such that
(1
if = an + elsewhere
va, b);
We use the notation 6, (long one-hot) to denote one-hot vector in R. with a-th position being to avoid confusion.
Definition G.1 (Matrix Representation of Sequence).
For a length-L imput sequence 8; = (., S, let its matrix representation be Mat(s,) & V, E Rn" that
V = v (s1,1, (81,3,8,4) v ( Si.L)
For each j E [L/2]. we use ( to denote the j-th column of V. We also denote the above conversion from a sequence 8; to matrix form as V, Mat( (x)) and assume that V serves as the input the surrogate model.
Note that the matricization operation is invertible by construction: for each column V:U) let arg max V:U), may read off the two characters the original alphabet by computing Mat = ([z/ |n], x%m).
At each level of translation, we assume the surrogate model will perform the following operations to V;: (a) Circular shift from to V.
Definition G.2 (Circular Shifting Operator Shift).
Given a matrix representation V E RD x/2 of sequence 8, the circular shifting operator Shi ft acts on V as Shift(V) V E x where for all j E [I/2].
かい Qv Q (0-1)))
where (I, 8 1,) (1, & I.) 1, pax1 is the all-ones vector, and O is the Hadamard product.
Lemma G.3 (Equivalence of Shi and circular shift). For any sequence A4, let be the circular shifted s, then
Mat (8) Shi Et (Mat (s)).
43

---

On the Power of Context-Enhanced Learning in LLMs
Proof of Lemma G.3. In this proof we will show that Mat (8) and Shi (Mat (s)) agrees on every column.
Fix column j [n/2]. without loss of generality let the imput sequence be (a, b, d) A starting from the (2) 1)-th position to the (2j + 2)-th position (wrapped around when necessary). By construction each output column of VG) is dependent on at most Vú) and V which corresponds to 4 characters in the sequence 8.
By definition V we then have V v(a. b) = ea @ Eb and Vi+18% v(c, d) = e @ ed. It follows that
マい Qv о Q v((++)%))
((hn I,) (1, 8 (e. @ (LI, @ I,) (1, @ 11,) (e, @ ex))
= -(In (1,)1,))1, 81 @ (L,, 11, (4 @ 1e))
-((h 8
((1, (e.
(ea, ed are one-hot, 1 1 e = 1)
= ((1n @ 1.) @ 1 ((1, @ 1.) (16e.)
= (es @ 1..) @ (1, @ 0.)
= Eb & E..
(by definition of Kronecker product)
Note that ε₁ @ e, is just Mat(8) since the (2j - 1)-th and the 2j-th character of the shifted sequence & is now (b, c).
This concludes the proof.
(b) Translation from V. to V+11
With Shift effectively completing Merge, Circular shift, and Split, what remains is the translation leveraging ,, Since A 1++1 [n]]. there natural bijection between the space of binary all possible (not necessarily bijective) mappings between 2-tuples from A. Concretely
matrix and the space of
Definition G.4 (Colunn-Stochastic Matrix Representation of phrasebook).
Given phrasebook æ : A22 + A2, its matrix representation defined to be Matrix(F) E R=*xn" that for i [n²],
Matrix(F)); = { [[<([)/)) ([i/n],j%m)
elsewhere
Let Vi+1 MatrixV, where Matrix x is a binary column-stochastic matrix, we can show that, a complete translation process for one level can formally expressed as follows:
Lemma G.5 (Equivalence of Mat.rix x and Translate). For any sequence 8; A4, let my a bijective phrasebook A - A defined Section 2.2, then
V411 = Matrix(m; Shift (Mat (s:)) = Mat (Tx, Mat($141)-
Proof of Lemma G.5. Fix any column j E [T/2]], let (a, b) be the (2j- - 1) and the 2ý-th character of the shifted sequence
5e Let (c, d) = n, (a, b), by the translation construction we know the (2j 1)-th and the 2j-th character 8:41 just c and d. Hence Mat(s, )0) d(c.d).
From Lemma G.3 know that Shi (Mat(s))))) know that Matrix(m) Shi ft(Mat(s))) is a one-hot vector at the (en + d)-th position, i.e. Matrix(m,) Shift (M))(s))))) = e, @ e. = d(c.d). Since Mat(s vσ for all j. we have V41 Mat(s.+1).
v(a, By construction of Matrix(m:) above,
v(i
H+1
44

---

On the Power of Context-Enhanced Learning in LLMs
Parameterization of the Translation Matrix P.
With the two key operations in place, we can introduce the surrogate model in its full detail. Since the multi-level translation task is naturally sequential operation, we model the surrogate operations as multi-layer network as well (which also matches with the solution found by Llama-based models when trained on real data as in Section 3.3).
For each level, the surrogate model needs to present both in-context learning capability at the initialization and in in-weight capability toward the end of context-enhanced learning on a certain set of phrasebooks IT", The in-weight capability requires certain parameter to store II" on its own that is independent of the context.
To capture both capabilities at the same time, parameterize the translation matrix P. combination of in-context information C₁ I=" xn² and in-weight memory W; R=""="
Definition G.6 (Effective Translation Matrix). For in-context information C; "="""" and in-weight memory W; R=*xx" at level the corresponding effective translation matrix P. is defined as
P: Ha rdMax (C; W.)
where HardMax the column-win wise hard-max function converting C, + W; binary column stochastic matrix.
Note that column P., which we denote (!()) is equal the one-hot vector at arg max (C + W (!
Surrogate Model with In-Context Capability Now can formally introduce the surrogate model.
Definition G.7 (Surrogate Model for MLT).
The surrogate model for MLT (d, n) can be represented by the recursive expression
V HardMax(C: W.) Shift(V.))
(18)
The learnable parameters for the surrogate model are the weight matrices { (Wijd = (Wi, W...., Wa). We
denote the surrogate model parameterized by ( (Wi) as SURR-MLT, (wije_) (.) which maps the input and the in-context information the output as
Va+1 MURT [Wi] (C1, Cz..... Cd, Vi)
A HardMax Wa)
(19)
Shift (HardMax(Caas + Wa-1) Shift HardMax(CC (arraaaaa + Wi) W
In this surrogate model, the in-context descriptive text DESC is just (Ci...., Ca), where each matrix is directly being passed into the corresponding layer. When providing information about set of phrasebooks = { we have C, Matrix(m;) where Matrix(T:) is the column-stochastic matrix representation of π. as defined in Definition G.4.
When partial phrasebook information is provided (corresponding dropping certain ab->CD entries in the language model context), we zero-out the corresponding column in C, When no in-context information is provided for level i, we just have C; be the all-zero matrix containing no information (assuming zero prior).
For simplicity of presentation, in Definition G.7 the in-context information C,'s are provided directly to the corresponding layers. We note that the equivalent operations can be exactly re- p--armettrrized such that C, 's are provided in-context (in concatenation with V1). The reparameterization of the surrogate model is provided below:
45

---

On the Power of Context-Enhanced Learning in LLMs
Definition G.8 CContext-Augmentine Surrogate Model for MI.T). Let the context- -ugemeted input be
x=dx((+++L))
X = -. Ca. Vi
then rewrite the same surrogate model
Iane
O1xl ]01]+ + ((x:
c ② In2
[nndxx
Ir. 1)
Odn
X = X;
w., shift
X;
0LXn²
0
(20)
= ................................ Ca. Onax1 ,o,...,,............ (C + W.) Shift (Vi)
|CI Ca. V++1]
With the reparameterization, the model is capable of generating [C1,.... Cd, Va+1l with input (C Cd, V1). with all information provided in-context in the input.
Now let us check what does Definition 2.2 (ICL-capable) and Definition 2.1 (specific task-capable) mean in the context of the surrogate model. make things more rigorous we introduce two stronger notions capabilities:
Definition G.9 (Strongly MLT n) -CC--caabble surrogate model). We say a surrogate model SURR-MLT, (w (-) strongly MLT(d, mz-IIC-capablle if for any of phrascbooks II [m)) in MLT(nn, d), for any input sequence 81 A where even, we have
SURR-MLT (Wi (aatida) Mat(s))) Mat(MLTT(ss))
Definition G.10 (Strongly MI.T ccapable surrogate model). For fixed set of phrasebooks II" ) in MLT(nn, d), say surrogate model SURR-MLA (W. (.) is strongly MLT if for any input sequence E A where L is even, we have
SURR-MLT (W. (0. 0, Mat(s) Mat(MLT (s1)).
Now we can show the following properties of the surrogate model SURR-MLT (Way (.):
Lemma G.11. When | } for all e [d]. SURR-MLT (wat strongly (L) -CCC-cppblle.
Proof. Fix any set of phrasebooks { and its corresponding matrix representations (Matrix(m:) Since [W.llo 1. for any column k. entries in (((A) flip the argmax of Matri (at)) + w(()) away from being arg max Mat Mat:i
Therefore we have P; HardMax(Matrix(t) + W,) Matrix(m: for all layers E [d].
By Lemma G.5, for all i [d] we have P. Matrix (m:) recovering TG Hence for any input sequence 81 A4, have SURR-MLT (w.y.. (((T)))..... {(ma)} Mat(s:)) Mat(MLTT(ss)).
Lemma G.12. Fix target set of phrasebooks II = { { when HardMax (W,) = -Matrix(;;)) for all E (d) SURR-MLT (wo._, (.) is strongly -cTan-capaale.
Proof. By Lemma G.5, for all i [d] we have P. HardMax W, + 0) Matrix(m;) recovering Tæy Hence for any input sequence 81 A4, we have SURR-MIL (w:)._, (O,..., (a)(s Mat(MLT (a:)).
46

---

On the Power of Context-Enhanced Learning in LLMs
Lemma G.11 suggests that when the weight matrices have small initializations, the model has perfect ICL capability, Meanwhile Lemma G.12 suggests that when the weights W; recover Mat Matrix(m;) in the column-wise hard-max sense, then the surrogate model can perform MT when context is being provided (C; 0).
G.2. Learning II" in ML) with Heuristics Search
In this section, provide brute-force algorithm that can learn any target set of phrasebooks II in MLT n) using single "short" sequence whose length is not exponentially dependent on d.
Before proceeding to the details, let us first investigate more on the nature of MLT and the surrogate model. For simplicity of notations, given IT = { Ta). we denote the general translation operator Matrix as P, and de- note the translation operator Matrix (m1) as W Also, with slight abuse of notations we use MLTTI (Vi) to denote Mat (MLTn (Mat" (Vi)))
First, we characterize the input sequence that is good for providing learning signals.
Definition G.13 (II' -coveabble input).
Fix a target set of phrasebooks II" {......, .ii.....,T. m in MLT(d, n) and an imput matrix Vi RT*x1, let Vi' V22 Vi,..., Vā, Vát1 E R= xL be the intermediate outputs when applying MLTTT- on Vi. We say Vi is -ovvevveebee if for all levels [d]. Vi is of rank-n?
Note that as a matrix with only one-hot columns, Vi- being rank-n² suggests that for A [n²]. there exists some column j E [L] such that V.+()) Ex. In the context of the translation process, it means that the correct translation process of covverabee input V1 would require all entries of all phrasebooks in II'.
Next we will show that if we use the context to condition all translation operators Pt of the surrogate model to be p(a for all one level E [d]& {i} as the unconditioned level), then for the surrogate model to correctly perform MLT the operator for the unconditioned level must also be equal p
Lemma G.14 (Uniqueness single P. when conditioning all other levels).
Fix target of phrasebooks II" = -- mr) in MLT n) and IT'-coveaable input V1. For any level i [d], consider surrogate model SURR-MLI (W:yd. (C1, Cz...., Ca. -) with certain context C1, C2...., Ca such that Pt W for all I [d] except i. Then SURR-MLJ (witg_ (C1, C2...., 22,...,C.,, Cd, VT) MLTT- (Vi) if and only if
P. = HardMax (C₁ + W.) = W
Proof. This lemma is a direct consequence of the bijective property of the translation process shown in Lemma E. Let Vi, V2, V ..... Vj Vā+ E R=²xA be the intermediate outputs when applying MLT on V1, and let V1, V2, Ve,..., Va, Va+i Bn= xL be the intermediate outputs when applying SURR-MLT (wa) SURR-MLT(Wa). (CI, Cz...., Ca. as described. Since we assume P.- p(7))) for all | we have V, V,* and therefore V V,".
On the other end, since SURR-MLTT (wart_ (C1, C2,..,, by the invertible property of the translation process (Lemma E.1) must have V V and thus V41 V Combining both ends we know that
... Ca, Vi) MLT (V1) Vi+1 and P, p(a)) for all | >> i,
P.V. = VA+1 = Va = W. = W.V.
(21)
Since V. V is rank n² by the II" coverabbee assumption and P, R xn? it must be that P: W
If we further condition on the held-out level P: such that we only leave one column of p. (k) not necessarily equal to (k) we have the following corollary
Corollary G.15 (Uniqueness of a single P, column when conditioning everything else).
Fix a target set of phrasebooks II (mi....,,TT) T } in MLT n) and a c'-coverable input V₁. For any level [d] and translation entry k E [n²], consider a surrogate model SURR-MLT (W.) (C1, Ci,.,,,..,CC Ca,)
(L..(
47

---

On the Power of Sentext-Enhaneel Learning in LLMs
with certain context C₁ .z..... Ca such that p(u) W for all (1,1) [d] [n²]] ((1,)))
Then
SURR-MLT (w.) (C1, C2,.... Ca. Vi) MLTTI (V1) if and only if
p(k) HardMax (c + w!
(g(+) W
This suggests that if we condition everything else except for column of the translation operator, then to match the final output on a II coverrabee sequence, the model must recover the held-out column to be the same the ground truth in the set of phrasebooks.
Now introduce the search algorithm, which simply enumerate over all translation columns W learning target, generate a contextual information that only leaves that column unconditioned, and search over all possible one-hot vectors for w until the output matches with ML Once the output matches, by Corollary G.15 we know Ha MardMax( ccMaa(M.a)) recovers (* and we move on to the next learning target. The algorithm can be formalized as follows:
Algorithm 2 Context-Enhancer Searching Algorithm for MLT n)
1: Input:
2: input V12 E RR"'LL, label V E 1="'x, descriptive text W Wi.... Wa E B='xn"
3:
4: Initialize WT, W..... Wa -
Start with zero initialization
5: for i 1 to d do
6: for k = to n do
7: Initialize C(())
W (I,a diag(ex))
# Create masked context matrix
8: # Search Loop
9: for a 1 to n do
((kk) < Ca
10:
Search over one-hot columns Va+1 < SURR-MLT W. (Wi... ..., W C(R), W W Vi)
11:
V4+1 Vo+1 then
break 14: end if
# Break when found the right column
15: end for
16: end for
17:
end for
18:
Return WI,..., Wa.
Theorem G.16 (I (Learning IT" with context-enhanced search with IT coverrabe input). For any target set of phrasebooks II {mī.... ........ corresponding ground truth label Vi+1 MLT (V1), Algorithm 2 terminates with W; W for all i e [d] with O d) forward passes through the surrogate model.
......TT)) T. in MLT(d, n), given an -cvvrrveell imput V1 and the
Proof. The statement can be proven by simple induction.
Let the inductive hypothesis such that when the enumeration goes to k-th column of the i-th layer, if Har HardMax(W (W))) + w, ( for all (1,j) [d] [n²] and w7(1) w-()) for all (1,j) such that
i A k. then the search loop (Algorithm 2: line 13) breaks with (+(A)) (*) while HardMax + w- w,(h) for all (1,j) [d] x (x²) is preserved.
The base case is satisfied with zero initialization, we have HardMax W (1)
(w HardMax(O+ w.
for all (I,j) [d] x [n²] and there are no requirements for ( = W yet.
For the induction step, we note that with the condition of rddd (adddMxxxXW (1()) wrll)) = w. for all (l.j) [d] x [n2).
48

---

On the Power of Context-Enhanced Learning in LLMs
WI.... Wi-1 Caae):
Wa will correctly condition all columns of P's except for the P(() since
Ci = W (I,² - (i)g(E)))) = 0.
(22)
Thus by Corollary G.15, we know that the search loop will terminate when it finds W W! column provides the correct inductive hypothesis on W (i) -W for the next enumeration step.
W, The newly added
By induction to i d and n will be able recover W; W for all i e [d].
Given that we can learn W; effectively with coveevvee input, how should we construct such inputs? It turned out that with high probability, short random strings suffices.
Lemma G.17 (Distribution intermediate sequences). Fix a target set of phrasebooks II- { a. MLT n.). Let V1 Ra X be random input matrix (L)(d, such that each column v is iid. sampled from v² (the uniform distribution over one-hot vectors (ex yn² the columns of intermediate random sequences Vj....,V V.. uniform distribution.
U((ex)k)))
xL. obtained by passing the input V, through MLTTI- also follow the same i.i.d.
Proof. We will prove the claim by induction on depth i. Let inductive hypothesis be that columns V, independently follow uniform distribution over the one-hot vectors {ex Fe Note that the base case is just the assumption.
Now we prove for the inductive step. For any j E [Z], we can write V ((n]). Intuitively this means each 2-tuple in the random input sequence is formed from two i..d. uniformly random characters, which is straightforward by construction.
Ca @ Caaan where a(s) 's follows i.i.d.
Now by Lemma G.3 we have V- Caran & Capa+aa also following )= Note that there is total independency of (n(()) with respect to the set of any other columns of V; since a(2)) and a(2++1) are independent from the generative process of any other columns in V,.
With columns in V, i..dd. following ((((e))) permuting the indices via p( does not change the distribution by symmetry of the uniform distribution. Therefore we have columns of Vi+1 = p V; also i.i.d. distributed as U({x; nn² and this complete the inductive step. By induction on we have the desired statement proved.
Since each column in V, i.i.d. follows U( sampling V, to be rank n² becomes identical to the classic coupon collection problem (see Lemma G.29 from Motwani (1995)). Thus we have the following bound:
Lemma G.18 (Short II -coveaable random sequence).
A random sequence V1 of length L 2 2n log 5 is II* coverrall with probability at least 1 5.
Proof. Let the event A; denote that V, is not rank n²a By Lemma G.17, each column of V, is i.i.d. distributed following U({es) Thus making V, being rank n² is equivalent coupon collecting problem (Motwaii, 1995) with set size n By Lemma G.29 we know that with = 2n log nd P [A++ < g Thus by a simple union bound the probability that V, being not II' -coveraaee is
5
d
P
A
P [.J]
d = 6.
(23)
d
==1
Now we can apply the above result and extend Theorem G.16.
49

---

On the Power of Context-Enhanced Learning in LLMs
Corollary G.19 (Learnnng IT with random input using heuristics search).
For any target set of phrasebooks II" f T2) MLT(d, n), with probability at least 1- ô over uniformly random input V1 of length terminates with W, W; for all i [d] with Onn' d) forward passes through the surrogate model.
2n log 8. Algorithm 2 provided with ground truth label Vä+1 MLTTT-
nd
G.3. Learning II" in M2)(2, with Surrogate Gradient Descent
In this section take the analysis one step beyond the heuristics searching regime. We will show that any set of phrasehooks II" { m2 can be samplc-efficienly learned by a gradient-descent based algorithm. In this particular case, the surrogate model is parameterized by
333 SURR-MLT (W:) (C,, C2. V1)= A HardMax(C, + W2) Shift (HardMax(CG + Wi) Shift (K)). (24)
We start with any weight initializations ( w E R=""x" satisfying || < the initialization strongly (LT cCpp-cae....e.
(0)
1 2 [ || <
1 2- by Lemma G.11
For simplicity, we denote the ground truth permutation matrix induced by कं as WT & p(rj) and similarly the ground truth permutation matrix induced by n as W & p(*2), From Lemma G.12 we know that the learning is successful if we have HardMax (W1) Wi and HardMax (W2) W (i.e. the maximum index of each weight column agrees with that of the ground truth).
We employ a wise gradient descent algorithm for the learning process. The algorithm takes in a single fixed sequence 81 with matrix representation Vi and corresponding ground truth label
Mat(MLTTT- (81)) = SURR-MLL (w() (Wi, W2, Vi).
(25)
Given the input and label, we employ the following gradient descent based algorithm to update the weights:
The training happens in layer-wise and column-wise fashion: We first freeze W2 and set W1 as the trainable parameter. For each entry & E [n²"], create context matrix Ci(R) A Wi (I,, diag(ex)) which essentially creates a copy of Wi except of setting the k-th column to be zero. Then we take a forward pass through the surrogate model with the one-column (w) (Ci(k)- W Vi). Here the subscript '(1,k) denotes the final output when the i-th column of the first context matrix is being dropped.
dropped-out context and get output V3(1,K) A SURR-MLT, w,))))((
The weight update follows a surrogate gradient update scheme where we use the MSE loss: L == V 112 Since it difficult to take gradient through the hardmax function, we instead compute the gradient of the loss with respect the translation matrix P1 = Ha rdMax (W1 C1(1)) and apply the update w <- --(()) ac We apply such gradient update twice for each dropped column k.
ap
For the second layer, we freeze the first layer W, and apply one-column dropouts to the Cz. Similarly we apply the ac but we only need one gradient step per column.
surrogate gradient update W2 W2 aP.
We claim the surrogate gradient descent update can correctly recover P and P similar the heuristics search case.
Theorem G.24 (Learning II with context-enhanced surrogate GD with -coveraeee input).
For any initialization W 1(0): W2(0) E "Rn2xnn xn² such that [| llo S 1 and || llo S for any target set of phrasebooks II. = { na } in MLT(2, n), given an II coveraabe input V1 and the corresponding ground truth label V MLT (Vi), Algorithm 2 terminates with HardMax (W1) = W and HardMax (W2) -W;-
To prove for Theorem G.24 we will carefully analyze the learning of the first layer and second layer respectively, and provide a similar induction argument as in the proof for the heuristics search case.
"Note that this is not the unique solution to attain strongly MI. n cappable model
50

---

On the Power of Context-Enhanced Learning in LLMs
Algorithm 3 Conheex-Enhanced Layerwise Gradient Descent
1: Input: input V1 Rx*xL label Vig E RT*x descriptive text Wi, W22 g=*xa" init w W( gn"x""
2:
3: # Train the first layer
4: for k = 1 to n² do
5: CI(a)) _ Wi (I,² - diag(es))
# Create context matrix with k-th column dropped.
6: fort = 1 to 2 do
7:
V3(1..) < SURR-MLT (URR-MLTTW........ (w..w)(CC(()) W Vi)
Forward pass
8:
+ Va(1.)) 12
( W(() at
9:
op(2)
# Surrogate gradient update
10: end for
11: end for
12:
13: # Train the second layer
14: for
k == to n² do
C2(2) - Wi (I diag(ex))
Create context matrix with k-th column dropped. Forward pass
V2(2,1) < SURR-MLT (W. C2(R), VT)
L+] (3 V
18:
(kk)) + 2 "P
8
# Surrogate gradient update
19: end for
20:
Return W1, W₂
G.3.1. I. LEARNING THE FIRST LAYER
To study the learning process we first need to compute the closed-fone gradient ac which requires the following lemma: ap())
Lemma G.20 (Gradient with respect to incorrect column in P1).
When only the k-th column of translation matrix p(() is not equal to P (:())) let P( = Ea @ ED and Pr = ea- @ Eb for some a, b, a², b* E [n]. If p(()) is used in the forward pass, there exists Z4 and B E N such that the gradient of L with respect to ((k) is of the form
(2a + 23) (1, @ eb) - (2a + 23) (1, 8 ev)) + 23 (e.- @ 1n)
ifa" a, # if # a. =
0C
(2a 23) (e, 1,,) (2a + 28) (e., & 23 (1, @ ee.))
(2a + 23) (e, @ 1,, + (2a + 23) (1,, @ ev) - 20 (ea 8 1x) 2a (1 8 ee))
ifa # a + b.
Proof of Lemma G.20. In this proof we use é, (long one-hot) to denote the one-hot vector R=² with on the a-th index and use e, (short one-hot) to denote the one-hot vector in R with on the a-th index. We use V,, V2, V2 and V₂ to denote the intermediate sequences attained with translation matrix (()) and Vi, V V and V to denote the counterfactual intermediate sequences should the forward pass is done with the ground truth translations p)((k) W
Now we can proceed to the gradient calculations. First note that with L(P)) = V Vi by chain rule we have
Oc
Σ
ov (1)
()) - V
= 2 Σ
(v:5)
(v,0) vol1)).
(26)
OP
aP
ovl))
(
OP
1=1
Ø=1
Specifically for each column / E [n of P we have
(v
Oc 0P())
(v,10) - -v,(1).
=
(27)
OP!!)
==1
51

---

On the Power of Context-Enhanced Learning in LLMs
Let us fix particular column j E [L] in the output V3 and compute the gradients. With MD) as the j-th column of the matrix M, the computation graph for the forward pass is of the form:
V
A
P
V
P
V
2
v(+11)
1
PS
v(i+1)
2
(28)
v(5+2)
We can see that (5) is only affected by Pi through ( and V22 (2(+1)
Let us first compute ov (1) /0Pt. Assume V = Cp and V v Eg for some P. E [n²], Vi(i is computed as
W
()) W
(Q
5) v2 v W (QP(") Q P(
(29)
We may express the Hadamard product in the following two ways: P("))
Vg Wi (QP()) p)
didi (QP(()) Q @
(30)
W Q p(o) QP(())
W didi
@Q T p(o)) (QTP(() QP!(
Therefore when & we have
av,() (1)
(i)
ov
W diag (QTP((0) Qi
0:
3
W diag (QP(-)) QT W (p.a)
0.
(31)
0P(P) (P)
0P(1)
aP(")
When p = q. the expression would be
av
oV (5)
W diag P(m) Q + Widdia (Qp(")) QT:
vI + p
ap() = 0.
(32)
ap(r)
Now we move on to compute (V-1)) used when computing V2 and (5+1)) Let us go over these cases one-by-one.
V) There are in total four cases to consider: depending on whether ((() is being
1. Case 1: V.(0) * ég and (((+1) + ēx.
Assume v-(1) e and ((++1) E, for some p, + k. Since V- Pi (i) while P1 equals Pf for all columns not equal to A by assumption, we have v2(1 = e,ep
V p(r) = pr(2) Píé e v2 With an identical argument
we have v-ό+1)= = V2 ()+) Thus in this case v v-(() and hence
3
()) V3(1)) = 0.
(33)
aP:
2. Case v.(0) ēy and ((+1) # Ex.
Assume v(i+1) e, for some 9 # k. from case we know V 6+1) Vg (++)
However since v-6)
-6 we have
v, P,a.- = ((k), which is not equal v=(() P.
Therefore
v( Vo(()) Wi QQP( R))
Q p(1))
W. (PP)(k) P(())
W
Q p (4)
Q p(()) -P- (p((k)
(34)
diggga
Q p(1)
Q (p(k) (-((k))
52

---

On the Power of Context-Enhancd Learning in LLMs
For the gradient with respect to P(A) combining Equation (34) with Equation (31) (substituring p = k) we have
Vv3 v2(
W diag (p(p()) Q)
T (iidiag (Q p( Q (P(h) -P!((t)
@P(A)
Q diag
Q (Q p(() W2 Widiag @ [P(())) Q (p(A) - -Pi(())
(35)
Q diag
(Q p(.)
diag (TTP(T)
p(o)) Q (p(h) Pi P1(())
((1,11)) @ In)
(P( Pi(M)
(by Lemma G.27)
where dropped W W in the third step since W is permutation matrix and W W In-
Now plugging in Pl = ea @ ey and P. = Ea- @ eb. we have
Vv."" ())
V/3 vo(
((1,, X In) (e, & eb) ((1,1)) @ In) (ea- 2 @er)
0P!(9
(11,, ea) @ I.e,) (11n1 ea-) @ kher))
(36)
1,, @ C 1, er
1n @ (es ex.)
3. Case 3: V-00 + ee and ((0+1) = ex-
This is a symmetric case with respect to case 2. Assume v,00) Ey for some p * k, from case we know
v = vi However since v ē₁ we have ((5+1) = Piex (()) which is not equal to V2 v*(++) = p. (k)
Therefore similar to case we have
v -V--Va-0
3 V-(
Wiidiag (QP(C)) QT
(37)
Combining Equation (37) with Equation (32) gives
) ap(R)/
T (v50)
v2(1)) (W diag (QP!(")) )²")²=""'""")))) (Widiag diag (QP( QT (p(N)-
- -P((()))
Qdiag (QP(1)) W Wødiag (QP((") Q (p(*)- - -p(()))
(38)
Qdiag (QP(()) diag (QP("") QT (p(h) - pi(kk)
(4n @ (1n11)) (p(()) - pi(1t)
(by Lemma G.28)
Now plugging in p(() ea es and p. c.. @ Es- we have
ov
(v)) - -3 = (In 2 (1n1.)) (ea 8ex) (In @ (1n111) (e.- @er)
0p(k]
(Inea 8 (I, e>))
(Iner- @ (In
(39)
e, & 1n - e @ 1,
= (e. ea.) & 1n
4. Case 1 v(it1) ck.
53

---

On the Power of Context-Enhanced Learning in LLMs
This is the most complicated case since the loss is contributed by different paths. We first decompose the negative residual
6 Vi W2 (QP)
о Q
p(())
W. (QP) (k)
2 Q
Q'Pi' (h))
W [QP)
p(())
W (QP( (k)
PT (k)
(40)
+ Wi (PP(t) (k)
② Q T
P (k) W2 QQP" (k)
① P (k)
Wiidiag
(QP(())
Q (p(h) - pi(li))
W2 diag Q T p!(r))
Q (P(l))
- - (p(4))
Combining with Equation (32), we have
(1) -V3
V3 vi(
OP
(W2 diag
(Q p(((()))
Q + W diag (QP(k)) ²T)T
(Widiag (P((1) (k) T (P(A) -P((A) digi (Q pP:))())) @ -P(())
QQ = diag
(Q_P) P diag QP(()) Q (P(N) (P(k)_Pp -PP(()))
(a)
+ Qdiag (QP((t)
diag (P((k) Q (P/(t) -P
(b)
+ QT diag
Q
T (1)
P
diag Q T (()(t)
(p(L)) -P pi(()))
(c)
Qdiag
(QP((t) diag
Q (TTPP(())
p-(())
((()) -P P.lkk)
(d)
Now let us analyze the four cross terms term-by-term.
(a) With (kk) & eb. by Lemma G.26 we have Q (P(()) 1n @ ea and QP(() eb 1n. Therefore
Q T p( QP
(k) = es @ ea and hence
diag (p)((*) diag (QP((t) diag (es @ e,) = diag (es) @ diag (ea)
(41)
It follows that
Q diag
(k)) (QT)({) diag (QP(k)) Q
(1n 8 4n) (4] 1 (diag (es) diag (ea)) (1, @ la) (4 x 1.))
(1n 8 In) [IT @ 1 (e, @ diag (ea)) (In @ 1!))
(42)
(1, @ In) [T] @ 1
(es )e,))
(1,, In) [I] 1 (ese) @
=1,, @ In) (ebe) @ 1
Plugging into (a) we have
(a) Q diag Q P((M diag (QP(*) Q -pp
(1, @ In) ( T
2 1.1 (e, @ en) - (1, 8 (n) (esez 8 1111) (ea- S
= (1n @ In) (ere ea 8 1) (1, @ In) (ese) 81)
(43)
(1n @ I.) 2 e e.))
(1, @ Im)) (1 & -
1, & when a # otherwise
54

---

On the Power of Context-Enhanced Learning in LLMs
(b) We have seen the in case 3, by Lemma G.28 have
(b) Qdiag (QP(()) diag (QP(()) Q² (P( - P!(())) = (ea -a-) 0 In.
)
ea-)
(44)
(c) With p(k) = @ and Pi(K)) Ca- @ Cb., we have
diag (1n @ ca)
if a" otherwise
diag
Q P (b) diag (QTP(())) diag (1, & ea) diag (1, & ex-)
(45)
0
Thus by Lemma G.27 we have
(Inl 2 Ia ifa=
Q Q"diag
@QT p(A) diag (Q (PTP)
(46)
otherwise
When a", then have
T diag @ P diag Q Pi(())) Q (P((t) - PP
Q
((1n1)) @ In) (ea @ es)) ((1n1n) @ In) (ea & err)
(47)
= 1n & (es, - eb-)
Thus in summary
(c) Q diag Q T p(k)) diag ( (PTP((()) Q (p(k) pi(())
1n @ (es -0--
if a" otherwise.
(48)
(d) Now for the last term, note that
Qdiag (QP((t) diag Q T P p/ Q
Qdiag (es @ 1n) diag (1, @ )
Qdiag (es @ ea-) Q
(Jn 6 1,) (1 @ T (diag (ex) @ diag (ea-)) (In @ I.) (1, Φ (
(49)
(4n @ 1x) (1]] @ 4.T) (diag (es) & e..) (1 8 1.T
(In @ 1,) (c) @ ca-) (1 @ I
(In @ 1n) e, 2 @1) (1 @ (
e, @ 1n) [1] @ IIT))
Plugging in (P(*) - ((())), ). we have that
Qdiag (QP(()) diag (Q p(((k) Q (P(1t)
-P(())) _
p(kk)
(ea- eg
8 1,) (1, 8 [,] (ea @ ea) - es 8 1m) (1 8 11,) (ea 8
(e. e
T
1n) (1 ex) - (ea e T @ 1n) (1 2 e..)
(ea eg
3 1,) (e, < 1) (ea eg x lm)
(er 3 1)
(50)
(ea- eb
es X 1m) (e.- eg es- & 81n)
Ca- 1n b b*
0
otherwise.
Summing the four terms together, we then have
when a* = a, D² =
(v
(v50) - v)((0)
1, 2 (e, e--) + ea. In
when a, b" # b
(51)
dP
1 @b + (e. ear) @ 1,
when a" + a, b" = b
1, eb + ea & 1,
when a" + a, b" + b.
55

---

On the Power of Eonhace-Enhaneed Learning in LLMs
Now we ready to provide the gradient expression for the loss over the entire sequence. Observe that for every consecutive sequence of m columns (V,*)), V v* V v*0+m--1) that all equals to Ck. will result in one incorrect column
V (f-1) in case one incorrect column
V V(itm-1) in case 3, and m incorrect columns (V v in case 4.
v(-11)
( - Ep ex)
p(o
(case 2)
(=
(= Ex)
P(N)
(
W
Vg
(case
7
2(5tm-2)
f (= (=ex)
pl V ((itm22)
(
W
Vs
v(ttm22
(case 4)
((+m-11)
- (= 2k)
p
v (itm-1)
vutm-1)
W
v 1)
(case 3)
v((+m)
v(+m) (= Eg + Ex)
plu)
(--))
(++m
W
v
(case or 2)
Figure 12. Error propagation of p(A)
For illustration, one can refer to the computation graph in Figure 12. In the graph, green entries agrees with the counterfactual
values with correct P (R) Red and pink entries are incorrect entries where red entries are consequence solely dependent on P(() (case 4) and pink entries depend on other correct columns (case 2.3).
Assume that in total there columns in V, under case columns in V/ under case 3, and 3 columns in Va under case.
4, then by Equation (27) the total gradient can be expressed as follows:
When a = a", b b*:
2C
2a (1, (e, - e5-) 2a (e. - ea-) & 1.)) + 23 (1, & (ee--z)) + e.- 2 1,)
ap(k)
(52)
(2a 28) (1, ev) - (2a 23) (1, @ eb-) ) 23 (c.- 1,).
When a + a", a²,b== bb::
dc
2a (1n 8 (es eT-)) )) + 2a ((e. e..) & 1,) + 2 (1, @ es + (ea - - 2 1n)
(53)
(2a + 2 (e., & 1n) (2a + 28) (e.. @ 1.) +3 (1, @e.).
56

---

On the Power of Eonhace-Enhaneed Learning in LLMs
When a + a b":
ac
2n (1n & (es eb-)) + 2a ((ea ea-) X 1.) 23 (In es ea 00 1n)
(P(h)
(54)
(2a + 28) (ea & 1n) (2a 28) (1, ea) 2a (ea @ 1n) -2 (1n @ @eb-).
This gives the desired expression of gradient//
Now with the gradient expression, are ready to prove for the learning of single missing column.
Lemma G.21 (Learning Column of W1).
Fix an input sequence Vi E x and any column index E [n²), HardMax (C2 W2) W and HardMax (Ci(s) W1) equals to Wi everywhere except for the k-th column and there exists non-empty subset of indices I C [L] such that V,+()) ēx for all j E for any initialization W R such that [W (18 Woo0llo llo 1. Taking two surrogate gradient updates on Wilk) as deseribed in Algorithm 3 gives W such that HardMax ( (()) = (1(
1(2)
Proof of Lemma G.21.
Without loss of generality, assume the initialization P earo: Chgo and Pi(A)) ea- & eyt for some a(u): b(o), a², [n]. We first note that the conditions specified the lemma meets the assumptions required by Lemma G.20, namely there is only one incorrect column in W, missing and that column is being used the forward pass at least one time (since J is non-empty).
To prove HardMax
W1 it is sufficient to show that arg max W 1(2) = a" n + +
Now we will leverage the gradient expressions Lemma G.20. We will dive into three different cases:
When ((o) # a and b(;) + b".
By Lemma G.20 we have for some a E E and E that
x 1n) + (2a + 2,3) (1n @ ebwo, 2a (ea- @ 1n) 2xx (1, @ eb-). dc = (2a + 2,3 (Ea
(55)
(p(k) 1(0)
It not hard to verify that the unique smallest entry is of index a" + b* with value 44a This entry contributed by the intersection of -2x (e.- & 1,,) 2a (1,, e.,) ). the remaining smaller entries are of value 2a contributed by non-intersecting entries in the same expression above.
Thus we know arg min ac/Op(k) (0) a" n b² with margin of least -2 (since 1).
Now apply the gradient step the weight initialization. Since the largest margin in the initialization (which is 1), we have
W (k) 1(0) llo
1. the margin of a"n++"" dominates
arg max ( (W10)
arg max
((()) 1(0)
0c
=
(56)
ap(h)
Therefore p(k) = P. (1()) with the first step. We will reach zero loss after the first step and hence the second step is static, so we have shown arg max (w(l)) 1(2) a'n + b* as desired.
When ((o) = and bi; # b*.
57

---

On the Power of Context-Enhanced Learning in LLMs
By Lemma G.20 we have for some 24 and [3 N that
ac
0P(k) (k) (2a + 23) (1n 2 ebjo)) - (2a + 23) (1, @ + (ea @ 1m).
(57)
1(0)
longer have an unique smallest entry, the of negative entries is of index b* where [n]. The values 2a for the case of x a and -2(a + 3) for other æ's. These negative entries contributed by the -(2x + 23) (1n es-). and all other entries are at least O. Thus we also have negative margin of at least -2 since 2 1. Therefore after taking the first gradient step, we know that there exists some x E [x] such that
In this
arg max
(Wr(k) 1(1)
arg max
W (4) 1(0)
OC
nEE +
(58)
ap
(R) 1(0)
Let [(k] I(1) arg max max(w(()) (1)))
are in the of b(1) bo since the negative margin of -2 dominates any entry-wise difference in the initialization as || 1 llo S 2. If 3 = 0 and it happens that a(1) = a",
1(1)) Ca()) Eb(sy then
then we have zero loss after the first step and we are done as the second step will be static. If b()) # b", then by Lemma G.20 the second step gradient of the form
0C (k)
(2a + 28) (Cars) @ 1n) -2a + 23) (e.- @ 1,)) 28 (1,, @ eb.).
(59)
(P(k) (1)
This is a bit tricky to analyze directly since we longer have the small entry-wise difference from the initialization in the weights, but may note that the sum of the update steps is of the form
dc
dc
+
ap(k)
ap(R)
1(1)
1(0)
(2a 2,8) (ears 1n) (2a + 28) (ea- 1,,) 23 (1,, e
(60)
+ (2a + 2,3) (In @ @bro,) - (2xx + 23) (1, 8 es-) + 2/ (ea- @ 1,))
(2ax + 23) (ear1) @ 1n) (2a + 23) (1, & Ebro)) 2a (e.- @ 1,) 2a (1, & es-).
This is identical to the single step gradient as in Equation (55) in the first case, so follow the identical argument we have
arg max (w((k) 1(2) = arg max
(x(k) 1(0)
Oc
ann
(61)
OP(4 1(0)
dp(
1(1)
as desired.
When ((o) # a" and b(o) = b*
Note that all expressions are symmetric with respect to and b with an additional swap of the Kronecker products, so we may follow the exact same argument as in the case of ((o) = a" and bço) * b* and arrive at the same conclusion.
Thus for any initializations, after two surrogate gradient steps on W I have Ha HardMax HardMax(WW()) ((2) w
G.3.2. LEARNING THE SECOND LAYER
Now for the second layer, will similarly first derive the gradient (which is much simpler) and show that with only one gradient step, one can learn the correct entry.
Lemma G.22 (Gradient with respect to incorrect column in P2).
When only the k-th column of translation matrix ((a) not equal to P (k))
P,P (k) is used in the forward pass, there exists
a E4 such that the gradient of C with respect to p! is of the form
0C
2a (P())
pk()))
aP( kk
58

---

On the Power of Enhaneee Learning in LLMs
Proof. We follow the same of notations as used in the proof for Lemma G.20. In partieular, we use V, V2, V2 and 3/3 to denote the intermediate sequences attained with translation matrix p( and Vi, V2, V and V to denote the counterfactual intermediate sequences should the forward pass is done with the ground truth translations P2 WT.
Since we assume that Pi = PT. we have V22 = Vz- Therefore for all column such that V o Val)) must be so that V. V.(() εR, and the residual can be written as
3
V33 (3) V3
P2V. P.V;10 P2ex P P P2(R)
(62)
Assume that p(() has been used times in the forward pass, follow the chain rule then have
ov
0P (k)
OC
2 Σ j=1
aPpK))
(v3 - -v2 = 2a (P -pp- p2(h))
(63)
Lemma G.23 (Learning W2).
Fix inpul sequence V1 E x and any column index k. E [n²), if Ha HardMax (C1 + Wi) = Wi and HardMax (C2(k) + W2) equals to W22 everywhere except for the k-th column and if there exists a non-empty subset of indices I C [L] such that V2 = êx for all j E J. for any initialization W E Rn such that (k) [W 2(0) llo 1. Taking one surrogate gradient updates W (k) (W 2(1))
2(0)
described in Algorithm 3 gives w (k)
2(1) such that
2(0)
(2(
HardMax
Proof. Without loss of generality, assume at the initialization P(R) Earo) Ebro) and P. ((u): b(o); ,",b" b= [n]. We first note that the conditions specified the lemma meets assumptions required by Lemma G.22, namely there is only one incorrect column in W2 missing and that column is being used in the forward pass at least one time (since I is non-empty). To prove HardMax (W 2(1))
ea. X es- for some
(x()) it is sufficient to show that
arg max w-(h) 2(1) a b.. By Lemma G.22, have
0C
2
Σ j-1
L
Vv
(
T
(i
Vo())) = 2a (p p2(1)) -
20 & Ehio 200 @ e
(64)
OP(()
(
Since 2 1, the negative margin of the a"n + -- entry dominates the initial difference in the initialization which is bounded by (=(A) Ilo S 1 Thus we have
2(0) 2'
arg (w(A) 2(1)
w (R)
dc
arg
2(0)
aP(h (R)
a n b*
(65)
2(0)
as desired.
G.3.3. PROOF FOR THEOREM G.24
Now we are ready to prove for the main theorem restated below:
Theorem G.24 (Learning II" with context-enhanced surrogate GD with II" -ooveraee input).
For any initialization Wn()) W2(0 E RRn2xx xx² such that [| llo & 2 and 22 S for any target set of phrasebooks II { my in MLT(2, n),, given an II -coeerable input V1 and the corresponding ground truth label V MLT (V1), Algorithm terminates with HardMax (W₂) Wi and HardMax (W2) W2
59

---

On the Power of Context-Enhancd Learning in LLMs
Proof. The statement can be proven by similar induction as in the proof for Theorem G.16.
Let the induction hypothesis be such that when the enumeration goes to k-th column of the i-th layer, if HardMax W, +
wr))) = w7()) for all (1, E [2] x [n²]] and HardMax (w/") =
Wi for all (4,5) such that < i or 1 iAj < k.
then the gradient update on the k-th column of the i-th layer ends with w(()) w while HardMax (W) W, w, for all (4,5) E [2] x [n²] is preserved.
The base case is satisfied as with initialization of || | lo < and 2 llo S 2 by Lemma G.II, we have
HardMax (W + w,b)) HardMax(O + w, = W, (1) for all (1,j) E [d] x [n²] and there are no require- ments for w W w. yet.
For the induction step, we note that with the inductive hypothesis of HardMax: W/ ( + w/ = w for all (1.j) E (2) x [n²]. (Ci(R), WT) (when = 1) or (WI, C2(1)) (when i = 2) will correctly condition all columns of P's except for the p(() since
( = W (T,2 - diag(es) (ex)) = 0.
(66)
Thus by Lemma G.21 (when i 1) or Lemma G.23 (when = 2), we know that after updating ((( we have Ha HardMax w(th) w,((h) The newly added column provides the correct inductive hypothesis HardMa HardMax(W (W,)) ()))) w,"h) for the next enumeration step.
By induction to i and n we will be able to recover HardMax (W,, = =W, for all i E [d].
Similarly we may use the coupon collecting argument to generalize the input to uniformly random strings as follows:
Corollary G.25 (Learning II in MLT(2, n) with context-enhanced surrogate GD with random input).
For any initialization W W1(0): W 2(0) Rud xn² such that (W10))l llo 1/b and || [2 llo S 1/2 for any larget set of phrasebooks II { T2) in MLT(2, n). with probability least - ô over uniformly random input V, of length L 2n log 5. HardMax (W1) = Wi and HardMax (W2)= = W2.
2n Algorithm 2 provided with the ground truth label V MLTT (V1) terminates with
G.4. Auxiliary Lemmas for Learning Surrogate Models
Lemma G.26. For any long one-hot vector Ca @ eb. with Q (In 1,,) (In @ IT) we have
Qu = c & 1
Q v = 1, @ ea.
(67)
Proof. Note that with v Ca ② Cb. we have
Qr (In 1,, (1, @ I (e. & &c,) (1
1.) (eb 1) c 2 1n;
(68)
Q 0 = (1, & In,) (T, 8 1 (ea 8 e.) = (1, @ Ix) (ea 01)= (1, (1) ea) = 1, e..
Lemma G.27, For any one-hot vector v = -ea eb RR=
(Q) Q T diag (Q T v) diag (Q v) Q = (1.,,] @ Tn-
Proof. By Lemma G.26, Q v 1, 2 e. Therefore diag (QT v) -diag (1,) @ diag (ea) = In @ diag (ew). Thus
diag (QTv) Q (In @ diag (e.)) (In & 1m) [1 Ø 1 (Ln 8 en) (1, 1 (1,))
(69)
60

---

On the Power of Context-Enhaneee Learning in LLMs
and therefore have
Q diag (QTv) diag (Q'v)
(In @ I.) (4 6 e (1, 8 01,)
(In @ 4.) (4, 8 1) (.)
(since e ea =
(In 69 [.) (1 I.) (1, x
(70)
= (1, @ (,1 8 (,))
(1,,1,) @ In
Lemma G.28. For any one-hot vector Ea Eb I
Qdiag (Qv) diag (Qv) Q I, @ (1
Proof. This proof is very similar to the proof for Lemma G.27. By Lemma G.26, Q = 1n @ e.. Therefore diag (Qv)
diag (et) & diag (1n) diag (ex) @ In- Thus
diag (Qu) Q" (diag (ex) @ [,) (I, 6 I.) (T, & 1 (es 60 I.) (7, @ 11.)
(71)
and therefore we have
Qdiag (Qv) diag (Qv) Q
(In & 1n) (e, &1) (es (T) 0 1
(Ln @ 1,) (1 (4, 11.)
(since e, eb = 1)
(In @ 1,) (4, 8 1) 1
(72)
(In 8 1,)(4, 01)))
In (1,11.)
Lemma G.29 (Tail Bound for Coupon Collector Problem (Motwani, 1995)). For a set of size n, with probability at least 1 - one can cover all unique elements of S in n nog " independent uniformly random sampling trials from S.
G.5. Learning II in MLT with Gradient Descent (Empirical Evidence)
In this section we provide more details on empirically optimizing the simple surrogate model SURR-MLT (w; which was only briefly discussed in the main text by the end of Section 5.1. We will first introduce the approximations made to the surrogate model to make gradient-based optimization easy and stable, then will present empirical results on the model learning target sets of phrasebooks ML. in MLT(5,, 10), MLT(10, 10), and even MLT(20, 10).
G. APPROXIMATED LATENT MODEL FOR GD
Recall that with input sequence represented by V1 E R=²xI, the surrogate model for a depth-d translation is being recursively defined by the translation + shifting operations
V411 = HardMax(C, + W.)
(73)
until we reach V4+1- While this model captures the essence of transition from ICL capability to memorization specific set phrasebooks, HardMax is making it not directly differentiable and hard to optimize. To address this issue, we approximate with an column-wer wise softmax function with very low temperature (T' 1/25). The recursive definition in the approximated model is then
V. S25(Max + Wi)) Shift(V.).
(74)
61

---

On the Power of Context-Enhanced Learning in LLMs
We denote the recursive surrogate model with the softmax substitution as SURR-MLT (W.).. (C1, Cz...., Cd. Vi) where
Va+1 = SURR-MLT (W. (C1, C2, ......, Cd, Vi)
A ScftMax(25C. + 25Wa)
(75)
Shift (SoftMax(25Ca11 + 25W 25Wa-i)Shift Shift (..SSoftMax(55C) + 25 (.)
We define the objective function as the column-wise cross-entropy loss between the final output and the input. Namely for input V1 with prediction Va+1 and ground truth label Vd-1+ the loss is computed
CrossEntropy (V( Vali))
(76)
We follow the same masking (dropout) curriculum as described in Appendix G.2 and Appendix G.3, that at each step we zero-out single column from a single context matrix C. We experiment on two gradient update schemes:
Layer-wise Training: at each step, we are masking a column on C. we only compute the gradient with respect to W; and update This training is more akin to the theoretical analysis described in Appendix G.3.
Full Parameter Training: at any step, compute the gradient with respect to each of the weight matrices and update all parameters. This is more akin to the real gradient-based training do not have the heuristics for localized update.
To allow for fast and stable training, we adopt a very large learning rate of 100 and apply parameter clipping between [0, 1] after each update. The complete algorithm is described as follows:
Algorithm Layerwise Gradient Descent with Conaext-Enhanced Learning For Optimizing SURR-MLT
Input:
2: input Vi E x label Vi+ R descriptive Wi...., W E = learning rate 77, total steps T
3:
4: Initialize W...., Wa <- 0
# Start with zero initialization
5: for = 1 to 7 do
6: <- [(1 - 1)/22 %d +
# Get the layer to be masked Get the column index to be masked Create masked context matrix W Vi)
7: ++ ((t 1)%?² +1
8: Initialize CHuy W (I,2 diag((k))
9: Ve+1 + BURR-MLI [W.) (Wi .... W Ci()), W
10: L+ CrossEntropy V441, Vār1)
11: W. W. w.
Update the weight for the layer with mask
12: end for
Return W W....., Wa.
62

---

On the Power of Context-Enhancdd Learning in LLMs
Algorithm 5 Full Parameter Gradient Descent with Context-Enhanced Learning For Optimizing SURR-MLT
1: Input:
2: input Vi E RRn'xL, label V E IR" descriptive text W
W E x="" learning rate y, total steps T
3:
4: Initialize W...., Wa <
# Start with zero initialization
5: for t = 1 to T do
6: i- (t 1)/n2 %d +1
Get the layer to be masked
7: k: ((t- 1)% +1
# Get the column index to be masked
8: Initialize ((R) W (I, diag(Ex))
# Create masked context matrix
9: V4+1 + SURR- MLT [W:) (Wi ..., W Caay, W
W V/)
10: 4 CrossEntropy( (V V4+1, 1)+
for 1 to d do
W. W W.
# Update the weight for all layers
13: end for
14: end for
15: Return WT,... Wa.
63

---

On the Power of
Learning in LLMs
Layer-wise GD
Full Param GD
1.0
1.0
0.8
0.8
W1
0.6
0.6
W2
W3
0.4
0.4
Wa
Ws
0.2
0.0
0.0
2000
4000
2000
4000
Steps
Steps
(a) (L, 10)
Layer-wise GD
Full Param GD
W,
1.0
1.0
W2
0.8
0.8
W3
W4
0.6
0.6
Ws
W6
0.4
0.4
W1
0.2
0.2
Ws
W9
0.0
0.0
W10
2500
5000 Steps
7500
10000
2500
5000
7500
10000
Steps
(b) (10 10)
Layer-wise GD
1.0
W1
0.8
W2
Wa
0.6
Wa
0.4
Ws
Ws
0.2
W7
0.0
Wa
Wg
2500
5000
7500
10000 Steps
12500
15000
17500
20000
W10
W11
Full Param GD
W12
1.0
W13
W14
0.8
W1s
0.6
W16
W17
0.4
W1a
Wi
0.2
Wzo
0.0
0
2500
5000
7500
10000
12500
15000
17500 20000
Steps
(c) MLT 10)
Figure 13. We perform Layer-wise and Full Parameter Training with Gradient Descent (defined in Appendix G.5.I) on the SURR-MLT (Definition G.7) designed for MLT in MLT(5,, 10), MLT(10,, 10, 10), M20 10) respectively (alternately, depth d = 5, 10, 20 re- spectively, while number of characters fixed ters, which after Ha HardMax application (HardMa rdMax (W,)) align with the corresponding stochastic matrices of the phraschooks (Matrix x (T)) [d - We observe that under both algorithms, the trainable parameters quickly learn the relevant stochastic matrices.
10). Here, we report the portion of columns from the trainable parame-
64

---

On the Power of Context-Enhancee Learning in LLMs
HardMax(W,)
Ground Truth Phu
20
20
20
0.6
60
80
80
80
0.0
5000
10000
25
50
75
Steps
0.25
0.50
0.75
1.00
0.00
0.25
0.50
1.00
0.00
0:25
0.50
0.75
1.00
W3
HardMax(Wg)
Ground Truth prù
1.0
HordMaxwW)))
40
40
40
0.4
60
60
60
5000
10000
25
50
75
25
50
25
50
75
Steps
0.00
0.25
0.50
0.75
1.00
0.75
1.00
0.00
0.25
0.50
0.75
1.00
W/3
HardMax(Ws)
Ground Truth pri
HantMoxiW)) A4c5
0.0
20
20
40
40
0.2
5000
10000
50
75
Steps
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
HardMax(W,)
Ground Truth pii
0
20
20
HaneMoxtW.)
0.6
60
60
0.2
80
80
80
0.0
5000
10000
25
50
75
25
50
75
Steps
0.00
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Wg
HardMax(Wa)
Ground Truth phin
1.0
20
FardMoriik:
40
40
40
0.4
60
60
60
5000
10000
25
50
75
25
50
25
50
75
Steps
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
0.00
0.25
0.50
0.75
1.00
Figure 14. Detailed analysis on the Full parameter training behavior of the trainable parameters in SURR-MLT from Figure 13 for ML.T MLT(10, 10) (i.e. depth d = and number of characters n = 10). We report the behavior of odd-index parameters Wi, Wa... We. (left to right) first, we show the number of columns of the trainable parameter, which after Ha HardMax (W;) align with the corresponding columns of Mat matrix (m²) .. Second, third and fourth visualize the matrices, Wi, HardMax (W;), and Matrix (x))) respectively. All matrices learn match Mat r 1x (x7) the end of training with HardMax operation.
65

---

On the Power of Context-Enhanced Learning in LLMs
H. Construction of a Transformer that can Simulate the Latent Model
H.1. Useful definitions and lemmas
Definition H.1 (Relative self-attention with 1 head). For a set of matrices { quuery Wkey W Wvatue with each matrix Rick for some > and set of (t + 1) biases {bilesigo. the self-attention computation on an input sequence 11 , a with each p=p2 IR* is given by the output sequence "/.".", y where for all p1 [1,L]
E01 kpz t0p-
yp
ap Oj, where apip
if p2 S P1, otherwise
Epd Spi E R
tb v2'
4p2 W Tur
kp = W Wkey Op2 W Valueepps
for all p2 E [1, L].
For relative self-attention with H heads, we will simply add the output of the H heads as the final output.
Definition H.2 (MIP). For of matrices Wouter E R W Winner R for some k, H 0 and activation function o, the output of the MLP layer an input sequence T1," Ey with each a; E R is given by the output sequence y/," yr. where for all i
yi W outer o W Wi)er
Lemma H.3. For GELU (Hendrycks & Gimpel, 2016) activation function, which takes E R as input and returns x (x) as output, with ((x) representing the standard Gaussian cumulative distribution function, for any two variables x. y E R, the following holds true:
2 + +) GELU(x) GELU(y)) -x 0(x2))))
The above lemma has been taken from Akyurek et (2023).
H.2. Transformer construction
Recall that a translation task in MLT ni) involves two primary operations each step: Circular shift and Translate. will refer to the surrogate model to use notations for different operations in the translation task. Recall from Equation (18), the surrogate model for MLT(d, n). denoted by SURR-MLT (w:)._, (+) with trainable parameters (W;) = (W..... Wa).
can be represented by the following recursive expression
V HardMax(C; + W;) Shift = HardMax W.).
(77)
for 1 i Here Shift represents Circular shift operation, and HardMax(C, + Wi) represents Translate operation, where the operation can be done either using the relevant in-context information C2 the in-weights memory parameters W. when in-context information isn't provided in form of C. HardMax is the column-wise hard-max function converting C, + W; to binary column stochastic matrix.
Lemma H.4. For the family of translation tasks MLT(d, (d,, n),, there exists transformer model with embedding size 2n² + 2d 4, 2d relative self-attention layers (containing either or 3 heads), and 2d MLP layers that can simulate the surrogate model, SURR-MLT (W:).=. (.)-
For each input sequence $1 of length L, the input sequence of embeddings to the transformer will be of length ndd + L/2 + d, where the last d embeddings are padding tokens (<P>) and given as Os, and the length of output sequence of the transformer model will be nd²d + + L/2, where the last L/2 output embeddings will be used for loss computation. The middle embeddings in output are represented by <THINK> tokens."
"In our experiments, we used output sequence length rd + d dL/2, where (d + 1)2 tokens the output sequence were represented by <THINK> tokens. Instead, we only use d <THINK> tokens in our output construction for simplicity, the proof can be easily modified to align with the experiments.
66

---

On the Power of Conhexe-Enhanced Learning in LLMs
Outline of the construction: Our argument will be for any general MLT(d, n). To create the transformer, will create similar transformer modules that handle Ha MardMax((; + W;) Shi (t()) for each i. We will refer to the constructed modules as MLT-MODULE and will mention the specific step i as an argument when we attempt to use the module to perform translation step HardMax(( C + Wi) Shift(Vi) W.l.og.,, we will assume we are building MLT-MODULE to perform translation step Ha HardMax C + W.) (V()). MLT-MODULEE will contain two modules, CIRCULAR SHIFT-MODULEE and TRANSLATE-MODULE which simulate Circular shift and Translate operations, and have been outlined in Algorithms 6 and 7.
Next, we explain the structure of the embeddings in the transformer architecture. Our embeddings will be built on the context representation matrices (Ci..... Ca) and the matrix representation V1 from input sequence 81 ($1,1,'* S1,L) and subsequent intermediate representation of the translation task. Furthermore, our embeddings will contain additional information like indices of the context matrices when utilizing them in-context, segment indicators that represent whether embeddings represent context matrices, or the input sequence tokens, and start and end indicators that indicate the start and the end embeddings representing the input sequence. This information can be extracted from the input sequence using a few input processing layers, though we do not delve into the specifics.
H.3. Structure of input embeddings to MLT-MODULE
Our embeddings will be split into 4 components: token, context matrix indicator, start and end indicator, and segment indicator. To maintain simplicity in our discussion, we will present them as 4 separate embeddings.
For module that will represent Ha HardMax(C, W,) SVi, assume the input will be provided as 2 major segments.
1. First segment: in-context matrices We will give the in-context information (Cp, 1d as follows.: each context matrix Cp, will be fed as n² token embeddings, {ley; Cpt ey] E R2n")² زグ=ı dimensional vector that contains 1 in position j and (ey; Cp. e,] represents a concatenation of ej and Czey. Thus, the in-context information will look as follows
" where ej represents a one-hot n²
(Le;; Cielly " [lej; Cae;l)
Additional embedding: context matrix level indicator In order to differentiate the different context matrices, we will use an additional d dimensions in the embeddings to represent a one-hot vector that indicates the index of the corresponding step they will be used for. For simplicity, we will represent these dimensions separately as separate embeddings: 1p Rd. which are one-hot vectors that contain in position P1 and otherwise.
2. Second segment: input query For a length- input sequence s; (s1,1 S1,2...... St,L), we will use the sequence of columns of its matrix representation V, E g=²xL/2, appended by Os to match embedding sizes, as token embeddings for the input query. A padding embedding <P>, containing Os, follows this sequence as an end of sequence embedding.
Additional embedding: start and end indicator embedding We will have 2 additional dimensions representing whether a token represents the start or the end of the input query sequence (first or second dimension activated respectively). Start of the input query sequence is determined by the first embedding in the input query, while end of the input query sequence is determined by the first padding embedding containing Os after the input query sequence. We will represent these dimensions separately as separate embeddings: (b1, bb, 0 E R2) heLL, which are one-hot vectors. by contains 1 in dimension 1 if the embeddings represents start of the sequence, by contains 1 in dimension 2 if the embeddings represent end of the sequence. Other embeddings have Os.
segment embedding We will differentiate the input embeddings in the two segments using 2 dimensions that represent one-hot vectors indicating segment indices. We will represent these dimensions separately as separate embeddings: 91, 92 E I2², where both are one-hot vectors, with gi containing 1 in dimension 1 and g2 containing 1 in dimension 2.
Additional embedding:
All the notations have been summarized in Table 8.
67

---

On the Power of Context-Enhanced Learning in LLMs
Additional optional imputs: There might be additional input embeddings, represented as <THINK> in the first segment, which are null inputs and ignored during self-attention computation. As discussed next, will right shift the sequence by 1 at each step, in order to handle Circular shift operation with causal masking in transformers.
H.4. Structure of the output embeddings from MLT-MODULE
All the embeddings in the first segment are kept intact. In the second segment, the module outputs null output <THINK>> followed by L/2 output embeddings that represent the columns of V,. We will require <THINK> to represent Circular shift with causal self-attention in transformers. <THINK> will be ignored in self-attention computation and so, we will ignore their discussion for simplicity of presentation. Other embedding values are kept intact for the generated output, except start and end indicator embeddings b1, b2, which need to be right shifted at each step. The right shift operation can be handled similar to our computations on the token embeddings and so, we ignore them in our construction below. We summarize these in Table 9.
Embedding Name
Dimension size
First segment values (In-context information)
Second segment values (Input sequence embeddings)
Token
2n²
fle;; Gee, Ale;; Caellip=
[V_-] 0
(L/2),
00. <P>
Context matrix index indicator Start End indicator Segment
d
(4)) p........ (La)p.a) 0,",,
2
by 0, b2
2
91. .1
92- .2
Table Input embeddings to MLT-MODULE that simulates HardMax + W.) Shi Shift((V) ej indicates a one-hot n² dimensional vector that contains in dimension j.
Embedding Name
Dimension size
First segment values (In-context information)
Second segment values
(Input sequence embeddings)
Token
2n²
[le;: -1 Ca
<THINK>, (V
Context matrix index indicator Start and End indicator
d
(h 12 1 ... (lahi.
0,0,
2
0, b 0,
0, b
Segment
2
. 91
0, g2, 92
Table Output of the transformer module MI.T-MODULE that simulates rrdaa((C Wi) Shi (t. <THINK represents null output and won't be attended to in the future modules. We ignore this symbol for simplicity, when analyzing any module. ej indicates a methottre n² dimensional vector that contains 1 dimension j.
Constructing the M.T-MODULE: The MLTTMODULLE consists of 2 self-attention layers and 2 MLP layers. We use one self-attention layer and an MLP layer to represent Circular shift operation, one self-attention layer to represent C, Shifft(V., and one MLP layer to represent HardMax(C + W,) (V;)tt(VV....... We name the two modules for Circular shift and Translate as CIRCULAR SHIFT-MODULE and TRANSIATT-- respectively. We have outlined their constructions in Algorithms 6 and 7.
H.S. Step (CIRCULAR SHIFT-MODULE): Represent Circular shift using a self-attention and an MLP layer
As Circular shift only focuses on the input query sequence and not the in-context matrices, we will simply focus the module's operation on embeddings in the second segment. The effect of the operation on embeddings in the first segment can be removed using gated residual connection. From Lemma G.3, we have that the output of the Circular shift operation on any sequence, represented by its matrix representation Vi, can be written as
V. QV oQ VTVVV(ia ((++1))) for all SjSLPP L/2.
where Q = (In & 1,, (1, 8 T 1n E Rax1 is the all-ones vector, and O is the Hadamard product.
In order represent the operation, we will first use self-attention layer to compute (1, 2 In) v and
(In & 1, 1() at each column j. Because the computation of ((0) requires the model to look forward to f (1) we need shift the computation of (1)) to position j + 1, as a causal attention mask is involved in self-attention computation.
68

---

On the Power of Context-Enhanced Learning in LLMs
After right shift operation, we will represent the output of the self-attention computation as <THINK>, (02; 0] (oy; .)..... .., 2 0), and we will ignore the <THINK> embedding. Note that the second half of the out- put embeddings will still contain Os and we will ignore them in the current computation. Then, the above computation can be rephrased as
o - -V-VVV111 ② V for all 2 S S L/2.
01/2+1 = QV, @ V.VV
Self-attention layer: The computation of oj for 2 S j L/2, requires the computation of (1,, @ I,,) v(--)) and
(In @ 1,,) v,6), This will require 2 attention heads, one head that attends to itself, and another that attends to previous embedding at each position. We will outline both below. We will require one additional head, as computing 01/2+1 will require the model to compute (In & 1n) V, V.
1. Attention Head 1 computes (1,, @ In) V ( at position for all 2 S L/2+1 This can be done using self-attentin head (Definition H.1) that sets query and key matrices W query. W key. and biases (bilesioo such that the attention score between embeddings at any two positions Pl, P2 is given as follows:
apr.p= p. P1 -1,
0 otherwise
W value is set such that for any input æ, the output of W Vellue is given by
(W value 2) pp
In-ttp
j=0
In simple words, this operation simply adds up the values in dimensions Pr.P1 + n, 2n, and stores them at position P1 for all 1 < p1 S n.
2. Attention Head 2 computes (4n @ 1n) v. at position j for all 2 j < L/2. This can be done using a self-attention head (Definition H.1) that sets W Wenery, Wkey Wintue and biases {4i such that the attention score between embeddings at any two positions P1, P2 is given as follows:
ap 1. Pı 0, 0 otherwise
Watue is set such that for any input x, the output of W Wontue is given by
W watuee 2 Pi
I P-P--
In simple words, this operation simply adds up the values in dimensions + (p1 - 11m,, 2 (p1 - 1)m..... and stores them at dimension p1 for S pr S n.
3. Attention head will compute (I, 1,,) v(() and store in 01/2+1. This can be done by a self-attention layer which activates only between positions Pi and P2 that represent the start and the end tokens of the sequence, i.e. contain b and b2 start and end indicator embeddings, and is 0 otherwise. Waatue is set same attention head 2.
The output of the three heads are simply added up. Hence, at each position 2 < S L/2 + 1, the output 0, has (I, @ I.)" (-(-1) in [1, n] dimensions and (T, @ 1,,) V. (1%L) in [n + 1, 2m) dimensions.
MIP layer: The objective with the MIP layer (Definition H.2) will be to multiply (1,, 6 T.)) dimensions and (I, & 1,,) v. (1%L) present in [n + 1, 2n] dimensions in each position j. This can be done by using an MLP layer with GELU activation by using Lemma H.3. The weights of the MLP layer are set as follows: W Winner is set such that
V present in [1, m] 1
69

---

On the Power of Enhaaeee Learning in LLMs
for all input æ, we have
1
( i:) N Bi
for all 1 1 n,
N
1
(W Winnera)it
N di:
for all 1 S S n.
W snner +)22
anti:
for all 1 s n.
N
W outer is set such that for all 3m
W W) = N2; Zitn zi42n), for all S S n.
All other coordinates in these matrices are set Os. N is set large number (say 100). By Lemma H.3, the output of the MLP layer will be <THINK>,02 02, 01/241, with o, containing
-レー +OON = (1n 2 In) vi--1) o OI (L.31,)
v(0% v, +O(N
at each position 2 L/2.
Embedding Name
Dimension size
First segment values (In-context information)
Second segment values
(Input query sequence embeddings)
Token
2n²
(le;; Gee,ly" 1.
Ale;; Ca
<THINK>, [V 0].
(4/2),
Context matrix index indicator Start and End indicator
d
(1)), (72 p... (la)t..
0, 0.
2
0.
0, b, 0, ... ,0, b
Segment
2
91 91
0, g2. ............. .
Table 10. Output of CIRCULAR SHIFT-MODLLE in MLT-MODULE that simulates Circular shift, i.e. computes Shi (V;) at second segment token embeddings. <THINK> represents a null output and won't 't attended to in the future modules. We ignore this symbol for simplicity, when analyzing any module. e, indicates one-hot n dimensional vector that contains in dimension j.
H.6. Step (TRANSLATE-MODDULE): Translate as a module containing a self-attention and an MLP layer
Our current token embeddings are given <THINK>, Joz; 0]. (o;; Oj.... /or/2-1: 0J, where each o, contain v-Ú-11 Other embeddings have been kept intact. The in-context information are given as {[eş; Cie,lly {le; Cae We will first use a self-attention layer to compute [C, (d-1), V (111) at position 2 < S L/2 + 1. We then use an MLP layer to represent HardMax(Ci + (.)))
Self-attention layer to represent (C, U 1) V. ( ) We will use two attention heads.
1. The first attention head computes C, (1-)); Matrices W Wyuery, W key are set such that the attention score between an token embedding in first segment [e,: Cte,] and a token embedding second segment oy is given by
(er, (.0-1)), if = and 0 otherwise.
for any j E [2, L/2 information. The condition can be set using the Context matrix index indicator vectors l which is present in each in-context information embedding.
+ 1],r E d]. The condition requires the model to attend to C, and ignore other in-context
The attention between any two input sequence embedding and o,r is computed as Os. The distinction between the attention scores of pairs of embeddings in second segment, oy and ojr, v/s attention scores between a token embedding second segment and token embedding in first segment, oj and [e,; C.ez], can be done by using the segment indicator embeddings gı and ga used to differentiate token embeddings in first segment and the input sequence embedding vectors.
Matrix W valuue set such that the columns of each Cys are picked from the token embeddings in the first segment: (fle; nn id
70

---

On the Power of Context-Enhanced Learning in LLMs
The second attention head simply copies the input 1701"); This can be done with an attention head that attends to itself at each position j and copies V. (ij-1) to output.
2.
The output of the two attention heads are simply added up. The output embeddings will now look as follows: <THINK> [CCVV (-1), V0-DD L/2
1=
MLP to represent HardMax C + WoV Our current token embeddings at any position j contain both C, V,V((-1) and V ( The first layer of MLP can be used to compute (C; + W;) (.(-1) by setting the weights of the layer using W., We simulate HardMax operation follows:
/ |65112: where o; = GELU(((C; W.)
The & normalization is equivalent to R.M.norm operation (Zhang & Sennrich, 2019). This is an approximation of the HardMax function, which are equivalent only under the following conditions: for each column j
1. either (()) or w
()) are all Os.
2. cl and W are both hot vectors and they match at the corresponding activated dimension.
71

---

On the Power of Context-Enhanced Learning in LLMs
Algorithm 6 CIRCULAR SHIFT-MODULE: Self-atiention and MLP layers for Circular shift
Require: Input embeddings (Token, Context matrix index indicator, Start and End Indicator, and Segment embeddingg split into 2 segments) (Table 8). Important (for the current module)
1. Token embeddings: First segment contains ([ey; Cie,,)s Ale; Cae;l)i= and the second segment contains [Shift [Shi .. [
2. Start and End indicator: First segment contains all Os and second segments contains bj and by at first and last embedding, while containing all Os everywhere else.
Step a: Using a self-attention layer with 3 attention heads, change the token embeddings in second segment as <THINK>, ]o 2 s.t. Oj has (1, @ T,) v-5-1) in [1, m] dimensions and (Tn & 1,,) V. (1%5)) in nn 1, 2n] dimensions. Primarily,
iff pa Pr =- and otherwise. Value matrix Wattue is set such that for any input æ, the output of walue is given by (for all p1 E n1)
Attention head I: Computes attention score between any two positions pt, P2 as aps
n-1
W Vellue
Σ
1-0
Attention head 2: Computes attention between any two positions Pr. P2 ap dr = 1 ift p2 Pr =0 and 0 otherwise. Wedlue set such that for any imput æ. output W Valuee is given by (for all pa E n1)
( Woalucw:pptt
Σ HttDinnn
Attention head Computes attention between any positions P1; P2 as 0p1.P2 1 iff by and are present positions Pı and P2 respectively and otherwise. Wetuue is set such that for any input a. the output W Wvatue is given by (for all E m))
W Walluell piin
Kjtpin-n-
Sum the output of the three heads.
Step b: Use ML.P layer to change token embeddings in second segment as <THINX>, small error. Primary computation at each position j is given as (for a large N)
[0: 1//41 s.t. og has v.!11) with some
22/mN² (GELU 1 (In & In)
v((-1)
(Ln In) (((%))
N
GELU N (1, 2 In) v(リ-1) GELU 1 N (In 2 1.)) ((.81I).
which will return
マリ ワリー O(N : (1, @ Jn) Tv ") (7, (1.11) (1, +ONN4 vil-"
Return the output embeddings (as given in Table 10).
72

---

On the Power of Context-Enhanced Learning in LLMs
Algorithm 7 "TANSSLATE-MODULE Self-attention and MLP layers for Translate
Require: We will require an index, and input embeddings as input:
Index (indicating index of the M..TMODULLE it is part of),
Embeddings (Token, Context matrix index indicator, Start and End Indicator, and Segment embeddings split into 2 segments) from the output of its preceding CIRCULAR SHIFT-MODULE (Table 10). Important ones (for the current module) are
1. Token embeddings: First segment contains {[e;; Cielli "M Ale5; Cae,lli= and the second segment contains <THINK>, [V ; ;0....,,IV [V. 0)
Context matrix index indicator: First segment contains (41),,0)) (12) ... (La)p, and second segments contains all Os vectors.
Step a: Using self-attention layer, change token embeddings in the second segment as <THINK>, [[C.V 100-"1) 2/2
Primarily, the self eeff-attention score between embeddings that contain a second segment token embedding (V.(=)); 0] and first segment token embedding [ep1 Cpep1] (for any E d). pi E [L,n"].pr [1, LJ) computed
(Ep1 V((2)) (4,4.):
where I, is the corresponding context matrix index indicator for the first segment embedding under consideration and 1, is constructed using the index i.
Crep, used value vector from each first segment embeddings.
Step b: Using an MLP layer, change token embeddings in the second segment to contain <THINK>, {05]
L/2+1 where
0j oj/ 1|, where o; GELU((C, W.)V(--1)
Return the output embeddings (as given in Table 9).
73

---

On the Power of Context-Enhaneed Learning in LLMs
H.7. Trainability of the constructed transformer
How does our constructed transformer perform on the MLT task? To do so, we hand-construct the designed transformer and train the transforer model random translation task in MLT(5, 10), that has depth 5 and number of characters 10 in each translation level. We train only parameters (Wi Wii= using Layer-wiss wise and Full Parameter optimization algorithms from Algorithms 4 and 5 respectively.
However, we make two changes. First, we use Adam optimizer instead of SGD, as we found SGD to get stuck frequently at bad minimas. Next, instead of masking in-context representation of layers in rotating curriculum fashion, where each layer i's in-context representation were masked in intervals [j (i 1)n²:J j in for all integers j 2 0, we mix in the masking of all layers together by randomly picking layer and mask random column from C. We found that mixing in masking of all layers helps train the model faster.
We vary the peak learning rate as {10-2, (10-3,10-3, 1,0-3,10-4 We use cosine decay learning rate schedule with warmup steps 100 and minimum learning rate at end of training as 0.1x the peak learning rate. We train with cross entropy loss, and use total of 200 training sequences length 20 during the course of training. We use a batch size 32 per gradient update step. We also use a small & regularization, with strength 10-4, on the rows of the parameters to help optimization.
Observations: We report the performance of model trained with both Layer-wise and Full Parameter optimization algorithm with peak learning rate 10-3 and & regularization with strength 10-4 in Figure 15. We observe that under both algorithms, the trainable parameters quickly learn the relevant matrices that represent the true set of phrasebooks to minimize the training loss.
Layer-wise Adam
Fullparam Adam
1.0
W1
0.5
W3
W4
0.0
0
500
1000
1500
0
500
1000
1500
Steps
Steps
Figure 15. We perform Layer-wise and Full Parameter Training with Gradient Descent (defined in Appendix G.S. 1) the handconstructed transformer (Lemmaa H.4) that represents SURR-MLT (Definition G.7) designed for MLT in MLT 10) Here, we report the portion of columns from the trainable parameters, which after HardMax application {Ha (Har (W.) align with the corresponding stochastic matrices of the phrasehooks {Matrix (m²)) (m]))==1 We observe that under both algorithms, the trainable parameters quickly learn the relevant stochastic matrices.
74

---

On the Power of Context-Enhanced Learning in LLMs
I. Additional Experiment Setups
I.1. Format of Training Text
Here we present examples training data used for experiments Section 3
INFUT:
<</begin_of.text/i <b blbeinnof.text <istart.hoader_ddl user <|end_hoader_dd]>
You are performing a spectal translation ion cask called language.taak.2 The subset of dictionaries used are as
follows:
Dictionary used from lznguage to language
I
T2 B J 3; EA > J A J, > BI F P >
Dictionary used from lznguage to language 3:
F -> x;
RR; + x 8 R -> R M. ~ 1 31 T 1 - V.
> 0 9:
Is x W, 0 Q0UN 330 20
8
I
2 21
1
JL 0 Ri - R R
NX
MI 0 P
K
Dictionary used from language to language
f;; x
5
2 2 VK
R 17
4
0 >> El W
8 -> 0
Dictionary used from language 4 to language 5:
a >> m 31 - Sanger -> mm c mi 2d hr
>>
1 m²
Y -> 11 b 11-11
e
Dictionary used from language to language
-> 15 9 "...
i h 5 1m
g > v h. u 8 1k > o h > t 1 11 > k > 1 u
New please perform language.task. translation from the following sequence in language to language 6. Do not use code! You must only reponse in the form: "Sequence in language 1: (the sequence in
language 1]: Sequenca language [the sequance languaga 2]; Sequance in language 3: [the sequence language 31; Sequence in language
[the sequence language 4]; Sequence language 5: [the sequence in language 5]; Sequence in language 6: [the sequence language El", The sequence need to translate from language 1 is: CREFEBDE B.R.B.CAACEEEEEEEE.EE.EEEECCECE. C PAACAAEPSADPOODEEEE
E.<lest.id] wtw..ehoooe.....aaaaaaaaaat... andhheeerii
LABEL:
<<bbginnoffttxtt I>Tha translation result Saquence language
CAE
G Sequence language
Sequerce language 3: Sequence in Language d b
S
N
Z
[ Y
b
Y [ 2
šequence in language 5: 11kaka n h
1
h
h
h
h à
:; Sequence in language 6: govttu
S
P
a t Pj </ 615"
Figure 16. Input with complete context and explicit chain-of-thought. We use this data format at the begining of stage 1 training.
75

---

65 On the Power of Context-Enhanced Learning in LLMs
INPUT:
b<lgnn of text |> <| blggin of text <| begin of text |> stattrt header idi> user <|end header id|>
are perferming special tronslation cailed Language.tosk.2.
dictianarias
follows:
Dictionary used from language to language 2: D J 1
D A -> N di
o c ~> R
E M B
C -> K 21
M K 1
P7 D SIDE-S 8
H a 3y
X d
L G.G. K2
Dictionary used from language 2 Lo language 3:
x
62 8
V7
M
Dictionary used from language 3 to language
U -> Y d; R 2 Y; >> 2 2
1 58 8 8 b. b.
5 s +> c v 10 2 ag -> 8,
- + 1
Q Ei R
VR a X
0
2
Y
Dictionary used from language to language 5:
a ii
-> b
31
0 1 11 Y
Yb
1:
Dictionary used from language 5 to language 6:
kg > as 1. > 5E j > 200 i 9 1 > > 1 2 1 > karke km > 9 > v wtt.pppmmmmmmmmerccc............
-> v 1 u -25 u 0 t - k -r 3.
please perform Language.task. translation only reponse in the form: "Sequence in language 1: (the sequence language 11; Sequence language 2: [the sequence in language 2]; Sequence in language 3: [the sequence in longuage 3]; Sequence in language 4: [the sequence in language 4); sequance in language 5: [the sequence in language 5]; Sequence in language 6: [the sequence in laaguage 61", The sequence
following sequence language language
6.
code!
you need to translate from language 1 is: CEEFESBOCCCAA
start http::/:ss..ttccceee..eeaaeeeee. ind
LABEL:
*<lbegin of text |She translation result is: Sequence in language 1: C
CCAD
F.G.B. C. 8y Saquence in language 2: < <T>
eD
<T>
<D> <
<T <T> <T>
<>
<T> <1> <>> <D>
<D> <t> <D>
KEM Sequence in
language 3: <T> <5> <T>
<TD <T>
<T>
<TD <TD <T <
<3 635
42 Sequence in language 4:
Sequence in language
< <T9 Sequence in language
idd
Figure 17. Input with complete context and internalized thain-of-thought. We this data format by the end of stage training well the base format for stage training (before dropout)
76

---

On the Power of Context-Enhanced Learning in LLMs
INPUT:
"<lgnnin of text |> <| begin of text <| begin text |> start header idi> user <|end header idj>
are perferming special tronslation cailed Language-tesk.2.
dictianarias
follows:
Dictionary used from lenguage to language 2:
Dictionary used from language 2 to language 3:
Dictionary used Erom language 3 to language 4:
bi at lonary used from language to language 5:
Dietionary used from language te language 6:7
code! You only reponse the following sequence language to language "Sequence in language (the sequence language Sequence language [the sequanee in languaga 2); Sequance in language 3: [the sequence in language 31; Sequence in language 4: [the sequence in language 4]; Sequence in language 5: [the sequence in lenguage 5]/ Sequence in language 6: [the sequence in language 6l", The sequence you need to translata from language 1 is: CREFER
please perform language.tass 2 translation
<</begin.of.toxt/ (>iha translation result Saquence in language
CBE
F B B, Sequence in language 2: <t <T> e et c dt <T> d
<T
< <T>
<2> <T> <ss <>> <
<T>
<T> d
<
d d d
Sequence in
language 3: <T> 619 10 do 68 d
To CT
63 C < <
<> < <g> 40> 48>
<1> 62 48>
<1>1 Sequence in language 4: 4
<8> 62
Sequence language
< Sequence in lenguage 6: 100111100000011111111111
Figure 18. Input with completely masked context and internalized chain-of-thougee We this data format evaluate the model's capability on conducting translation without anything (useful information) in context.
77


Reviews:
[{"review_id": "im2bHCYGod", "reviewer": "Reviewer_6pqm", "summary": "The authors study the impact of augmenting context with additional data for learning in LLMs. Notably, on the added data no autoregressive gradients are computed. The authors consider a stylized problem of multi-layer text translation over finite alphabet where text from one language is translated to another in $d$ steps, where each step uses a phrasebook and is invertible.  The paper differentiates the learning difficulty between when all the phrasebooks are augmented in the context vs absence of such context enhancement. Notably, these phrasebooks are not provided at test time. Therefore, explicit CoT is avoided and the model operates in silent/internalized CoT mode adding it's own CoT before producing the results. \n\nThe authors tune the Llama 3.2B model with to expose the LLM to the specific translation task using SFT with random phrasebook. Next, they train the model using the context-enhanced learning technique with Dropout and Annealing to force the model to not memorize. The experiment show without dropout or annealing the test time performance suffers. This is augmented with interesting mechanistic observation of the learned model.\n\nThe authors next establish exponential sample complexity improvement with context-enhanced learning using a stylized learning model. They show without context enhancement the learning task has a $n^{\\Omega(d)}$ with $d$-layer translation over $n$-sized alphabet (under statistical query model). However, with context enhancement a heuristics search algorithm exists with $O(d n^6 )$  sample complexity, and for $d=2$ a gradient based algorithm exists with sample complexity $O(n^4)$.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your comments and insightful questions!\n\nPlease find our responses to your questions / concerns below.\n\n---\n\n### Generalization in Context-Enhanced Learning\n\n>\"The MLT task does not shed light into the generalization capability of the technique ... the setup only shows quicker memorization (not generalization capabilities).\"\n\n**Response:** Our main message goes beyond quicker memorization of the phrasebooks. We also highlight that:\n\n1. Even when **no next token prediction loss** is not computed on the phrasebook tokens, their simple presence in context can still accelerate learning.\n\n2. 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.\n\nWithin 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).\n\nPlease 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.\n\n---\n\n### What is the Motivation of Having no Context at Test Time\n\n**Response:** Our motivation is twofold: philosophical and theoretical explanation of existing empirical works leveraging context-enhancement.\n\nFirst, 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.\n\nSecond, 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.\n\n---\n\n### How Phrasebook Information are Provided to Surrogate Model\n\n>\"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.\"\n\n**Response:** The surrogate model is motivated by our mechanistic experiments in section 3.\n\nFirst, 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.\n\n>\"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.\"\n\n**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:$\n\\def\\R{\\mathbb{R}}\n\\def\\bm#1{\\begin{bmatrix} #1 \\end{bmatrix}}\n\\def\\p#1{\\left(#1\\right)}\n\\def\\x{\\times}\n$\n\nLet $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 \n\n$\\\\begin{align*}\nX_{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}\\texttt{Shift}\\p{X_i\\bm{0_{n^2d\\times L}\\\\\\\\I_L}}}\\bm{0_{dn^2\\times dn^2} & 0\\\\\\\\0 & I_{L}}\\\\\\\\\n&=[C_1,\\dots, C_d, 0_{n^2\\x L}] + [0,\\dots, 0, \\p{C_i + W_i}\\texttt{Shift}\\p{V_i}]\\\\\\\\\n&=[C_1,\\dots, C_d, V_{i+1}].\n\\\\end{align*}$\n\nThis 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.", "claims_and_evidence": "The authors claim that for multi-layer translation task context-enhanced learning boosts learning of LLMs, even when the context-enhancement is not present at test time. They prove their claims both experimentally, and on a surrogate model for learning.", "methods_and_evaluation_criteria": "This paper studies the utility of context enhancement in learning for LLM. The authors study this on a specific task, namely multi-layer translation. The methods seem logically sound. The mechanistic interpretation of the learning dynamics provide further evidence that context enhancement helps with multi-layer translation.", "theoretical_claims": "The authors claim that context enhancement provides exponential improvement in the learning of multi-layer translation in a surrogate learning model. They also show that the surrogate model can be a simulated using transformers.", "experimental_designs_or_analyses": "The designed experiment is non-standard but seems suited for understanding the effectiveness of context enhancement for multi-layer translation. The loss evolution with various different annealing/dropout method shows the effectiveness of the context enhancement.\n\nThe mechanistic explanations also provide some structural insights into the learned models.", "relation_to_broader_scientific_literature": "This paper provides new insights into the utility of contexts in LLM learning, albeit in a very stylistic setup.", "essential_references_not_discussed": "I have basic, but not in-depth, knowledge of the literature, and I did not find any important references lacking. But I may have missed some references given I lack in-depth knowledge of the literature.", "other_strengths_and_weaknesses": "Strengths\n- The results, theoretical, mechanistic, and experimental, presented in the paper seems insightful in showcasing the effectiveness of context enhancement.\n- The exponential improvement due to context enhancement in the surrogate learning model seems \n\nWeakness/Questions\n- The multi-layer translation task presented does not shed light into the generalization capability of the technique. It seems the model learns the rules of translation through context enhanced learning. But due to the assumed bijection of the translation rules (phrasebooks) the setup only shows quicker memorization (not generalization capabilities).  Does that capture the scope of the work?\n- The authors consider the setting where context is not enhanced at test time. What is the motivation behind? \n- 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. This choice is somewhat unsatisfactory. The transformer model that mimics this surrogate learning also doe snot include the phrasebooks at the input place, so the claim that standard transformer model can replicate the learning seems misleading.", "other_comments_or_suggestions": "See Weakness part.", "questions_for_authors": "N/A", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "jmCBQp8ub4", "reviewer": "Reviewer_evJV", "summary": "This paper investigates the power of context-enhanced learning in large language models. The authors propose a synthetic machine translation task that utilizes phrasebooks to transform initial strings. Experimentally, the paper demonstrates that if the base model is MLT(d, n)-ICL-capable, context-enhanced learning enables more sample-efficient adaptation to MLT_{\\Pi_}-capability. Furthermore, the study establishes that there is no information leakage regarding \\Pi_* in the SFT model. Using a mechanistic interpretability approach, the authors analyze the trained model’s behavior to distinguish between an ICL-capable model and an MLT_{\\Pi_*}-capable model. Theoretically, the paper shows that standard SFT requires an exponential number of samples to learn an MLT_{\\Pi_*}-capable SURR-MLT, whereas context-enhanced learning only requires a polynomial number of samples.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your comments and insightful questions!\n\nOur responses to your questions are as follows:\n\n---\n### What accounts for the OOD generalization to empty context?\n\n**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.\n\n**More details:**\n\n**Mechanism of the OOD Generalization**\n\nWe 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):\n\n1. The sequential structure of translation in the model, in which each step can take information from context or the weights.\n\n2. Phrasebook information at different translation stages are internalized in a localized manner, compensating for the corresponding dropped phrasebook information in the context.\n\nDuring 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.\n\nOur 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.\n\n>“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”\n\n**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.\n\n>”Is the generalization driven by the base model’s capabilities (Llama 3.2-3B), or by its MLT(d, n)-ICL capability?”\n\n**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.\n\n---\n### SQ dimension and Sample Complexity of SGD\n**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.\n\n**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.\n\nHowever, 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.\n\n---\n### Presentation\nThank you for your advice on the presentation! We will add illustrative figures in future versions of the manuscript.", "claims_and_evidence": "The paper provides extensive experimental results, a mechanistic interpretability analysis, and theoretical justifications to support its claims. The presented evidence appears sufficient to substantiate the claims.", "methods_and_evaluation_criteria": "The methodology and evaluation criteria are sound.", "theoretical_claims": "I did not verify the proofs in detail, but they appear sound at first glance.", "experimental_designs_or_analyses": "The experimental results appear robust.", "relation_to_broader_scientific_literature": "The paper addresses training methods for LLMs, a topic that is closely related to the broader scientific literature.", "essential_references_not_discussed": "N/A", "other_strengths_and_weaknesses": "The proposed framework of context-enhanced learning and the design of the synthetic MLT task is novel. The contribution is solid. \n\nThe paper was somewhat challenging to follow on the first read due to the dense mathematical formulations and technical descriptions of the task and methods. The second read was more manageable. Including a figure to illustrate the concept of context-enhanced learning, the MLT-ICL task, and the training/testing methodology for context-enhanced learning would greatly improve readability. However, I acknowledge that space limitations may make this difficult.", "other_comments_or_suggestions": "N/A", "questions_for_authors": "1.\tThere appears to be an intriguing OOD generalization effect in the SFT and testing phases of the downstream MLT_{\\Pi_*} task: the training prompt includes part of the phase codes, whereas the test prompt does not. What accounts for this OOD generalization? Standard empirical risk minimization cannot explain this behavior. Is the generalization driven by the base model’s capabilities (Llama 3.2-3B), or by its MLT(d, n)-ICL capability? In other words, 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?\n2.\tI did not fully understand how the “SQ dimension of MLT(d, n)” implies that “any algorithm attempting to learn an MLT_{\\Pi_*}-capable SURR-MLT requires at least n^d sample complexity.” Could the authors provide further clarification on this reasoning?", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "BDfwupVAUw", "reviewer": "Reviewer_vWg8", "summary": "This paper introduces a new concept called context-enhanced learning for LLMs, where they add context related to the task and training time step t in addition to the training data of x and y, and no loss is applied to the context text. To study how this will help LLM learning, they introduced a task called multi-level translation, which is a bijection task that maps input to output according to phrasebooks. Using this task, they prove in a simplified setting that context-enhanced learning can be exponentially more sample-efficient than standard learning when the model is capable of in-context learning. They also demonstrate mechanistically that the benefit arises from more accurate gradient learning signals and show that it's difficult to detect or recover learning materials used in the context during training.", "questions": "", "limitations": "", "rating": 3, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your constructive comments (especially regarding section 3.3) and questions!\n\nOur responses to your questions and concerns are as follows:\n\n---\n\n### Definition of Necessary Rules and Unused Rules (line 165)\n\nFor 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.\n\n*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).\n\n---\n\n### Proper Control in Section 3.3\n\nTo 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](https://ibb.co/ym187Skp)). 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.\n\nFrom 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.\n\n---\n\n### Real-World Tasks\n\nWe have conducted additional experiments using context-enhanced learning to improve math capabilities. Please refer to [our response to reviewer PGoK](https://openreview.net/forum?id=Gn6L4QRKf7&noteId=qbF30B1xu4).\n\n---\n\n### FLOP Improvement\n\nEmpirically, 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.\n\nTheoretically, 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.\n\n---\n\n### Training with Loss on Context\n\nFollowing 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).\n\n|# Samples | 10,000 | 25,000 | 50,000 | 100,000 | 250,000 |\n|---|:--:|:--:|:--:|:--:|:--:|\n| No Loss on Context | 12.7% | 20.0% | 100.0% | 100.0% | 100.0% |\n| Loss on Context (new) | 15.6% | 31.6% | 98.1% | 97.1% | 98.9% |\n\nIn general there is no qualitative difference in sample efficiency when we train with loss on context.\n\nTo understand what other difference does the additional loss incur, we conduct the following additional experiments:\n\n* 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. \n* We [reproduced figure 4 for the model trained with loss on context](https://ibb.co/ccn6NsNk) trained with 250k samples and observe similar localized storage similar to the case of having no loss on context.\n* 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.\n\nThe 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.\n\n---\n\n### Bad Performance of No Dropout\n\nWe would like to clarify that our central claim (both empirical and theoretical results) all require proper dropout in the helpful context.\n\nIn 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.\nWhen no context is dropped, we would expect no internalization.", "claims_and_evidence": "I have concerns about the experiment in section 3.3 on sequential processing:\nWhen replacing STR(pi i) with STR(pi hat_i), the experiment doesn't control for the positional effects of the replacement. Transformers process information based on both content and position, so the observed changes in representations could simply reflect the model's sensitivity to any change at that position, rather than specifically processing the phrasebook.\nA proper control would include replacing non-phrasebook content at the same positions to test whether the observed pattern is specific to phrasebook processing or a general property of how transformers process sequential information. Without this, it's premature to conclude that \"the first layer showing a significant difference is identified as the layer where the model begins processing the phrasebook.\"", "methods_and_evaluation_criteria": "please see below.", "theoretical_claims": "no i didn't check fully.", "experimental_designs_or_analyses": "please see below", "relation_to_broader_scientific_literature": "The paper's concept of context-enhanced learning extends the Learning Using Privileged Information, and one of there method of annealing relates to curriculum learning.", "essential_references_not_discussed": "no", "other_strengths_and_weaknesses": "Pros:\n1. The paper introduces an interesting concept and clean formalism along with an interesting synthetic task to test the idea.\n2. The experiments suggest that context-enhanced learning significantly improves sample efficiency on the MLT task.\n3. Extensive mechanistic experiments are conducted.\n4. Strong theoretical analysis showing an exponential gap in sample complexity.\n5. The paper offers privacy implications for data security and copyright through the difficulty of recovering context materials.\n\nCons:\n1. The MLT task is artificial and synthetic with explicit rules. In real-world tasks like math or coding, there aren't many explicit rules that are universal and guarantee the mapping from a coding question to a block of code. It would greatly improve the paper if the authors had experiments on how to apply context-enhanced SFT in real-world tasks.\n2. i have concern in sec 3.3 regarding proper control is needed to test whether the observed pattern is specific to phrasebook processing or a general property of how transformers process sequential information.", "other_comments_or_suggestions": "1. the \"No Dropout (ablation)\" curve is invisible from figure 2.\n2. In the paper, they have showed \"Context-enhanced learning from an ICL-capable model greatly improves training sample efficiency.\", what about overall training efficiency in terms of training FLOPs? Since adding the context introduces additional computation in the attention calculations, although no loss is taken on the context.\n3. What will happen if you take loss on the context, will that leads to even better sample efficiency?", "questions_for_authors": "1. In the line 165 \"• Annealing Dropout: A better strategy: for s1, select the necessary rules from Π∗ plus 25% unused rules.\", how are the necessary rules selected and what's the unused rules?\n2. Have you considered applying this approach to more realistic tasks like math or coding problems?\n3. In figure 2, i'm assuming no dropout is overlapping with no-context, why is this happening? Why is using all the useful context not improving sample efficiency. If \"No Dropout\" is not improving sample efficiency over \"No Context,\" This might contradict the paper's central claim that context enhances learning.", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "JppMuD7LEN", "reviewer": "Reviewer_PGoK", "summary": "The paper introduces *context-enhanced learning* as a training setup where an LLM is fine-tuned on a task with extra context provided, but no gradient is directly taken on that context. \n\nThe main claims of the paper are:\n\n1. Context-enhanced learning improves training sample efficiency on reasoning tasks. This is presumably because the phrasebook rules are internalized atomically (*i.e.* the model is learning to apply intermediate steps, as opposed to memorizing end-to-end shortcuts).\n    1. Context-enhanced learning only works with models that are ICL capable. \n    2. Context-enhanced learning only works if you dropout parts of the context during training.\n2. Phrasebooks are processed sequentially layer-by-layer. \n3. It is challenging to prompt the model to regurgitate the rules of the intermediate phrasebooks without token filtering.", "questions": "", "limitations": "", "rating": 4, "confidence": null, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you so much for your constructive comments and questions!\n\nOur responses to your questions and concerns are as follows:\n\n---\n\n### Applying Context-Enhanced Learning to Real World Data\n\nContext-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).\n\nFollowing 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.\nWe experimented with three strategies for selecting in-context learning (ICL) examples:\n\n* **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.\n\n* **Random ICL**: Five in-context examples are randomly sampled for each training example individually.\n\n* **Skill-based ICL**: In-context examples are selected based on skill similarity, following the strategy proposed in [1].\n\nIn 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.\n\n|Pre-trained Model| Baseline (SFT) | Fixed-ICL CEL | Random-ICL CEL| Skill-based ICL CEL|\n|:--:|:--:|:--:|:--:|:--:|\n|Qwen-7B base|85.5|87.4|87.9|88.4|\n|Llama-3.1 8B base|76.5|77.4|76.5|77.2|\n\n---\n\n### Difficulty of Curating In-Context Curriculum\n\n>\"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.\"\n\nThank 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.\n\nPrior 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.\n\nWhile 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.\n\n---\n\n### Alternative Ways for Eliciting In-Context Rules\n>\"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.\"\n\nWe 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).\n\nThat 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.\n\n---\n\n### Other suggestions on related work and presentation\nThank 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.\n\n---\n\n#### References\n[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.", "claims_and_evidence": "Overall, the claims made in this paper are well-supported by empirical evidence on **synthetic tasks**. However, in my view, the main limitation of this work is that there are no experiments on real data and even the synthetics are limited to a single synthetic task. \n\n1. This claim is supported by the results in Figure 2 as well as the theoretical results in Section 5. It would be nice to see a demonstration of this phenomenon on real data or other synthetic tasks. \n2. I found the mechanistic analysis in Figures 3 and 4 to be quite compelling for this claim — and the presentation was clear. \n3. There seems to be some evidence on synthetic data supporting this claim in Section 4. However, I’m concerned that this result may not transfer when moving to real language — 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.", "methods_and_evaluation_criteria": "Overall, the authors make clever use of synthetic datasets and mechanistic analyses to illustrate their point. \n\nAs I discuss above, the paper would be greatly strengthened if it included experiments on real language data.", "theoretical_claims": "I did not carefully check the correctness of the proofs, but read the theoretical claims closely. The theoretical tools the authors use make sense and their theorems align nicely with their empirical claims.", "experimental_designs_or_analyses": "Yes see discussion of claims above.", "relation_to_broader_scientific_literature": "The paper studies an important open problem in the literature: how do we get models to perform sequential chain of thought without explicitly supervising the model to do so. Context-enhanced learning is a creative alternative to explicit reasoning supervision that I have not seen explored in recent literature.", "essential_references_not_discussed": "There are a number of prior works that study sequential reasoning tasks with synthetics and theory that could be important to cite given the mechanistic analysis of sequential reasoning. Consider mentioning them in related work. \nhttps://arxiv.org/abs/2310.07923\n\nhttps://openreview.net/forum?id=NikbrdtYvG#discussion\n\nhttps://openreview.net/pdf?id=2N3CtUdoB0", "other_strengths_and_weaknesses": "The paper’s presentation quality is excellent. \n\nPerhaps the most significant limitation is how well these findings translate beyond the neat confines of synthetic experiments. Real-world tasks rarely come with a perfectly curated “phrasebook” of rules that can be placed in the context. **One concern is the availability and quality of context**: 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.", "other_comments_or_suggestions": "**Clarification and typos.**\n\n- In Figure 2, I can’t see the pink or brown lines. Presumably they are underneath the random guess line? Is there someway this can be made more clear.", "questions_for_authors": "None", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Meta-review: This paper provides a systematic study of context-enhanced learning in LLMs. The paper considers a synthetic task of multi-level translation (MLT) when phrasebooks (mappings between tuples in two languages) are provided as additional context to language models. Even though the additional context is available during training, the goal is to improve the model performance on the MLT task without access to the additional context during test time. The paper empirically demonstrates that simply training with dropout (on additional context) can significantly speed up the model training compared to the SFT without additional context during training. The paper also explores the mechanistic understanding of how the model internalizes the additional context available during training. By studying a surrogate model, the paper establishes a clear separation between vanilla SFT and context-enhanced training in terms of their sample complexity. 

The reviewers agreed that the paper studies an interesting problem while making significant contributions toward developing a systematic understanding of the problem. Some of the reviewers raised questions about OOD generalization which the authors successfully addressed. The reviewers also pointed out the lack of treatment of real-world tasks in the initial submission. The authors' response provided new experiments on the GSM8K dataset to showcase the value of additional context during training. The authors also addressed reviewers' questions about including loss on additional context tokens. The authors are encouraged to incorporate the additional discussion/results from the author-reviewer discussion phase in the final version of the paper.
Decision: Accept (spotlight poster)