What Makes a Good Toy Example for Diffusion Models?

Toys help us see diffusion algorithms at the level of distributions—but the moment we design a toy, we may also (without noticing) choose a side by baking in preferences. code: https://github.com/851695e35/Leaves_Toy

Why We Build Toys

Diffusion and flow matching models learn and generate data by gradually transforming noise into samples, many of their most important behaviors are easiest to understand at the level of distributions. From this view, we can directly see how probability mass moves and what structure is kept or lost along the way.

This is why toy examples appear so often in diffusion research. They give us a small, visual, and controllable setting where training and sampling behavior can be observed directly, without the extra complexity of high-dimensional image data or heavy evaluation pipelines. In a good toy problem, effects that are hard to notice in realistic datasets can become clear at a glance, which makes toy examples useful for introducing ideas, testing intuitions, and showing empirical phenomena in a direct way.

But this simplicity comes with a price: what we conclude from a toy example might depend on how the toy is set up, and those choices can quietly encode preferences.

When Simple Toys Mislead

Consider a simple toy distribution: two-component Gaussian mixture, with Class A centered at $(-2, 0)$ and Class B centered at $(2, 0)$. We train a toy diffusion model conditioned on the class label. In the visualization below, you can see:

• Light green dots: Class A ground truth samples
• Dark green dots: Class B ground truth samples
• Orange dots: Generated samples for Class B using classifier-free guidance (CFG)

Notice the effect of increasing the guidance scale. When $w = 1.0$, that is, without additional guidance beyond the conditional prediction, the orange samples closely match the dark blue Class B distribution. As $w$ increases beyond $1.0$, the behavior changes. The orange samples begin to drift away from the dark green target, and the deviation grows with the guidance strength. At first glance, this seems at odds with practical image generation, where much larger guidance scales are often preferred, such as $w=7.5$ in Stable Diffusion .

Why does this happen?

To understand this phenomenon, let’s start with the ideal case. In flow matching (given the equivalence of flow matching and diffusion, we use the former for brevity.), assume an Optimal Transport probability path:

\[\mathbf{x}_t = (1-t)\mathbf{x}_0 + t\boldsymbol{\epsilon},\qquad \mathbf{x}_0 \sim p_{\text{data}}(\cdot \mid c), \quad \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})\]

Under this path, the optimal conditional velocity field is the conditional expectation of the endpoint displacement given the intermediate state:

\[v^*(\mathbf{x}_t, t, c) = \mathbb{E}_{\mathbf{x}_0 \sim p(\cdot | c), \boldsymbol{\epsilon} \sim \mathcal{N}(0, \mathbf{I})} [\boldsymbol{\epsilon} - \mathbf{x}_0 \mid \mathbf{x}_t, t]\]

Which is analytically:

\[v^*(\mathbf{x}_t, t, c)=\frac{\sum_{i \in \mathcal{N}_c}(\mathbf{x}_t-\mathbf{x}_0^i)\,\mathcal{N}\!\left(\mathbf{x}_t;\,(1-t)\mathbf{x}_0^i,\;t^2\mathbf{I}\right)}{t\sum_{i \in \mathcal{N}_c}\mathcal{N}\!\left(\mathbf{x}_t;\,(1-t)\mathbf{x}_0^i,\;t^2\mathbf{I}\right)},\qquad \mathcal{N}_c=\{\,i\mid \mathbf{x}_0^i \in c\,\}.\]

If the model learns this optimal field, sampling follows the ground true marginal distributions and can reconstruct the data points. In our two-dimensional toy setting, this is not very hard. The space is low-dimensional, the structure is simple, and training can get very close to the optimum. As a result, when $w=1.0$ (i.e. no guidance), the learned conditional model already matches the ground-truth distribution well.

Real images and videos are very different. Their distributions are much more complex, so the optimal velocity field is much harder to learn. The predicted field is therefore only approximate, and these small errors can build up during sampling. Over many steps, the generated trajectory may drift away from the true conditional path.

Now consider classifier-free guidance, written in velocity form as

\[v_{\mathrm{cfg}}(\mathbf{x}_t, t, c)=v(\mathbf{x}_t, t)+w\bigl(v(\mathbf{x}_t, t, c)-v(\mathbf{x}_t, t)\bigr).\]

In 2-Gaussian toy setting, this extra amplification is not needed. The conditional model already separates the classes well, so increasing $w$ adds little useful information. Instead, it over-pushes the velocity field and moves samples away from the true target distribution.

AutoGuidance Toy: Richer Structure, Specific Context

AutoGuidance (AG) takes toy experiments a step further by constructing a richer example with more complex in-class distributions. Instead of simple unimodal Gaussians, each class contains multiple internal modes, giving the toy enough structure to reveal differences in how guidance strategies handle within-class diversity.

