Uncertainty Quantification for Deep Regression using Contextualised Normalizing Flows

AI summary

Overview

Research area: Uncertainty quantification (UQ) for deep regression models, combining Monte Carlo Dropout with conditional normalizing flows.

Technical level: Advanced. The paper assumes familiarity with normalizing flows, Monte Carlo Dropout, conformal prediction, quantile regression, and the distinction between epistemic and aleatoric uncertainty.

Scope: The paper introduces MCNF (Monte Carlo Normalizing Flow), a post hoc, distribution-free UQ method that produces both prediction intervals and a full conditional predictive distribution on top of an already-trained deep regression model, and benchmarks it against MCD, CQR, MCCP, MCQR and DQR across nine datasets.

What This Paper Is About

Deep regression models are increasingly used in high-stakes settings, but many existing UQ methods return only prediction intervals and discard distributional shape, while full or approximate Bayesian methods require modifying and retraining the architecture. The authors present MCNF, which takes an already-trained dropout-equipped regression model, draws Monte Carlo Dropout samples from it, and trains a small normalizing flow after the fact to model the distribution of prediction errors conditioned on those samples. The goal is to obtain a well-calibrated, arbitrarily complex predictive distribution without retraining the underlying predictor.

Key Contributions

1. The MCNF method — a post hoc UQ technique whose estimates take the form of a distribution-agnostic predictive distribution, rather than only intervals.
2. A comprehensive evaluation of MCNF against state-of-the-art UQ methods (MCD, CQR, MCCP) on a range of standard benchmarks plus a physicochemical (solubility) dataset.
3. An open-source prototype MCNF tool and case-study repository, released at https://github.com/alexiatoumpa/MCNF.
4. Implicitly, a demonstration that the approach transfers to deep learning architectures other than feed-forward regression networks, and that it can adapt its estimates to the quality of the underlying predictive model.

Main Findings

• Smallest prediction error: MCNF overall yielded the smallest MAE, closely followed by MCD. MAEs for the other methods were typically larger by roughly one order of magnitude compared to MCNF and MCD. The authors attribute the larger CP-based errors to CQR and MCCP adjusting upper and lower intervals homogeneously, which only suits near-Gaussian distributions.

• Competitive, near-nominal coverage: Apart from MCD, all methods produced marginal coverage values close to the theoretical 90%, showing similar capability and some conservativeness. This was especially true for the conformalized methods (CQR and MCCP) on the smallest datasets (Boston Housing and Concrete). MCD, which does not account for aleatoric uncertainty, was highly non-conservative — for example coverage of 0.726 on Boston Housing, 0.601 on Concrete, 0.341 on Abalone, 0.354 on Protein, 0.239 on Romano-Original and 0.537 on Solubility.

• Better coverage than MCD at similar interval size: MCNF outperformed its MCD counterpart, and MCNF is computationally more efficient than MCD, especially as the number of samples used to approximate the predictive distribution grows.

• Narrowest intervals at maintained coverage: Although all methods were sensitive to the actual uncertainty in each dataset, MCNF yielded the smallest interval sizes (~Δ) while maintaining the expected 90% marginal coverage. CQR and MCCP, in achieving 90% coverage through conformalization, produced conservative intervals reflected in larger interval sizes.

• Best coverage/interval trade-off: MCNF provided the best trade-off between coverage and interval size, giving smaller intervals for similar marginal coverage.

• Captures multimodality: On the synthetic Romano-Mod dataset, which exhibits heteroskedasticity and varying distributions of the predicted variable for different values of x, MCNF effectively captured the multimodality of y for small x, transitioning to a unimodal distribution as x increases. MCNF also outperformed CQR, MCCP and DQR on marginal coverage, interval size and MAE on both Romano-Original and Romano-Mod.

• Robust to an underfitted base model: Coverage between well-trained and underfitted predictive models was similar and around 90% for most UQ methods (except MCD), while a well-trained predictive model gave narrower intervals. Even with an underfitted predictive model, MCNF yielded narrower intervals than the state-of-the-art UQ methods and comparable coverage, showing MCNF adapts its estimate to the quality of the predictive model.

• Epistemic uncertainty propagation: The authors report results for MCNF with the epistemic uncertainty propagation through MCD sampling removed in the last column of Table 1 (labelled NF). The description of those results is cut off in the provided content, so the specific conclusions drawn from the NF-only configuration are not reported here.

Methodology in Plain English

The idea rests on splitting "predicting the value" from "saying how unsure we are."

• A normal predictive model does the prediction. The base model is a Deep Quantile Regressor with a batch-normalization input layer, two fully connected layers with ReLU nonlinearities and dropout at rate 0.1, and an output layer with three linear units for the quantiles q = {0.05, 0.5, 0.95}. It is trained for 100 epochs with Adam, a custom pinball loss, batch size 32, learning rate 5e-4, weight decay 1e-6, and an 80:20 train/test split.

• Monte Carlo Dropout provides a rough picture of uncertainty. At inference, the model is run many times with dropout active. The paper fixes n_MCD = 50 prior samples to keep overhead low. These samples give a sample mean and log-variance.

• A context vector summarises the situation. Rather than feeding raw inputs to the flow, MCNF builds a context c = { ȳ_MCD, log s²(ȳ_MCD), h(x) }, where h(x) is an internal representation taken from the trained regression model at a chosen layer depth (a tunable hyperparameter affecting context dimensionality), and the first two terms are sample mean and log-variance of the MCD samples.

