Rethinking Approximate Gaussian Inference in Classification

We are grateful to the reviewer for their careful evaluation of our paper. We thank them for acknowledging the principled nature of our method, the theoretical analyses that support our claims, and the extensive set of experiments.

Eqs. (10) and (12). We agree that the approximation in Eqs. (10) and (12) is the key approximation in our work. Therefore, it should be better introduced. It consists of bringing the expectation inside the fraction: the expectation of the fraction is approximated by the fraction of the expectation. Such an approximation is sometimes referred to as a mean-field approximation [1], although in this work, we specifically use the term to describe a method proposed for approximating the Gaussian softmax integral [2]. We discuss the quality of this approximation later in the paper, both theoretically (Theorem 2.1) and empirically (in the synthetic experiment, Figure 1). As the variance across the different classes decreases, the approximation becomes more accurate (in the limiting deterministic case, the approximation is exact). The aim of Theorem 2.1 is to quantify this limiting behaviour.

To improve the readability of the paper, we propose to explain in words after Eqs. (10) and (12) what the approximation consists of, and point the reader toward Theorem 2.1 to give the reader a glimpse of the analysis that is to come.

Related work. A Related Work section indeed greatly helps contextualise the present work in the literature. To avoid repetition, we have included the full text of the new Related Work section in our response to Reviewer LBev.

Sections 3 and 4. Section 3 addresses how to obtain Dirichlet distributions, which offer more expressivity than a single predictive value. These can be used, for instance, for uncertainty disentanglement or out-of-distribution (OOD) detection, as we showcase in the experiments.

Section 4 discusses the advantageous training opportunities that the element-wise sigmoid and normCDF activations provide. The goal is to demonstrate that these activations are not only beneficial for approximate Gaussian inference but also for training the underlying neural network.

We appreciate the reviewer's point about connecting these sections to the broader framework. To address this, we have expanded our introduction with a detailed list of contributions that explicitly shows how each section builds toward our overall framework. We have pasted it in full in our response to Reviewer LBev.

Used metrics. While Appendix H provides a formal overview of all the metrics we use, we agree with the reviewer that a more intuitive description would aid readability. We provide such an overview below and also give more context in the updated paper.

• The negative log-likelihood (NLL) metric measures how well the trained model fits the test datasets (i.e., whether it gives a high probability to the correct label). Given the model's predicted probability $p(y \mid x)$ of label $y$ for input $x$, we compute $- \log p(y \mid x)$. As the negative of the log probability scoring rule, NLL penalises both over- and underconfidence, reaching its minimum only when the model outputs the ground-truth class probabilities. Lower values indicate better performance (indicated by the arrows in each plot's $y$-axis).
• The expected calibration error (ECE) measures whether a model's confidence aligns with its actual accuracy. We define confidence as the maximum predicted class probability (e.g., for predicted probabilities $(0.3, 0.2, 0.5)$, the confidence is $0.5$). Test samples are binned into fifteen equal intervals over $[0, 1]$ based on their confidence scores. Within each bin, we compute the L1 distance between the average confidence and the actual accuracy, then average this distance across all bins. A well-calibrated model has similar confidence and accuracy in each bin.
• The correctness prediction task is a binary problem to determine whether the model can predict its own correctness. For each input $x$, the model's predicted class $\hat{y}(x)$ is the argmax of the predicted probability vector. We define the correctness indicator as $\ell(x) := 1[\hat{y}(x) = y]$, which is one if and only if the model correctly predicts the label for input $x$. These serve as the binary labels for the correctness prediction task. Given the model's uncertainty estimate $u(x) \in [0, 1]$ for each input $x$ (e.g., its confidence, defined in the previous point), we evaluate how well $u(x)$ predicts $\ell(x)$ using the binary cross-entropy loss over a test dataset. The negative of this value is often referred to as the log probability strictly proper scoring rule.
• Like correctness prediction, the OOD detection problem is also a binary prediction task, where a mixture of in-distribution and out-of-distribution samples is provided. The model's uncertainty estimate $u(x)$ (see above) is tasked to separate in-distribution and out-of-distribution samples. To measure the model's performance, we use the Area Under the Receiver Operating Curve (AUROC) metric, which measures how separable the in-distribution and out-of-distribution samples are based on the model's uncertainty estimate (where 1.0 indicates perfect separation and 0.5 indicates random performance).

