Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Revisiting Diffusion Model Predictions Through Dimensionality

Qing Jin, Chaoyang Wang

TL;DR

The paper tackles the question of how to choose the optimal diffusion prediction target as data dimensionality grows. It introduces a unified framework with a continuous target $ m{u}=km{x}-(1-k)m{n}$ and proves the optimal choice is $k^*=\frac{D}{D+d}$, linking target selection to intrinsic and ambient dimensions. To avoid explicit dimension estimation, it proposes k-Diff, a learnable target where $k=\mathrm{sigmoid}(w_k)$ (and optionally time-varying $k(t)$), enabling end-to-end optimization. Empirical results in both latent and pixel spaces show k-Diff matches or exceeds fixed-target baselines, automatically adapting to representation space and data scale with minimal overhead.

Abstract

Recent advances in diffusion and flow matching models have highlighted a shift in the preferred prediction target -- moving from noise ($\varepsilon$) and velocity (v) to direct data (x) prediction -- particularly in high-dimensional settings. However, a formal explanation of why the optimal target depends on the specific properties of the data remains elusive. In this work, we provide a theoretical framework based on a generalized prediction formulation that accommodates arbitrary output targets, of which $\varepsilon$-, v-, and x-prediction are special cases. We derive the analytical relationship between data's geometry and the optimal prediction target, offering a rigorous justification for why x-prediction becomes superior when the ambient dimension significantly exceeds the data's intrinsic dimension. Furthermore, while our theory identifies dimensionality as the governing factor for the optimal prediction target, the intrinsic dimension of manifold-bound data is typically intractable to estimate in practice. To bridge this gap, we propose k-Diff, a framework that employs a data-driven approach to learn the optimal prediction parameter k directly from data, bypassing the need for explicit dimension estimation. Extensive experiments in both latent-space and pixel-space image generation demonstrate that k-Diff consistently outperforms fixed-target baselines across varying architectures and data scales, providing a principled and automated approach to enhancing generative performance.

Revisiting Diffusion Model Predictions Through Dimensionality

TL;DR

The paper tackles the question of how to choose the optimal diffusion prediction target as data dimensionality grows. It introduces a unified framework with a continuous target and proves the optimal choice is , linking target selection to intrinsic and ambient dimensions. To avoid explicit dimension estimation, it proposes k-Diff, a learnable target where (and optionally time-varying ), enabling end-to-end optimization. Empirical results in both latent and pixel spaces show k-Diff matches or exceeds fixed-target baselines, automatically adapting to representation space and data scale with minimal overhead.

Abstract

Recent advances in diffusion and flow matching models have highlighted a shift in the preferred prediction target -- moving from noise () and velocity (v) to direct data (x) prediction -- particularly in high-dimensional settings. However, a formal explanation of why the optimal target depends on the specific properties of the data remains elusive. In this work, we provide a theoretical framework based on a generalized prediction formulation that accommodates arbitrary output targets, of which -, v-, and x-prediction are special cases. We derive the analytical relationship between data's geometry and the optimal prediction target, offering a rigorous justification for why x-prediction becomes superior when the ambient dimension significantly exceeds the data's intrinsic dimension. Furthermore, while our theory identifies dimensionality as the governing factor for the optimal prediction target, the intrinsic dimension of manifold-bound data is typically intractable to estimate in practice. To bridge this gap, we propose k-Diff, a framework that employs a data-driven approach to learn the optimal prediction parameter k directly from data, bypassing the need for explicit dimension estimation. Extensive experiments in both latent-space and pixel-space image generation demonstrate that k-Diff consistently outperforms fixed-target baselines across varying architectures and data scales, providing a principled and automated approach to enhancing generative performance.
Paper Structure (24 sections, 2 theorems, 55 equations, 5 figures, 4 tables, 2 algorithms)

This paper contains 24 sections, 2 theorems, 55 equations, 5 figures, 4 tables, 2 algorithms.

Table of Contents

  1. Introduction
  2. Related Work
  3. Low-Dimensional Diffusion Models
  4. Pixel-Space High-Resolution Diffusion
  5. Learning and Generalization Dynamics of Diffusion Models
  6. Learning and Generalization Dynamics of Neural Networks
  7. Theoretical Analysis of the Optimal Prediction
  8. Problem Definition
  9. Learning Dynamics
  10. Remark.
  11. $k$-Parameterized Diffusion Prediction Target
  12. Learnable Diffusion Prediction Target
  13. Experiments
  14. Latent Space Diffusion
  15. Pixel Space Diffusion
  16. ...and 9 more sections

Key Result

Theorem 3.1

Under the loss and learning dynamics defined in Eq. (eq:loss) and Eq. (eq:learning_dynamics), the optimal weight $W^*$ is achieved at the equilibrium of both modes, given by: The corresponding optimal loss is obtained by substituting $W^*$ into Eq. (eq:loss):

Figures (5)

  • Figure 1: FID-50k convergence dynamics of $k$-Diff in latent and pixel spaces. (a) Performance comparison in latent space using a LightningDiT-XL/1 architecture on ImageNet-256 over 64 epochs. $k$-Diff maintains a convergence profile nearly identical to the original Flow Matching ($\bm{v}$-prediction) framework. (b) Convergence comparison in pixel space using a JiT-B/16 architecture over 200 epochs. While $\bm{x}$-prediction exhibits a slightly faster initial descent, $k$-Diff recovers this gap as the learnable target stabilizes, indicating that the dynamic optimization of $k$ does not compromise training efficiency or final generative performance across disparate architectures.
  • Figure 2: Long-term FID-50k dynamics for $k$-Diff on ImageNet-256. Training was conducted in latent space using a LightningDiT-XL/1 architecture. The model reaches optimal generative performance at epoch 384 with a peak FID of 1.22. As training extends to 800 epochs, a marginal degradation occurs, with the final FID converging to 1.34. This suggests a subtle shift in the model's generalization-memorization balance during the latter stages of extended training.
  • Figure 3: Evolution of $k$ in Latent vs. Pixel Space. We plot the trajectory of $k$ against normalized training progress (total epochs: 800 for latent, 600 for pixel). In pixel space (JiT-B/16), $k$ exhibits a sharp ascent, converging to $k=1.0$ ($\bm{x}$-prediction) within the first 5% of training. Conversely, in latent space (LightningDiT-XL/1), $k$ climbs gradually to a steady-state value of $0.66$. This stark difference in optimization dynamics validates that $k$-Diff autonomously adapts to the representation space: the high ambient dimensionality of pixels strongly favors $\bm{x}$-prediction, while the compressed latent manifold identifies an optimal target between $\bm{v}$- and $\bm{x}$-prediction.
  • Figure 4: Qualitative Results. Selected examples on ImageNet 256$\times$256 from LightningDiT-XL/1 trained with $k$-Diff.
  • Figure 5: (a) The learned time-dependent $k(t)$ for 128 time bins. (b) Probability distribution of the logit normal distribution with $\mu=0.0, \sigma=1.0$.

Theorems & Definitions (2)

  • Theorem 3.1
  • Theorem 3.2