How Compositional Generalization and Creativity Improve as Diffusion Models are Trained

We thank the reviewer for their feedback and address their specific concerns below.

Scope of the analysis. Our theoretical results are limited to the RHM, which is a form of PCFG. As suggested, we will modify the writing to stress this point in the abstract, introduction, and conclusion.

Modeling choices and random production rules. We agree that a more detailed discussion of the RHM's applicability will improve the paper. The RHM is designed as a minimal yet analytically tractable framework that captures some phenomenology of real data as highlighted in previous work. Its simplifying assumptions - such as fixed tree topology and random production rules - allow us to compute all relevant correlations and derive theoretical predictions, which can then be empirically tested. In a more realistic setting with non-random rules, these correlations would be different, but our conclusions, formulated in terms of such correlations, would hold. Importantly, our experiments on OpenWebText, which contains richer syntactic structure, suggest that insights from the RHM, such as the fact that generated sentences achieve progressively larger coherence lengths as learning proceeds, extend to more realistic domains. We will revise the manuscript to articulate better both the strengths and limitations of the RHM, and to outline promising directions for future work, including more expressive models with context dependences or varying topology.

Relation to previous work. We have added a discussion of these papers to the related work section. Park et al. (2024) study compositional generalization in concept space for vision diffusion models on a toy dataset, where data are shallow compositions of several features (shape, size, color, etc.). Li et al. (2024) consider hierarchical Gaussian mixtures and study when diffusion sampling localizes on a sub-mixture. Neither of these two works study sample complexity. Allen-Zhu and Li (2020) show that deep networks trained with a variant of SGD can learn a class of multivariate polynomials with a time/sample complexity polynomial in the dimension. None of these papers considered synthetic data models based on trees.

Limits of validity of the theory. Thanks for the feedback. We will emphasize the set of assumptions made in Section 4.

Answers to questions:

1. How realistic is the RHM? Please see the paragraph above on modeling choices and the answer to Reviewer SjY7.

2. Alternative models. We believe the local-then-global coherence predicted by our theory is not specific to the RHM, but rather reflects a broader phenomenon that would persist for more general context-free grammars. We expect our results to hold if some limited amount of context-dependence is present (corresponding to a graphical model that is not perfectly tree-like). While a complete characterization of which classes of graphical models exhibit this behavior is beyond the scope of the present work, we see this as an important direction for future research. Importantly, our findings stand in contrast to those predicted by simpler models that rely only on second-order statistics - such as Gaussian models, which dominate much of the theoretical literature in machine learning. In those settings, spectral bias theory predicts that modes are learned in order of decreasing variance, meaning that low-frequency (global) components are acquired first (e.g., Wang, 2025). This would lead to an opposite phenomenology, where global structure appears early in training, and local detail is learned later. The fact that the pattern predicted by our theory is also observed in real data (e.g., in Fig. 4-5) adds empirical support to the modeling assumptions behind the RHM and the broader relevance of our theoretical framework. We will clarify this distinction and its significance in the revised manuscript.

Wang, B., 2025. An Analytical Theory of Power Law Spectral Bias in the Learning Dynamics of Diffusion Models. arXiv preprint.

3. Autoregressive models. Studying the sample complexity required by autoregressive models to learn data generated by a hierarchical grammar such as the RHM is a compelling direction. We view this as an important next step in understanding generative language models and plan to explore it in future work.

We thank the reviewer for their feedback and address their specific concerns below.

Role of branching factor $s$. The number of samples required to learn the full hierarchy in the absence of weight sharing is $m^{L+1}$, which is independent of $s$. In the presence of weight sharing, we expect an additional prefactor $1/d=s^{-L}$. Expressing the sample complexity as a polynomial of the dimension, one gets respectively in the two cases $md^{\alpha}$ and $md^{\alpha-1}$ with $\alpha = \log m / \log s$. When $\alpha$ is large, both sample complexities are dominated by the same effect. We will clarify this point in the manuscript.

Network's capacity. The generalization ability of the architecture is not due to a lack of capacity: the experiments are always run in the heavily overparametrized regime, and the networks eventually memorize their training data at the end of training. For this reason, we perform early stopping and avoid overfitting. We believe that not matching the network to the RHM structure does not change the dominant factor in the sample complexity as long as the architecture is deep enough. This is the case for both classification and last-token prediction with the RHM (Cagnetta et al., 2024; Cagnetta & Wyart, 2024), where fully connected networks and transformers (agnostic to the RHM structure) show the same dominant factor (i.e., polynomial in the dimension) in the sample complexity as CNNs architectures matched to the data structure. We will clarify this point in the paper.

