Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks?
Authors: Amirhesam Abedsoltan
Huaqing Zhang
Kaiyue Wen
Hongzhou Lin
Jingzhao Zhang
Mikhail Belkin
Venue: ICML 2025 poster (2025)
Keywords: In Context Learning
Chain of Thought
Parity problem
Composition Generalization
Primary area: deep_learning->theory

Abstract: Large language models (LLMs) exhibit remarkable task generalization, solving tasks they were never explicitly trained on with only a few demonstrations. This raises a fundamental question: When can learning from a small set of tasks  generalize to a large task family? In this paper, we investigate task generalization through the lens of autoregressive compositional structure, where each task is a composition of T operations, and each operation is among a finite family of D subtasks. This yields a total class of size~D^T. We first show that generalization to all D^T  tasks is theoretically achievable by training on only \tilde{O}(D) tasks. Empirically, we demonstrate that Transformers achieve such exponential task generalization on sparse parity functions via In-context Learning (ICL) and chain-of-thought (CoT) reasoning. We further demonstrate this exponential generalization in arithmetic and language translation, extending beyond parity functions.

Task Generalization With AutoRegressive Compositional Structure:
Can Learning From D Tasks Generalize to D Tasks?
Amirhesam Abedsoltan 1
Huaqing Zhang 2 Kaiyue Wen
Hongzhou Lin Jingzhao Zhang?
Mikhail Belkin 15
Abstract
Generalization to unseen tasks
Large language models (LLMM)) exhibit remark- able task generalization, solving tasks they were never explicitly trained on with only few demon- strations. This raises fundamental question: When can learning from a small set of tasks gen- eralize to a large task family? In this paper, we investigate task generalization through the lens of autoregressive compositional structure, where each task is a composition of T operations, and each operation is among finite family of D sub- tasks. This yields total class of size D We first show that generalization to all D tasks is theoretically achievable by training on only Ō(D) tasks. Empirically, we demonstrate that Trans- formers achieve such exponential task general- ization on sparse parity functions via In-context Learning (ICL.) and chain-of- soning, We further show generalization in arith- metic and translation, beyond parity functions.
Test Accuracy (%)
90
80
With CoT
No CoT
70
Chance Level
60
50
10
15
20
25
30
35
Number of Training Tasks
Figure 1. We train a Transformer learn parity functions through In-Context Learning (ICL), where all demonstrations and the query concatenated into single input prompt. Each prompt corre- sponds to one specific target function f, and contains binary input- output pairs (x1, fx(x) )), ..... (x.. I: (xx.)) followed by query Xquery with x; E [0, 133d. The functions fs are drawn from the Parity(10, 2) family (d 10, A 2). yiclding 45 possihle tasks. For testing, 10 such functions were withheld entirely, and different subsets the remaining ones were used for training. Consistent with prior work (Bhattamishra et al., 2024), we observe that stan- dard ICL fails generalize across tasks. In contrast, incorporating Chain-of- -hiiughhTThhhttt (CoT) reasoning significantly improves perfor- mance on unseen tasks.
(CoT) rea-
1. Introduction
Large language models (LLMs) demonstrate a remarkable ability to solve tasks they were never explicitly trained on. Unlike classical supervised learning, which typically as- sumes that the test data distribution follows the training data distribution, LLMs can generalize to new task distributions with just few demonstrations in-context learning (ICL. (Brown et al., 2020; Wei et al., 2022; Garg et al., 2022; He et al., 2024; Lee et al., 2023). Recent studies suggest that trained Transformers implement algorithmic learners capable of solving various statistical
lasks-such as linear regression at inference time in con- text (Li al., 2023; Bai al., 2023).
phenomenon known as
Despite their success in tasks such as learning conjune- tions linear regression, Transformers relying solely on ICL struggle with other problems, particularly those re- quiring hierarchical reasoning. A notable case where Transformers fail is out-of-domain learning of sparse parity functions, as examined (Bhattamishra et al., 2024). During training, a Transformer is exposed to data (x1, fx.f(()))) ...., ( f(x))) generated from multiple tar- get functions fr, E I, each defining distinct sparse parity task. In testing, data from a different task Sj.j & I is given in-context and evaluated on new data point. The results in (Bhattamishra et al., 2024) show that Transformers fail to generalize to these unseen functions, even when each tasks is simple XOR operation on two bits of the binary input
'qqual contribution 'Department Computer Science and Engineering, UC San Diego "Institute for Interdisciplinary In- formation Sciences, Tsinghua University. 'Stanford University. "Amazon. This work independent of and outside of the work at Amazon. Halicioglu Data Science Institute, UC San Diego Corre- spondence to: Amirhesam Abedsoltan aaabedsoltan @ edu."
Proceedings of the nd International Conference Machine Learning, Vancouver, Canada. PMLR 267, 2025. Copyright 2025 by the author(s).

---

