Quantifying Uncertainty in the Presence of Distribution Shifts

Dear reviewer,

Thank you for your thoughtful review and your clarity regarding score reevaluation. Please see our responses.

Question 1:

To the best of our knowledge, this is indeed the first work to introduce a covariate-dependent prior. Below we briefly review the literature you mentioned and highlight key differences from our approach with respect to the prior treatment. We will address these references in the paper as well.

1. Conditional Neural Processes (Garnelo et al. 2018)

   This paper addresses the problem of predicting outputs for a new set of inputs (the target set) given a set of input-output pairs (the context set). The model uses an encoder to map each labeled pair $(x_i, y_i)$ in the context set to a representation $r_i$. These representations are then aggregated into a summary vector $r$, which, together with a new target input $x^*$ is passed to the predictor to produce the output.

   Key difference : The proposed model is trained by maximizing the conditional likelihood, following a purely frequentist approach with no prior. Conditioning occurs only with respect to training covariates (and labels) and the model remains fixed afterward. New covariates are treated as standard inputs to the prediction function. In contrast, our model is fully Bayesian and incorporates an adaptive prior: not only does it produce a different posterior for each new covariate $x^ *$, but the prior itself changes with $x^ *$. This dependence is learned by optimizing a variational objective, where both the prior and posterior are explicitly conditioned on $x^ *$.

2. Bayesian Model-Agnostic Meta-Learning (Yoon et al. 2018)

   The paper aims to learn a parameter initialization $\theta$ such that a few gradient updates on a new task $T_j$ lead to good performance. To capture task-specific dependencies, it places a particle-based prior over the task-specific parameters.

   Key difference: The particle-prior is $\{\theta_1, … , \theta_m\}$, corresponding to tasks $T_1, …, T_m$. The prior is not covariate specific. It uses training data for different tasks, but does not depend on individual covariates.

3. Hypernetworks have a similar dependence structure: a hyper-network $H$ generates the parameters $\theta$ of the prediction network $f_\theta$. Rather than learning a fixed $\theta$, the model learns a generator $H$ that maps task embeddings $z_j$ to task-specific parameters $\theta_j = H(z_j)$. This constitutes a form of amortization—task-specific parameters are not learned independently, but inferred via a shared network. In the deterministic case, no sampling is involved and no prior is placed. In the latent-variable formulation a prior is placed over task embeddings $z_j$, which induces a prior over parameters $p(\theta) = p(z) {\vert \frac{dH(z)}{dz} \vert}^{-1}$. This prior captures variability across tasks, but does not condition on individual covariates.

4. In amortized inference, we can learn a function that produces different posterior parameters for each input. However, such models (e.g., Variational Auto-Encoders) are trained with a shared, fixed prior $p(\theta)$ across all inputs and only varying posterior $p(\theta \vert x)$.

Question 2:

ACE and ECE refer to the same metric, differing only in whether the expectation is taken over the true data distribution (in the case of ECE) or over a finite test sample (in the case of ACE). In practice, evaluations are based on empirical estimates, and thus ACE is a slightly more precise term. That said, the term ECE is commonly used for both (so we have no objection to changing the name).

Question 3:

Please see the following table summarizing the average running time of each method across all experiments.

Our method differs from others in that it pre-trains a representation and performs variational inference only on the prediction layer. As a result, for smaller models (e.g., linear regression), our approach can be more computationally expensive due to per-environment optimization. However, for larger models, our method has an advantage since competing methods must optimize many more parameters (see the reported times for the CelebA and CIFAR experiments).

Dear authors, thank you very much for your quick response. I fully agree that emphasizing the conditioning on individual covariates in the Related Work section will be helpful, as will the runtime comparisons in the Appendix.

I also appreciate your clarification regarding ECE/ACE. Indeed, this is the metric that I had in mind, and there seems to be some inconsistency in terminology across the literature. So I do think laying this out briefly in the paper does make sense.

