Feature Information Dynamics in Diffusion

AI summary

Overview

• Research area: Information theory applied to generative diffusion models (stat.ML), specifically the interplay between mutual information, minimum mean-squared error (MMSE), and the denoising trajectory.
• Technical level: Advanced. The paper builds on the I-MMSE identity, Gaussian channels, MMSE curves, and flow-matching parameterizations; readers need prior familiarity with diffusion models and information theory to follow the derivations.
• Scope: The paper introduces "feature information dynamics," a framework that localizes when along the diffusion noise schedule a given feature (frequency band, class label, segmentation mask, Canny edges) is generated, and applies it to compare four image representations.

What This Paper Is About

Diffusion models are widely observed to generate coarse structure before fine detail, but this intuition has been mostly empirical and qualitative, and it does not transfer cleanly to models operating in latent representations rather than pixels. The authors build an information-theoretic account of when a feature's information enters the denoising trajectory, expressed as the gap between optimal unconditional and feature-conditional denoising losses. They then use it to test spectral autoregression in pixel diffusion, and to compare how pixel space, SDVAE, VAVAE, and RAE order semantic, structural, and textural information under an identical conditioning chain.

Key Contributions

1. Feature information dynamics via I-MMSE. Starting from the classical I-MMSE identity, the authors define the SNR-resolved information between a specified feature and the noisy image, and relate its derivative to the gap between optimal unconditional and feature-conditional denoising losses. They derive a family of equivalent estimators for this "feature information density" (single-feature MMSE gap, denoiser gap, score gap, and further posterior-mean, Fisher, and likelihood-score forms) and give a practical procedure for obtaining the information profile of a general feature, compared in detail in Appendix A.
2. Chained information decomposition. For a hierarchy of features, computing each single-feature density separately double-counts shared information. The authors instead group features into nested cumulative bundles and define the chained information density of each level as its incremental information conditioned on the preceding bundle, verified experimentally on frequency bands and on a class → mask → Canny hierarchy.
3. Representations shape feature information dynamics. Applying the chained framework to a general visual feature hierarchy on ImageNet-256, with the conditioning chain and evaluation protocol held fixed, they show that information dynamics differ significantly across pixel space, SDVAE, VAVAE, and RAE, recasting representation choice as a controllable knob over feature information dynamics rather than a black box on top of FID.
4. A matched-recipe training-speed test. They run a representation-only comparison in which training noise shift, sampling noise shift, and learning rate are selected on a smaller LDiT-B model and then frozen across all four XL trajectories, isolating the representation as the remaining variable.

Main Findings

• Pixel diffusion approximates spectral autoregression. On a controlled MNIST frequency hierarchy, single-band conditioning produces heavily overlapping densities because shared low-frequency information is counted in every band, while forward chaining (conditioning on the cumulative bundle) removes already-explained information and reveals a clean coarse-to-fine progression along log-SNR. The same chained construction on CIFAR-10 shows a monotone spectral progression, with low-frequency bands concentrating information at lower log-SNR and higher-frequency bands shifting toward cleaner denoising states.
• Structure can enter before semantics in pixel space. In pixel-space diffusion, the mask density peaks at lower log-SNR than the class density, so mask information enters earlier and more strongly than class identity — challenging a simple semantic-first account of the denoising trajectory.
• Convergence ordering under a matched recipe. Best FID over 200 XL epochs: RAE 8.39 (epoch 180) < VAVAE 21.16 (epoch 200) < SDVAE 30.83 (epoch 200) < pixel 78.51 (epoch 200). The ordering pixel < SDVAE < VAVAE < RAE is consistent across the whole 200-epoch budget.
• Feature ordering tracks convergence speed. VAVAE recovers the semantic-first half (class peak at lower log-SNR than mask), but mask and Canny overlap rather than separating. Only RAE, the fastest-converging representation, achieves a strict class → mask → Canny ordering, with its three density peaks in textbook high → mid → low order along the log-SNR axis. The authors report this as a correlation and explicitly do not claim the ordering causes faster convergence.
• Order and size checks. With approximately 132M-parameter RAE-S checkpoints, peak log-SNRs stabilize at −3.75 (class), −2.75 (mask), −1.75 (Canny) for 60- and 80-epoch runs, matching the approximately 842M-parameter DiTDH-XL model (−3.75, −2.75, −1.75). Using the alternative mask → class → Canny order shifts the class and mask peaks to −4.00 and −3.25 but preserves their separation on RAE.
• Estimator caveats. At very high SNR, denoising errors and gradients are small, so a finite-capacity denoiser can underfit the clean-data limit; the authors fit the single-feature MMSE gap with a smooth double-pole Lorentzian family on a reliable interval and extrapolate the tail analytically.