How realistic is the RHM? We think of the RHM as the simplest yet tractable model that captures hierarchical compositionality, which is a property shared by many real-world data modalities. Making the model more realistic for specific datasets would involve introducing variable tree topologies, more complex non-random production rules, and possibly some context dependence - but all this would come at the cost of theoretical tractability. Despite its abstract form, the RHM captures key behaviors observed in real data, such as the staged learning and increasing-range correlations. These results cannot be obtained using simpler models (e.g., Gaussian random fields with long-range spatial correlations). This alignment between theory and empirics underscores the value of the RHM to study images and text.

Synonyms. We thank the reviewer for pointing this out. We have added to point iii) (L148) the following sentence: "We call the $m$ strings produced by any given symbol synonyms".

Compositionality. In our work, we adopt the formal notion of compositionality from theoretical linguistics (e.g., Montague, 1970; Partee, 2008), where it refers to the principle that the meaning of a complex sentence is determined by the meanings of its parts and the rules used to combine them. This principle underlies the ability to understand and generate an unbounded number of novel sentences by knowing the meanings of words and the grammar rules that govern their composition (Chomsky, 1957). In our setting, we interpret compositionality as the ability of diffusion models to represent and learn such structured, rule-based combinations of components. The RHM provides an idealized but expressive formalism for studying this phenomenon: it defines hierarchical, symbolic structures whose generative rules allow for compositional generalization. Regarding hierarchical structure specifically, it mirrors the hierarchical organization of syntax in language and the layered processing observed in biological vision systems. We will clarify our usage of the term "compositional".

Montague, R., 1970. Universal grammar.

Partee, B.H., 2008. Compositionality in formal semantics: Selected papers. John Wiley & Sons.

Chomsky, N., 1957. Syntactic Structures. De Gruyter Mouton.

We thank the reviewer for their feedback and address their specific concerns below.

Paper claims

1. Creativity. By creativity, we refer to the well-established phenomenon in linguistics (e.g., Chomsky 1976) in which a finite set of rules can be learned by an infant so as to generate an infinite number of novel and well-formed sentences. This technical notion of creativity aligns with the behavior we observe in our models: the U-Net demonstrates compositional generalization on the RHM by generating novel, structurally valid samples that were not seen during training. We will revise the manuscript by recalling this definition.

2. Hierarchical learning. We respectfully disagree with the reviewer's assessment of the novelty of our first contribution. The works by Sclocchi et al. and Li et al. (which we will cite) involve no training and focus exclusively on the sampling behavior of diffusion models that have already been pre-trained up to their best performance. In contrast, our work is the first to theoretically describe the hierarchical learning process of diffusion models: how they learn to compose low-level elements together, with how many data, and how generated data depend on training set size. These aspects are entirely novel with respect to the cited works and constitute the key contributions of our paper, as we will emphasize further.

3. Compositional generalization in real data. In our view, the compositional abilities of diffusion models have already been established in various domains. Language diffusion models, for instance, are able to create new sentences using previously seen words. Regarding images, for example, the work of Sclocchi et al. (2025) demonstrates compositional abilities using forward-backward experiments: by adding a small amount of noise and then denoising an image, some of the low-level features are combined in a different way (e.g., a dog might open its mouth and bend its ears). Once again, our goal is not to prove that diffusion models compose but instead to study how this skill is gradually learned as models are trained on more and more data.

Given these points, we kindly invite the reviewer to reconsider their score.

Figure 1. We believe there may be a misunderstanding: our theory does indeed predict the cascaded learning behavior across hierarchical levels. Specifically, it shows that rules at level $\ell$ require a number of samples scaling as $m^{\ell+1}$ to be learned. This prediction is precisely what we observe in Figure 1. In the inset of the figure, we explicitly rescale the horizontal axis by the theoretical prediction for each level, resulting in a collapse of the curves across layers and highlighting an excellent agreement between theory and experiments.

Other comments

1. We thank the reviewer for pointing out the typo.

2. Discrete diffusion. We have added references to Hoogeboom et al. (2021) and Austin et al. (2021) for Equation 2. As suggested, we will also update the notation to follow Austin et al., using $Q$ for the forward transition matrix in the discrete setting.

Hoogeboom et al., 2021. Argmax flows and multinomial diffusion: Learning categorical distributions. NeurIPS 34.

Austin et al., 2021. Structured denoising diffusion models in discrete state-spaces. NeurIPS 34.

