Information-Guided Noise Allocation for Efficient Diffusion Training

Abstract: Training diffusion models typically relies on manually tuned noise schedules, which can waste computation on weakly informative noise regions and limit transfer across datasets, resolutions, and representations. We revisit noise schedule allocation through an information-theoretic lens and propose the conditional entropy rate of the forward process as a theoretically grounded, data-dependent diagnostic for identifying suboptimal noise-level allocation in existing schedules. Based on these insight, we introduce InfoNoise, a principled data-adaptive training noise schedule that replaces heuristic schedule design with an information-guided noise sampling distribution derived from entropy-reduction rates estimated from denoising losses already computed during training. Across natural-image benchmarks, InfoNoise matches or surpasses tuned EDM-style schedules, in some cases with a substantial training speedup (about $1.4\times$ on CIFAR-10). On discrete datasets, where standard image-tuned schedules exhibit significant mismatch, it reaches superior quality in up to $3\times$ fewer training steps. Overall, InfoNoise makes noise scheduling data-adaptive, reducing the need for per-dataset schedule design as diffusion models expand across domains.

• The paper introduces InfoNoise, which reallocates training updates using the conditional entropy rate to focus on the most informative noise window.
• It employs online MMSE estimation from denoising losses to dynamically adjust noise sampling, enhancing efficiency across various domains.
• Empirical results demonstrate up to 2.8x reduction in training steps for target FID levels, highlighting benefits in both training and inference.

Information-Guided Noise Allocation for Efficient Diffusion Training

Motivation and Problem Setting

The training of diffusion models, particularly in the variance-exploding (VE) Gaussian setting, traditionally relies on hand-crafted or empirically tuned noise schedules. These schedules determine the allocation of learning signal along the noise (corruption) path but frequently misallocate computational effort when shifted across datasets, resolutions, or data representations. Such misallocation manifests as over-sampling noise levels with minimal information or under-sampling those with highest learning leverage, resulting in suboptimal convergence and poor cross-setting transferability. This paper introduces a principled framework for adaptive, data-driven noise scheduling based on information-theoretic diagnostics, with a focus on the conditional entropy rate of the forward process as a direct measure of denoising informativeness.

Theoretical Foundations: Conditional Entropy Rate and I–MMSE

A central observation is that the reduction of uncertainty about clean samples $x_0$ as noise is removed is highly inhomogeneous along $\sigma$ (the noise scale): most of the informative learning signal is concentrated in a narrow "informative window" at intermediate noise levels. Formalizing this, the paper leverages the I–MMSE identity for Gaussian channels, which connects the conditional entropy $H[x_0 \mid x_\sigma]$ to the Bayes-optimal denoising mean-squared error (MMSE):

$$\frac{d}{d\sigma} H[x_0 \mid x_\sigma] = \frac{\mathrm{MMSE}(\sigma)}{\sigma^3}.$$

This relation provides an information gradient (entropy-rate profile) along the noise path. The optimal strategy then becomes: allocate training updates proportionally to the local entropy-rate along the path, directly targeting regions where the model can most efficiently reduce uncertainty.

Figure 1: Uncertainty about $x_0$ (the clean data) collapses rapidly in an intermediate range of noise, delineating a “decision window” where the posterior transitions from ambiguous to confident.

InfoNoise: Adaptive Schedule Construction

Building on the above, the paper introduces InfoNoise, an online scheduling algorithm that periodically estimates the entropy-rate profile during training. The methodology is as follows:

• Online Estimation: MMSE is estimated directly from the denoising losses already produced during SGD, avoiding additional forward passes.
• Normalization and Regularization: The raw entropy-rate is regularized to avoid degenerate concentration near the low-noise boundary and then normalized to form a probability density $\rho(\sigma)$.
• Sampling Schedule: The sampling density $\pi(\sigma)$ is set such that the effective emphasis (the product of schedule and objective-specific loss weighting) matches the normalized entropy-rate, thus shifting computational mass into the most informative region.

  Figure 2: Construction of the InfoNoise schedule via online estimation and normalization of the entropy-rate, and realization of the corresponding training sampler.