• A normalizing flow models the errors. The flow is trained to represent the distribution of the prediction error δ = y − y_MCD conditioned on that context, rather than modelling y directly. The flow is a sequence of two Neural Spline Flows with a 3-layered MLP of 64 hidden units producing the 16 support vectors of the spline transformation and their inner derivatives, with a factorized Gaussian base distribution with trainable parameters ψ. Training uses batch size 32, Adam, and a 0.001 learning rate, with the same partition as the base model.

• Training objective and regularisation. The flow is trained by minimising the forward KL divergence, equivalent to minimising negative log-likelihood. Because plain likelihood minimisation can overfit and produce deformed, uncalibrated distributions — especially for low-uncertainty data corrupted by large outliers — the authors weight each observation by w_n = σ(−log p_MCD(y_n|x_n)/τ; τ) ∈ (0,1), where σ is a softmax and τ is a temperature hyperparameter. They set τ = 1e10, giving all observations in a mini-batch the same weight. Equation (8) reduces to the unweighted likelihood when τ → ∞.

• Inference is hierarchical. For each input, run n_MCD forward passes, collect and aggregate hidden states, build the context, feed it to the flow, draw n_NF samples from the base distribution and push them forward, then correct a randomly chosen prior sample with the sampled error: y_n = y_{n,j(k),MCD} − δ_{n,k} where j(k) ~ U(1, n_MCD). Likelihoods are computed in the same pass by solving the change-of-variables equation, or by running the flow in reverse.

• Comparisons. MCNF is benchmarked against MCQR (1000 MCD resamples averaged at q = {0.05, 0.95}), MCD (1000 resamples from the median q = 0.5), CQR and MCCP (both using 20% of the test set for calibration, targeting 90% marginal coverage). For all MCD-sampling methods the aleatoric term of the original formulation is left out, propagating epistemic uncertainty only. Results are reported across 20 independent runs, with a fresh train/test partition per run.

• Datasets. Boston Housing (506 observations, 14 attributes), Concrete (1030, 9), Abalone (4177, 11), Tertiary Protein Structure (45730, 10), wave energy (63600, 149), superconductivity (21263, 81), the solubility physicochemical dataset, plus two synthetic sets: Romano-Original (univariate predictor with few large outliers) and Romano-Mod (an extension with a multimodal distribution).

Why This Matters

Impact on research. MCNF shows that a full conditional predictive distribution can be recovered post hoc, without retraining the base model and without the architecture modifications Bayesian approaches require. It also addresses a gap the authors identify: CQR and MCCP deliver intervals but cannot represent arbitrarily complex uncertainty distributions, whereas MCNF can. Interval-only outputs ignore the effect of multimodal or asymmetric distributions on decision-making, which MCNF does not.

Real-world applications (domains where deep regression is already used, as cited in the paper):

• Drug discovery and molecular property prediction (the physicochemical solubility dataset used here).
• Credit scoring.
• Energy forecasting.
• Medical diagnostics and other safety-critical, high-stake decisions.

Industry relevance. The method is deliberately non-invasive: any deep learning model with dropout layers can be paired with MCNF after the fact, so organisations with existing trained regressors avoid retraining costs. It is also reported to be more computationally efficient than MCD, particularly as the number of predictive-distribution samples grows. Because it produces a full density rather than just an interval, downstream decision systems can price risk asymmetrically instead of treating the interval as uniform.

Future Directions

• Characterising the flow-only (NF) variant. The paper reports an ablation in which epistemic uncertainty propagation through MCD is removed, but the conclusions drawn from it are not contained in the supplied content; understanding when the epistemic pathway is worth its cost remains open.

• Extending beyond the evaluated architectures. The authors note applicability to deep learning architectures other than feed-forward regression networks is "also showcased," suggesting broader architectural coverage (for example sequence or vision-based regressors) as a natural extension.

• Robustness to the base model's quality. Results show MCNF adapts to underfitted predictors, but the paper's underfitted model was trained for only 6 epochs; how MCNF behaves with severely misspecified or systematically biased predictors is not established.

• Hyperparameter sensitivity. Several design choices — the depth at which h(x) is extracted, the context dimensionality, the prior sample count n_MCD = 50, and the temperature τ = 1e10 — are set rather than studied, leaving open how much performance depends on them and how they should be selected in new domains.

Target Audience

Machine learning researchers and practitioners working on uncertainty quantification for regression, particularly those in safety-critical domains; engineers who need calibrated prediction intervals or full predictive densities from models they cannot afford to retrain; and readers already familiar with Monte Carlo Dropout, conformal prediction, quantile regression and normalizing flows who want a practical, post hoc alternative that preserves distributional shape.

Authors’ abstract

Quantifying uncertainty in deep regression models is important both for understanding the confidence of the model and for safe decision-making in high-risk domains. Existing approaches that yield prediction intervals overlook distributional information, neglecting the effect of multimodal or asymmetric distributions on decision-making. Similarly, full or approximated Bayesian methods, while yielding the predictive posterior density, demand major modifications to the model architecture and retraining. We introduce MCNF, a novel post hoc uncertainty quantification method that produces both prediction intervals and the full conditioned predictive distribution. MCNF operates on top of the underlying trained predictive model; thus, no predictive model retraining is needed. We provide experimental evidence that the MCNF-based uncertainty estimate is well calibrated, is competitive with state-of-the-art uncertainty quantification methods, and provides richer information for downstream decision-making tasks.