Figure from AutoGuidance: the toy experiment comparing CFG and AG on multi-modal in-class distributions.

In this setting, CFG tends to concentrate samples in the high-density region of each class, reducing mode coverage. AG, by contrast, preserves coverage within classes. This pattern is consistent with the empirical finding of AutoGuidance on ImageNet, where AG outperforms CFG in generation quality.

However, other studies have found that CFG remains more robust, and in some cases better than AG, in certain settings. Our interpretation is that the difference arises from context. AG has a clear advantage when the underlying conditional model is already well trained, such as EDM2 on ImageNet. In that context, class information is already captured reliably, so the main remaining issue is mode coverage in the class, which is precisely where AG helps.

CFG performance across different models on ImageNet512.

Model	FID (w/o CFG)	FID (w CFG)	CG/CFG scale w
ADM-U	9.96	3.35	4.0
DiT-XL	12.03	3.04	1.5
EDM2-S	2.56	2.23	1.4
EDM2-XXL	1.91	1.81	1.2

This is also the context encoded by the toy design. The example contains only two classes (A and B), and each class is modeled in a setting that is relatively easy to fit. Under these conditions, CFG mainly exaggerates separation between 2 classes and can therefore degrade the learned distribution. The key assumption is thus not in AG itself, but in the experimental setup: the toy implicitly assumes a near-well-fitted unguided conditional model. In jointly trained conditional models with shared parameters and limited capacity, conditional fields may overlap or blur, and in that context the stronger contrast induced by CFG can serve as a useful separation mechanism. From this perspective, AutoGuidance toy is better viewed as a richer but still context-specific benchmark: it introduces meaningful within-class structure, but under assumptions that unguided conditional model is well fitted.

A Multi-Class Toy: Leaves

So far, we have moved from a two-Gaussian mixture to the more structured multi-modal toy used in AutoGuidance. Each step makes the toy more realistic, but also brings in new assumptions, such as a small number of classes, well-fitted models, or a fixed context. To study behaviors like guidance in diffusion model more systematically, we need a toy that can vary the main sources of difficulty in a controlled way.

Two factors matter in here. The first is class-condition complexity, which describes how many classes the model must separate. The second is within-class complexity, which describes how much internal structure each class contains. Real datasets vary along both axes at once. ImageNet, for example, contains many classes, and each class also has large variation in pose, background, and subtype. Text-to-image settings go further by introducing additional modalities and richer conditions. A useful toy should let us adjust these factors separately and study how they affect guidance behavior.

This is the motivation behind the Leaves dataset, a parameterized family of toy distributions designed for diagnosis rather than realism. It is controlled by two parameters, num_classes and max_depth. Each class grows from a shared trunk into a branching structure. Increasing max_depth adds recursive splits and makes the structure within each class more complex. Increasing num_classes adds more branches and makes class separation harder. We denote each setting as num_classes/max_depth. For example, 4/3 means four classes with depth-three branching, while 24/1 means twenty-four classes with minimal internal structure.

These two parameters mainly control between-class crowding and within-class branching. At the same time, like any designed toy, they also affect other properties such as geometry, support overlap, local curvature, and optimization difficulty. We do not fully isolate these secondary effects, and we acknowledge them as part of the limitation of the setup.

4 classes / depth 3

12 classes / depth 2

24 classes / depth 1

Ground-truth distributions for three Leaves configurations. From left to right: fewer classes with deeper branching (richer in-class structure) to many classes with shallow structure (harder class separation, simpler in-class modes).

The above three configurations illustrate the trade-off. In 4/3, each class forms a deep recursive tree, so the main challenge is to model rich within-class structure while separating only four classes. In 24/1, each class has simple geometry, but the number of classes is much larger, so the main challenge shifts to between-class discrimination. The 12/2 setting lies between these extremes. By sweeping across this space, we can ask when CFG helps, when AG performs better, and when both become inadequate.

One minor distinction to AG is that we train a single conditional model over all classes jointly, which is closer to standard diffusion training in practice. This means the model shares capacity across classes, and guidance is studied in a more crowded label space rather than in a setting where each class is handled separately.

Ground Truth Unguided CFG AG SGG

4 classes / depth 3, T=2^15

12 classes / depth 2, T=2^15

24 classes / depth 1, T=2^12

Generated samples (black dots) overlaid on ground-truth distributions across different guidance strategies for each Leaves configuration. (class numbers, recursive depth and training iteration (T))

4/3: Preserving within-class structure

In the 4/3 T=2^15 regime, the main challenge is to preserve fine within-class structure rather than to separate classes. Because there are only four classes but each class has deep branching, the model must capture detailed conditional geometry. In this setting, unguided sampling gives reasonable class separation but misses some branch detail. CFG tends to over-amplify the class signal and pushes samples toward the main trunk, which reduces structural fidelity. AG performs better, it refines the conditional prediction more locally, helping preserve the branching structure without introducing unnecessary outliers.