Timings. We agree that including more reports of timings in the experimental section helps to underscore the efficiency of our method. Regarding your specific question about the 7% overhead, this measurement was from preliminary experiments on different hardware (MacBook Pro M4). For more systematic comparisons, we measured the time to obtain the predictives from the logit-space Gaussians compared to the time of a forward pass on an NVIDIA RTX 2080 Ti GPU (with similar results on Tesla A100 GPUs). As shown below, ResNet-50 with 1000 samples shows 6-8% overhead depending on the batch size, confirming our claim. In contrast, the runtime cost of computing our closed-form predictives is negligible compared to a forward pass, especially for larger batch sizes.

ResNet-26 (C=10) Runtime Overhead (sampling time as % of total inference time)

Wide-ResNet-26-5 (C=100) Runtime Overhead (sampling time as % of total inference time)

ResNet-50 (C=1000) Runtime Overhead (sampling time as % of total inference time)

These results demonstrate that even with a diagonal covariance structure, MC sampling yields a nontrivial overhead, especially for large numbers of classes and samples. In comparison, our closed-form approach has negligible computational cost.

Thank you again for your constructive feedback. We have incorporated all your suggestions to improve the paper's readability: adding a Related Work section, a list of contributions to explain the organisation of the paper, clear explanations of the key approximations, timing results, and intuitive descriptions of all metrics. We believe these changes significantly strengthen the presentation of our work.

[1]: Friedman & Koller, (2009). Probabilistic Graphical Models -- Principles and Techniques

[2]: Lu et al., (2021). Mean-field approximation to Gaussian-Softmax integral with application to uncertainty estimation

We are very grateful to the reviewer for their thorough read of our paper, including the technical appendices, and for acknowledging the novelty of our work, the clarity of the presentation, the empirical and theoretical arguments that support the claims made in the paper, and the extensive set of experiments.

Below, we address the points raised by the reviewer one by one.

Logit-space Gaussian assumption. As the reviewer points out, our work concerns methods that output Gaussian distributions over the logits. While this setting might seem restrictive at first, it encompasses many state-of-the-art methods for uncertainty quantification, making it arguably one of the dominant paradigms in uncertainty quantification for deep learning. Our empirical comparison of approximate Gaussian inference methods against evidential deep learning methods also supports this claim (see our response to Reviewer LBev). In the revised manuscript, we include a Related Work section which, among other things, describes the significance of such methods. To avoid repetition, we have pasted this new related work section in full in our response to Reviewer LBev.

Computational benefits. We agree that providing more timing reports in the experimental section emphasizes the benefits of sample-free inference. We compare the time to construct Gaussians with diagonal covariance compared to the time of a forward pass. Below we show results on an NVIDIA RTX 2080 Ti. NVIDIA Tesla A100 GPU gave similar results. In reality, the overhead of MC sampling may be even greater, as each of these logit Gaussian samples needs to be mapped through the softmax function. In contrast, the runtime cost of computing our closed-form predictives is negligible compared to a forward pass.

Resnet-26 (C=10)

Wide-Resnet-26-5 (C=100)

Resnet-50 (C=1000)

Pre-trained models. Thank you for raising the important point about pre-trained models. While Appendix K.1 discusses the limitations of applying our approximation directly to off-the-shelf softmax models, we have an encouraging observation from our experiments: In our ImageNet setup, we started with pre-trained ResNet-50 and ViT-Little backbones and fine-tuned them with sigmoid/normCDF activations. We observed that these models adapted remarkably quickly, achieving near-optimal loss and accuracy after just 4-5 ImageNet epochs. This suggests that our approach remains practical in the era of large models, requiring only brief fine-tuning rather than training from scratch. Given the strong uncertainty quantification benefits demonstrated throughout the paper, we believe this short adaptation period is a worthwhile trade-off. This finding also opens interesting directions for future work on efficient uncertainty transfer between pre-trained models and our framework.

