How Compositional Generalization and Creativity Improve as Diffusion Models are Trained

TL;DR

This paper explores how diffusion models learn hierarchical compositional rules, showing that their ability to generate coherent data improves with more training data and time, supported by theoretical and empirical analysis.

Contribution

It introduces a theoretical framework linking diffusion models' learning process to hierarchical clustering of features, revealing how sample complexity scales with feature context size.

Findings

Diffusion models learn composition rules with sample complexity similar to word2vec.

Higher-level features require more data to identify, leading to hierarchical learning.

Generated data coherence increases with training time and dataset size.

Abstract

Natural data is often organized as a hierarchical composition of features. How many samples do generative models need in order to learn the composition rules, so as to produce a combinatorially large number of novel data? What signal in the data is exploited to learn those rules? We investigate these questions in the context of diffusion models both theoretically and empirically. Theoretically, we consider a simple probabilistic context-free grammar - a tree-like graphical model used to represent the hierarchical and compositional structure of data such as language and images. We demonstrate that diffusion models learn the grammar's composition rules with the sample complexity required for clustering features with statistically similar context, a process similar to the word2vec algorithm. However, this clustering emerges hierarchically: higher-level features associated with longer contexts require more data to be identified. This mechanism leads to a sample complexity that scales polynomially with the said context size. As a result, diffusion models trained on an intermediate dataset size generate data coherent up to a certain scale, but lacking global coherence. We test these predictions across different domains and find remarkable agreement: both generated texts and images achieve progressively larger coherence lengths as the training time or dataset size grows. We discuss connections between the hierarchical clustering mechanism we introduce here and the renormalization group in physics.