Source-linked AI summary

Being Bayesian about Categorical Probability

Taejong Joo, Uijung Chung, Min-Gwan Seo

arXiv:2002.07965v2cs.LGstat.ML

TL;DR

Softmax classification lacks uncertainty representation and can be poorly calibrated, while Bayesian neural networks face costly approximate inference and difficult prior specification. The paper instead applies Bayesian inference to categorical probabilities, reporting gains in generalization, uncertainty estimation, and calibration.

  • Problem

    Softmax cannot represent predictive uncertainty and can be poorly calibrated, despite its standard use for neural-network classification.

    Source: §1 · p. 1
  • Method

    The method treats categorical probability as a random variable and uses a prior with Bayesian inference to construct the classifier’s target distribution.

    Source: §1 · p. 1, §3.1 · p. 3
  • Results

    Across experiments, the method improves generalization, uncertainty estimation, and calibration; on CIFAR-10, ECE improves from 3.82 with softmax to 1.66 with BM.

    Source: §1 · p. 1, §5.2 · p. 7, §6 · p. 9
  • Takeaways & Limitations

    BM can be applied to standard deep models by replacing only softmax and cross-entropy, with negligible added computational complexity.

    Source: §1 · p. 1, §6 · p. 9

    Bayesian neural networks require approximate posterior inference and can make meaningful weight-prior specification difficult in highly nonlinear networks.

    Source: §1 · p. 1

Abstract

from arXiv · show

Neural networks utilize the softmax as a building block in classification tasks, which contains an overconfidence problem and lacks an uncertainty representation ability. As a Bayesian alternative to the softmax, we consider a random variable of a categorical probability over class labels. In this framework, the prior distribution explicitly models the presumed noise inherent in the observed label, which provides consistent gains in generalization performance in multiple challenging tasks. The proposed method inherits advantages of Bayesian approaches that achieve better uncertainty estimation and model calibration. Our method can be implemented as a plug-and-play loss function with negligible computational overhead compared to the softmax with the cross-entropy loss function.

1. Introduction

Softmax classification lacks uncertainty representation and can be poorly calibrated, while Bayesian neural networks bring practical inference and prior-specification obstacles. This paper applies Bayesian inference to categorical probabilities to improve generalization, uncertainty estimation, and calibration with minimal changes to existing models.

1. Introduction

  • Softmax is widely used for classification but cannot represent predictive uncertainty and can behave poorly in calibration.The passage also notes that confidently wrong predictions and overfitting are concerns.Source: §1 · p. 1
  • Bayesian neural networks can address softmax’s issues, but approximate posterior inference is computationally expensive and meaningful weight priors are difficult to determine.They may also require substantial changes to baselines or result in degraded performance.Source: §1 · p. 1
  • The method treats categorical probability as a random variable and uses Bayesian inference to construct classifier target distributions.The prior belief regularizes the target, with its impact controlled by the number of observations.Source: §1 · p. 1
  • Experiments show improvements in generalization, uncertainty estimation, and calibration, with negligible added computational complexity.The approach replaces the softmax and cross-entropy loss in existing neural networks.Source: §1 · p. 1
  • Figure 1 contrasts direct label-to-target conversion under softmax cross-entropy with combining the label and a prior Dirichlet distribution before updating the target belief.Source: Figure 1 · p. 2

2. Preliminary

The preliminary section formulates classification learning through softmax cross-entropy and distribution matching, then describes weaknesses in empirical target estimation and regularization alternatives.

2. Preliminary

  • A classifier maps inputs to labels, and the training data comprise input-label samples with labels represented as one-hot vectors.Source: §2 · p. 2
  • Softmax transforms network logits into a categorical distribution, and cross-entropy measures its mismatch with the one-hot label.The softmax output is treated as the categorical distribution’s parameter.Source: §2 · p. 2
  • Minimizing softmax cross-entropy can be expressed as matching the network’s categorical distribution to an empirical target formed from label counts at each input.The empirical risk sums these input-specific matching terms, up to a constant with respect to the network weights.Source: §2 · p. 2
  • Because most inputs are unique or rare, the empirical target distribution may be estimated from only one or a few samples.Source: §2 · p. 2
  • Label smoothing mixes empirical counts with a uniform distribution, but its fixed mixing coefficient may limit use of empirical counts when needed.Source: §2 · p. 2
  • Weight-prior Bayesian neural networks balance target fit with model complexity, but weight norms may poorly proxy complexity in modern networks.The passage attributes this concern to strong weight-output nonlinearity and batch-normalization weight-scale invariance.Source: §2 · pp. 2–3

3. Method

The method treats categorical probabilities as random variables, uses a Bayesian prior to construct target distributions, and trains neural networks to match them. Its Dirichlet outputs retain distributional information beyond the softmax mean, while the analytical loss remains computationally scalable.

3.1. Constructing Target Distribution

  • The categorical probability is treated as a random variable, with a Dirichlet prior producing a posterior target whose mean is smoothed by prior belief and adapts its relative strength to each data point.The prior acts like adding pseudo-counts; small β0 limits prior dominance when most inputs are unique.Source: §3.1 · p. 3

3.2. Representing Approximate Distribution

  • BM models a Dirichlet density over the simplex, retaining information beyond the softmax-equivalent mean, including predictive uncertainty such as mutual information.The density can represent how concentrated predictions are around their mean.Source: §3.2 · p. 3