Softmax likelihood alternatives. We appreciate the pointer to alternatives to the softmax likelihood and have clarified this in the paper. Of the cited works, robust-max [1], logistic-softmax [2], stick-breaking [3], and GP-Tree [4] do replace the softmax likelihood, whereas one-vs-each [5] is a training surrogate for softmax (not a normalised likelihood) and the Dirichlet/tempered perspective for BNNs [6] retains softmax while re-specifying aleatoric confidence. All of these approaches require sampling or one-dimensional quadrature at inference. By contrast, given Gaussian logits from any back-end, our method provides closed-form, sample-free predictives and second-order Dirichlet quantities, with a theoretical error analysis. To better contextualise our work in this significant body of literature, we now include a Related Work section in the paper. To avoid repetition, we have pasted it in full in our response to Reviewer LBev.

Technical Questions

T1: Appendix C1 performs an informal theoretical analysis specifically under the setting of the synthetic experiment from Figure 1. In Eq. (43), we use that the logit means and variances are sampled i.i.d. for each class; hence, $\mathbb{E}[Q_c] \\approx \\mathbb{E}[Q_1]$ for all $c$, and they are of the same order in the number of classes. The reviewer is correct in saying that this is not exact, but an approximation. Thank you for spotting this. We have replaced the equality sign with an approximation sign and explained the source of this approximation. Note that this approximation is valid in this informal analysis because we only care about the order of the quantities as a function of $C$, and in this case, $\\mathbb{E}[Q_c]$ will be the same order as $\\mathbb{E}[Q_1]$, i.e., constant in $C$.

T2: Thank you for catching the missing summation in Eq. (58). We have now corrected this, and the rest of the analysis carries through as before. For completeness, we demonstrate this here:

$= \sum_{c=1}^C p_c\left(\frac{p_c - \tilde p_c}{p_c} + \frac{(p_c-\tilde p_c)^2}{2\omega_c(\tilde p_c)^2}\right)$

$= \\sum_{c=1}^C p_c - \\sum_{c=1}^C \\tilde p_c + \\sum_{c=1}^Cp_c \\frac{(p_c-\\tilde p_c)^2}{2\\omega_c(\\tilde p_c)^2}$

$= \sum_{c=1}^C\frac{p_c}{2\omega_c(\tilde p_c)^2}\cdot O\left(\operatorname{Var}\left(\sum_{c'=1}^C Q_{c'}\right)^{1/2}\right)^2$

$= O\left(\operatorname{Var}\left(\sum_{c=1}^CQ_c\right)\right)$

as $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right)\to 0$, where $\omega_c(\tilde p_c)\in [p_c,\tilde p_c] \text{ or }[\tilde p_c,p_c]$, where we used that $p_c$, $\tilde p_c$, and hence $\omega_c(\tilde p_c)$ is bounded away from $0$ by compactness of $\mathcal K$.

T3: In practice we observe that $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right)$ is very small (between $0.001$ and $0.01$). Meanwhile, the individual variances $\operatorname{Var}(Q_c)$ remain meaningful for uncertainty quantification, as we see in the various experiments that the approximate Gaussian inference methods significantly outperform the neural networks' point estimates (Appendix K.5 for NLL). Moreover, there is a bound of the form $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right) \leq C \sum_{c=1}^C \operatorname{Var}(Q_c)$. On the other hand, $\sum_{c=1}^C \operatorname{Var}(Q_c)$ cannot be bounded by $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right)$ in general. This shows that, for our approximation to be good, we do not necessarily require the individual uncertainties in the classes to be small, a small $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right)$ is sufficient.

To demonstrate this, we run a synthetic experiment where we sample $200$ logit means and covariance matrices such that $\sum_{c=1}^C \operatorname{Var}(Q_c) > 0.1 \cdot C$, take the minimum $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right)$, and show that it is much smaller than $\sum_{c=1}^C \operatorname{Var}(Q_c)$, see the table below. This provides an existence proof that our approximation can work well even with significant individual class uncertainties, as long as $\operatorname{Var}\left(\sum_{c=1}^C Q_c\right)$ is small.

Minor

