Revisiting Diffusion Model Predictions Through Dimensionality

TL;DR

This paper develops a theoretical framework explaining why data prediction becomes optimal in high-dimensional diffusion models and introduces k-Diff, a data-driven method to select the best prediction target.

Contribution

It provides a formal analysis linking data geometry to prediction target optimality and proposes k-Diff to automatically learn this target without explicit dimension estimation.

Findings

Theoretical justification for x-prediction superiority in high ambient dimensions.

k-Diff outperforms fixed-target baselines in image generation tasks.

Demonstrates the importance of data geometry in diffusion model predictions.

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.