Task Generalization With AutoRegressive Compositional Structure
string.
In contrast, we found that incorporating Chain-of-Thought (CoTT) reasoning - adding intermediate reasoning steps to the prompt allows Transformers to easily generalize to unseen tasks. This illustrated in Figure 1. Consistent with (Bhattamishra et al., 2024), we observe that Transformers without CoT perform only slightly better than chance level, no matter how many training tasks are presented the model. Conversely, as the number of training tasks increases, Transformers with CoT achieve near- pear-perfect generalization on the held-out set of unseen tasks. We see that the extra information provided by CoT enables the model to exploit the compositional structure of the parity problem.
Pol, 1/1
Poz, 1/2
2, 1/1
Por b yr
D. Vi,..........................y Vr
We will call this function class AutoRegressive Composi- tional structure. Assuming that the cardinality of the set of possible values for each parameter 0g is finite and is equal to D, we will use the notation F = ARC(T, D). The cardinality of this class D".
Motivated by this example, we aim to systematically ana- lyze how models can leverage autoregressive compositional structures extend their capabilities beyond the training tasks. Conventionally, learning involves approximating target function f drawn from function class F using examples from a training distribution over the input space X: generalization is then measured by testing f on new examples. In contrast, our focus is on task generalization, where training is restricted to subset of functions "tasks" Fuain& F, leaving the remaining functions, unseen during training. Our goal is to investigate whether a model trained on tasks from Firain (with inputs from X²) can generalize to all tasks, including unseen tasks.
For the sparse parity problem with & secret keys in this framework, the output sequence has length T' = k. As an example, consider n 5 and k = 3 with the secret key set S = {1, 4). Here, the inputs a = (2.......... Fn) are binary sequences of length n, and the target function is fa 22, I3, 13, I| @ 22 @ 24, where @ denotes bitwise XOR. The computation can be expressed step-by- step as:
91
1/2 II
22
1/3
I1
22 4 24
Step
Stop 2
Step 3
That is, each 1/: recovers the XOR of the first secret coor- dinates. For example, if the input sequence is 10101, the corresponding thain-of-thought will be 111. In this case, the output distribution at each step is deterministic, assign- ing probability 1 to the correct XOR value and to other values.
This notion of task generalization goes beyond the stan- dard -fitt--distriuution (OOD) settings (see, e.g,, (Zhou al., 2022) for review) by shifting the focus from adapting to new input distributions to learning entirely new tasks. Specifically, we ask:
We can now address the following fundamental question:
How many tasks in Farain must a model be trained on to generalize to all tasks in F, including those it has not seen? In particular, can a model trained on ( tasks generalize across the entire set of DJT tasks?
How can we quantify the number of tasks model must be trained on to generalize to the entire class F?
To analyze task generalization, we consider a finite set of functions F, where each function maps an input x E X' to a tuple of random variables y = (y1..... yr)). We assume each function can be characterized by a parameter tuple 0 (01,02.... .2.....,7)) 0r). The outputs are generated autoregressively: first, y1 is produced from æ; then y/2 is generated from x and y1; and then 1/3 is generated from æ, yı and y2; and this process continues until yr is produced. Specifically, the sequence is generated sequentially as:
Our main contributions are as follows:
We define AutoRegressive Compositional structure and introduce a framework to quantitatively analyze task generalization when the function class follows an Au- toRegressive Compositional structure. (Sections 3.2 and 3.3)
We establish that under this structure, task general- ization to all DT tasks is theoretically achievable by training on Ó(D) tasks up to logarithmic terms (Sec- tion 3.4).
yt ~ Pa, (yr | ær y<t), for 1. = 11,11.... 1,...,T,
where yęr- (u ...,.... Wt 1) denotes the previously gener- ated outputs, and Ps. is some conditional probability distri- bution that is parametrized by f2 and is conditioned on Yer and x. This structure can also be interpreted as a sequence of compositions,
We demonstrate how the parity problem aligns with our framework and empirically show that Transformers trained on i.i.d. sampled tasks exhibit exponential task generalization via thain-of-thougeee (CoT) reasoning, consistent with theoretical scaling (Section 4).
2

---

Task Generalization With AutoRegressive Compositional Structure
Finally, we show that the selection of training tasks significantly impacts generalization to unseen tasks. If tasks are chosen adversarially, training on even nearly all D of the tasks with CoT may fail to generalize to the remaining tasks (Section 5.1).
2021) defines variation and informativeness between differ- ent environments based on common representation, while (Arjovsky et al., 2019) addresses the problem by designing specific training objectives. Earlier studies on linear and kernel models also explore this direction (Du et al., 2017; Lei et al., 2021).
2. Related Works
Another perspective considers the testing environment as a distribution shift, where the model may sample during infer- ence to achieve domain adaptation. (Mansour et al., 2009) analyzes generalization error when a model is trained on distribution P but tested on a different distribution Q, intro- ducing an error bias dependent on the distance d(P.Q). To mitigate this bias, (Cortes et al., 2010) proposes reweighting training samples when test samples from Q are available.
2.1. Composition and Generalization
The role of composition in reasoning for language models has been widely studied. (Saparov et al., 2023) explores various out-of -itoo-dditributionn (OOD) generalization formats, including compositional generalization, showing that neu- ral network's s ability to generalize compositionally is highly dependent on both architecture and task properties. Similar conclusions have been drawn in prior works (Lakee & Baroni, 2018; Keysers al., 2020). Further, (Bhattamishra al., 2024; Dziri et al., 2023; An et al., 2023; Xu et al., 2024) examine compositional generalization in in-context learning (ICL) and find that generalization to composing multiple steps is in general hard for LLMs. One notable observation is that LLMs succeed in compositional generalization for clause satisfaction problems but not for parity problems.
A related line of research investigates scenarios where both training and test samples accessible. Notable setups include covariate shift (Kpotufe & Martinet, 2021; Ma et al., 2023) and domain adaptation (Sugiyama et al., 2007; Ben- David & Umerr 2014). When direct sampling from the test distribution is not feasible, alternative strategies focus on training robustly against worst-case shifts. This can be achieved through adversarial perturbations or min-max optimization formulations (Madry et al., 2018; Raghunathan et al., 2020; Duchi et 2023).
Another line of research investigates composition mech- anism underlying emergent abilities in language models. (Arora& & Goyal, 2023) demonstrates that language model- ing can lead to learning tuples of skills, which are small compositions of fundamental capabilities. Building on this idea, (Kaur et al., 2025; Zhao et al., 2024) leverage compo- sitional structures to generate supervised fine-tuning (SFT) data, leading to improved language model performance.
In this work we impose an AutoRegressive Composi- tiRRal(ARC) structure on the function class and propose a new framework to study task generalization. This compo- sitional structure decomposes the function class into atomic subtasks, enabling modular approach to learning. By leveraging this structure, we establish a quantitative un- derstanding of how many training tasks are required for generalization. Our results provide a theoretical foundation for structured learning and demonstrate how models can efficiently generalize beyond the training distribution.
Beyond sequential composition, other forms of composi- tionality in neural networks have been explored. (Song et al., 2025) investigates layer-wise composition in trans- formers, while (Schug et al., 2024) proposes a modular neural architecture for learning hidden compositional repre- sentations. Additionally, (Wiedemer et al., 2024) examines compositional structures in image reconstruction, and (Lippl & Stachenfeld, 2025) provides theoretical analysis of com- position in kernel and linear models.
3. Theoretical Framework for Task Generalization
In this section, we present a theoretical framework to study task generalization with autoregressive compositional struc- ture. When the compositional structure holds, we show that there exists learning algorithm that is only trained on Ö(D) different tasks, but can generalize to exponentially many unseen tasks.
While prior work has largely focused on qualitative insights into compositional generalization, our work takes a quantita- tive approach: studying how many training tasks are needed to achieve task generalization over an entire function class.
3.1. Preliminaries and Notations
2.2. Learning and Testing with Multiple Distributions
For a positive integer n. denote [r] = 21, .... For a finte set S, we denote by A(S) the probability simplex over with support S. Given i sets S2,,,,..., S, their Cartesian product is defined as
Our work aims to analyze generalization when the train- ing and testing distributions differ. This problem has been studied from various perspectives in the statistical learn- ing community. One approach is to frame it as learning shared representation across multiple tasks. (Ye et al.,
X S {(11, 82,)..,8 8: E St for all E [1]- im
3

---

Task Generalization With AutoRegressive Compositional Structure
We further denote S X butions P and Q over discrete space S, the total variation distance is defined as
S. For two probability distri-
Modularity. The generation process is decomposed into T sequential steps, each governed by an independent con- ditional distribution Pa, Pe, This modular structure allows tasks to be constructed by combining different components at each step.
==
TV(P,) = Σ P(s)) -
SES
Exponential Growth. The task class size grows expo- nentially in T as |J] = D despite each step having only D choices. This reflects the combinatorial nature of task construction, where variations each step lead to an exponentially large set of possible tasks.
We let the bold letter y denote a sequence, and the sub- scripted y denote the Jrh sequence / example. Within each sequence y² (y1,....yr) the regular letter yt denote the tn token in the sequence.
3.2. AutoRegressive Compositional Structure
3.3. Task Generalization
In the following definition, we formally introduce the Au- toRegressive Compositional (ARC) task class, which mod- els structured sequence generation through composition of conditional distributions.
Under the autoregressive task learning setup, there are two levels of generalization:
1. Generalizing to unseen inputs within task.
Definition 3.1. (AutoRegressive Compositional task class). Let X' and Y denote the finite input and output spaces, re- spectively. The AutoRegressive Compositional (ARC) task class consists of sequential generation processes:
2. Generalizing to unseen tasks in the class F.
We focus on the latter one, referred as task generalization.
(fo (Po.. Po,) Po. E Pe, for all t e (T)},
Training Phase During training, the model can only cess to a small subset of tasks Ftrain (Jot...., S F with ny = |Ftrainl- For each task for E Firains we observe ny i..d. demonstration samples:
where each task fe E F for any input æ E X' generates an output sequence y (yr,....9r) E y through an autoregressive sampling process:
D. ( (
د iid. Po v).
Vt Po. x. yct,, for all [T].
where Pg: is defined as in Equation (1). The full training dataset the union of D, denoted by Dusin (Dr
At each step the conditional probability distribution Po. drawn from a subtask family Pe.. parametrized by 0s:
We assume the learner does not know the true subtask con- ditional distribution families (Pe, Y priori. Instead, it accesses to larger hypothesis class:
Pe. = { Pe (-1 v, y<t) X' yr11 + A(D) 0, E 0.).
Here, O, represents the parameter space at step and the overall task parameter space is
X" -- O, Assum- ing each step has a finite number of possible subtasks, i.e., [0.| d for all [T]. the RutoRegres Compositional task class ARC T') consists of [F| |+] d" tasks.
Pz, = (Pc x. yc.) Se Ee Pe..
where E, parameterizes the learner's model class at step f.
The goal of training is to identify the true subtask families Pe, from Pz, through D Duriin
Given any input X and sequence y E y", the joint distribution for task fo = (Po..... )o,) E F is:
Inference Phase At test time, the learner is given & inference-time demonstration samples:
Pola. P(œ) II Po. (wa a, Neo),
(1)
D Pinfer {(0), 8' iid Pa(x.))
and for partial sequences up to any & E [T]:
where fa F unseen task. The leamner must identify the true conditionals {Pa, 1/ from Pe, for each step t. Formally, the learner A that is trained on Dtrain and given Dinfer as input, produces an output sequence of conditional distributions:
Pen. (2. Via) = (x: II Po. (v. x. V.s)
(2)
At a high level, an AutoRegressive Compositional task class ARC(D, T) is characterized by two key properties:
A (Dinfor; Derain) E ((PG..... P&r) 1 E Pa,, E [T]).

---

Task Generalization With AutoRegressive Compositional Structure
3.4. Main Result: Exponential Task Generalization
10101
1
00100
00111 ?
Promot (no CoT)
We now establish our main theoretical result: with the com- positional structure in Definition 3.1, learner can achieve exponential task generalization with only (D) training tasks. This demonstrates how compositional structure fun- damentally reduces the sample complexity of task learning from exponential to polynomial in D. Our results hold under the following mild assumptions:
Example
Fxample 2
abba
Quary
Prompt (with CoT)
10101
111
00100
000
00111 ???
Exemple
Label
E 2
Labe2
Query
Figure 2. Prompt Format Comparison: With vs. Without CoT.
Assumption 3.2 (Compositional Identifiability). The au- toregressive task class F satisfies:
Remark 3.4. If we additionally assume that every subtask distribution is separated from any incorrect hypothesis by fixed variatiai margin, i.e. for all t [T], 61:-1
1. Finite Subtask Families. For each t E [T]. the hy- pothesis class Ps, is finite and the subtask conditional distribution family Pe, S Pz, has size [Pa,1 = D.
s. and 0g E O2, Se E E, with Pc. # Po..
TV(Po Os Po 2 2r>4 0r>0
2. Task Identifiability. For any t E [T] 01..01111 E x and 0g e <e E E., Ps. F Po. the induced distribu- tions stasify:
s-1 ...
then one replace the MLE procedure used in the training stage with the same distribution distribution-disciinnnee -iscrimiinaioon approach from the inference stage (Lemma A. Under this condition, we can derive a non-asymplotic bound on the n2 needed per task for accurate identification. See Appendix for details.
TV (Por.-1,01: Pona-n.s) 0.
Furthermore, for any E [T]. 01.1-1 E x== O., and 0g + 02 2₂ the induced distributions satisfy:
3.5. Example: Sparse Parity Problem
TV (Po... Por P 2 c <> 0.
To illustrate the role of the AutoRegressive Compositional structure, we use sparse parity problem as an example.
Under these conditions, we establish our main theorem:
Theorem 3.3 (Exponential Task Generalization), Let F be an AutoRegressive Compositional(ARC) task class sat- isfying Assumption 3.2. Then there exists a learner A with the following property: if during training. one samples me D In (100 D DT) tasks uniformly and independently from F. each provided with n2 i..dd demonstration samples as the training dataset, and if at inference one observes
Sparse Parity Problem. Given d binary variables x (b,, b2..b....... ... ........,,)) a sparse parity function selects k secret indices S = {in ..,..., , i and outputs I if the sum of the corresponding variables odd, and 0 otherwise:
paritys be..... ba) = b @ b/2 .... + bi.
where @ denotes the XOR (exclusive OR) operation. We define Parity(d, k) as the set of all parity functions with d variables and ls secret indices, yielding total of
2 In (1000T ne) i..dd demonstration samples from a 2
previously unseen task Pg E FF, then
lim Pr Pr[AA(aateri Derain) # (P S 0.02,
d ( = O(d").
|aarity(d, k))
=0
ny 400
where Dieain and Dinter denote the training dataset and inference-time demonstration samples respectively, and the probability is taken over the random selection of train- ing tasks Fersin F, the training data D Diraius and the inference-time demonstration samples Dinfer-
Representation Matters. Without Shain-of-Thoughh (CoTT), the sparse parity problem Parity k) is of ARC(()) 1). but with CoT, it becomes an autoregressive compositional structure of ARC k).
* No CoT + ARC (0).).
In other words, Theorem 3.3 shows that the learner can generalize from only Ö(D) tasks an exponentially large number of unseen tasks, on the order of DJ. The learn- ing algorithm A operates in two stage. In the training stage, it applies a maximum-likellhoon estimation (MLE) procedure to D; in order to identify the subtasks of the th training task (Poç... ... Po.)) In the inference stage, it then uses a wwwaaaatonbbaaseee distribution discrimination test (Lemma A. inference-time demonstration sam- ples Dinfor to recover Sā = (Pa.." ".. Pay) The proof is deferred to Appendix A.1.
With CoT ARC(d,
Indeed, without COT, the model maps input (b,, (bi....,ba) ba) directly to output y = b,, + b ..... O b,, in single step, hence we have length T = 1 and the breadth D = |Parity(d, k) O(d*). Such exponential depen- dency on the breadth suggests that one would need to train on Odd*) tasks, explaining what we are observing unsuc- cessful task generalization in the introduction figure.
5

---

Task Generalization With AutoRegressive Compositional Structure
In contrast, with CoT (Abbe et al., 2024; Wen al., 2025), the parity computation is decomposed into small steps:
ing steps. In this case, the loss is the average on all these intermediate steps.
y = (bi, bi, + bi....... b,, + + by
Training Data Generation. We split both the task space and input space into training and testing. In other words, the parity tasks are split into Fiain and Fiest; and the binary sequences are split into Xrain and Kest- The split in the input space helps us monitor the in-distribution training process while as the split in the task space aims to measure generalization to unseen parity functions.
enabling structured representation that reduces the breadth. More precisely, at step i, the class of subtask is exactly defined by the XOR operation with previous token and one secret index:
(u yee) = 1gye = 1 + b..l-
To construct the ICL data, we sample points 2i,...., Fm uniformly at random from X Similarly, we sample a function I uniformly at random from Fimin- and generate the sequence (#1 f(z1)....,.) T..... fzm))) This process is repeated N times, each time with a fresh sample of m points and a new function f.
where i, is the 1-th secret index, by default lying E [1,d] Therefore, the breadth D at each step is exactly d.
This said, the CoT effectively reduces the breadth from (dd) ) to d. According to Theorem 3.3,
Corollary 3.5. For the sparse parity problem described above, we can show that the parameter e in Assumption 3.2 is 1 thus when no 2 d'In(100kd), E 2 8 $11(100kna). it holds that
Evaluating Task Generalization. For evaluation, sam- ple ʃ randomly from the held-out tasks Frest and sample inputs uniformly at random from Visa. We report the ac- curacy of the prediction f(zm) given the demonstration (x1, fz11)..... #m 1. f( Øm). This setting chal- lenges the model to generalize to novel tasks beyond those encountered during training.
2.
lim Pr (Dintorr Dirain) & (Pa..)) P(ra))) S 0.02. -0-000
In other words, the family of Parity(d, k) with CoT is learnable with (dd log(d) training tasks.
4.2. Experimental Results
As discussed in Section 3.5, introducing CoT transforms the parity problem class Parity (d, k) into an AutoRegressive Compositional structure ARC(D,T) with D d and T k. A key empirical question is:
4. Experiments: Parity Problem Case Study
As we have shown, there exists a learning algorithm that is only trained on ((d) different tasks to fully generalize on all the tasks in Parily(d, k). However, from practical standpoint, it is not clear whether a Transformer can actually match this task complexity. In this section, we present em- pirical evidence demonstrating that a standard Transformer can indeed learn the sparse parity function with CoT using (d) training tasks. The paper's GitHub repository can be found online.
how does the number of training tasks scale with d and T to achieve fixed target accuracy on unseen task sets?
To investigate this, we conduct experiments on:
1. Scaling T(= k)), i.e. the length of secret indices.
Scaling D(= i.e. the ambient dimension of input.
4.1. Experimental Sctup: In-Context Learning
Scaling T for a Fixed D. We examine how the number of training tasks affects test accuracy while keeping the am- bient dimension fixed at d = 15. Specifically, we evaluate test accuracy for k = {3, 4, 4,5, 6, 7} under varying numbers of training tasks. As & increases, the size of the parity class grows significantly- from approximately 500 for k = 3 to around 6500 for k: = 7.
Our empirical setup is a refinement of the theoretical frame- work presented in Section 3.3, and closely follows that of (Garg et al., 2022; Bhattamishra et al., 2024). In this setup, a sequence model M (such as Transformers) is trained using N sequences, each sequence consisting of m demonstration samples (I1, V/...., Im, y.)) The model trained for the next token-prediction task, except that we only consider y₁ in loss optimization: for each context (x1, V1..... zz 11 Z,), the model predicts Új. and the loss is given by 1 Σグ= ((ūj,y)) In our experiment, we use cross-entropy loss to measure the discrepancy be- tween the predicted output أj and the true output yj-
Remarkably, despite this increase, the test accuracy follows similar trajectory. With just 3d In(d) 122 training tasks, the model generalizes to unseen cases with high accuracy (> 95%). For A = 7, this means training on 122 tasks enables generalization to about 6,400 unseen ones! This empirical results suggests that the required number of training tasks remains roughly the same, regardless of k, consistent with the theoretical scaling of Ö(d) tasks.
When Chain of Thought (CoT) is used, each yy is itself a sequence of length i representing intermediate reason-
6

---

Task Generalization With AutoRegressive Compositional Structure
Fixed dimension (d=15)
Fixed number of secret keys (k=3)
100
A
100
B
Test Accuracy (%)
90
80
Ambient dimension
# Secret keys
a' 10
70
k= 3
k=
80
d= 15
60
k 5
25
70
3 Number of Training Tasks
2
1
Number of Training Tasks
D In(O)
x In(D)
Figure 3. Test accuracy on unseen tasks. For parity task: D = d as the ambient dimension and T' = k as the number of secret indices. We show that the empirical scaling closely follows the theoretical scaling of Dln(D). (A) For fixed D 15, increases, the test accuracy on unseen tasks remains similar, even though the total number of tasks (~> D²) grows exponentially with T. (B) For a fixed secret length 3, as increases, number of tasks grows polynomially with D. yet the number of tasks required to generalize reasonably to unseen tasks remains in x D log D.
Scaling for a fixed T. We examine the effect of increas- ing d E (10, 15, 20) while keeping h: 3 fixed. For each d, total number of tasks proportional d (n(),,
the i-th token in the Chain-of of-Thought, aiming to predict the i-th secret index. Only the linear classifier is trained, while all Transformer parameters remain frozen. Table 2 shows that for d E (10, 15, 20). k = 3, the linear probe consis- tently achieves high accuracy at all the coordinates. This suggests that at each CoT step, the model indeed first infers the relevant subtask (the secret index) from the in-context examples and then executes that subtask to generate the output token-an an ability acquires during training. Further experimental details appear in Appendix D.
we train up to 4 x d [nd). Figure 3, Panel B, shows similar task gen- eralization performance across different ambien dimension of d, providing further evidence that (dd i.i.d. training tasks are sufficient for generalization to unseen tasks on praity functions with COT.
Scaling D and T together. We examine the ef- feet of jointly increasing d and h, with kd E {(10,5) (15, 7), (20, 10). (25, 12), (30, 15)]- For each d, the model is trained on total of (d) i.i.d. tasks. Table I shows that generalization performance remains con- sistent across these settings, providing further evidence that Ódd training tasks are sufficient to generalize to unseen parity tasks using CoT prompting. As a concrete example, for = generalization to approximately 155 million unseen tasks.
Token 1
Token 2
Token 3
10
100.00%
99.83%
91.08%
15
95.42%
99.97%
98. 73%
97.38%
95.54%
91.92%
Table 2. Validation accuracy (%) of linear probes trained to predict each secret token position from the final hidden state.
30, 15, training on just 306 tasks enables
5. Task Generalization Beyond i.i.d. Sampling and Parity Functions
d
k
#Training Tasks
#TotalTasks
Accuracy (%)
10
5
69
252
98.51
15
7
121
6,400
99.12
In this section, we extend our experiments beyond i.i.d. task sampling and parity functions. We show an adversarial ex- ample where biased task selection substantially hinders task generalization for sparse parity problem. In addition, we demonstrate that exponential task scaling extends to non- parity tasks including arithmetic and multi-step language translation.
10
185,000
98.67
25
12
241
3,200,000
98.60
30
15
306
155,000,000
98.10
Table 1. Task generalization performance as d and k increase. Training on only (d) tasks enables generalization exponen- tially many unseen tasks in the parity function family.
5.1. Task Generalization Beyond i.i.d. Task Sampling
Subtask Identification via Linear Probing. Finally, we probe the hidden representations of the Transformer (see (Alain& & Bengio, 2017)) to see if it identifies and then ex- ecutes subtasks at the inference time, consistent with the framework of Section 3. Specifically, we add a linear classi- fier the final attention layer's hidden state when producing
In previous sections, we focused on i..dd. settings, where the set of training tasks Ferain were sampled uniformly at ran- dom from the entire class F. Here, we explore scenarios that deliberately break this uniformity to examine the effect of task selection on distt-iiitiiitiiio (OOD) generalization.

---

Task Generalization With AutoRegressive Compositional Structure
How does the selection of training tasks influence model's ability generalize to unseen tasks? Can we predict which setups are more prone to failure?
evaluates the expression by applying the operations se- quentially from left-to-right order: for example, if S (+, +)) then the arithmetic function would compute Arithmetic ((1 b2. by, ba) = ((b) + b2) x by) + b1.
To investigate this, we consider two cases parity problems with d = 10 and k = 3. where each task is represented by its tuple of secret indices (51, $2, 83):
By introducing step-by-step CoT, arithmetic class belongs to A d d11); 1):
this is because at every step, there is only D |0,| = 2 choices (either + or x) while the length is L, resulting a total number of 2d tasks.
Generalization with a Missing Coordinate. In this setup, we exclude all training tasks where the second coordinate takes the value $2 5, such as (1, 5, 7). At test time, we evaluate whether the model can generalize to unseen tasks where = 5 appears.
T
Task generalization for the arithmetic task with CoT. It has d 2 and T d 1 as the ambient dimension, hence D Pn 2 In lnT)). We show that the empirical scaling closely follows the theoretical scaling.
2. Generalization with Missing Pair. Here, we remove all training tasks that contain both 4 and 6 in the tplle
Arithmatic task
Test Accuracy (%)
(51, 82, 83), such as (2, 4, 6) and (4, 6). At test time, we assess whether the model can generalize to tasks where both 4 and 6 appear together.
75
50
Amblent dimension d= 10
25
d 20
If you had to guess. Which scenario is more challenging for generalization to unseen tasks? We provide the experimental result in Table 3.
0-
2
8 Number Training Tesks
10 12
Notably, when scaling with T. we observe in the figure above that the task scaling closely follow the theoretical O log (DT) dependency. Given that the function class grows exponentially as 2", is truly remarkable that train- ing on only few hundred tasks enables generalization an exponentially larger space-on the order of 225 > 33 Mil- lion tasks. This exponential scaling highlights the efficiency of structured learning, where a modest number of training examples can yield vast generalization capability. 5.2.2. MULTI-STEP LANGUAGE TRANSLATION TASK
In the first scenario, despite being trained on all tasks except those where 82 = 5, which is of size O(D"), the model struggles to generalize to these excluded cases, with predic- tion at chance level. This is intriguing as one may expect model to generalize across position. The failure suggests that positional diversity plays a crucial role in the task gen- eralization of Transformers.
In contrast, in the second scenario, though the model has never seen tasks with both and together, it has encoun- tered individual instances where appears in the second position (e.g., (1,4,5) or where 6 appears in the third po- sition (e.g.. (2, 3, 6)). This exposure appears to facilitate generalization to test cases where both 4 and 6 are present. As result, when the training tasks are not i.i.d, an adver- sarial selection such as exclusion of specific positional con- figurations may lead to failure to unseen task generalization even though the size of Ftrain is exponentially large.
In this task, we study a sequential translation process across multiple languages (Garg et al., 2022). Given a set of D languages, we construct a translation chain by randomly sampling a sequence of T languages with replacement:
L2, Lg,..,, LT, where each L, is a sampled language. Starting with word, we iteratively translate it through the sequence:
LT. For example, if the sampled sequence is EN - FR - DE FR, translating the word "butterfly" follows:
5.2. Task Generalization Beyond Parity Problems
5.2.1. ARITHMETIC TASK
butterfly + papillon -+ schmetterling + papillon.
We introduce the family of Arithmetic task that, like the sparse parity problem, operates on d binary inputs bi,bb,..,bu The task involves computing a structured arithmetic expression over these inputs using sequence of addition and multiplication operations.
This task follows an AutoRegressive Compositional struc- ture by itself, specifically ARC(D, 1). where at each step, the conditional generation only depends on the target language, making D the number of languages and the total number of possible tasks is DJ' This example illus- trates that autoregressive compositional structures naturally arise in real-world languages, even without explicit CoT.
Formally, we define the function:
Arithmetics: : {0, 1]2 - {0. 1, I.....d).
where S = (op,. UP2. 1 is a sequence of d - operations, each op, chosen from (+, x). The function
We examine task scaling along D (number of languages) and T (sequence length). As shown in Figure 5, empiri-
8

---

Task Generalization With AutoRegressive Compositional Structure
Table 3. Generalization Results for Scenarios I and 2 for 10, 3.
Scenario
Tasks excluded from training Generalization accuracy {4. 6} C {81, 82)
Generalization with Missing Pair
96.2%
Generalization with Missing Coordinate
$2 5
45.6%
Fixed number of languages languages)
Fixed number of steps (T
100
100
A
B
Test Accuracy (%6
75
50
Values
50
Number of languages
-
8
25
-2-
5
+
9
10
01
5
10
15
20
5
10
15
20
Number of Training Tasks
Number of Training Tasks
x In(D)
Figure 5, Task generalization for language translation task: D is the number of languages and T is the length of steps.
cal D-scaling closely follows the theoretical O In DT). However, in the T --scaling case, we observe a linear depen- dency on T rather than the logarithmic dependency OIIn 7). A possible explanation is error accumulation across sequen- tial steps-longer sequences require higher precision in intermediate steps to maintain accuracy. This contrasts with our theoretical analysis, which focuses on asymptotic seal- ing and does not explicitly account for compounding errors in finite-sample settings.
context learning (ICL) on parity tasks but achieve exponential generalization when Chain-of-Thought (CoT) reasoning is introduced. Our results provide the first empirical demonstration of structured learning enabling efficient generalization in this setting.
Scaling Laws in Arithmetic and Language Transla- tion. Extending beyond parity functions, we demon- strated that the same compositional principles hold for arithmetic operations and multi-step language transla- tion, confirming that structured learning significantly reduces the task complexity required for generaliza- tion.
Despite this, the task scaling is still remarkable training on a few hundred tasks enables generalization to 1 x 10 tasks!
6. Conclusions
Impact of Training Task Selection. We analyzed how different task selection strategies affect generalization, showing that adversarially chosen training tasks can hinder generalization, while diverse training distribu- tions promote robust learning across unseen tasks.
In this work, we quantitatively investigated task general- ization under the autoregressive compositional structure, demonstrating both theoretically and empirically that expo- nential task generalization to D" tasks can be achieved with training on only ÓD) tasks. To summerize:
We introduce a framework for studying the role of composi- tionality in learning tasks and how this structure can signifi- cantly enhance generalization unseen tasks. Additionally. provide empirical evidence learning tasks, such as the parity problem, demonstrating that transformers follow the scaling behavior predicted by our compositionality-bun based theory. Future research will explore how these principles extend to real-world applications such as program synthesis, mathematical reasoning. and decision-making tasks.
Theoretical Framework for Task Generalization. We introduced the AutoRegressive Compositional (ARC) framework to model structured task learning, demonstrating that a model trained on only Ó(D) tasks can generalize to an exponentially large space of T tasks.
Formal Sample Complexity Bound. We established fundamental scaling law that quantifies the number of tasks required for generalization, proving that expo- nential generalization is theoretically achievable with only logarithmic increase in training samples.
By establishing principled framework for task generaliza- tion, our work advances the understanding of how models can learn efficiently beyond supervised training and adapt to new task distributions. We hope these insights will in- spire further research into the mechanisms underlying task generalization and compositional generalization.
Empirical Validation on Parity Functions. We showed that Transformers struggle with standard in-

---

Task Generalization With AutoRegressive Compositional Structure
Acknowledgements
complex skills in language models. arXiv preprint 1XXv::307.119366 2023.
We acknowledge support from the National Science Founda- tion (NSF) and the Simons Foundation for the Collaboration on the Theoretical Foundations of Deep Learning through awards DMSS2031883 and #814639 as well as the TILOS institute (NSF CCF-21 12665) and the Office of Naval Re- search (ONR N00014240 1000142412631). This work used the pro- grams (1) XSEDE (Extreme science and engineering discov- ery enviromment) which is supported by NSF grant numbers ACI-1548562, and (2) ACCESS (Advanced cyherinfrastrue- ture coordination ecosystem: services & support) which is supported by NSF grants numbers #2138259, #2138286, #2138307, #2137603, and #2138296. Specifically, used the resources from SDSC Expanse GPU compute nodes, and NCSA Delta system, via allocations TG-CIS220009.
Bai, Y., Chen, F., Wang, H., Xiong, C., and Mei, S. Trans- formers as statisticians: Provable in-context learning with in-context algorithm selection. In Thirty-seventh Conference on Neural Information Processing Systems, 2023. URL http:::://///.. htt..tttts//oppeeeviieeee http////epeeeeeeeeeeeeoee wttpeeewwwwwweeeeeeeeee..eeeeeee n forum? id-1iM8qUuVg3.
Ben-David, S. and Urner, R. Domain adaptation-can quan- tity compensate for quality? Annals of Mathematics and Artificial Intelligence, 70:185-200 2014.
Bhattamishra, S., Patel, A., Blunsom, P., and Kanade, V. Understanding in-context learning in transformers and LLMS by learning to learn discrete functions. In The Twelfth International Conference on Learning Represen- tations, 2024. URL https: //pprr/iiee/eeeeeeeeeeeee net// forum?id=ekeyCgeeti
Impact Statement
This paper advances the theoretical understanding of task generalization, investigating how learning on small sub- set of tasks enables generalization to much larger set of unseen tasks. Our study is primarily theoretical, supported by synthetic experiments designed to isolate key structural properties of task composition. By analyzing autoregressive compositional structures, we provide insights into sample complexity and task transferability. While our work lays foundational groundwork for broader applications in ma- chine learning, its immediate impact remains theoretical, with no direct societal implications requiring emphasis.
Brown, T., Mann, B., Ryder, N., Subbiah, M., Kaplan, J. D., Dhariwal, P., Neelakantan, A., Shyam, P., Sastry, G., Askell, A., et al. Language models are few-shot learners. Advances in neural information processing systems, 33: 1877-1901, 2020.
Cortes, C., Mansour, Y., and Mohri, M. Learning bounds for importance weighting. Advances neural information processing systems, 23, 2010.
Du, S. S., Koushik, J., Singh, A., and Póczos, B. Hypothesis transfer learning via transformation functions. Advances in neural information processing systems, 30, 2017.
References
Duchi, J., Hashimoto, T., and Namkoong, H. Distribution- ally robust losses for latent covariate mixtures. Opera- tions Research, 71(2):649-664, 2023.
Abbe, E., Bengio, S., Lotfi, A., Sandon, C., and Saremi, O. How far can transformers reason? the locality barrier and inductive scratchpad. Advances in Neural Information Processing Systems, 2024.
Dziri, N., Lu, X., Sclar, M., Li, X. L.., Jiang, L., Lin, B. Y., Welleck, S., West, P., Bhagavatula, C., Le Bras, R., et al. Faith and fate: Limits of transformers on compositionality. Advances in Neural Information Processing Systems, 36: 70293-70332, 2023.
Alain, G. and Bengio, Y. Understanding intermediate layers using linear classifier probes. In 5th International Con- ference on Learning Representations, ICLR 2017, Toulon, France, April 24-26, 2017, Workshop Track Proceedings, 2017.
Garg, S., Tsipras, D., Liang, P. S., and Valiant, G. What can transformers learn in-context? case study of sim- ple function classes. Advances in Neural Information Processing Systems, 35:30583-30599 2022.
An, S., Lin, Z., Fu, Q., Chen, B., Zheng, N., Lou, J.-G., and Zhang, D. How do in-context examples affect composi- tional generalization? In Proceedings of the ólst Annual Meeting of the Association for Computational I Linguistics (Volume 1: Long Papers), pp. 11027-11052, 2023.
He, T., Doshi, D.,, Das, A., and Gromov, A. Learning to grok: Emergence of in-context learning and skill compo- sition in modular arithmetic tasks. Advances in Neural Information Processing Systems, 37:13244-13273, 2024.
Arjovsky, M.. Bottou, L., Gulrajani, I. and Lopez- Paz, D. Invariant risk minimization. arXiv preprint 197700190, 2019.
Kaur, S., Park, S., Goyal, A., and Arora, S. Instruct-skillmix. A powerful pipeline for Ul instruction tuning. In The Thirteenth International Conference on Learning Repre- sentations, 2025.
Arora, S. and Goyal, A. theory for emergence of
10

---

Task Generalization With AutoRegressive Compositional Structure
Keysers, D., Schärli, N., Scales, N., Buisman, H., Furrer, D., Kashubin, S., Momchev, N., Sinopalnikov, D., Stafiniak, L., Tihon, T., Tsarkov, D., Wang, X., van Zee, M., and Bousquet, O. Measuring compositional generalization: A comprehensive method on realistic data. In International Conference on Learning Representations, 2020.
deductive reasoning capacity of large language models using ood examples. Advances in Neural Information Processing Systems, 36:3083-3105, 2023.
Schug, S., Kobayashi, S., Akram, Y., Wolezyk, M., Proca, A. M., Von Oswald, J., Pascanu, R., Sacramento, and Steger, A. Discovering modular solutions that generalize compositionally. In International Conference Learn- ing Representations, 2024.
Kingma, D. P. and Ba, J. Adam: A method for stochastic optimization. In International Conference Learning Representations, 2015.
Song. J., Xu, Z., and Zhong, Y. Out-of-distribution general- ization via composition: lens through induction heads in transformers. Proceedings of the National Academy of Sciences, 122(6)e24171821222 2025.
Kpotufe, S. and Martinet, G. Marginal singularity and the benefits of labels in covariate-shift. The Annals of Statistics, 49(6))3299-33333 2021.
Lake, B. and Baroni, M. Generalization without systematic- ity: On the compositional skills of sequence coaecce-tt-seeee te-senuence recurrent networks. In International conference on ma- chine learning, pp. 2873-2882. PMLR, 2018
Sugiyama, M., Krauledat, M., and Müller, K.-R. Covariate shift adaptation by importance weighted cross validation. Journal of Machine Learning Research, 8(5), 2007.
Lee, N., Sreenivasan, K., Lee, J. D., Lee, K., and Papail- iopoulos, D. Teaching arithmetic to small transformers. arXiv preprint ar arXiv:2307.03381, 2023.
Wei, J., Tay, Y., Bommasani, R., Raffel, C., Zoph, B., Borgeaud, S., Yogatama, D., Bosma, M., Zhou, D., Met- zler, D., Chi, E. H., Hashimoto, T., Vinyals, O., Liang, P., Dean, J., and Fedus, W. Emergent abilities of large language models. Transactions on Machine Learning Research, 2022. ISSN 2835-8856. URL https: openreview net forum? wwarmrddddzxksxsommm Sur- vey Certification.
Lei, Q., Hu, W., and Lee, J. Near-optimal linear regression under distribution shift. In International Conference Machine Learning, pp.
PMLR, 2021.
Li, Y., Ildiz, M. E., Papailiopoulos, D., and Oymak, S. Transformers as algorithms: Generalization and stability in in-context learning. In International Conference Machine Learning, pp. 19565-19594. PMLR, 2023.
Wen, K., Zhang, H., Lin, H., and Zhang. J. From sparse dependence to sparse attention: Unveiling how chain- of-thought enhances transformer sample efficiency. In The Thirteenth International Conference on Learning Representations, 2025.
Lippl, S. and Stachenfeld, K. When does compositional structure yield compositional generalization? a kemnel theory. In The Thirteenth International Conference on Learning Representations, 2025.
Wiedemer, T., Mayilvahanan, P., Bethge, M., and Brendel, W. Compositional generalization from first principles. Advances in Neural Information Processing Systems, 36, 2024.
Ma, C., Pathak, R., and Wainwright, M. J. Optimally tack- ling covariate shift in rahs-based nonparametric regres- sion. The Annals of Statistics, 51 2023.
Wolf, T., Debut, L., Sanh, V., Chaumond, J., Delangue, C., Moi, A., Cistac, P., Rault, T., Louf, R., Funtowicz, M., and Brew, J. Transformers: State-of-the-art natural lan- guage processing. In Proceedings s of the 2020 Conference on Empirical Methods in Natural Language Processing: System Demonstrations, pp. 38-45, 2020.
Madry, A., Makelov, A., Schmidt, L., Tsipras, D., and Vladu, A. Towards deep learning models resistant to adversarial attacks. In International Conference on Learn- ing Representations, 2018
Mansour, Y., Mohri, M., and Rostamizadeh, A. Domain adaptation: Learning bounds and algorithms. In Annual Conference on Learning Theory, 2009.
Xu, Z., Shi, Z., and Liang, Y. Do large language models have compositional ability? an investigation into limitations and scalability. In Conference on Language Modeling, 2024.
Raghunathan, A., Xie, S. M., Yang, F., Duchi, J., and Liang, P. Understanding and mitigating the tradeoff between robustness and accuracy. In International Conference Machine Learning, pp. 7909-7919. PMLR, 2020.
Li, Z., and Wang. L. Towards Ye, H., Xie, C., a theoretical framework distrriissti generaliza- tion. Advances in Neural Information Processing Systems, 34:23519-2355 2021.
Saparov, A., Pang, R. Y., Padmakumar, V., Joshi, N., Kazemi, M., Kim, N., and He, H. Testing the general

---

Task Generalization With AutoRegressive Compositional Structure
Zhao, H., Kaur, S., Yu, D., Goyal, A., and Arora, S. Can models learn skill composition from examples? Advances in Neural Information Processing Systems, 37:102393- 102427, 2024.
Zhou, K., Liu, Z., Qiao, Y., Xiang, T., and Loy, C. Do- main generalization: A survey. IEEE Transactions Pattern Analysis and Machine Intelligence, 43(4) 4415, 2022.
12

---

and denote the KI. divergence as Task Generalization With AutoRegressive Compositional Structure
A. Proofs in Section 3
Notations For distribution over finite class S, denote supp(P) C S as the support P. For two distributions P, Q over finite set S, denote the cross entropy as
H(P = Σ
P(x) P(r) log Q(z),
xES
KL(P.Q) =- - P(x) log Q(x)
P(x)
FES
A.1. Proof of Theorem 3.3
To identify the true distribution from θ at inference time, we need the following lemma to provide a non-asymplotic bound for distribution diserimination.
Lemma A.1 (Non- -symtssyioo Discrimination of Two Distributions). Let S be any finite set and P.Q E A(S) be two 0. Suppose we observe independent samples XI.... Xn from an unknown distribution Y, where Y is either P or Q. Then, there exists testing algorithm that identifies Y with probability at least 8 provided that
distributions with total variation distance TV Q)
n 2mm(1/8)
Proof. Testing Procedure. Define the testing statistic @ as follows:
n1
Equivalently, & can be written in terms of the empirical distribution Yn:
Φ = Σ
(n(x) (-1)))(()))))G)
#ES
where n E x].
Compute the expected values of under Y = and Y = =Q:
Σ P(x) (-1))((())))) pQ = Q(x)) (-111)(1)))))))
=
DES
#ES
The algorithm reports y = if 16 - upl < HQl and V = Q otherwise.
Proof of Correctness, Without loss of generality, assume Y P. We analyze the probability of error:
Pr ( P) = Pr (l0 prl
2 |ob pQl)- -
Note that E[Ø) pr under Y P. From the definition of total variation distance:
pQ
[(Q(x) P(z)) (-111(()00a))) = 2.dr((P..))
2c.
ZES
Thus, the error probability can be bounded as:
Pr (lo - pel 2 10b - xel) < Pr (Φ 2 pr + c).
Since & is the average of n independent random variables taking values in {-1, by Hoeffding's inequality:
n
Pr (6 2 up + < exp
Setting exp (_2n < ô gives the required sample complexity >
2la(1/8) 64
13

---

Task Generalization With AutoRegressive Compositional Structure
To prove Theorem 3.3, introduce the following lemma:
Lemma A.2. Consider a compositional task class satisfying Assumption 3.2. Suppose the training tasks Fuain (fo V satisfy the per-component coverage condition: for every timestep E (T).
1 -Pe
Then there exists a learner such that when trained on Drrsin = (D.Yi= with Dr Ltd. Po. where Dr consists ng i.idd samples, and given e > Dinter then it holds that: 6²
21n(1007me) inference-time demonstration samples i..d sampled from unseen task Pg denoted by
lim Pr [A(D Dirain) + (P....... ....,.. pe,), 0.01,
+00
where the probability is over the randomness Dirain and Dinfer.
Proof.
Step 1: Training Stage.
We construct a learning algorithm that recovers
for (Poj (Poj....,, Poz.
from
D. (( داار itd (o
which is part of the training set Drrain. The procedure shown below:
Algorithm 1 Training Stage
Require: Training set Derain
1: for 1 to ny do
2: for 1 to T do
3: Observed Support,
4: Pzy - Ps. E Ps,
supp(Pes (Pojn- Observed Supporty
log Po 5: 0g +- arg max
( (2t1, vl
Eo E 3
6: end for
7:
end for
8: return Pe.- (P&)) 1m for each i E [T].
We now show that, as ntz -
Pr Pr|Pe. Pe, [T + 1.
Assume 0 81 As n2 +
Pr [[00bberved.Suporrt # supp(Pa:.))
Σ
Pr[( vid (x. vra) Vj [na]] 0.
(m.wiir E supp (rot.)
Conditioned on LnsservedSSuportt = supp (Pop)(Po..)) Pot..) any Pe, in P'd must share the same support:
supp (Pa P (ppp(Ppp. (Pot....)))
Thus the cross-entropy is finite. By the law of large numbers and Pinsker's inequality, for Ps. # Poş we have
14

---

Task Generalization With AutoRegressive Compositional Structure
lim
log Pot
o (2 311r) ==1
log Por e (xt .y(())
KLLPo 1.0f Pof.. i.6))
2TTV 11.-1.00 Poj 1 )
>0
almost surely. Thus.
ni log Po of (xt), 81/2) + 0.
Pr[Po
Po:] s
Σ
Pr
log
P in
virt)
Pi
Rx
Pa
Recursively applying this argument from to T and taking union bound over all i [no] and E [T).
lim (r( Po. * Po:] = 0.
Finally, by the assumption that ( .... Poro = -e, for each f, we conclude
Pr Pr[Po. Pe, VLE
1 as 72 +
Step 2: Inference Stage
Assume that, in the training phase, the learner identifies the true distribution sets
Po. = PPo pgro} Pe. for all ( = I....,
T.
la Ia(1000 no)
We now construct an algorithm that, given
independent samples from the unseen composite task
f8 (Po. Por
outputs the same distribution tuple (Pa..... Po with probability at least 0.99.
To achieve this, we leverage the distribution discrimination technique of Lemma A.1 to distinguish between candidates in Pa. given sufficiently many demonstration samples.
15

---

then Task Generalization With AutoRegressive Compositional Structure
Algorithm 2 Inference Stage
Require: Inference-time demonstration samples Dinfer
((0), 1') i..dd from Pg. and the identified sets Pe.
{Pa 0 ................................ Pg.. (go
for each t E [T].
1: for f = to T do
2: Initialize Pc. +- Pa:
3: for: to ng do
4: Compute
1
1 P tha) < < c:n-1: 1. (((" sla)
(-1) Σ
==1
5: if
Pu (a, (-1) ( Pe qa (a.11.))
Σ (z...))XXx))
(, (1:) (-1)2(pan, (a
Σ
(XXx)
6:
Update Pco +- Poc.
7: end if
8: end for
9: end for
return (Pa...., Pct)
Error Analysis. Let Pc, be the chosen distribution at step Given Pa.. Po Lemma A.1 ensures that any incorrect candidate can be detected with high probability, long Pa. E Po. Specifically,
Pr Ps. + Po. Ps Phrr-.)
Si = 0g or 0 =6,] s no expl -xp(-c22)
< Pr distribution testing fails at step i
2
union bound over ..... T then implies
Pr[(Pa..... P(r) 4 .Pa....., Poy) S T exp( -
ha T1) no)
Since l
we get
Tno exp (-22) - 0.01.
Step 3: Combining Training and Inference
From the results of Step 1 and Step 2, have
lim Pr Pr[AADDiterr: Derain) *
ns 000
< lim
(Pr[(Pa...... Pcx) + (Po, ... Pon) ) Pa. PP, Vi er(3 Ps. + #Poil)
Ry 000
= 0.01 + 0 = 0.01.
16

---

Task Generalization With AutoRegressive Compositional Structure
Proof of Theorem 3.3. By Lemma A.2, it suffices verify that per-component coverage condition holds for each timestep with high probability when no D In (100
For each timestep & E (T), observe
Pr[{Pop)1m1 + Po.l S
Σ
Pr Pr(Pot + Po. Vie [me]]
PopEPo:
D ( 1
(1)
1 n0
Hence,
T
e (Pa # Po.]
- Σ Pr (P lhg
* Pe.l
1=1
S DT 1 (1- 1 m
DT DTe-D 0.01.
This completes the proof.
A.2. Proof of Corollary 3.5
Proof of Corollary 3.5. It suffices to verify that for the sparse parity problem, the constant i Assumption 3.2 2-
We denote P..... d, (x, ye1) = 1 2 and it holds that
1 [y11 = Nú, I 1 Wy = Vo-1 + x;,]. For any in. .. ,4-1 E [d]. and i, # i E [d]
TV(P do P. Pap. d 118()=
Σ
122 + ..... Di._ D4, 2i, @ .... zi.- + zzl
2n
(0,0)
1| + x:c
n
ZE (
This completes the proof.
B. Non-Asymptotic Analysis of Training-Time Demonstration Sample Complexity
In Theorem 3.3, we established only an asymptotic result, showing that as the number of demonstration samples per task at training time 7l co, the probability of correctly identifying subtask families Pe, tends one. However, by imposing an additional assumption the total variation gap between the true distributions and any other hypotheses, it is possible to derive a non-asymptotic guarantee on how large ntz must be for accurate subtask identification.
Although maximum-likelihooe estimation (MLE) does not directly yield such a non-asymptotic bound in this setting, we can use the same distribution discrimination approach introduced in the inference stage (Lemma A.
Assumption B.1 (Compositional Identifiability with fixed tv marigin). The autoregressive task class F satisfies:
1. Finite Subtask Families: For each t E [Z'] the hypothesis class Pz, is of size at most H and the subtask conditional distribution family Pe. S PPz has size iPe. = D.
2. Task Identifiability: For any [T], 01a-1 x O., and 0, E On E Er, Pc. + Po., the induced distributions stasify:
TV (Po Saae,PPacce PonA-1ki) 2 2r>0 0.
Furthermore, for any timestep & E [T'. 011-1 E x-=1 O. and 0, + 02 E O, the induced distributions satisfy:
TV (Pon.-100 Porc-1,9 2 <> 0.
17

---

then Task Generalization With AutoRegressive Compositional Structure
Theorem B.2 (Exponential Task Generalization). Let F be autoregressive compositional task class satisfying Assump- tion 3.2. Then there exists learner A with the following property: during training, one samples no tasks uniformly and independently from F, each provided with ne i..dd. demonstration samples as the training dataset, and if at inference one observes e i..dd demonstration samples from previously unseen task Po E F, then
Pr A([inerr Dirain) * (Po....,,PP. Po.)) S DDeDD + noTe 21/2 + no Ho -
where Derain and Dinter denote the training dataset and inference-time demonstration samples respectively, and the probability is taken over the random selection of training tasks Firain S F, the training data Drrain and the inference time demonstration samples Dintor-
Proof. Denote the hypothesis class Pz,- = (Pe..:" P J. we present the training stage of the learner.
Algorithm 3 Training Stage with Distribution Dislimination
Require: Training set Derain = (D
1: for to ng do
2: for to T do
3: Initialize Pa: P
4: for 3: = 2 to E do
5: Compute
(-1) 2-1)=(Pam p wil) < Po 1.be (00 (0' w1
Σ
Th.
(wtt JED,
6:
Po (x, v1.) P (e.0))))))))))))) Poso 1. (2.
Σ
(a (xx
Pen o (xx y1:e) (-1)'(pa Pin (onema))))) < S0, (mmna))
Σ
7:
Update Po. <- PEL..
8: end
9: end for
10: end for
11: end for
12:
return Pe. (Pa Mi for each / E [T].
Using the same approach as Step 2 of the proof of Theorem 3.3,
Pr((p .., + ... .oo)l s THe Pn./2
By union bound,
Pr E [ Po, # Pe.. S PP(3(())) Pa: + Po:] S noTH Te - -
The remainder the proof then proceeds exactly as Theorem 3.3.
18

---

Task Generalization With AutoRegressive Compositional Structure
C. Extra experiments
Effect Context Length. The theory assumes access to an infinite number examples for each training task but does not require infinite demonstrations during inference, However, in practice, we camnot train on infinite number of examples. Figure shows that providing sufficient context length during both training and inference crucial for strong performance. Empirically, we observed that context length of 40 works reasonably well across all experiments with dimensions up to d = 20.
d 10
3
# secret keys (k)
Test Accuracy (96)
96
10
20
10
Context length
Figure 6. The effect of context length on performance.
ICL with no CoT fails in even in-distribution generalziation. We observe in Figure 7 that transformers with ICL and no CoT struggle generalize even in simpler in-distribution settings the number of tasks increases. In the parity task, refer to in-distribution generalization evaluated on the same set of tasks Ferain but with entirely new sequences Steat that were not seen during training.
setting where the model is trained on Ftrain tasks and Strain sequences, and then
Here, the setting is the same as in (Bhattamishra et al., 2024) for Parity(10, 2), but we used the same tasks during both training and testing. trained on half of the total sequences, 2 and tested on unseen sequences while keeping the tasks unchanged.
In-Distribution generalization
100
Test Accuracy (%)
90
70
W/O CoT
60
W CoT
Chance Level
50
1
4
7
Number of Training Tasks
Figure 7. ICL without CoT even fails to generalize in distribution.
D. Experiment Details
Model and optimization. We used the transformers library from Hugging Face (Wolf et al., 2020) to instantiate and train our GPT-2 model from scratch. In all experiments, we used a 3-layer, 1-head configuration. We used the Wadam optimizer (Kingma & Ba, 2015) with learning rate of 10-5 and a batch size of 64.
19

---

Task Generalization With AutoRegressive Compositional Structure
Parity and arithmetic. In all experiments shown in Figures 3 and 4 for both parity and arithmetic tasks, we used context length of 40.
For the arithmetic problem, across all dimensions, we used total of 25,000 training examples, equally distributed across the training tasks.
For the parity problem, we used 20,000 training samples, equally distributed across the training tasks for dimensions up to 15. For dimension 20, we increased the total number of training samples to 50,000.
At lesting time, we always randomly select the minimum between 200 subsets and all remaining tasks, each containing 500 different sequences with the same context length of 40.
Language experiments. For the translation experiments, train a 2-layer Transformer with 3 heads and embedding dimension 768. We use an Adam optimizer with betas being 0.9, 0.95 and learning rate 3e-4. We will keep the number of total training samples to be 1e5 and train for 1 pass for 6250 steps. We choose the languages randomly from the following set (Enllish, French, Spanish, Chimese, German, Italian, Japanese, Russian,
Portuguese, Arabic} and meanings (in English) from {cat, dog, house, apple, sky, car, road tree, bed, water, sun, moon). We use GPT-2 tokenizer and in our demonstrations, we will prepend the language of the corresponding word before each word in the following format like "English: cat".
Linear Probing We append linear classifier the checkpoints of models of "Increasing D for fixed 7" tasks, trained on the hidden states of the final attention layer when generating the i-th token in the Chain-of- Thoughh,, with the goal of predicting the i-th "secret index." The models are trained on a total number of of 20, 000, 20, 000, and 50, 000 training samples for d = 10, 15, and 20, respectively. The tasks used for training and validation are disjoint. Only the linear classifier trained, while the parameters the transformer are frozen. We use the Adam optimizer with a learning rate of 4 x 10-5, and the batch size is set to be 32.
20


Reviews:
[{"review_id": "5AKgbEGbMu", "reviewer": "Reviewer_B7oU", "summary": "This paper investigates task generalization in large language models (LLMs) through the lens of AutoRegressive Compositional (ARC) structure. The central question explored is: When can learning from a small set of tasks enable generalization to a much larger task family? The authors propose that LLMs, particularly transformers, can achieve exponential task generalization when tasks are structured compositionally.", "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": "We thank the reviewer for their positive and constructive review. Below, we address each point in detail.\n\n> Scope is limited to synthetic tasks\n\nOur goal is to establish a foundation where the compositional structure of tasks is clearly defined and generalization to unseen tasks is fully controllable. To this end, we deliberately focus on synthetic tasks to isolate the core principles underlying ARC-style generalization, explicitly specifying seen tasks and rigorously evaluating generalization on truly unseen tasks, free from leakage or contamination from pretraining.\n\nWe also note that several influential works studying transformer behaviors similarly focus on synthetic or simplified settings, such as parity functions [1] and modular arithmetic [2,3], which have proven effective for uncovering key insights. In our case, the key insight is *exponential task generalization*: the model generalizes to a combinatorially large set of unseen tasks after training on only a small subset.\n\nWhile extending this framework to broader NLP tasks is important, it incurs substantial computational cost, which is beyond the scope of this paper. We view our current results as a necessary first step and are actively exploring such extensions.\n\n[1] Bhattamishra et al., *Understanding in-context learning in transformers and LLMs by learning to learn discrete functions*, ICLR 2024.  \n[2] Power et al., *Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets*, 2022.  \n[3] Nanda et al., *Progress measures for grokking via mechanistic interpretability*, ICLR 2023.\n\n\n>The range of $d$ and $k$ is narrow\n\nWe agree that a broader range of $d$ and $k$ provides additional support for our findings. However, both $d$ and $k$ appear in the base and exponent when determining the total number of tasks, leading to exponential growth. For instance, in the parity task, even with $d=25$ and $k=10$, the number of distinct parity functions is $\\binom{25}{10} \\approx 3$ million. We have strengthened our empirical support by adding this new experiment on parity tasks with $d=25$ and $k=10$, \n\n| $d$ | $k$ | \\# Training Tasks | \\# Total Tasks | Accuracy (%) |\n|-----|-----|------------------|----------------|--------------|\n| 10 | 5 | 69 | 252 | 98.51% |\n| 15 | 7 | 121 | 6,400 | 99.12% |\n| 25 | 10 | 241 | 3.2M | 98.60% |\n\nIt is still a valid question to ask why can’t we scale to significantly larger $d$. The main challenge comes from context length. Unlike prior works focusing on learning a fixed parity function where $d$ simply reflected input dimensionality, we are in an ICL setup where the model needs to learn a general algorithm that infers *any* parity function drawn from a large function family. For the problem to be identifiable, the number of in-context examples $N$ must scale linearly with $d$. For example, when $d=100$ and $k=50$, $N=100$ examples result in an input length of roughly $N(d+k) \\approx 15,000$ tokens -- which surpasses our available computational resources. This quadratic scaling of context length is the main bottleneck limiting us from larger $d$ and $k$.\n\n\n>Lack of model scale experiments\n\nIn early experiments, we varied the number of layers and heads. We observed that while increasing layers from 2 to 3 improved generalization, going beyond 3 layers yielded no additional benefit. Similarly, increasing attention heads did not improve performance. This motivated our choice of a 3-layer, 1-head model for faster optimization without sacrificing generalization. A supporting plot from our early experiments is available in this [link](https://files.catbox.moe/x1r11m.png). \n\nWe also conducted a new set of experiments comparing models with width 192 vs. 384, 3 layers vs. 6 layers, and 1 head vs. 4 heads, across different numbers of training tasks (120 and 160), with fixed input dimension \\(d = 15\\) and \\(k = 3\\). The results were consistent, showing comparable performance across configurations, in this table [Link](https://files.catbox.moe/dlsa4k.png). \n\nImportantly, our model already exhibits the near-linear task complexity predicted by theory. Larger models would not meaningfully improve this scaling.\n\n\n>Would pretraining on structured reasoning datasets improve ARC generalization?\n\nThis is an excellent question. Recent work on R1 distillation has shown that strong generalization can be achieved from a relatively small number of SFT samples [4,5]. Inspired by this, we hypothesize that pretraining on structured reasoning datasets could further enhance ARC generalization.\n\nHowever, this direction requires significant additional work and resources, as pretraining large models on diverse reasoning datasets is far more demanding than our current ICL setting. Nonetheless, we believe this is a promising avenue for future research.\n\n[4] Ye et al., *LIMO: Less is More for Reasoning*, 2025.  \n[5] Muennighoff et al., *s1: Simple test-time scaling*, 2025.\n\n\n>Typo in line 933\n\nThank you for pointing this out. We will fix it.", "claims_and_evidence": "The paper makes several claims that are appropriately supported by empirical results. \n\nClaim: Generalization to exponentially large task sets is possible with limited training tasks when tasks follow ARC structure.\n- Empirical results in sparse parity and arithmetic tasks show that transformers\n\nClaim: Task selection significantly impacts generalization.\n- When tasks are sampled adversarially (e.g., omitting tasks with a specific coordinate value), transformers fail to generalize even when trained on an otherwise large task set.", "methods_and_evaluation_criteria": "Yes, their proposed framework and the evaluation criteria make sense and support their theroem. \nMore specifically, \n- the ARC framework is clearly defined and analyzed mathematically\n- the experimental setup isolates core factors influencing generalization (number of secret indices, ambient dimension, task selection, CoT usage)\n- uses linear probing for subtask identification.", "theoretical_claims": "Yes. I checked Theorem 3.3 and the proof in the supplementary material seems correct.", "experimental_designs_or_analyses": "The experimental setup is well-designed, focusing on synthetic benchmarks that allow controlled evaluation of task generalization.\n(See Methods and Evaluation Criteria part).", "relation_to_broader_scientific_literature": "The paper connects well to prior work in:\n- Compositional generalization in neural networks\n- In-context learning and task transferability\n- CoT reasoning and its impact on LLM expressiveness", "essential_references_not_discussed": "I am not aware of essential references not discussed", "other_strengths_and_weaknesses": "### Strengths\n- Theoretical contributions are rigorous, providing a well-defined framework for studying task generalization.\n- Scaling laws are empirically validated, demonstrating consistency between theory and practice.\n- Task selection ablations are insightful, highlighting the importance of diverse training examples.\n\n### Weaknesses\n- Scope is limited to synthetic tasks, making it unclear whether the results translate to broader NLP applications.\n- Lacks model scale experiments\n- The range of k (number of secret indices), and d (ambient dimensions) is narrow.", "other_comments_or_suggestions": "- (minor) Typo in line 933 Wadam -> adamW? \n- model scale experiments (rather than with 3-layer, 1-head model) could enhance the empirical results.", "questions_for_authors": "- How does task generalization change with model scale?\n- Would pretraining on structured reasoning datasets improve ARC generalization?", "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": "768KXyLddL", "reviewer": "Reviewer_vSqF", "summary": "This paper investigates when models trained on a small set of tasks can generalize to a much larger task family. The authors approach this through the lens of \"autoregressive compositional structure\" (ARC), where tasks are composed of T sequential operations, each chosen from D possible subtasks, creating a total task space of D^T. The key theoretical contribution is proving that generalization to all D^T tasks can be achieved by training on only O(D log DT) tasks. Empirically, the authors demonstrate that Transformers can achieve this exponential generalization on sparse parity functions and arithmetic operations.", "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": "We thank the reviewer for their comments and valuable suggestions. Below, we address each point in detail.\n\n>Dependency on $d$ and $k$ in Experiments\n\nWe agree that a broader range of $d$ and $k$ provides additional support for our findings. However, both $d$ and $k$ appear in the base and exponent when determining the total number of tasks, leading to exponential growth. For instance, in the parity task, even with $d=25$ and $k=10$, the number of distinct parity functions is $\\binom{25}{10} \\approx 3$ million. We have strengthened our empirical support by adding this new experiment on parity tasks with $d=25$ and $k=10$:\n\n| $d$ | $k$ | \\# Training Tasks | \\# Total Tasks | Accuracy (%) |\n|---|---|-|-|-|\n| 10 | 5 | 69 | 252 | 98.51 |\n| 15 | 7 | 121 | 6,400 | 99.12 |\n| 25 | 10 | 241 | 3.2M | 98.60 |\n\nThese results clearly indicate that the number of training tasks grows approximately linearly while the total number of tasks grow exponentially. \n\nIt is still a valid question to ask why can’t we scale to significantly larger $d$. The main challenge comes from context length. Unlike prior works focusing on learning a fixed parity function where $d$ simply reflected input dimensionality, we are in an ICL setup where the model needs to learn a general algorithm that infers *any* parity function drawn from a large function family. For the problem to be identifiable, the number of in-context examples $N$ must scale linearly with $d$. For example, when $d=100$ and $k=50$, $N=100$ examples result in an input length of roughly $N(d+k) \\approx 15,000$ tokens -- which surpasses our available computational resources. This quadratic scaling of context length is the main bottleneck limiting us from larger $d$ and $k$.\n\n>Theoretical Claims and Autoregressive Specificity\n\nWhile the proof may appear general, our framework is specific to the autoregressive (ARC) setting. The distribution of $y_t$ explicitly depends on $y_{<t}$, and our constructed learner solves each subtask autoregressively, only using $(x,y_{\\leq t})$ from the input sequence. Setups without ARC structure can be seen as a special case (when $T=1$), but compositionality, which is central to our analysis, only emerges through ARC composition when $T>1$.\n\n>Assumption of Infinite Demonstrations\n\nThe assumption $n_x\\to\\infty$ in the main text is for simplifying exposition. We also provide a non-asymptotic guarantee in Appendix B (Remark 3.4, Theorem B.2). Specifically, if each subtask distribution is separated from incorrect hypotheses by a total-variation margin $\\epsilon>0$, then $n_x=\\widetilde{O}(1/\\epsilon^2)$ demonstrations suffice. Thus, our analysis covers the finite-sample regime.\n\n>Visualization of Experimental Results\n\nThe visualization suggested by the reviewer—plotting the number of training tasks needed to reach a given accuracy—is a valid alternative. However, we believe there is no single best way to present the results. Our current presentation is chosen to align closely with the theoretical prediction on task complexity. It also provides both lower and upper bounds for a given accuracy. For instance, Figure 1 shows that in the parity task, all setups achieve ≥98% accuracy within $[2D\\log D,4D\\log D]$ training tasks.\n\n>Relation to Prior Literature on Compositional Generalization\n\nOur framework differs from prior work in two key ways. First, most existing work is largely empirical, focusing on assessing compositional abilities of pretrained LMs, while we explicitly study how compositionality emerges during training when structured task families are presented. Second, we shift the focus from standard sample complexity to *task complexity*, aiming for generalization to unseen tasks where training on a small number of tasks generalizes to exponentially many unseen tasks.\n\n>Experimental Details: Arithmetic Task\n\nThe arithmetic task fully follows the in-context learning (ICL) setting, consistent with the parity task. We will add more experimental details in the final version.\n\n>Use of Synthetic Tasks\n\nWe chose synthetic tasks to ensure that the compositional structure is fully controlled and interpretable. The goal is to establish a foundation where the compositional structure is clearly defined, and generalization to unseen tasks is fully controllable. We also note that several influential works studying transformer behaviors similarly focus on synthetic or simplified settings, such as parity functions [1] and modular arithmetic [2,3], which have proven effective for uncovering key insights. In our case, the key insight is *exponential task generalization*.\n\nExtending this framework to real-world NLP tasks is important and an active direction we are exploring. However, doing so requires substantial more work, which we believe deserves a dedicated follow-up study.\n\n[1] Bhattamishra et al. (https://arxiv.org/abs/2310.03016 )\n\n[2] Power et al. (https://arxiv.org/abs/2201.02177)\n\n[3] Nanda et al. (https://arxiv.org/abs/2301.05217)", "claims_and_evidence": "1. The theoretical claims are supported by the proof. \n2. The empirical claim on the dependency of the function space (D) and the number of compositions (T) are not sufficiently convincing. First, for each task, only a handful number of D and T choices are tested, which is still relatively small to verify the bound.", "methods_and_evaluation_criteria": "Yes the proposed methods and/or evaluation criteria (e.g., benchmark datasets) make sense for the problem or application at hand.", "theoretical_claims": "1. I briefly read the proof but didn't check the details. One confusing thing is that it seems to be more like a general approach rather than being specific to the autoregressive prediction, which is probably also the reason that the authors observe a different dependency on T in the language translation task.\n2. In the paper, it assumes the number of demonstration (nx) goes to infinity which seems not very realistic.", "experimental_designs_or_analyses": "1. About the experimental setup: There is no detail on the arithmetic task. Is it still following the ICL format and what does it look like?\n2. About the empirical results: Besides the number of tested D and T is small, the visualizations to showcase the dependency of the number of training tasks are also relatively confusing. The current plots have test acc (the y axis) vs num of training tasks/Dlog(DT) (the x axis) which is hard to see whether it scales in a linear/log way. More natual visualizations would be number of training tasks used to achieve ~100% test acc vs Dlog(DT)), with the actual linear/log line as a reference.", "relation_to_broader_scientific_literature": "The key theoretical contribution is proving that generalization to all D^T tasks can be achieved by training on only O(D log DT) tasks. I fell like this result is relatively new but it would be great if the author could clarify that how the proposed frame work and the theoretical result is different from the compositional generalization literature. For the empirical work, though the authors claim that it verifies the theory, I still feel like they're not very convincing for the reasons mentioned in Claims And Evidence & Experimental Designs Or Analyses.", "essential_references_not_discussed": "Most recent works in this area that I am familar with are cited.", "other_strengths_and_weaknesses": "Strengths:\n1.  The paper introduces a structured mathematical framework that quantifies compositional complexity through parameters D and T and provide theoretical result for the sample complexity.\n2. The paper designs ICL experiments to verify the theortical results.\n\nWeaknesses:\n1. The dimensions tested are relatively small compared to the parameter spaces of modern LLMs. The generalization properties might behave differently at much larger scales.\n2. The experiments focus on relatively synthetic tasks with clear compositional structure. It remains unclear how these findings might transfer to more complex, natural language tasks.", "other_comments_or_suggestions": "1. I recommend making the experimental section more comprehensive with clearer visualizations for the linear/log dependency.\n2. The paper would benefit from more details on experimental setups for each task.", "questions_for_authors": "See above.", "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": "l1CyueCJO3", "reviewer": "Reviewer_7pHd", "summary": "This paper presents the simple but useful and intuitive idea that *chain-of-thought generation reduces the theoretical complexity of compositional problems, and therefore leads to theoretically and empirically less required samples for strong generalisation* to unseen functions. Specifically, prior work has shown that transformers fail with this type of generalisation, but the authors of this work show they can in fact do it. The question the authors investigate is when can transformers generalise to a large number of unseen tasks from a smaller number of tasks? They show for an autoregressive compositionally structured task formulation how many tasks a model needs to theoretically see to generalise to all heldout tasks, and empirically demonstrate that this holds for three examples of such tasks. Even when task space gets larger (e.g. through more steps in the sequence), the amount of tasks needed to generalise follows the theoretically found scaling law. The authors further show that this is only the case when the model is trained on i.i.d. tasks and examples of these tasks, and has worse generalisation performance when specific examples/tasks are structurally held-out.\n\n## Update after rebuttal\n\nMy points have been addressed after the rebuttal and I am still in favour of accepting this paper.", "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": "We sincerely appreciate the reviewer's support and valuable suggestions. Below, we address each point in detail.\n\n> The sequential translation task is not a real-world language task.\n\nThank you for pointing this out. We agree the translation example is synthetic and will clarify that our goal was merely to illustrate how decompositional structure can emerge without explicit CoT, as long as compositional structure holds, rather than to suggest that this task reflects real-world language use. We will clarify this in the final version.\n\n>  A failure case where the compositional identifiability assumption is violated so that the scaling law fails.\n \nThank you for the suggestion. At a high level, the compositional identifiability assumption is required to ensure that any subtask is identifiable regardless of the other subtasks it is composed with. To illustrate what happens when this assumption is violated, we present an extreme failure case where a subtask is identifiable only in the presence of specific preceding subtasks. In this setting, task complexity grows exponentially with $T$:\n\nLet the vocabulary be $\\mathcal{Y} = \\{0, 1\\}$. Consider the following autoregressive compositional task class where the output is deterministic and independent of the input $x$.\n\nFor timesteps $1 \\leq t \\leq T-1$, there are two possible tasks:\n\n- $\\theta_t^0$: Always output $y_t = 0$, regardless of $(x, y_{<t})$.\n- $\\theta_t^1$: Always output $y_t = 1$, regardless of $(x, y_{<t})$.\n\nAt timestep $t = T$, there are also two tasks:\n\n- $\\theta_T^0$: Output $y_T = y_1 \\land y_2 \\land \\cdots \\land y_{T-1},$ where $\\land$ denotes the AND operation.\n\n- $\\theta_T^1$: Always output $y_T = 1$, regardless of $(x, y_{<T})$.\n\nIn this case, identifying the subtask $\\theta_T^0$ as distinct from simply outputting $y_T = 0$ requires that the specific task $(\\theta_1^1, \\theta_2^1, \\dots, \\theta_{T-1}^1, \\theta_T^0)$ appears in training. The probability of sampling this task decreases exponentially with $T$, making the task complexity exponentially large for successful learning. We will briefly discuss this in the final version.\n\n\n> Why the compounding of errors does not show up in the other tasks than the word translation one\n\nThis is a great and challenging question. One plausible explanation is that the implicit biases of Transformers vary significantly across tasks. In the word translation task, the model may adopt undesirable shortcuts—such as only partially solving the intended task (e.g., learning to translate from Chinese to French without generalizing from English to Chinese)—which increases the likelihood of errors at each generation step. In autoregressive generation, an error rate of $p$ per step compounds over time, meaning that achieving correct sequence generation with constant probability requires $p = o(1/T)$. In contrast, algorithmic tasks like parity and arithmetic appear to induce stronger biases toward learning the correct underlying computation, potentially mitigating this issue. That said, we stress that this compounding error argument is merely a hypothesis and does not constitute a rigorous explanation. We do not fully understand why these tasks exhibit different scaling behaviors, and we leave a complete explanation for future work.\n\n> No real-world tasks.\n\nOur goal is to establish a foundation where the compositional structure of tasks is clearly defined and generalization to unseen tasks is fully controllable. To this end, we deliberately focus on synthetic tasks to isolate the core principles underlying ARC-style generalization, explicitly specifying seen tasks and rigorously evaluating generalization on truly unseen tasks, free from leakage or contamination from pretraining.\n\nWe also note that several influential works studying transformer behaviors similarly focus on synthetic or simplified settings, such as parity functions [1] and modular arithmetic [2,3], which have proven effective for uncovering key insights.\n\nWhile extending this framework to broader NLP tasks is important, it incurs substantial computational cost, which is beyond the scope of this paper. We view our current results as a necessary first step and are actively exploring such extensions.\n\n[1] Bhattamishra et al., *Understanding in-context learning in transformers and LLMs by learning to learn discrete functions*, ICLR 2024.  \n[2] Power et al., *Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets*, 2022.  \n[3] Nanda et al., *Progress measures for grokking via mechanistic interpretability*, ICLR 2023.\n\n> Typos and Related Papers in the paper.\n\nWe thank the reviewer for the careful feedback and suggestions. We will correct the typos and include the suggested reference in the final version of the paper.", "claims_and_evidence": "All claims are supported by clear and convincing evidence. The authors derive a theoretical result and show it empirically holds for three separate tasks for which the critical assumption in the theory holds.\n\nThis is not central to the thesis of the work, but the claim that is not supported is the following. In Section 5.2.2. the authors show results on a toy word translation task, and claim that *\"This example illustrates that autoregressive compositional structures naturally arise in real-world languages, even without explicit CoT.\"* But this is a bit of a stretch, given that sequential translation of words to different languages is not a real-world language task.\n\nOne suggestion I would make to make the evidence stronger is to show that the scaling law does not hold for a task that violates the assumption about compositional identifiability.", "methods_and_evaluation_criteria": "All methods and evaluation criteria make sense.", "theoretical_claims": "I did not check the correctness of any proofs and theoretical claims, but the empirical results validate them and they make intuitively perfect sense. E.g. if a problem is compositional, in that each next step to tackle is independent of the previous steps if you keep track of a state, chain-of-thought can break down the problem and leads to a lower complexity task. Generalisation then is independent of the overall task complexity in parameters like sequence length cause one can compositionally apply the algorithm to each step.", "experimental_designs_or_analyses": "The experimental design is sound and valid.", "relation_to_broader_scientific_literature": "The authors successfully demonstrate how transformers can show strong generalisation in cases that prior work said they struggle with, and come up with a new theoretical result that explains why and how, which can potentially connect to generalisation properties of LLMs.", "essential_references_not_discussed": "Some papers the authors could consider citing are the following (just a suggestion, I leave it up to the authors if they are worth discussing):\n- https://arxiv.org/abs/2307.03381 (related in that they show CoT is crucial for strong generalisation)\n- https://arxiv.org/abs/2406.02550 (related in that they focus on task generalisation)", "other_strengths_and_weaknesses": "**Strengths**\n- Very clear writing\n- Neat connection of theory to multiple empirical results\n- The idea presented here contributed to my understanding of generalisation in transformers\n\n\n**Weaknesses**\n- No real-world tasks\n- As mentioned above, would be interesting to show a failure case (different number of tasks required than predicted by theory for a task that violates compositional identifiability)", "other_comments_or_suggestions": "List of typos:\n- 117r typo finte -> finite\n- 315r typo ambien -> ambient\n- 346-348l needs a rewrite (mostly in terms of verbs and determiners)\n- 365-366l as well\n- 111r typo hide -> hides\n- 193l \"can only access to\" -> \"can only access\"\n- 202l \"it accesses to\" -> \"it has access to\"\n- 211r \"two stage\" -> \"two stages\"\n- Typo title figure 3 arithmatic -> arithmetic", "questions_for_authors": "- Do you have any idea why the compounding of errors does not show up in the other tasks than the word translation one?", "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": "tIDHxd47mb", "reviewer": "Reviewer_w4XA", "summary": "This paper demonstrates that for Boolean tasks with an AutoRegressive Compositional Structure (e.g., sparse parity), using chain-of-thought (CoT) to break down the tasks into simpler sub-problems—while predicting the outcomes of intermediate steps—significantly improves generalization in the GPT-2 model.", "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": "We thank the reviewer for their constructive comments. Below, we address each point in detail.\n\n> **(1) Tasks that were previously unlearnable without CoT (2) The method also appears to have an advantage...** \n\nWe thank the reviewer for raising this question regarding whether the difference between CoT and non-CoT models is primarily an optimization issue or a generalization issue. To clarify, in our experiments, models trained without CoT are able to reduce the training loss to near-zero and achieve nearly 100% training accuracy. For clarity, we include an additional plot showing the training dynamics corresponding to Figure 1 of our main paper: [Link](https://files.catbox.moe/95rhrv.png). As shown, while the training accuracy approaches 100%, the test accuracy remains at random chance. This indicates that the key challenge is not optimization, but generalization: models trained without CoT fail to generalize to unseen compositions of subtasks, despite perfectly fitting the training data. In contrast, models trained with CoT exhibit strong generalization.\n\nThe core contribution of our work is a shift in perspective from traditional sample complexity to task complexity. We show that for tasks with compositional structure—where CoT prompting serves as an example—training on a small set of tasks enables generalization to an exponentially large number of unseen tasks. Prior work [Wen et al (2024)] shows that transformers can, in principle, learn sparse parity functions without CoT, but doing so requires exponentially more data. In our setting, which focuses on in-context learning (ICL), we observe that even when the model fits the training tasks without CoT, it fails to generalize across tasks. In contrast, CoT enables generalization to an exponentially large space of unseen compositions, in line with our theoretical predictions. We will make this distinction clearer in the revision.\n\n> **Multirun experiment.**\n\nWe appreciate the concern. It is worth noting that as we are focusing on task generalization, each time we change the number of tasks we need to rerun the experiments from scratch. In other words, each data point in our plot corresponds to a completely new random sample of tasks, which already provides an implicit robustness check. Nevertheless, in the revised version, we further strengthen our results by including additional experiments with repeated runs (at least 5 random seeds) for selected settings—specifically for $d=15$ and $d=20$—to directly assess the stability of our results. A plot summarizing the variance across seeds for these settings is available [here](https://files.catbox.moe/acfo1o.png).\n\n> **Essential References Not Discussed**\n\nWe thank the reviewer for highlighting Task Diversity Shortens the ICL Plateau, which shows transformers can learn sparse parity when trained on a diverse task mix (e.g., parity with linear or quadratic regression). However, it's important to note that their setting differs from ours as they do not separate training tasks from testing tasks. Specifically, the model is trained continuously in an online fashion, where all tasks are repeatedly visited during training. In contrast, our work explicitly evaluates out-of-distribution task generalization by testing models on entirely unseen tasks. The key message of our work is that the ARC structure enables generalization to an exponentially large space of unseen compositions. Moreover, in their Figure 1, sparse parity (red line) isn’t learned on its own with transformers—it only becomes learnable when mixed with other task.\n\n> **How is the input prompt structured when training with CoT?**\n\nThe input and overall setting are exactly the same as in the paper you mentioned, *“Understanding In-Context Learning…”*. The only difference is that, instead of providing only the final parity label, each label now includes the intermediate XOR computations (see the figure in the right column of page 5). For example, if $d = 15$, context length is 40, and $k = 3$, then the input size becomes $40 \\times (15 + 3)$, compared to $40 \\times (15 + 1)$ without CoT—just as in the original paper. As a result, zero-padding is not needed.\n\n> **Permutation-invariant**\n\nThat's an interesting question. As the task has the same compositional structure, our theoretical framework still applies. However, as you noted, the number of possible tasks increases by a factor of $k!$, so the constants in the theory may differ. We ran experiments with $d = 15$ and $k = 3$, varying the number of training tasks. The results are summarized in the table in this [link](https://files.catbox.moe/63l1hh.png). (The number of training tasks is $n \\cdot 15 \\cdot \\ln(15)$, where $n = 2, 3, 4$, and the total number of possible tasks is ${15 \\choose 3} \\cdot 3! = 2730$.) As expected, the generalization behavior roughly follows the pattern observed in the non-permutation case. However, we found that training was more difficult and required more steps to converge.", "claims_and_evidence": "**\bWeakness 1**\n\nI believe the paper’s most significant contribution is demonstrating that by leveraging chain-of-thought (CoT) to train tasks with an auto-regressive compositional structure:\n\n(1) Tasks that were previously unlearnable without CoT (in the sense that the training loss could not be reduced) can now be learned.\n(2) The method also appears to have an advantage in terms of generalization (i.e., the sense that the test loss is small).\n\nHowever, regarding point (1), I am not particularly surprised, because it seems the authors have used prior knowledge of the ARC structure to explicitly break the tasks down into smaller units via CoT, thereby artificially lowering the complexity and making it more learnable.\n\nAs for point (2), I am skeptical of the claim that this approach truly confers a generalization advantage. The paper’s primary example, Sparse Parity, is well-known to be essentially unlearnable (and certainly not generalizable) by Transformers. Thus, from my perspective, the experiments mainly show that CoT makes learning possible in the first place, but do not convincingly demonstrate how much of a generalization benefit is gained over a No CoT approach.\n\nFor a fair comparison of generalization performance, I believe it is necessary to examine tasks that can be learned both with and without CoT. Without such a baseline, it seems to me that the paper merely shows how reducing task complexity via CoT enables successful training, which, while potentially useful, is not particularly surprising.", "methods_and_evaluation_criteria": "I believe the methods and evaluation criteria proposed by the authors are reasonable.", "theoretical_claims": "Under the stated assumptions, the proofs for the theoretical claims are logically coherent and employ standard techniques appropriately. While some assumptions (e.g., infinite samples, fixed TV gaps) might be limited, I think that they are acceptable for a theoretical analysis.", "experimental_designs_or_analyses": "**\bWeakness 2**\n\nDespite the fact that the experiments are neither time-consuming nor require extensive GPU resources, the authors did not provide results from multiple repeated trials. As a result, it is difficult to assess the consistency of the findings.", "relation_to_broader_scientific_literature": "This work appears to contribute to the research area of using synthetic tasks to better understand language models. The parity task is known to be difficult for Transformers to learn, but other researchers could gain insights from this paper on how to make it learnable.", "essential_references_not_discussed": "The author mentions that, so far, Transformers have been unable to learn sparse parity. However, while reviewing recent studies, this paper* claiming that sparse parity can be learned when multiple tasks are trained together.\n\n*Task Diversity Shortens the ICL Plateau, arXiv 2025. (https://arxiv.org/abs/2410.05448)", "other_strengths_and_weaknesses": "**Strength**\n\n1. Providing experimental and theoretical explanations of how Transformers can effectively learn and generalize by leveraging CoT.\n2. In the sparse parity and arithmetic tasks, it is interesting that the results closely align with the theoretical findings.\n3. Overall writing and structure were easy to read!\n\n**Weakness**\n\nPlease see weakness 1 and weakness 2.", "other_comments_or_suggestions": "The multilingual translation experiments are conducted in a relatively simplified setting, where a single word (e.g., \"cat\" or \"dog\") is translated through a language chain using a 2-layer, 3-head Transformer. This simplified experimental environment may not adequately capture the diverse linguistic and contextual factors present in real-world translation scenarios, making it difficult to substantiate the claim that the ARC structure naturally arises in actual language translation tasks. Therefore, there are no experimental results that sufficiently bridge the gap between the paper's contributions and real-world scenarios.", "questions_for_authors": "1. How is the input prompt structured when training with CoT? I'm curious about how zero padding is applied. For example, when using 10101 -> 111 as input, do you add two zero paddings to 111, making it 11100? It would be helpful if you could provide a detailed explanation of the overall process. (I understand the setting used in the paper \"Understanding In-Context Learning in Transformers and LLMs by Learning to Learn Discrete Functions\")\n\n2. I’m curious whether the experimental results are permutation-invariant. There are multiple ways to apply CoT (technically, there are k! possible ways). In the current paper, the model processes the input autoregressively in order, starting from the smallest index (such as figure in page 5). What happens if a randomly permuted order is used instead of ascending order? If you have experimented with this, I’d love to know the results. Additionally, I’d be interested in hearing the authors' insights on this matter.", "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 investigates transformer models' ability to generalize from a small set of tasks to an exponentially larger task family. The authors prove that generalization to D^T tasks can be achieved by training on only O(D log DT) tasks under the task identifiability assumption. (Note that O(D log DT) tasks imply O(DT log DT) 'subtasks', so one can learn the D subtasks for each time step.) The theoretical results are validated empirically across both synthetic tasks that satisfy the assumptions and those that don't.

The paper merits acceptance due to its rigorous theoretical foundation supported by consistent empirical validation. Reviewers appreciated the clear mathematical framework for the compositional structure, the demonstrated alignment between theoretical predictions and experimental results, the excellent writing quality, and the insightful ablations that highlighted the importance of diverse training examples.
Decision: Accept (poster)