Revisiting Diffusion Model Predictions Through Dimensionality

papers.abstract

Jüngste Fortschritte bei Diffusions- und Flow-Matching-Modellen haben einen Wandel beim bevorzugten Vorhersageziel aufgezeigt – weg von Rauschen (ε) und Geschwindigkeit (v) hin zur direkten Datenvorhersage (x) – insbesondere in hochdimensionalen Settings. Eine formale Erklärung, warum das optimale Ziel von den spezifischen Eigenschaften der Daten abhängt, bleibt jedoch bislang unklar. In dieser Arbeit stellen wir einen theoretischen Rahmen vor, der auf einer verallgemeinerten Vorhersageformulierung basiert und beliebige Ausgabeziele umfasst, wobei ε-, v- und x-Vorhersage Spezialfälle darstellen. Wir leiten den analytischen Zusammenhang zwischen der Geometrie der Daten und dem optimalen Vorhersageziel her und liefern eine rigorose Begründung dafür, warum die x-Vorhersage überlegen wird, wenn die Umgebungsdimension die intrinsische Dimension der Daten signifikant übersteigt. Während unsere Theorie die Dimensionalität als bestimmenden Faktor für das optimale Vorhersageziel identifiziert, ist die intrinsische Dimension von manifold-gebundenen Daten in der Praxis typischerweise nicht greifbar zu schätzen. Um diese Lücke zu schließen, schlagen wir k-Diff vor, einen Rahmen, der einen datengesteuerten Ansatz verwendet, um den optimalen Vorhersageparameter k direkt aus den Daten zu lernen und so die Notwendigkeit expliziter Dimensionsschätzung umgeht. Umfangreiche Experimente sowohl in der Bildgenerierung im Latent-Space als auch im Pixel-Space demonstrieren, dass k-Diff feste Ziel-Baselines über verschiedene Architekturen und Datenskalen hinweg konsistent übertrifft und somit einen prinzipienbasierten und automatisierten Ansatz zur Steigerung der generativen Leistung bietet.

English

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.