3.3. Distribution Matching

  • Maximizing the ELBO matches the network’s approximate distribution to the target, with an analytically computable loss and O(K) computation per sample, equal to softmax.The resulting distribution-matching loss can be applied directly to logits as a plug-and-play loss.Source: §3.3 · p. 4

3.4. On Prior Distributions

  • Because small concentration parameters create large ELBO gradients, the method sets β = 1 and scales the KL term by λ to stabilize optimization without changing a local optimum.The paper notes that a fully Bayesian hierarchical model for β is left for future research.Source: §3.4 · p. 4

4. Related Work

Related work includes Bayesian neural networks and other approaches to uncertainty and calibration. The paper contrasts these methods with alternatives that modify softmax or model Dirichlet concentrations for different purposes.

Related work

  • Bayesian neural networks use approximate posterior inference, but may not scale to state-of-the-art architectures on large datasets or may reduce generalization performance in practice.The cited approximation schemes include variational inference and stochastic-gradient MCMC.Source: §4 · p. 5
  • Other approaches include MC dropout and sequential weight-posterior estimation; deep kernel learning combines deterministic neural networks with Gaussian processes for uncertainty representation and robust extrapolation.The passage cites MC dropout, sequential estimation methods, and deep kernel learning (Wilson et al., 2016a;b).Source: §4 · p. 5
  • Non-Bayesian alternatives use ensembles or temperature scaling for uncertainty and calibration, and auxiliary losses or label-noise models to address overconfidence.The passage cites ensembles, temperature scaling, entropy penalties (Pereyra et al., 2017), and label-noise approaches.Source: §4 · p. 5
  • Prior Dirichlet-concentration methods target uncertainty on OOD or adversarial samples using auxiliary losses, unlike BM’s stated distribution-matching objective.Sensoy et al. (2018) minimize prediction variances on training samples; Malinin & Gales (2018) and Chen et al. (2018) use auxiliary losses on OOD or adversarial samples.Source: §4 · p. 5

5. Experiment

Across image-classification, uncertainty, transfer-learning, and semi-supervised evaluations, the experiments test BM against softmax and Bayesian baselines. They report improved generalization and uncertainty representation, alongside sensitivity to β and gains from distribution matching in semi-supervised learning.

5. Experiment

  • The experiments use residual-network benchmarks and compare BM with softmax, MC dropout, and deep ensembles; BM training uses warm-up and gradient clipping for stability.The paper notes difficulty training deep models without warm-up and clipping because of initial gradient explosion.Source: §5 · p. 5

5.1. Generalization Performance

  • BM achieves the best generalization performance across CIFAR configurations and improves softmax test errors on ImageNet.The CIFAR comparison reports test classification error rates; the ImageNet comparison uses validation error rates.Source: §5.1 · p. 6, Table 1 · p. 6, Table 2 · p. 6
  • Without the KL term, ResNet-50 test errors are 5.68% on CIFAR-10 and 24.69% on CIFAR-100, while distribution averaging still outperforms softmax.The authors attribute the remaining improvement to considering all possible categorical probabilities; the prior contributes a separate regularization effect.Source: §5.1 · p. 6
  • Generalization remains robust for β ∈ [exp(−1), exp(4)], while extreme β values can trigger gradient explosion or substantially increase error.With only β changed, the error rate reaches 7.37 at β = exp(8); jointly tuning λ and β changes the extreme-value behavior.Source: §5.1 · p. 7, Figure 3 · p. 7

5.2. Uncertainty Representation

  • BM lowers softmax ECE from 3.82 to 1.66 on CIFAR-10 and from 13.48 to 4.25 on CIFAR-100.Deep ensembles reach ECEs of 1.04 and 3.54, respectively, with 5× more training and inference computation.Source: §5.2 · pp. 7–8, Figure 4 · p. 7
  • For OOD samples, BM produces clear high-predictive-entropy peaks, unlike softmax’s relatively flat uncertainty; BM also assigns more high uncertainty to in-distribution samples.The paper notes that some in-distribution samples warrant uncertainty because test accuracy is not perfect.Source: §5.2 · p. 8, Figure 5 · p. 8

5.3. Transfer Learning

  • In transfer learning, BM improves softmax test error rates and uncertainty representation after fine-tuning only the last layer, but confidently predicts Cars examples.The authors relate the Cars result to CIFAR-10 containing the automobile category; BM shows clear high-uncertainty peaks on SVHN and Food-101.Source: §5.3 · p. 8, Table 3 · p. 8, Figure 6 · p. 9

5.4. Semi-Supervised Learning

  • On CIFAR-10 semi-supervised learning, BM’s distribution-matching consistency metric reduces error more effectively than softmax’s moment-matching metric.The experiment uses 4K labeled, 41K unlabeled, and 5K validation samples.Source: §5.4 · p. 9, Table 4 · p. 9

6. Conclusion

The conclusion presents BM as Bayesian target construction that treats categorical probability as random rather than fixed by the training label. It reports gains in generalization, uncertainty estimation, and calibration, while identifying richer distributions and input-dependent priors as future directions.

6. Conclusion

  • BM constructs its target distribution by treating categorical probability as a random variable rather than fixing it to the training label.The method can replace the softmax and cross-entropy loss in standard deep-learning models.Source: §6 · p. 9
  • The authors report consistent improvements in generalization, uncertainty estimation, and calibrated behavior from this construction.Source: §6 · p. 9
  • Future directions include more expressive distributions for class-label similarities and input-dependent priors that can control student prediction smoothness.The paper connects these directions to large-class classification, imbalanced datasets, and multi-domain learning.Source: §6 · p. 9
Loading 2002.07965v2…