Methodology in Plain English

The authors treat the forward diffusion process as a Gaussian channel that adds noise to clean data at a controllable signal-to-noise ratio. For any feature of interest (a frequency band, a class label, a mask, an edge map), they compare two denoisers: one that sees only the noisy image, and one that additionally sees the feature. The best achievable squared error in each case is the MMSE. Because the feature-conditioned denoiser has strictly more information available, its error can never be larger, and the difference between the two errors measures how much that feature helps at that noise level. A classical result (I-MMSE) says the derivative of mutual information with respect to SNR equals half the unconditional MMSE; applying the same argument conditionally yields the derivative of the feature's mutual information as half the MMSE gap. This "feature information density" plots against log-SNR shows where along the trajectory the feature's information is concentrated, and the location of its peak indicates when that feature is resolved. Because features share information, the authors chain the conditioning: each level is measured as its incremental contribution beyond all coarser levels, so the density of level k is half the SNR-weighted difference between the MMSE conditioned on the first k−1 levels and the MMSE conditioned on the first k. Diffusion models are instantiated as Gaussian flow matching with linear interpolation, so the standard training loss is the MMSE up to known linear relations. In practice, all MMSE terms are estimated by Monte Carlo using trained denoisers, with the previously mentioned Lorentzian tail fit to handle the high-SNR regime.

Why This Matters

The paper converts a qualitative folk intuition about coarse-to-fine generation into a measurable quantity that works for arbitrary features and arbitrary representations, which matters because latent representations are now the mainstream compromise in image generation and the choice of representation affects training efficiency.

• Representation selection before expensive training runs. Researchers could compare candidate representations by asking whether semantic, geometric, and local information separate along the SNR axis, rather than committing to a full diffusion training run first.
• Timestep and conditioning schedules. The location and width of density peaks identify the SNR ranges in which particular features are most active, and could inform candidate timestep schedules or conditioning orders.
• Conditional sampling and guidance. Samplers could apply feature- or SNR-dependent guidance strengths instead of a single guidance weight applied throughout the trajectory.
• Diagnosing entangled information. Comparing profiles across architectures could help identify where information becomes entangled or remains unresolved, motivating targeted representation or conditioning designs.

Industry relevance: the framework offers a cheap diagnostic signal for representation choice and schedule design in image generation pipelines, and suggests that ordered generation could be beneficial for training diffusion models. The authors frame the observed association between clean hierarchical ordering and convergence speed as suggestive, not causal.

Future Directions

• Determining why a particular representation makes diffusion training easier or harder under a matched optimization recipe, which the authors state is an open problem and which their training-speed test does not explain.
• Extending the analysis beyond additive Gaussian channels with independent coupling to settings with non-independent endpoint coupling, which is outside the present scope.
• Scaling the direct chained estimator to larger feature hierarchies, where additional cumulative conditional models become a practical bottleneck.
• Testing whether density profiles can serve as a representation-selection signal or a predictor of convergence speed, and whether feature- or SNR-dependent guidance and candidate timestep schedules derived from peak locations actually improve generation.

Target Audience

Readers with a background in information theory or diffusion modeling who want a quantitative, representation-agnostic account of when features are generated — particularly researchers working on representation learning for generative models, latent-diffusion design, timestep scheduling, and conditional or guided sampling. The empirical cross-representation comparison is also relevant to practitioners deciding between pixel-space and latent diffusion pipelines.

Authors’ abstract

Diffusion models generate data through a continuum of denoising problems, and are widely observed to reveal coarse structure before fine detail. Yet, this intuition is mostly empirical and qualitative. We introduce feature information dynamics, an information-theoretic framework for localizing when a feature is generated during diffusion. Using the I-MMSE identity, we connect the rate of feature mutual information change to a gap between optimal unconditional and feature-conditional denoising losses, yielding practical estimators for feature information density. We further develop a chained decomposition that separates shared from incremental information in a feature hierarchy. We use this framework first to quantitatively confirm spectral autoregression in pixel diffusion, and then to extend the analysis beyond frequency: under a class $\to$ mask $\to$ Canny conditioning chain, the per-feature information densities differ across pixel, SDVAE, VAVAE, and RAE, exposing fundamental differences between these representations and suggesting that ordered generation could be beneficial for training diffusion models. Our code is available at https://github.com/AI4Science-WestlakeU/feature-information-dynamics.