Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks?

We thank the reviewer for their constructive comments. Below, we address each point in detail.

(1) Tasks that were previously unlearnable without CoT (2) The method also appears to have an advantage...

We 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. 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.

The 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.

Multirun experiment.

We 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.

Essential References Not Discussed

We 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.

How is the input prompt structured when training with CoT?

The 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.

Permutation-invariant

That'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. (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.

We sincerely appreciate the reviewer's support and valuable suggestions. Below, we address each point in detail.

The sequential translation task is not a real-world language task.

Thank 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.

A failure case where the compositional identifiability assumption is violated so that the scaling law fails.

Thank 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$:

Let 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$.

For timesteps $1 \leq t \leq T-1$, there are two possible tasks:

• $\theta_t^0$: Always output $y_t = 0$, regardless of $(x, y_{<t})$.
• $\theta_t^1$: Always output $y_t = 1$, regardless of $(x, y_{<t})$.

At timestep $t = T$, there are also two tasks:

• $\theta_T^0$: Output $y_T = y_1 \land y_2 \land \cdots \land y_{T-1},$ where $\land$ denotes the AND operation.

• $\theta_T^1$: Always output $y_T = 1$, regardless of $(x, y_{<T})$.

In 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.

Why the compounding of errors does not show up in the other tasks than the word translation one

This 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.

No real-world tasks.

Our 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.

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.

While 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.

[1] Bhattamishra et al., Understanding in-context learning in transformers and LLMs by learning to learn discrete functions, ICLR 2024.
[2] Power et al., Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, 2022.
[3] Nanda et al., Progress measures for grokking via mechanistic interpretability, ICLR 2023.

Typos and Related Papers in the paper.

We 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.

We thank the reviewer for their comments and valuable suggestions. Below, we address each point in detail.

Dependency on $d$ and $k$ in Experiments

We 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$:

These results clearly indicate that the number of training tasks grows approximately linearly while the total number of tasks grow exponentially.

It 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$.

Theoretical Claims and Autoregressive Specificity

While 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$.

Assumption of Infinite Demonstrations

The 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.

Visualization of Experimental Results

The 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.

Relation to Prior Literature on Compositional Generalization

Our 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.

Experimental Details: Arithmetic Task

The 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.

Use of Synthetic Tasks

We 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.

Extending 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.

[1] Bhattamishra et al. (https://arxiv.org/abs/2310.03016 )

[2] Power et al. (https://arxiv.org/abs/2201.02177)

[3] Nanda et al. (https://arxiv.org/abs/2301.05217)

We thank the reviewer for their positive and constructive review. Below, we address each point in detail.

Scope is limited to synthetic tasks

Our 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.

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: the model generalizes to a combinatorially large set of unseen tasks after training on only a small subset.

While 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.

[1] Bhattamishra et al., Understanding in-context learning in transformers and LLMs by learning to learn discrete functions, ICLR 2024.
[2] Power et al., Grokking: Generalization Beyond Overfitting on Small Algorithmic Datasets, 2022.
[3] Nanda et al., Progress measures for grokking via mechanistic interpretability, ICLR 2023.

The range of $d$ and $k$ is narrow

We 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$,

It 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$.

Lack of model scale experiments

In 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.

We 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.

Importantly, our model already exhibits the near-linear task complexity predicted by theory. Larger models would not meaningfully improve this scaling.

Would pretraining on structured reasoning datasets improve ARC generalization?

This 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.

However, 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.

[4] Ye et al., LIMO: Less is More for Reasoning, 2025.
[5] Muennighoff et al., s1: Simple test-time scaling, 2025.

Typo in line 933

Thank you for pointing this out. We will fix it.