A warm-up phase ensures stable statistics collection, and the approach is agnostic to the underlying diffusion architecture and loss, making InfoNoise a drop-in replacement for previous heuristics.

Experimental Evaluation

Transferability to Discrete and Non-Image Domains

The paper demonstrates that heuristically tuned image schedules fail to transfer to discrete sequence domains, resulting in substantial misallocation of learning updates. InfoNoise, by contrast, dynamically shifts emphasis into the data-dependent informative window, leading to accelerated convergence and improved sample quality.

Figure 3: Fixed (EDM/Log-uniform) samplers misallocate training emphasis in discrete domains, whereas InfoNoise matches the diagnostic allocation and achieves higher efficiency.

Empirically, InfoNoise achieves up to $2.8\times$ reduction in training steps required to reach a fixed FID/Sei-FID quality metric on DNA and binarized image data.

Efficiency and Quality on Standard Image Benchmarks

On classical continuous image domains (e.g., CIFAR-10, FFHQ-64), InfoNoise matches or slightly outperforms the best hand-tuned fixed schedules (such as EDM log-normal). Notably, InfoNoise achieves $1.4-1.5\times$ faster convergence to target FID levels on CIFAR-10 (both unconditional and conditional) (Figure 4, Figure 5), demonstrating the practical benefit of information-adaptive allocation even in mature pipelines.

Figure 4: On CIFAR-10, InfoNoise aligns its training emphasis with that of optimized EDM schedules, but achieves target FID faster via adaptive reallocation.

Figure 5: In conditional CIFAR-10, InfoNoise reduces the number of processed examples needed to achieve EDM-level FID.

Online Emergence of Informative Window

InfoNoise’s adaptive profile rapidly stabilizes early in training and closely matches offline diagnostics computed from converged checkpoints, indicating that the informative band is reliably detectable from partial training dynamics.

Figure 6: Online emergence and stabilization of the InfoNoise schedule during training on a discrete DNA domain.

Inference-Time Utility: Information-Guided Discretization

Beyond training, the entropy-rate profile learned by InfoNoise can inform inference-time discretization of $\sigma$, allocating solver steps where uncertainty is resolved—resulting in more uniform progress in information space and improved final sample fidelity at fixed computational cost.

Figure 7: Qualitative class-conditional CIFAR-10 and FFHQ samples show that InfoNoise yields cleaner samples versus EDM baselines under matched NFE and grid.

Implications and Theoretical Significance

By re-framing noise scheduling as an information-allocation problem, this work solidifies the connection between data-dependent uncertainty reduction and the allocation of learning updates in diffusion training. The results offer several key implications:

• Reduced Reliance on Heuristics: InfoNoise makes per-dataset schedule tuning unnecessary, as the method dynamically adapts to the intrinsic information geometry of the corruption process and data distribution.
• Cross-Domain Transfer: The approach is robust across modality (continuous images, discrete sequences) and representation, eliminating the empirical failures of rigid, manually tuned schedules.
• Enhanced Theoretical Understanding: By revealing that the hallmark "intermediate noise window" is a necessary consequence of the information flow in Gaussian channels, it provides a normative underpinning to many successful ad hoc scheduling heuristics developed over the past years.
• Future Directions in Diffusion: Opens the prospect for schedule optimality with respect to compute budgets and objectives, as well as generalizations to non-Gaussian or more exotic corruption processes. Extending the method's theoretical guarantees to non-Gaussian settings remains a relevant open problem.

Conclusion

InfoNoise provides a theoretically principled, practical, and efficient solution to the longstanding problem of noise schedule design in diffusion model training. By leveraging the I–MMSE connection and using online statistics for entropy-rate estimation, it adaptively reallocates training updates to the most informative noise regimes. Empirical results in both continuous and discrete domains confirm that InfoNoise achieves superior training efficiency and obviates brittle manual tuning. The methodology also offers co-benefits for inference-time sampling strategies, strengthening its relevance for the ongoing expansion of diffusion models across modalities and data types.

Figure 8: InfoNoise-derived entropy-rate profiles across multiple datasets highlight the data-dependence and representation-specific structure of the informative window.