M1: Our $\mathrm f(x)$ is a Gaussian distribution, not a point estimate (see Eq. 5). So we believe our Eq. (6) is correct.

M2: We have now fixed the link in the appendix. Thank you for pointing it out.

Thank you again for your detailed review. We believe we have addressed your main concerns regarding the scope of our method, computational benefits, and technical questions. We are happy to clarify any remaining points.

[1]: Hernández-Lobato et al., (2011). Robust Multi-Class Gaussian Process Classification

[2]: Galy-Fajou et al., (2020). Multi-Class Gaussian Process Classification Made Conjugate: Efficient Inference via Data Augmentation

[3]: Linderman et al., (2015). Dependent Multinomial Models Made Easy: Stick-Breaking with the Pólya-gamma Augmentation

[4]: Achituve et al., (2021). GP-Tree: A Gaussian Process Classifier for Few-Shot Incremental Learning

[5]: Titsias, (2016). One-vs-Each Approximation to Softmax for Scalable Estimation of Probabilities

[6]: Kapoor et al., (2022). On Uncertainty, Tempering, and Data Augmentation in Bayesian Classification

We are grateful to the reviewer for their thorough evaluation of our manuscript and for including several important references. We thank them for praising the clarity of our work, the extensiveness of our experimental analysis, the versatility of our framework, and the detailed implementation provided. We address their two central questions/concerns below.

Significance

We agree that the probit approximation for the Gaussian sigmoid integral is well-known, and certainly not novel to our work. It dates back to [1, 2], and as the reviewer points out, is used in recent works for binary classification, such as [3]. As the reviewer writes, we generalise it to multi-class settings. Furthermore, this is the first work that provides a theoretical analysis of such an approximation, to the best of our knowledge. We also show empirically that it outperforms other approximations from the literature. We extend it further to allow for not just predictive approximations but second-order Dirichlet approximations. Finally, we leverage the advantageous structure of the activation to train with the binary cross-entropy loss. Combined with our closed-form predictives, we demonstrate the improved uncertainty quantification capabilities of this pipeline.

To clarify these points, we now include a List of Contributions section at the end of the introduction:

List of Contributions

Our work extends the classical probit approximation for the Gaussian sigmoid integral in binary classification [1, 2] in the following manner:

• Generalising it to the multi-class setting (Sections 2.1 and 2.2),
• Analysing it theoretically (Section 2.3),
• Analysing it empirically, and how it outperforms previous approximations (Figure 1 and Section 6),
• Generalising it to second-order distributional Dirichlet approximations (Section 3),
• Combining it with advantageous training (Sections 4 and 6).

These contributions are detailed in our revised manuscript, which also includes expanded Related Work and experimental comparisons.

Contextualisation

We agree that contextualising our work and comparing it to the existing literature is crucial, and we have spent a significant amount of effort in this rebuttal to do so.

Firstly, we compare our approach to evidential deep learning methods [4, 5] and present the results below.

ImageNet

CIFAR-100

CIFAR-10

Following our experimental setup (Section 6), we optimise the hyperparameters of both methods using a ten-step Bayesian Optimization scheme in Weights & Biases. Compared to the results presented in our paper, our approach outperforms the two EDL methods across all metrics, significantly so on NLL and OOD AUROC. The results on the latter metric indicate that our moment-matched Dirichlet distributions capture second-order epistemic information more effectively than the loss-based EDL methods, which aligns with the findings of Bengs et al. [6].

We also now compare to an additional softmax Gaussian integral approximation from [7].

Finally, we have written a new Related Work section that contextualises our work w.r.t. all references shared by the reviewer. We paste this verbatim below.

Related Work

In uncertainty quantification, methods that output logit Gaussians, which we refer to as approximate Gaussian inference methods, have become one of the dominant paradigms in uncertainty quantification [3, 8, 9, 10]. In classification tasks, such methods have the downside of requiring sampling at inference to transform Gaussian distributions into predictive probability vectors.

