Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Improved probabilistic regression using diffusion models

Carlo Kneissl, Christopher Bülte, Philipp Scholl, Gitta Kutyniok

TL;DR

This work extends diffusion-based regression by explicitly modeling the full distribution of the diffusion noise rather than only its mean, using strictly proper scoring rules to train a flexible noise distribution $p_{\\theta}^{\\epsilon}(\\epsilon_t \\mid \\bm x_t)$. By adopting Gaussian-mixture parametrizations (including univariate, mixture, and multivariate forms with efficient covariance structures), the authors obtain closed-form backward passes and improved probabilistic calibration across tasks. The approach yields superior or on-par predictive performance and provides estimates of epistemic uncertainty, addressing a key limitation of standard diffusion models in regression settings. The framework demonstrates versatility across UCI benchmarks, autoregressive tasks, and depth estimation, highlighting practical impact for calibrated probabilistic prediction in diverse domains.

Abstract

Probabilistic regression models the entire predictive distribution of a response variable, offering richer insights than classical point estimates and directly allowing for uncertainty quantification. While diffusion-based generative models have shown remarkable success in generating complex, high-dimensional data, their usage in general regression tasks often lacks uncertainty-related evaluation and remains limited to domain-specific applications. We propose a novel diffusion-based framework for probabilistic regression that learns predictive distributions in a nonparametric way. More specifically, we propose to model the full distribution of the diffusion noise, enabling adaptation to diverse tasks and enhanced uncertainty quantification. We investigate different noise parameterizations, analyze their trade-offs, and evaluate our framework across a broad range of regression tasks, covering low- and high-dimensional settings. For several experiments, our approach shows superior performance against existing baselines, while delivering calibrated uncertainty estimates, demonstrating its versatility as a tool for probabilistic prediction.

Improved probabilistic regression using diffusion models

TL;DR

This work extends diffusion-based regression by explicitly modeling the full distribution of the diffusion noise rather than only its mean, using strictly proper scoring rules to train a flexible noise distribution . By adopting Gaussian-mixture parametrizations (including univariate, mixture, and multivariate forms with efficient covariance structures), the authors obtain closed-form backward passes and improved probabilistic calibration across tasks. The approach yields superior or on-par predictive performance and provides estimates of epistemic uncertainty, addressing a key limitation of standard diffusion models in regression settings. The framework demonstrates versatility across UCI benchmarks, autoregressive tasks, and depth estimation, highlighting practical impact for calibrated probabilistic prediction in diverse domains.

Abstract

Probabilistic regression models the entire predictive distribution of a response variable, offering richer insights than classical point estimates and directly allowing for uncertainty quantification. While diffusion-based generative models have shown remarkable success in generating complex, high-dimensional data, their usage in general regression tasks often lacks uncertainty-related evaluation and remains limited to domain-specific applications. We propose a novel diffusion-based framework for probabilistic regression that learns predictive distributions in a nonparametric way. More specifically, we propose to model the full distribution of the diffusion noise, enabling adaptation to diverse tasks and enhanced uncertainty quantification. We investigate different noise parameterizations, analyze their trade-offs, and evaluate our framework across a broad range of regression tasks, covering low- and high-dimensional settings. For several experiments, our approach shows superior performance against existing baselines, while delivering calibrated uncertainty estimates, demonstrating its versatility as a tool for probabilistic prediction.

Paper Structure

This paper contains 44 sections, 2 theorems, 52 equations, 15 figures, 15 tables.

Table of Contents

  1. Introduction
  2. Background
  3. Probabilistic regression
  4. Diffusion models
  5. Our methodology
  6. Learning the full distribution in each denoising step
  7. Theoretical justification
  8. Parametrization of $p_\theta^\epsilon(\cdot \mid \bm x_t)$
  9. Numerical experiments
  10. UCI regression
  11. Autoregressive prediction tasks
  12. Monocular depth estimation
  13. Related work
  14. Conclusion
  15. Approximations of the noise distribution
  16. ...and 29 more sections

Key Result

Theorem 1

Let $p_\theta^\epsilon(\epsilon_t \mid \bm x_t)$ be given as in Equation eq:noise_parametrization. Then the reverse distribution in Equation background:diffusion:backward_formula admits the closed form where $\hat{\bm x}_0 = \frac{\bm x_t - \sqrt{1 - \bar{\alpha}_t} \bm\mu_{\theta,k}^\epsilon(\bm x_t, t)}{\sqrt{\bar{\alpha}_t}}$, $\lambda_t \coloneq \sqrt{1 - \bar{\alpha}_{t-1} - \sigma_t^2}$ and

Figures (15)

  • Figure 1: Method overview. At any given time $t = 1, \ldots, T$ in the diffusion process, we need a prediction of the noise $\epsilon_t$ given our current state $x_t$. Traditionally, this is achieved by a network that approximates the conditional mean $\mathbb{E}[\epsilon_t \mid x_t]$, which, when viewed under a probabilistic viewpoint, treats the distribution $p^\epsilon(\cdot \mid x_t)$ as a Dirac distribution. We propose to model this distribution by a Gaussian Mixture family.
  • Figure 2: Comparison of the $\bm \delta_\theta$ model (upper) and $\bm\Sigma_\theta^\mathrm{mix}$ model (lower) predictions for a selected sample of the dynamics of the 2-meter surface temperature.
  • Figure 3: Comparison of aleatoric and epistemic uncertainty estimates for a sample trajectory of the $\bm\Sigma_\theta^\mathrm{diag}$ model.
  • Figure 4: Comparison of the $\bm \delta_\theta$ model (upper) and $\bm\Sigma_\theta^\mathrm{mv}$ model (lower) predictions for the depth estimation task on DIODE vasiljevic2019diode.
  • Figure 5: Correlation matrices for different covariance parameterizations across several time steps $t$ for a selected sample of the Kuramoto--Sivashinsky equation.
  • ...and 10 more figures

Theorems & Definitions (6)

  • Theorem 1
  • proof
  • proof
  • Theorem 2
  • proof
  • proof