My questions have been resolved satisfactorily. I do not have any more questions at this time. I'm planning to adjust my score upwards accordingly.

Dear reviewer,

Thank you for your thorough feedback and for responding so quickly!

We are glad to hear that the distinction and novelty of our approach are now clearer. As mentioned, we will include a corresponding discussion in the Related Work section of the paper, emphasizing the conditioning on individual covariates.

Following your suggestion, we will also add runtime comparisons to the appendix.

Regarding ACE and ECE, thank you for pointing us to the relevant reference. We now see that there is some inconsistency in terminology across the literature. In any case, the metric reported in our paper (as can be seen in the code attached to the original submission) corresponds to what you referred to as ECE. We will ensure that the paper explicitly defines this metric to avoid any ambiguity.

Metric computation: We divide the predicted confidence scores (i.e., the predicted probabilities assigned to the predicted class) into 10 equal-width bins with edges at 0.0, 0.1, 0.2, ..., 1.0. For each bin, we compute:

• the average predicted probability $\hat{p}$
• the fraction of samples where the predicted label matches the true label We then calculate the absolute difference between the bin's accuracy and its average confidence. This difference is weighted by the proportion of total samples that fall into the bin. We report the sum of these weighted differences across all bins.

Dear reviewer,

Thank you for your review. We address the raised questions and weaknesses below.

• Smoothness in language models: The non-smooth conditional marginals in language models are the conditional distributions of the target (e.g., next word) given the context. However, our model operates in a continuous latent space where the conditional distributions are of the parameters $\theta$ (variational).

• Interpretation of calibration for uncertainty and application: Calibration, accuracy, and other metrics reported in the paper serve to evaluate how well the model captures the true conditional distribution $p(y \vert x)$. Among these, calibration is often viewed not just as an evaluation metric, but as a goal in itself—reflecting the reliability of the predicted probabilities rather than just the correctness of the most likely label. Intuitively, if for every $\alpha \in [0,1]$, for example $\alpha=0.7$, indeed 70% are positive, it reflects correct uncertainty, that can be used for decision making, selective prediction, and risk-sensitive applications.

• Computation of calibration: In classification, calibration is estimated via binning: predicted probabilities are grouped into discrete intervals (bins), and for each bin, the average predicted probability is compared to the empirical accuracy—that is, the proportion of correctly classified examples within that bin. We follow this standard estimation.

• NN classifiers vs. regression models: Our use of regression models was mainly since they allow a simple visualization of the model and the predictive variance. However, the rest of our experiments were on NN architectures (e.g., CIFAR, Celeb-A, UCI datasets). In fact, our method is agnostic to the architecture and can be applied to any neural network.

Dear reviewer,

Thank you for dedicating the time to write a thorough review. Please see below.

1. Runtime:

The reported run-times were for the entire experiment (all methods). Here is a table summarizing the average running time of each method across all experiments:

You are correct that using multiple environments increases computational cost. However, this overhead is significant only for small models with few parameters. For larger models (CIFAR, CelebA experiments), our method is more efficient: it pre-trains a representation using standard likelihood optimization and performs variational inference only on the final prediction layer. In contrast, competing methods involve more complex computations across all layers.

As for hyperparameter tuning, our method shares the standard set of hyperparameters used by all methods (e.g., batch size, learning rate). The additional hyperparameters relate to environment sampling: the number of environments, their size, and the penalty coefficients. We found that performance was largely insensitive to the number of environments, as long as there are more than a few (we used 10 in the CIFAR experiments). We did observe sensitivity to the penalty coefficient, with consistently better results for smaller values.

We will include these additional details in the appendix.

2. Effectiveness of synthetic environments:

The main limitation of our strategy lies in the support of the source and target distributions. If the supports overlap sufficiently—or, as in our regression experiments, if learning from a sub-region of the training support suffices—then the method is effective.