In parallel, prior work has developed ways to propagate probabilistic uncertainty through neural networks [11]. Several works develop analytic propagation of covariances up to the logits [12, 13]; our method complements these by addressing the propagation of the logit-space Gaussian distributions onto the probability simplex. Few works address this last step. The Laplace bridge [14] proposes a map from Gaussians to Dirichlets on the probability simplex. However, this approach offers no guarantees on the quality of its predictives, and it severely underperforms Monte Carlo sampling. Other approaches include the mean-field approximation to the Gaussian softmax integral [15] and the approximation of [7]. In this work, we propose approximate predictives that outperform existing closed-form predictive techniques and provide second-order Dirichlet distributions for other closed-form quantities of interest. Furthermore, this is the first work to provide a theoretical analysis of such an approximation.

Other works have also proposed various activations for classification to replace the softmax. Particularly relevant is [16], which composes element-wise sigmoids with normalisation and utilises auxiliary-variable augmentation to obtain closed-form ELBO terms during training; however, its predictives require Monte Carlo estimation. Related conjugate/augmentation-based objectives exist for stick-breaking and tree-structured multi-class GPs [17, 18], and for BNNs, a tempered-likelihood/Dirichlet perspective with a factorised-Gaussian surrogate has been explored [19]. Yet, all these approaches still require sampling or numerical quadrature at inference time. [1, 2, 3] use the probit approximation for the Gaussian sigmoid integral in the binary case, which we generalise to the multiclass setting. [20] approximates the softmax by providing lower bounds for point estimate training, but replacing the softmax with these bounds during training does not yield a normalised multi-class likelihood. The binary cross-entropy loss with the element-wise sigmoid activation has been reported to outperform cross-entropy in some large-scale settings [21, 22]. Our work is the first to use this advantageous loss while normalising the output to obtain proper multi-class probabilistic models.

Evidential deep learning (EDL) approaches [4, 5] offer sample-free outputs through Dirichlet distributions but have been shown to be unreliable for uncertainty estimation and disentanglement [6, 23, 24, 25]. We focus on making approximate Gaussian inference methods sample-free in a principled and performant manner; however, for completeness, we compare our approach with EDL methods in Appendix L.

Finally, multiple authors have explored continuous approximations to categorical distributions, which enable gradient-based optimization for discrete stochastic models through sampling, such as the Gumbel-max trick [26, 27, 28, 29]. These methods principally optimise a single categorical distribution, while the present work is concerned with constructing probability densities over the (continuous) simplex of such categorical distributions.

[1]: Spiegelhalter & Lauritzen, (1990). Sequential updating of conditional probabilities on directed graphical structures

[2]: MacKay, (1992). The Evidence Framework Applied to Classification Networks

[3]: Kristiadi et al., (2020). Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks

[4]: Sensoy et al., (2018). Evidential Deep Learning to Quantify Classification Uncertainty

[5]: Charpentier et al., (2020). Posterior Network

[6]: Bengs et al., (2022). Pitfalls of Epistemic Uncertainty Quantification through Loss Minimisation

[7]: Shekhovtsov & Flach, (2018). Feed-forward Propagation in Probabilistic Neural Networks

[8]: Collier et al., (2021). Correlated Input-Dependent Label Noise in Large-Scale Image Classification

[9]: Daxberger et al., (2021). Laplace Redux

[10]: Liu et al., (2020). Simple and Principled Uncertainty Estimation with Deterministic Deep Learning

[11]: Gast & Roth, (2018). Lightweight Probabilistic Deep Networks

[12]: Wright et al., (2024). An Analytic Solution to Covariance Propagation in Neural Networks

[13]: Li et al., (2024). Streamlining Prediction in Bayesian Deep Learning

[14]: Hobbhahn et al., (2022). Fast Predictive Uncertainty for Classification with Bayesian Deep Networks

[15]: Lu et al., (2021). Mean-field approximation to Gaussian-Softmax integral

[16]: Galy-Fajou et al., (2020). Multi-Class Gaussian Process Classification Made Conjugate

[17]: Linderman et al., (2015). Dependent Multinomial Models Made Easy

[18]: Achituve et al., (2021). GP-Tree

[19]: Kapoor et al., (2022). On Uncertainty, Tempering, and Data Augmentation in Bayesian Classification