24/1: Separating many classes

In the 24/1 T=2^12 regime, the situation is reversed. Each class has simple internal structure, but there are many classes packed into the same space, so the main difficulty is between-class separation. When the model is less underfitted $T=2^{12}$, unguided sampling often shows class confusion, with samples leaking across nearby class boundaries. CFG is more effective because its stronger conditional contrast sharpens class separation and helps keep samples aligned with the intended class. AG is less helpful in this context, since refinement offers limited benefit when the main problem is not subtle geometry but coarse discrimination among many classes.

Context-dependent behavior

Taken together, Leaves suggests that guidance methods work best in different context. CFG is most useful when the dominant challenge is separating many classes, while AG is more useful when the main challenge is preserving rich within-class structure. The intermediate 12/2 case lies between these extremes, and both methods show mixed behavior there.

We are Still Limited

Leaves makes context dependence visible, but it remains a two-dimensional distribution. Once the question moves beyond coordinate-space behavior to issues such as representation geometry, semantic abstraction, or feature collapse, a toy of this kind can no longer provide a faithful answer.

Beyond coordinates: where 2D toys fail

One can find such scenarios in training improvements. The following case studies mark the limit of what coordinate-space toys can represent. We compare three training variants on the Leaves dataset to illustrate phenomena that depend on high-dimensional representation geometry.

• Vanilla flow matching: the standard objective.

\[\mathcal{L}_{\text{vanilla}}=\mathbb{E}\left[\sigma^2 \left\| s_{\theta}(\mathbf{x},\sigma,c)-(\boldsymbol{\epsilon} - \mathbf{x}_0)\right\|_2^2\right]\]

• + Dispersive loss : adds a regularizer that encourages the model to spread its predictions, penalizing collapsed covariance in the denoised estimates.

\[\mathcal{L}_{\text{dispersive}}=\mathcal{L}_{\text{vanilla}}+\lambda_{\text{disp}}\cdot\frac{1}{\mathrm{Tr}\left(\mathrm{Cov}(\hat{\mathbf{x}})\right)+\varepsilon},\quad\hat{\mathbf{x}}=\mathbf{x}+\sigma^2 s_{\theta}(\mathbf{x},\sigma,c)\]

• + Deep supervision : provides auxiliary supervision at intermediate dit blocks using the same stochastic objective used in vanilla training.

\[\mathcal{L}_{\text{deep-supervision}}=\mathcal{L}_{\text{vanilla}}+\lambda_{\text{ds}}\cdot\mathbb{E}\left[\sigma^2 \left\| s_{\theta}^{\text{aux}}(\mathbf{x},\sigma,c)-(\boldsymbol{\epsilon} - \mathbf{x}_0)\right\|_2^2\right]\]

Wasserstein distance during training for three loss variants on the Leaves dataset.

The results show that some regularizers, such as deep supervision, can help in this low-dimensional setting, whereas others, such as dispersive loss, may have little effect or become unstable. The key point, however, is that methods such as dispersive loss and REPA are designed to shape representation geometry in high-dimensional latent spaces. In a 2D toy, such mechanism is largely absent. The model maps directly from coordinates to scores, leaving no meaningful feature space in which collapse, alignment, or diversification can occur.

This is not a weakness of the methods, but a limitation of the toy setting itself. Toy distributions can reveal distribution-level effects, such as mode coverage, class separation, and guidance drift, but they cannot faithfully capture perceptual or representation-level phenomena that depend on high-dimensional structure.

Inevitable preferences in toy construction

A second limitation is more subtle: every toy dataset requires many design choices, and these choices inevitably encode assumptions about what should count as realistic or important.

Take leaves for example, the construction depends on several decisions, including the number of classes, the branching depth, the variance at each node, the angular spread of branches, and the sampling density. Each choice changes the resulting distribution and can affect which guidance method appears more effective. A setting with many classes and an underfit model may favor CFG, whereas a setting with richer within-class structure and a well-fitted model may favor AG. As a result, outcomes are shaped not only by the method, but also by the preferences built into the toy design.

Other toy settings can carry similar assumptions. The two-Gaussian mixture assumes a simple distribution with clearly separated classes, while the AutoGuidance toy assumes a class-conditional model that is already well fitted. These choices are not neutral, they may influence the context before comparison begins.

Closing Thoughts

Perspective shapes both how we see the world and how we model it. Diffusion models are powerful generative tools that extend across images, videos, proteins, language, and simple toy settings, yet each domain requires its own adaptation and reveals its own phenomena. The assumptions and preconditions we adopt in defining algorithms and benchmarks may carry implicit or explicit human bias. For that reason, we view toy analysis as a way to study a given context rather than answer the full question. Its value lies in presenting the setup, the assumptions, and the behavior one observe.