Title: Rethinking Approximate Gaussian Inference in Classification
Venue: NeurIPS 2025
Accepted: yes
Keywords: Uncertainty quantification

Predictives

Predictive distribution

Sampling-free
Abstract: In classification tasks, softmax functions are ubiquitously used as output activations to produce predictive probabilities. Such outputs only capture aleatoric uncertainty. To capture epistemic uncertainty, approximate Gaussian inference methods have been proposed. We develop a common formalism to describe such methods, which we view as outputting Gaussian distributions over the logit space. Predictives are then obtained as the expectations of the Gaussian distributions pushed forward through the softmax. However, such softmax Gaussian integrals cannot be solved analytically, and Monte Carlo (MC) approximations can be costly and noisy. We propose to replace the softmax activation with element-wise normCDF or sigmoid, which allows for the accurate sampling-free approximation of predictives. This also enables the approximation of the Gaussian pushforwards by Dirichlet distributions with moment matching. This approach entirely eliminates the runtime and memory overhead associated with MC sampling. We evaluate it combined with several approximate Gaussian inference methods (Laplace, HET, SNGP) on large- and small-scale datasets (ImageNet, CIFAR-100, CIFAR-10), demonstrating improved uncertainty quantification capabilities compared to softmax MC sampling. Our code is available at https://github.com/bmucsanyi/probit.


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper addresses the problem of uncertainty quantification in classification tasks. Building upon previous work on Gaussian inference in logit space, the Authors propose multiple enhancements to the standard procedure, aiming at\n- Sidestepping MC sampling, via closed-form evaluations of the Gaussian integral and the use of tractable approximations of the pushforwards.\n- Generalizing the final Softmax activation to sigmoid and normCDF, enabling accurate and MC-free approximations.\n\nThese methods are tested on standard image classification tasks.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n\n- The proposed methods are principled, and they address the concrete problem of MC sampling for approximate Gaussian inference.\n- There are some theoretical analyses supporting the proposed methodology  \n- Experimental results tested a good variety of activations and methods on standard image classification benchmarks, showcasing reasonable performance.\n\n## Weaknesses\n\n- From a readability perspective, I think that the paper could be greatly improved. While there are some extensive appendices and many details are reported, while reading the paper, I had a hard time piecing together the information and getting an overall picture. Some concrete examples:\n    - One of the main ingredients of the recipe is the approximation in equations 10 and 12, which is introduced unceremoniously by a simple $\\approx$ sign, without any further explanation highlighting its importance.\n    - A \"related works\" or similar section is missing, and I find the broader literature around the problem not adequately introduced. The works that are discussed are instrumental in the description of the methodology, but a broader context would be appreciated.\n    - Sections 3 (and its subsections) and 4 describe an undifferentiated sequence of additional methods for which it is hard (at least to me) to define a hierarchy of importance, and no indication is given about their possible use in concrete scenarios.\n    - As reflected in my confidence score, I do not consider myself an expert in uncertainty quantification. That being said, I had some hard times understanding what key metrics were evaluated in the experiments. For most of them, the inexperienced reader is diverted to external references to get some intuition of the tested metric.\n- I think that the experimental section critically lacks a report on the timings. One of the motivations behind the proposed methods is computational efficiency, which is theoretically discussed in Section 5. One of the first things I was looking for while reading the experimental section is a report of the actual timings to compute the predictive for the various methods, but I couldn't find one.\n\nQUESTIONS:\nIn line 173 they write \"For example, on a ResNet-50, sampling one thousand logit vectors from the logit-space Gaussian means an approximately 7% overhead on the forward pass\". Where is this number coming from?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 2, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-11"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper proposes an alternative to Monte Carlo sampling when computing the predictive distribution in Bayesian NNs which model the logits using a Gaussian distribution per class in classification tasks.  The key idea is to decouple the nominator and denominator when normalizing the logits and compute each term using closed-form formulas involving Gaussian densities. The authors present a bound on the error invoked by the approximations made to derive the predictive distribution. The authors further fit a Dirichlet distribution over the simplex by means of moment matching. Empirically, the authors show comparable results in terms of accuracy to standard MC-based predictive distribution calculation while benefiting from increased uncertainty estimation and OOD detection.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n- The paper deals with a problem that is given less attention in the literature as far as I am aware. Commonly MC samples are taken which can be suboptimal. \n- Most steps in the method are supported either empirically or theoretically, for instance estimating the expectation of a ratio with the ratio between the expectation of the nominator and the expectation of the denominator. Or, using Beta distribution for relevant activation functions.\n- The paper is written clearly and the method has novelties in it.\n- The empirical results are reasonable, the method is comparable to the Softmax baseline in terms of accuracy and tends to improve in other metrics.\n\nWeaknesses/Questions:\n- The scope of the method has limitations as it assumes a Gaussian distribution over the logits per class which can be overly restrictive.\n- The authors present the computational benefit as one of the advantages of their approach. Yet, besides section 5 which presents it in terms of the number of operations I didn't notice any discussion. Perhaps the authors can show a comparison in terms of wall-clock time and memory to understand better when their method is preferred (e.g., as a function of the number of MC samples and classes).  \n- Perhaps I missed it, but do you have results for using your approximation for pre-trained networks with common cross-entropy loss and Softmax activation? If indeed models need to be trained or fine-tuned to accommodate the objective in Eq. 29 it can be a harsh restriction, especially in this era of large models. \n- In light of the Gaussian assumption over the logits, I believe that authors should have discussed, and compared where appropriate, to other alternatives for the Softmax likelihood. Several examples can be found in the Gaussian process literature. Non exhaustive list: one-vs-each [1], robust-max [2], Dirichlet observation model [3], logistic-softmax [4],  stick-breaking [5], and GP-Tree [6].\n- Technical questions:\n  - In Eq. 43, why do you assume that the expectation of $Q_c$ is equal for all classes?\n  - In Eq. 58, second line, shouldn't the last term be $\\sum_{c=1}^C p_c \\frac{(p_c - \\tilde{p}_c)^2}{2 \\omega (\\tilde{p}_c)^2)}$?\n  - Is there a justification to assume that $Var(\\sum_{c=1}^C Q_c) \\rightarrow 0$?\n- Minor:\n  - In Eq. 6 should it be $p(f(x))$ in the last term?\n  - Broken link in line 876 at the Appendix.\n\n[1] AUEB, T. R. (2016). One-vs-each approximation to softmax for scalable estimation of probabilities. Advances in Neural Information Processing Systems, 29.  \n[2]  Hernández-Lobato, D., Hernández-Lobato, J., & Dupont, P. (2011). Robust multi-class Gaussian process classification. Advances in neural information processing systems, 24.  \n[3] Kapoor, S., Maddox, W. J., Izmailov, P., & Wilson, A. G. (2022). On uncertainty, tempering, and data augmentation in Bayesian classification. Advances in neural information processing systems, 35, 18211-18225.  \n[4] Galy-Fajou, T., Wenzel, F., Donner, C., & Opper, M. (2020, August). Multi-class Gaussian process classification made conjugate: Efficient inference via data augmentation. In Uncertainty in artificial intelligence (pp. 755-765). PMLR.  \n[5] Linderman, S., Johnson, M. J., & Adams, R. P. (2015). Dependent multinomial models made easy: Stick-breaking with the Pólya-Gamma augmentation. Advances in neural information processing systems, 28.  \n[6] Achituve, I., Navon, A., Yemini, Y., Chechik, G., & Fetaya, E. (2021, July). GP-Tree: A Gaussian process classifier for few-shot incremental learning. In International conference on machine learning (pp. 54-65). PMLR.\n\nQUESTIONS:\nSee Strengths And Weaknesses section.\n\nLIMITATIONS:\nYes.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 3, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe authors use the fact that expectations over one-dimensional Gaussian distributed random variables\nare analytically tractable (or can be tightly approximated) for a number of transformations,\nsuch as the exponential function, or normal cdf, which often appear in classification objectives.\nThey use this to improve classification performance by predicting Gaussians in the logit space \nand transforming them using this result for a richer training signal. \n\nThe method is evaluated with several architectural variations on a range of image classification data sets\nand performance metrics, that demonstrate not only improved predictive performance, but also\nimproved calibration.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- The paper is well-written and can be followed quite easily\n- The method is properly motivated with an extensive set of experiments (however, see the weakness below)\n- The proposal is generic enough to be added to a wide variety of prior approaches\n- The authors provide a detailed implementation\n\n## Weaknesses\n### Significance\nThe closed-form probit solution to the expectation of a normally distributed random variable is well-known in the literature (for recent work see, e.g., Kristiadi et al., 2020; Haussmann et al., 2020; Li et al., 2025 among others and the references therein).\nAs such, I would consider this common knowledge in the field with the remaining primary contribution being the application in a multi-class setting, which is rarer. \n(At least, I can't point to a specific reference, although it might have been explored.)\nAs such, the paper's main contribution is its detailed discussion and evaluation of such closed-form approximations. \n\n### Prior work\nThe second main weakness of the paper is its lack of discussion of prior work and experimental comparison against such work. Given the small theoretical contribution, such a discussion and evaluation\nwould be the main strength of this work.\n\nThe two main areas of research that are (partially) omitted are\n\n- **deterministic/sampling-free BNN methods:** \nAs there is a large variety of approaches in that area of research, I will only refer to a few representative ones. These are, Gast et al. (2020); Wu et al., (2019); Haussmann et al. (2020) for older approaches and Wright et al. (2024); Li et al. (2025) and the references therein for current work.\nAll of them aim to infer BNNs in a sampling-free manner and apart from Wright et al. also discuss sampling-free classification approaches. \nGast et al. parameterize a Dirichlet distribution (see below). Wu et al. derive a Taylor approximation-based approach (but do not evaluate it and stick to regression). Haussmann et al. discuss the Probit approximation, but stick to evaluating an adaptation of Wu et al.'s approximation. Finally, Li et al. use the Probit approximation for their classification experiments, refering to Kristiadi et al. (2020) as their motivation. \nGiven their relevance such related work should be discussed and compared against.\n\n- **Dirichlet-based approaches:**\nApart from Gast et al. (2020) mentioned above, a major research direction is known as _Evidential Deep Learning_. See Sensoy et al., (2018) who introduced it and the diverse literature that builds on it.\nWhile Sensoy et al. (2018) are referred to in passing, their relation should be discussed and compared against in greater detail.\n\n\n\n_____\nGast et al. (2018), _Lightweight Probabilistic Deep Networks_  \nHaussmann et al. (2020), _Sampling-Free Variational Inference of Bayesian Neural Networks by Variance Backpropagation_   \nKristiadi et al. (2020), _Being Bayesian, Even Just a Bit, Fixes Overconfidence in ReLU Networks_  \nLi et al. (2025), _Streamlining Prediction in Bayesian Deep Learning_  \nSensoy et al. (2018), _Evidential Deep Learning to Quantify Classification Uncertainty_  \nWright et al. (2024), _An Analytic Solution to Covariance Propagation in Neural Networks_\n\nQUESTIONS:\nGiven the points mentioned above:\n- Q1: Can the authors provide a deeper discussion that highlights the strengths and innovations of their approach with respect to prior probit-based knowledge in the literature?\n- Q2: Can the authors discuss their approach with relation to the greater context of approaches mentioned above, and, if possible time-wise, a comparison against evidential deep learning or one of the Taylor approximation-based methods?\n\nLIMITATIONS:\nSocietal implications are not discussed. I don't consider there to be any that go beyond the general implications due to deep learning. As such this is not a weakness of the paper.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 5, "CLARITY": 4, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-26"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThe paper \"re-studies\" supervised classification problems from the perspective of softmax functions and other transformations from the *logit* space. Based on concerns around the way of modelling epistemic/aleatoric uncertainty, and the memory cost required for using Monte Carlo approximations to obtain predictive probabilities in the standard way --- the authors introduce closed-form predictives and other distributional approximations. Sadly, the work ignores a central line of research in ML, which is the Gumbel-Max trick, the Gumbel-Softmax trick, and all other contributions that in the period 2014-2020 (approx) did great advances in problems related to sampling from discrete distributions or dealing with their log-partition functions.\n\nSTRENGTHS AND WEAKNESSES:\nFirst of all, I have to say that I am surprised that the work seems to ignore and does not review classic methods to deal with classification. I was for instance wondering about the lack of reference to the following works:\n\n- Maddison, C. J., Tarlow, D., & Minka, T. (2014). A* sampling. NIPS.\n- Jang, E., Gu, S., & Poole, B. (2016). Categorical reparameterization with Gumbel-softmax. _arXiv preprint arXiv:1611.01144_\n- C. J. Maddison, A. Mnih, and Y. W. Teh, “The Concrete Distribution: A Continuous Relaxation of Discrete Random Variables,” 5th Int. Conf. Learn. Represent. ICLR 2017\n\nFrom my point of view, the authors identify certain problems in the standard use of the softmax function and how it may affect certain predictive quantities, which are obviously related to the problem of sampling from a distribution (this is somehow described in the first paragraphs of pp. 2). In addition to the reasoning about aleatoric/epistemic uncertainty, which I think does not completely fit well here (for instance, the sentence starting in L30), they observe the general property of the logits before being used in the one-hot encoding.\n\nThe situation, therefore, is the following: Eq. 7 and Eq. 8 are presented as the main discovery or method to use, where it is a Gaussian expectation of the function $\\exp(y)$. I see the advantage of the analytical result, but I would be very surprised if this has not been considered before or studied.\nThis result is then pushed forward and used to make the main approximation in Eq. 10, whose contribution I don't understand very well, or at least I question myself about the convenience of the expectation on top of the one-hot-encoding fraction. Are we really allowed to do this approximation? Seems extreme in my opinion, some discussion could be perhaps helpful.\n\nIn that direction, the absence of reference to log-partition functions, the log-sum-exp operator, and other (very) well-known methods makes me feel lost, or at least very disconnected from concepts and principles that (as a reader and practitioner) I know and have used quite often. So I cannot judge the importance of the method in a good way without such links.\n\nAfter this central result on pp. 3 is included, the same sort of technique is reused two to three times, and the Gaussian integral of the exp() function and other simple functions is similarly considered in the Section 3. I reviewed experiments as well, but I cannot really value them well without first understanding the connections and solving the previous concerns and issues first.\n\nQUESTIONS:\nCould the authors explain in some detail why the Gumbel-softmax trick or similar methods are not of interest here? Or why should they not be considered? Any connection or linking with them would be really helpful for this reviewer, who is actually a bit lost on \"why this work does not align with the references mentioned before\"\n\nIn this direction, I would also like to point out the first 3 paragraphs in Section 2 of Maddison et al. (NIPS 2014).\n\nThis TPAMI review can also help to find some central references:\n\n- Huijben, I. A., Kool, W., Paulus, M. B., & Van Sloun, R. J. (2022). A review of the gumbel-max trick and its extensions for discrete stochasticity in machine learning. _IEEE transactions on pattern analysis and machine intelligence_, _45_(2), 1353-1371.\n\nLIMITATIONS:\nN/A", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 3, "ORIGINALITY": 2, "DATE": "2025-07-04"}