Multivariate Latent Recalibration for Conditional Normalizing Flows

Thank you for the valuable feedback and for recognizing the extensiveness of our experiments.

On the extension to other models

We discuss this point in detail in our response to Reviewer 4tSF. While LR is not applicable to VAEs and standard diffusion models, it is directly extendable to Conditional Flow Matching (CFM). We showed that LR also leads to significantly improved performance with CFM.

On the comparison to other methods

We compared LR to the uncalibrated baseline and HDR-R because, to our knowledge, these are the only existing methods for post-hoc recalibration of multi-dimensional distributions. The authors of HDR-R [1] also noted that "methods which account for the joint multi-dimensional distribution in assessing calibration and recalibrating the prediction are generally lacking." Our work helps fill this important gap.

[1] Chung et al (2024). Sampling-based Multi-dimensional Recalibration. ICML.

Computational efficiency

We agree that reporting the computation time will strengthen the paper. The difference in computation time can be measured in two aspects:

• For calibration, the computational complexity of HDR-R is $O(M F n)$ and LR is $O(R n)$ where $M = 100$ corresponds to the number of samples of HDR-R per instance, $F$ the time for the forward mapping $\hat{T}$ and $R$ the time for the reverse mapping $\hat{T}^{-1}$.
• For inference, it is a bit more subtle. Given a test insance $x$, HDR-R requires to sample at least $M$ times ($O(M F)$) to obtain a recalibrated sample, which can be a weakness, e.g., if only one conditional sample is needed. LR only incurs a low fixed cost $C$ for evaluating the recalibration map ($O(C + F)$). Thus, the inference time is not directly comparable.

We report the calibration time of HDR-R and LR in seconds, using the convex potential flow model and averaged over 10 runs. To avoid the table being too large, we restrict the table to the largest datasets.

On CIFAR-10 with TarFlow, the time difference is larger and can be prohibitive for HDR-R:

Clarifications for the figures

Thank you for pointing out the ambiguity in our figures. You are correct that the green bars in Figure 3 represent the HDR-R method. The bars in Figures 2 and 3 show the mean metric value across 10 independent runs, and the black error bars depict one standard error of the mean. We apologize for the omission and will update the figure captions and legends to make this explicit.

Image generation quality

This is a great question. We clarify this point in detail in our answer to Reviewer qdRU.

Thank you for your valuable feedback and recognizing the paper's clarity and the experimental evaluation.

Clarifying the paper's goal and contributions

Our main contribution is to introduce latent calibration, a principled framework for multivariate calibration, and simultaneously propose Latent Recalibration (LR), a practical and efficient method that realizes this framework. The two are intrinsically linked: the framework provides the theoretical foundation, and the method demonstrates its practical value.

A key advantage of our approach is that LR produces a fully-specified, recalibrated PDF. This is a critical feature that distinguishes it from set-based or sampling-based methods and enables superior performance in important downstream tasks, as we demonstrate below.

1) Performance of LR on a synthetic experiment

To provide further evidence of the improved performance of LR, we compare the performance of LR and HDR-R on the synthetic dataset of Figure 1. The BASE model is a Masked Autoregressive Flow that has been misspecified by being only trained for 5 epochs. The dataset is divided into 2500 training points, 500 validation/calibration points and 2000 test points. The metrics, averaged over 200 runs, are displayed below with the standard error of the mean.

LR significantly improves both NLL and the Energy Score (ES). HDR-R, which is limited to ES evaluation, shows a much smaller improvement. This highlights that LR can correct model deficiencies that HDR-R does not.

2) Practical advantages of producing an explicit PDF

This is an excellent point. To make the benefits of a full PDF concrete, we have conducted a new experiment on a decision-making task, which we will add to the paper.

Motivation: Access to a full, calibrated PDF is crucial for any task requiring estimation of the probability mass within an arbitrary, non-standard region of the output space, a capability that set-based methods (like Conformal Prediction) or pure sampling-based methods do not provide. A direct application is anomaly detection, where low-density points are classified as anomaly [1, 2]. Other examples include risk assessment in engineering, targeted material design, or optimal control.

[1] Amit Rozner et al (2024). Anomaly Detection with Variance Stabilized Density Estimation. UAI.

[2] Lorenzo Perini et al (2024). Uncertainty-aware Evaluation of Auxiliary Anomalies with the Expected Anomaly Posterior. TMLR.

Experimental Setup: We use the SLUMP dataset (first dataset in Figure 3), where inputs are ingredients for producing concrete and outputs $Y = (S, F, C) \in \mathbb{R}^3$ are three concrete properties. A manufacturer must decide whether a given batch of ingredients is suitable for one of two projects (A or B), each with specific, non-negotiable requirement regions, or if it should be discarded. The decision has different financial utilities and risks.