3. Synonyms. We thank the reviewer for the question. We have added to point iii) (L148) the following sentence: "We call the $m$ strings produced by any given symbol synonyms". The value of $m$ controls how many valid strings can be generated by the model. Since the number of produced strings is $v m$ ($m$ for each of the $v$ values of the parent node) and the maximum number of distinguishable strings is $v^s$, when $m=v^{s-1}$ all strings of size $s$ can be produced at any level, implying that all possible $v^d$ input strings are generated, and the data distribution has little structure. Instead, when $m<v^{s-1}$, only a small fraction of all possible strings is generated by the production rules, and spatial correlations - enabling self-supervised learning - appear.

4. Low-noise limit of the score. This argument is presented in Section 4.1, in the paragraph "score function at low noise". There, we show that in the low-noise regime, the score function reduces to the conditional expectation of a token given the rest of the sequence, which is proportional to its correlations with the remaining input. As discussed, these correlations are invariant under the exchange of synonyms.

We thank the reviewer for their feedback and address their specific concerns below.

Models of images. We agree that the relation between pixel and latent space in images is subtle. While images are continuous, they admit high-level abstract representations that can be described with discrete hierarchies. This idea was formalized in pattern theory, where scenes are parsed into objects and parts. The RHM describes these semantic aspects of images, in contrast to approaches emphasizing more geometric aspects. We also agree that the hierarchical level of a latent does not map perfectly to spatial locality, as in the RHM, e.g., low-level visual features - like a large patch of uniform color - can exhibit long-range spatial dependencies. This motivates our use of CNN representations, encoding image semantics, rather than raw pixels in Sec. 5.2.

Shuffled RHM. With a CNN U-Net using non-overlapping filters, shuffling the RHM input is expected to significantly increase sample complexity because the model relies on spatial locality to infer the latent structure, and a shuffled input disrupts this property. In contrast, a fully connected network, invariant to input permutations, can overcome this issue. In that case, we expect a slight drop in performance compared to our results but still a polynomial scaling in the dimension, as observed in Cagnetta et al. (2024) for classification.

Clustering. We will add clustering details to the paper. Briefly, for a given (visible or latent) patch $h$, we fix it to one of its possible values $\nu$ and compute its mean context vector by averaging the one-hot-encoded tokens observed in the nearest neighboring visible patch $x$. In other words, we compute the empirical conditional expectations $v_{\nu}=E[x|h=\nu]$ for each of the possible values $\nu$, which are proportional to the correlations discussed in Sec. 4. We apply k-means clustering to these vectors. When the training set size is sufficiently large, synonymous patches $\nu$ will produce similar mean contexts and be clustered together.

Previous work. While Kadkhodaie et al. focus on geometric inductive biases, we focus on the compositional and hierarchical structure of data. Our framework is complementary: instead of framing sample efficiency in terms of how networks capture geometric structure, our results show that compositional generalization can be achieved with a limited amount of data. We will incorporate this reference and clarify our scope. Similarly, the work of Kamb & Ganguli, contemporaneous and already discussed, studies locality and translational equivariance in trained models, but not sample complexity and finite-size effects, which are central to our work.

Hallucinations. We thank the reviewer for pointing out these works. The findings that diffusion models have a local inductive bias and fine-grained rules have weaker signals in the loss are aligned with ours. We will cite these works as practical examples that can be quantitatively studied within our model.

Practical takeways. Fundamental research has intrinsic value: a deeper understanding of how diffusion generalizes is important for building a foundation that could lead to long-term practical progress. Moreover, our findings already suggest practical paths. Our results show that overparameterized models - having the capacity to fully memorize their data, resulting in issues such as copyright infringement - have an inductive bias for reaching generalizing solutions before memorization. Trying the early stopping used in our synthetic experiments in real settings is thus an interesting future direction. Moreover, we propose new tools, e.g., MMD across depths or correlation analysis, that could serve as metrics for evaluating model quality beyond standard ones. We will clarify the broader significance of the paper.

Unambiguity. While natural grammars have some ambiguity, we assume unambiguous rules to simplify the analysis. Yet, our theory indicates that limited ambiguity should not significantly affect results, as long as correlations among features exist. Extending the RHM to ambiguous rules and computing how correlations are affected is an interesting future direction. Note this assumption is not related to compositional generalization. In the example of objects with varying features, they can be thought of as different classes of the RHM data, i.e., different high-level concepts that can be composed with common low-level features. The example of generating two objects could be modeled by adding a production rule at the top, deciding which objects will appear on the left and right.

Synonymic invariance. Collapsing distinctions between synonyms is key to lower dimensionality and denoise high-level latents. In the U-net, it is achieved in the internal layers. Yet, residual connections preserve low-level information to ultimately denoise the data at the finest level.