[20]: Titsias, (2016). One-vs-Each Approximation to Softmax for Scalable Estimation of Probabilities

[21]: Wightman et al., (2021). ResNet strikes back

[22]: Li et al., (2025). Binary Cross-Entropy Loss Leads to Better Calibration than Focal Loss on Imbalanced Multi-Class Datasets

[23]: Jürgens et al., (2024). The Epistemic Uncertainty of Evidential Deep Learning is not Faithfully Represented

[24]: Pandey & Yu, (2023). Learn to Accumulate Evidence from All Training Samples: Theory and Practice

[25]: Mucsányi et al., (2024). Benchmarking Uncertainty Disentanglement

[26]: Maddison et al., (2014). A* Sampling

[27]: Jang et al., (2017). Categorical Reparameterization with Gumbel-Softmax

[28]: Maddison et al., (2017). The Concrete Distribution

[29]: Huijben et al., (2023). A Review of the Gumbel-max Trick

We thank the reviewer for their careful review of our manuscript.

Contextualisation. We agree that better contextualization is needed. To this end, we have added a new Related Work section and an explicit list of contributions. To avoid repetition, these appear in full in our response to Reviewer LBev.

The Gumbel-max trick. Thank you for highlighting this important related work. We have added all four of your suggested references of the Gumbel-max trick and related continuous relaxations to our Related Work section. While both our work and the Gumbel-max trick involve alternatives to standard softmax, they address different challenges. The Gumbel-max trick enables gradient-based optimization through discrete random variables by providing continuous relaxations of categorical sampling. Our approach instead provides fully sample-free inference by deriving closed-form expressions for both predictives and second-order distributions over the probability simplex through moment matching. Among other things, our analytic approximate push-forward is radically faster than sampling (see also the response to Reviewer a52w below). The key difference is the source of stochasticity. In Gumbel-max, the randomness originates from sampling discrete categories to enable gradient-based optimisation. In our setting, the randomness comes from Gaussian distributions over logits for uncertainty quantification, which we refer to as 'approximate Gaussian inference'. The prototypical example of an approximate Gaussian inference method is the Laplace approximation (e.g., [1]). This is an approximate Bayesian method that approximates the posterior over the neural network parameters with a Gaussian distribution. This Gaussian is then pushed forward to logit space, and the present work addresses how to further push this Gaussian forward onto the probability simplex without sampling.

Used approximations. Thank you for examining our approximations carefully. To the best of our knowledge, Eqs. (7) and (8) have never been used for analytically approximating predictives. Importantly, Eqs. (7) and (8) are presented as a bridge between traditional softmax-based approximate Gaussian inference and our proposed predictive approximation for the sigmoid and normCDF activations, which we present as the main method to use. While the probit approximation has been used in binary classification, we extend it to the multi-class case. We make this contribution clearer in our new list of contributions, which we display in full in our response to Reviewer LBev.

In Eq. (12), the expectation of the fraction is approximated by the fraction of the expectation. Such a type of approximation is well-known in the Bayesian literature and sometimes referred to as a mean-field approximation [2], though in this work, we refer to the mean-field approximation as a specific one that has been proposed to approximate the Gaussian softmax integral following Lu et al. [3]. We discuss the quality of our approximation later in the paper, both theoretically (Theorem 2.1) and empirically (in the synthetic experiment, Figure 1). Crucially, this approximation is highly accurate with our proposed normCDF and sigmoid activations, as shown in our experiments. As the variance across the different classes decreases, the approximation becomes more accurate (in the limiting deterministic case, the approximation is exact). The aim of Theorem 2.1 is to quantify this limiting behaviour, which we believe is the first error analysis for such a mean-field approximation in probabilistic deep learning.

We hope that our new Related Work section, the list of contributions, and clarifications above address the reviewer's concerns. We are happy to clarify any remaining points.

[1]: Daxberger et al., (2021). Laplace Redux -- Effortless Bayesian Deep Learning

[2]: Friedman & Koller, (2009). Probabilistic Graphical Models -- Principles and Techniques

[3]: Lu et al., (2021). Mean-field approximation to Gaussian-Softmax integral with application to uncertainty estimation