• Project A requirements: $7 \leq S \leq 20, 55 \leq F \leq 65, 25 \leq C \leq 40$.
• Project B requirements: $20 \leq S \leq 29, 70 \leq F \leq 100, 15 \leq C \leq 30$.
• The optimal action is chosen by maximizing the expected utility, which depends on the estimated probabilities $\hat{P}(Y \in \text{Region}_A | X)$ and $\hat{P}(Y \in \text{Region}_B | X)$. These probabilities are estimated either by (1) sampling or (2) numerical integration of the PDF. We use either 125 samples for sampling or a 5 x 5 x 5 grid for integration using trapezoidal rule.

Results:

The results show two key observations:

1. Using the PDF via numerical integration leads to better decisions (higher utility) than relying on a finite number of samples.
2. The improved calibration from LR provides a more accurate PDF, leading to a significant further increase in utility. HDR-R, which relies on resampling from the original uncalibrated density, actually harmed decision quality in this task.

This demonstrates a concrete scenario where an explicit, calibrated PDF is practically useful for optimal decision-making.

3) Computational efficiency

You are correct that we should substantiate our claim of computational efficiency with experimental results. As detailed in our response to Reviewer 54XG, we have now benchmarked the calibration time of LR against HDR-R. The results confirm that LR is substantially faster, particularly on larger datasets and computationally intensive models.

Energy score

Thank you for the insightful question, which touches upon a point we briefly noted in Section 5.3 regarding the Energy Score's (ES) weaker discriminative ability. The observation that the ES remains largely unchanged is an expected outcome that stems from the score's fundamental limitations in discriminative ability. As established in [1] and corroborated by [2], the ES is sensitive to shifts in the mean but notoriously insensitive to misspecifications in variance, correlation, and overall dependency structure. LR is a post-hoc procedure that primarily corrects the shape and spread of the predictive distribution. Therefore, the ES is fundamentally ill-suited to capture the specific improvements LR provides.

In contrast, the NLL is uniquely suited for this evaluation. As the only local strictly proper scoring rule, its value depends only on the probability density at the precise location of the observed outcome [3]. This locality makes it highly discerning of the very improvements LR makes to the distributional shape, which is why we observe significant and consistent NLL reductions. This result thus highlights an inherent property of the evaluation metric itself, not a shortcoming of our method, and we will clarify this important distinction in the revised manuscript.

[1] Pinson, P., & Tastu, J. (2013). Discrimination ability of the energy score. Technical Report.

[2] Alexander, C., et al (2022). Evaluating the discrimination ability of proper multi-variate scoring rules. Annals of Operations Research.

[3] Du, H. (2021). Beyond Strictly Proper Scoring Rules: The Importance of Being Local. Weather and Forecasting.

Image generation quality

This is a great question, and we appreciate the opportunity to clarify this point. To be perfectly clear, LR is not expected to improve the visual quality of generated samples. We will state this explicitly in the paper to avoid any potential confusion. The goal of our image data experiment is to understand the behaviour of latent calibration and recalibration in a very-high dimensional setting.

In the noisy setting, LR preserves the visual quality of the samples from the base model, with no perceptually visible changes, which aligns with the very small change we observed in NLL. In the noiseless setting, where the test data is out-of-distribution relative to the noisy training data, applying LR can cause samples to appear less realistic with more uniform colors. This is an expected outcome, as a simple post-hoc recalibration step cannot be expected to generalize to a domain shift and improve generation quality.

Samples from HDR-R are only resampled from an initial pool of samples. Thus, differences in visual quality are not expected either.

Dear Reviewer qdRU,

Thank you for your follow-up and for engaging with our rebuttal. We appreciate you clarifying your remaining concern regarding the ES results in Table 4.

We agree that the lack of significant improvement in the ES for LR in Table 4 warrants a detailed explanation. Our argument, supported by both the literature and the new experiment below, is that this observation stems not from a shortcoming of our method, but from a well-documented limitation of the ES itself. As we noted in our initial rebuttal, the ES is known to be relatively insensitive to changes in variance and correlation structure, which are precisely the aspects of the distribution that our post-hoc LR method is designed to correct.

To provide a clear, empirical illustration of this point, we designed a controlled synthetic experiment based on the dataset in Figure 1. The goal here is not to replace the real-world results of Table 4, but to isolate this specific property of the scoring rules in a setting free from the confounding variables of complex, real-world data.

We use an oracle predictor that knows the true data-generating distribution from Figure 1 for everything except the spread around the arc, which is controlled by a standard deviation parameter $\sigma$. We then evaluate the predictor's NLL and ES as we vary its estimate of $\sigma$. The true value is $\sigma=0.05$.

This experiment strikingly illustrates the issue:

1. The NLL shows a sharp, clear minimum at the true value of $\sigma=0.05$, correctly identifying the best model.
2. The ES remains almost completely flat for a wide range of $\sigma$ values (from 0.01 to 0.10). It fails to reliably distinguish a model with the correct variance from one that is substantially over- or under-confident.