To build intuition, we can consider the linear regression example. Since the noise variance grows with $x$, the region near $x=a$ is most informative for learning the heteroscedasticity structure. Given that training inputs are uniformly distributed, the probability that a sample falls within an $\epsilon$--neighborhood of $a$ is $\frac{\epsilon}{a}$. The probability $\rho$ that at least $k$ out of $m$ samples in a single environment fall in that region follows a binomial distribution. To ensure, with confidence at least $1-\alpha$, that at least one of $J$ sampled environments contains $k$ such points requires $J \geq \frac{\log \alpha}{\log 1-\rho}$. For example, with $\alpha=0.05, m=20, k=10, a=0.5, \epsilon=0.1$ this yields a requirement of at least 11 environments.

A similar argument can be developed for all experiments. For example, in the CIFAR-experiments we have a probability of 0.1 of a single sample to be corrupted, and then the same calculation applies only with the corresponding $\rho$.

3. Baseline comparisons:

We mention ensemble models in the related work section, as they are commonly interpreted as uncertainty estimators. However, their uncertainty estimates arise from averaging predictions across multiple models—effectively replacing the original prediction model rather than evaluating uncertainty of a given model. As a result, the quality of the uncertainty estimate depends heavily on the strength of the ensemble, making it an unfair basis for comparison.

Moreover, ensemble methods do not explicitly address distribution shifts, unlike many distance-based approaches. For these reasons, our comparisons focus on methods that are most closely aligned with our approach: those that integrate Bayesian and distance-based components.

4. Ablation: We did not include an ablation study since the components of our method are not designed to function in isolation. By sampling a sufficiently large number of synthetic environments, we are guaranteed to encounter one that approximates the unseen distribution shift. However, since we don’t know which one, the model must be constrained to perform well across all environments.

That said, we did run the proposed ablation. The results below show that without the penalty across environments, our method still outperforms competing baselines in most cases, but with a smaller margin than the complete version.

5. Minor clarifications:

• Tables 3 and 4: The main reason lies in the experimental design. For the classification datasets, we had access to additional annotations: corruption types in CIFAR and attributes in CelebA. This allowed us to define distribution shifts explicitly. In contrast, for the regression datasets, the subtypes were constructed via clustering, making the shifts less clearly defined.
• Normalization constant: Thank you for pointing this out. The implementation approximates the log of the normalizing factor $z$ using the log of the mean of exponential log-likelihoods: $\log v \approx \log (\frac{1}{n} \sum_i exp(v_i))$ where $v_i$ are the integrated log likelihoods. We will add this to the manuscript.
• Figure 2: Yes, the shift is one feature which is constant at training and no longer at test. We will explain this in the legend as suggested (the main text already mentions this).
• Batch sizes: the reported numbers account for the multiple environments.
• Notation: Throughout all the paper $M$ is the number of overall test examples, and $m$ is the number of environment test examples (within a synthetic environment).

Dear reviewer,

Thank you for your review. We address the raised questions below.

1. The prior is not conditioned on the true outcome $y$. It integrates over all possible values $y$. For example, in binary classification this would mean $\log p(\theta \vert x, y=0) + \log p(\theta \vert x, y=1)$, regardless of the true unknown $y$.

2. Aside from standard hyperparameters shared across all methods (e.g., batch size, learning rate), the additional hyperparameters in our method are for environment sampling: the number of environments, their size, and the penalty coefficients. We found that performance is insensitive to the number of environments, as long as there are more than a few (e.g., we used 10 in the CIFAR experiments). We did observe sensitivity to the penalty coefficient, with consistently better results for smaller values. Hence, as the tables in the appendix show, for each type of experiment (UCI datasets, 3 corruptions of CIFAR, 3 attributes in Celeb-A) share the same hyper-parameters.

3. Please see the attached running times:

Unlike other methods, our approach first pre-trains a representation and then performs variational inference only on the prediction layer. Thus, for large models, competing methods must optimize the full parameter set, and thus ours has an advantage. This is reflected in the reported times for the CelebA and CIFAR experiments. However, for small models our method is likely to be more computationally intensive due to computations across multiple environments.