Predictive Diffusion Regression Models

• Predictive regression models are statistical frameworks that estimate the full predictive distribution p(y|c) to capture uncertainty, heteroscedasticity, and multimodal outcomes.
• Diffusion-based methods recast regression as a sequential denoising process using proper scoring rules to nonparametrically learn the entire noise distribution.
• Enhanced parameterizations, including mixture and full covariance models, improve calibration and scalability, yielding competitive results in diverse tasks.

Predictive regression models constitute a foundational class of statistical and machine learning frameworks devoted to learning mappings from covariates to response variables, while providing quantification of uncertainty and full probabilistic characterizations of the prediction process. Recent advances, notably the introduction of diffusion-based generative architectures for regression, have extended model flexibility and expressiveness far beyond classical mean-based formulations, enabling robust probabilistic inference, multimodal output distributions, and highly calibrated uncertainty estimates in both low- and high-dimensional settings.

1. Mathematical Foundations of Probabilistic Predictive Regression

The general objective is to infer the conditional predictive distribution of a response $y \in \mathbb{R}^{d_y}$ given covariates $c \in \mathcal{C}$ and observed data $\mathcal D$: $p(y \mid c; \mathcal D) \approx p_\theta(y \mid c)$ Classical regression typically targets point estimation, i.e., $\mathbb{E}[y | c]$. Probabilistic approaches elevate this by modeling the full $p(y | c)$, capturing heteroscedasticity, non-Gaussian noise, and even multimodal behaviors critical for calibrated decision-making and uncertainty quantification.

Diffusion models reinterpret regression as a sequential denoising generative process:

• Forward process: For $x_0 = y$, iteratively add Gaussian noise:

$$p(x_t \mid x_{t-1}) = \mathcal{N}\left(x_t; \sqrt{\alpha_t} x_{t-1}, \beta_t I \right)$$

with a schedule $\{\beta_t\}_{t=1}^T$, $\alpha_t = 1 - \beta_t$, $\bar\alpha_t = \prod_{s=1}^t \alpha_s$.

• Marginalization yields:

$$x_t = \sqrt{\bar\alpha_t} x_0 + \sqrt{1 - \bar\alpha_t} e_t, \quad e_t \sim \mathcal{N}(0, I)$$

• Reverse process: Learn $p_\theta(x_{t-1} | x_t, c)$ via a parameterized mean and covariance.

Instead of learning just the mean $\mathbb{E}[e_t \mid x_t, c]$ (as in conventional DDPM/DDIM regression), the improved framework proposes full nonparametric modeling of $q_\theta(e_t | x_t, t, c)$.

2. Nonparametric Predictive Posterior via Diffusion Noise Modeling

The standard DDPM regression loss only fits the first moment: $$\mathcal{L}_{\text{simple}} = \mathbb{E}[\|e_t - \epsilon_\theta(x_t, t, c)\|^2]$$ where only the mean of the noise is regressed, and covariance is fixed and isotropic. The enhanced framework replaces this with a strictly proper scoring-rule based objective: $$\mathcal{L}_{SR} = \mathbb{E}_{t, x_0, e_t, c} \left[ S\left(q_\theta(\cdot \mid x_t, t, c), e_t \right) \right]$$ where $S$ is e.g. CRPS, energy score, or kernel score, enforcing that the predicted $q_\theta$ matches all aspects (not just the mean) of the true noise distribution.

3. Noise Parameterizations: Trade-Offs and Scaling

Three principal parameterizations for $q_\theta(e_t | x_t, t, c)$:

Parameterization	Model Capacity	Sampling/Comp. Complexity
Diagonal Gaussian ($K=1$)	Independent, unimodal	$O(d)$ per step
Diagonal Mixture ($K>1$)	Multimodal marginals	$O(Kd)$ per step
Full Covariance ($K=1$)	Arbitrary correlation	$O(d^2)$ (Cholesky), $O(dr^2 + r^3)$ (low-rank + diag)

• Diagonal Gaussian is efficient, suitable for weakly correlated noise.
• Diagonal mixtures capture multimodality in marginals.
• Full covariance (Cholesky or low-rank representations) is essential for tasks with highly structured uncertainty.
• Low-rank+diag is scalable for $d\gg 1$ and maintains expressive capacity.

Automated selection of parameterization remains an open challenge; post-hoc scaling of $\Sigma_\theta$ (covariance multiplier) can restore empirical calibration.

4. Algorithmic Workflow

Training: For each mini-batch:

1. Sample random timestep $t \sim \mathrm{Unif}\{1, ..., T\}$.
2. Draw $e_t \sim \mathcal{N}(0, I)$.
3. Form $$x_t = \sqrt{\bar\alpha_t} y + \sqrt{1 - \bar\alpha_t} e_t$$.
4. Predict mixture parameters $\{w_k, \mu_k, \Sigma_k\}$ with a neural network.
5. Compute loss $\ell = S(q_\theta(\cdot | x_t, t, c), e_t)$.
6. Backpropagate and update $\theta$.

Inference (Sampling):

1. Given covariates $c$, set $x_T \sim \mathcal{N}(0, I)$.
2. For $t = T, ..., 1$:
   • Predict $\{w_k, \mu_k, \Sigma_k\}$.
   • Sample $e_t \sim \sum_k w_k \mathcal{N}(\mu_k, \Sigma_k)$.
   • Compute $x_{t-1}$ via closed-form mixture reverse step.
3. Return $x_0$ as a sample from $p_\theta(y | c)$.

5. Uncertainty Quantification and Calibration

• Aleatoric uncertainty assessed via sample variance of $\{y^{(m)}\}$; CRPS and energy scores measure distribution calibration.
• Epistemic uncertainty quantified by the variance of predicted means $\mu_\theta$ or by second-order statistics over denoising steps:

$$\text{EU} \approx \sum_{t=1}^T \mathrm{Var}[ \mu_\theta(x_t, t, c) ]$$

This approach enables epistemic quantification not available in single-variance diffusions.

• Coverage: Empirical frequency of true $y$ within predicted quantile intervals; post-hoc scaling of covariances can be used to restore nominal coverage.

6. Comparison to Classical Predictive Regression Approaches

Model Type	Key Properties	Limitations
Gaussian Processes	Closed-form; calibrated	Cubic cost in $N$; single modality
Quantile Regression	Marginal quantile estimation	No joint distribution; monotonicity issues
Mixture Density Nets	Flexible multi-component	Sensitive to $K$ selection; MLE log-score may miscalibrate
Diffusion-Based (proposed)	Nonparametric; multimodal, heteroscedastic; scoring rule calibration	Scaling to multivariate mixtures remains open

Diffusion regression with noise distribution learning achieves:

• Nonparametric learning of predictive distributions
• Heteroscedasticity, multimodality, and improved calibration
• Scalability via U-Net backbones and proper scoring rules

7. Empirical Results across Task Families

A) Low-dimensional UCI regression ($d_y=1$):

• Emix (univariate mixture) and Ediag (diagonal variance) improve CRPS and energy score by $\sim$10–20% over CARD and deterministic diffusion baselines.
• Coverage at 95% matches nominal values.

B) Autoregressive PDE forecasting (Burgers’, Kuramoto–Sivashinsky, Weather):

• Ediag/Emix models reduce RMSE by $\sim$15% and halve CRPS; coverage is sustained.
• In chaotic PDEs, multimodal mixture bests RMSE/CRPS metrics; Ediag sometimes underconfident (improved via scaling).

C) Monocular depth estimation (multiple benchmarks):

• Emv (multivariate) achieves best AbsRel and CRPS, outperforming Marigold by 5–10%, providing calibrated uncertainty estimates.

8. Implementation Details

Typical deployment combines:

• U-Net variants with Fourier embeddings (32 frequencies)
• Timestep count $T=50$, linear beta schedule ($\beta_1=10^{-3}$ to $\beta_{50}=0.35$)
• Adam/AdamW optimizer, learning rate $10^{-3}$–$10^{-5}$, batch size 64–128, early stopping
• Scoring rule: CRPS or kernel energy score
• Mixture components $K=3$ suffice for most; low-rank $r=10$ for $d \sim 10^3$
• Covariance scaling ($\tau<1$) employed post hoc for calibration

Extensions under exploration:

• Automated parameterization selection
• Multivariate mixture modeling for highly structured output spaces
• Advanced noise schedules, stochastic contraction algorithms
• Rigorous covariance scaling theory
• Epistemic uncertainty via ensembles or Bayesian diffusion models

9. Outlook and Open Problems

Key challenges include:

• Adaptive selection/optimization of noise model structure (diagonal, mixture, full covariance) for diverse task domains.
• Scaling to multivariate Gaussian mixtures with full covariance for highly structured or correlated outputs.
• Theoretical analysis of calibration procedures, e.g., the effect of global covariance rescaling on predictive reliability.
• Bayesian or ensemble-based approaches for epistemic uncertainty modeling within sequential diffusion architectures.

The nonparametric diffusion-based predictive regression paradigm enables a unified framework for calibrated, uncertainty-aware probabilistic regression that is competitive with, or superior to, classical and neural baselines, and is extensible to arbitrary problem dimensions and output structures.