This insensitivity is so profound that detecting a statistically significant signal with the ES requires an impractically large number of samples. The table below shows that only with 5000 runs does the ES minimum align with the true $\sigma$, and even then the differences are minuscule:

This controlled experiment, therefore, explains the results in Table 4: LR is improving the shape and spread of the distribution, but the ES is not the right tool to measure this specific kind of improvement, highlighting the importance of using multiple, complementary evaluation metrics like NLL.

We will add this analysis to the appendix to provide context for our results and guide future work. We hope this addresses your concern and clarifies the interpretation of our findings. We are happy to elaborate further if anything remains unclear.

The core requirements for LR are (1) an invertible transformation between the response and latent spaces, and (2) a latent random variable with a known, tractable density. This makes Normalizing Flows (NFs) a natural and ideal fit for our method.

You are correct that this formulation does not directly apply to models like VAEs or Denoising Diffusion Probabilistic Models (DDPMs), which lack an exact, tractable inverse mapping. We will add a discussion of this limitation to the paper.

While it is not the main focus of our paper, LR is directly compatible with the more recent Conditional Flow Matching (CFM) methods, which also learn invertible maps and have a known latent distribution. We performed additional experiments that we will include in the appendix.

Details on Conditional Flow Matching: Given $x \in \mathcal{X}$, we model the conditional PDF $\hat{f}(y|x)$ using a transformation defined by an Ordinary Differential Equation (ODE), $\frac{d\tilde{y}}{dt} = \hat{v}(t, \tilde{y}, x)$, with a neural network-parameterized vector field $\hat{v}$. Trained via the Conditional Flow Matching (CFM) objective, the model learns to map samples from a latent variable $Z \sim \mathcal{N}(0, I)$ to the target conditional distribution. Forward numerical integration of this dynamic generates samples, while reverse integration performs encoding. To compute the likelihood for $y \in \mathcal{Y}$, we first obtain its latent representation $z = \hat{T}^{-1}(y|x)$. The log-likelihood is then given by the continuous change of variables formula, which integrates along the unique path $\tilde{y}(t)$ that solves the ODE with boundary conditions $\tilde{y}(0)=z$ and $\tilde{y}(1)=y$: $\log \hat{f}(y|x) = \log f_Z(z) - \int_0^1 \text{Tr}\left(\nabla_{\tilde{y}} \hat{v}(t, \tilde{y}(t), x)\right) dt$ The computationally intractable Jacobian trace is efficiently approximated using the unbiased Hutchinson's estimator, enabling practical likelihood computation.

Results: To avoid large tables, we compare the methods on the largest datasets. The CFMs results are aligned with the NFs results, with LR standing out particularly on the L-ECE and NLL metrics, with a good performance on HDR-ECE and ES.

Practical advantages of producing an explicit PDF

We designed a new experiment to show a practical scenario where having a calibrated, explicit PDF is beneficial. To avoid redundancy, you can find the details of this experiment in the answer to Reviewer qdRU.

LR only adjusts the length of latent vectors, not their directions

You are right, and it is a core design choice of our method.

LR intentionally adjusts only the magnitude of latent vectors, not their direction. We acknowledge that this means LR cannot fix miscalibration arising from errors in the orientation of the learned latent manifold.

However, this focused approach is a principled trade-off that yields significant benefits:

• Tractability and guarantees: It simplifies the difficult multivariate calibration problem into a tractable univariate one (calibrating norms). This is precisely what enables us to provide finite-sample guarantees and establish a clean connection to conformal prediction and existing Quantile Recalibration.
• Empirical performance: Our experiments demonstrate that this radial-only adjustment is remarkably effective, suggesting that correcting the norm distribution addresses a major source of miscalibration in many conditional normalizing flows.

We will add a discussion of this deliberate design choice and its implications to the limitations section of the paper.

Does the NLL performance gain come from the base or the jacobian term?

To investigate the source of the NLL improvement, we consider the decomposition of the recalibrated NLL: $-\log \hat{f}'(y \mid x) = -\log f_{Z}\left( z \right) - \log \left| \det\left( \nabla_z R(z) \right) \right|^{-1} -\log \left| \det\left( \nabla_y \hat{T}^{-1}(y; x) \right) \right|$ with $z' = \hat{T}^{-1}(y; x)$ and $z = R^{-1}(z')$. The third term is identical for both BASE and LR. We analyze the contributions of the first two terms across the largest tabular datasets to avoid the table being too large. All reported terms are averaged over the test set and over 10 runs.

The table reveals a clear pattern. The NLL improvement from LR is primarily driven by the first term. By radially transforming the latent codes $z'$ to new points $z$ that are more consistent with the base density $f_Z$, the latent density term $-\log f_Z(z)$ is significantly reduced. The recalibration Jacobian (second term) typically adds a small penalty (increases NLL), but this is almost always outweighed by the large gains from the first term. This confirms that LR works by finding more "plausible" latent codes for the observed data under the base latent distribution.