Density-Softmax: Efficient Test-time Model for Uncertainty Estimation and Robustness under Distribution Shifts

Thanks for your valuable feedback. We additionally address your concerns below.

1. The role of the Lipschitz constraint is not clear; The resulting feature extractor $f$ may not be 1-Lipschitz. As discussed in Apd.B.3, we note that the Lipschitz constraint is a necessary component to combine with the density function in our framework. Because it not only helps to improve robustness but also supports improving uncertainty performance by ablation results in Tab.5.

Regarding the 1-Lipschitz, we claim that our feature extractor $f$ aims to satisfy 1-Lipschitz by gradient-penalty regularization. We do not claim that $f$ is 1-Lipschitz. We agree that our feature extractor $f$ may not be 1-Lipschitz. We will remove the sentence "1-Lipschitz feature extractor" to avoid misleading. However, we note that our feature extractor $f$ satisfy $\sup_{x}||\nabla_x f(x)||_2 = 1$ in training data with the Rademacher Theorem.

Since $f$ is a neural network, we believe that there is no way to find the exact $L(f)$ at the moment. Even SeqLip [2] is only an approximation estimation method. SeqLip also mentioned that their upper bounds estimation is still extremely large for modern neural networks and probably overestimates the true $L(f)$. As requested by the reviewer, we additionally provide the upper bound of $L(f)$ by using SeqLip's algorithm with the Wide Renet 28-10 backbone on the CIFAR-10 dataset. We observe that the upper bounds of $L(f)$ are about $4.5 \times 10^6$ without gradient-penalty and $2.8 \times 10^4$ with gradient-penalty regularization. This does not show whether our true $L(f)$ is equal to 1 or not, but at least, it shows the gradient-penalty regularization can reduce the upper bound of $L(f)$. Regarding [1], this is only provable for a specific backbone, it is not true for the ResNet backbone we used in our experiments.

2. Comparison/discussion with joint probability estimation $p(y,x)$. Thanks for your suggestion. However, we are slightly unsure of what you mean by this point. Any further clarification would be appreciated.

Our goal is to improve the uncertainty quality of $p(y|x)$ by using the density model on feature space $p(z)$. Our motivation is not to use it for generative modeling so we only compare our method with models directly modeling the conditional probability. We also cannot formulate $p(y,x) = p(y|x) p(x)$, where $p(y|x)$ is softmax output and $p(x)$ is normalizing flow's output. Because we estimate feature space $p(z)$ not sample space $p(x)$.

In terms of modeling $p(y|x)$, we agree with the reviewer that we can compare Density-Softmax with JEM [3]. However, according to Tab.1 in JEM, $p(y|x)$ of JEM only achieves $92.9\%$, and has a lower classification accuracy than the discriminative Derterministic-ERM ($95.8\%$) on CIFAR-10. As also mentioned by JEM, this worse performance is due to the fact that JEM is a hybrid model and needs to jointly optimize $p(y|x)$ with $\log p(x)$. Additionally, training $\log p(x)$ with SGLD is very sensitive and needs a long time to converge. In contrast, Density-Softmax is a discriminative model, archives a higher accuracy than Derterministic-ERM (see our Tab.2-3-4), and does not require any sampling during training.

3.1. Theorem 4.3: the maximum value of distribution itself can be less than one. We agree with the reviewer that the maximum value of an arbitrary function $p(z; \alpha)$ can be less than 1. However, as discussed in Apd.B.1, we can re-scale the maximum output of $p(z; \alpha)$ to $1$ by dividing it by the maximum likelihood value in the training set. We also empirically show that the maximum value of $p(z; \alpha)$ equals $1$ with IID data in Fig.7 in our paper.

3.2. The normalizing flow cannot assign equal probability to all IID samples. We do not claim that $p(z; \alpha)$ assigns the same likelihood for all IID samples. Please note that in Theorem 4.3, we only consider the strong IID input, i.e., when $d(z_{iid},Z_s)\rightarrow 0$.

3.3. Possible flaws in proof of minimax optimality (corollary 4.4). We agree with the reviewer that the statement "Density-Softmax satisfy by Theorem 4.3 and Proof A.1 for IID data" in our proof in Apd.A.2 does not apply to every IID data. We will remove this sentence. However, the result "$\inf_{\mathbb{P}}\sup_{\mathbb{P}^*}\left[S_{iid}(\mathbb{P},\mathbb{P^*})\right]$ is fixed" still holds due to Density-Softmax model’s predictive distribution learned from IID data (see condition (a) and proof in Proposition 2 in [4]). Please note that our density function does not modify the order of entries in the logit vector.

We hope we have thoroughly addressed your concerns. We are willing to answer any additional questions.

[1] Singla et al., Improved techniques..., 2022

[2] Scaman et al., Lipschitz regula..., 2018

[3] Grathwohl et al., Your clas..., 2020

[4] Jeremiah et al., Simple and principled uncertainty estimation with deterministic deep learning via distance awareness, 2020

Thanks for your valuable feedback. We additionally address your concerns below.

1. Comparison with NatPN. We agree with the reviewer that we should add NatPN [1] to our baselines and we added them in our revised version. That said, NatPN does not consider improving the robustness and according to Tab.2-3 in NatPN [1], it even has a lower classification accuracy than Posterior~Net [2]. Additionally, different from our deterministic approach, NatPN is a conjugate-prior-based Bayesian Inference approach and requires predefined prior hyper-parameters for the model. This hyper-parameter is often sensitive and hard to choose in practice, which has been shown in the uncertainty and robustness baseline codebase with Posteror-Net. In contrast, our model does not require any pre-defined prior hyperparameter. Therefore, from the tables in our general response, we can see that our method is much more robust and also has a better uncertainty quality than NatPN.

2. Normalization of the $p(z; \alpha)$ to be in $[0, 1]$. As discussed in Apd.B.1, we can re-scale the output of $p(z; \alpha)$ to be in $[0, 1]$ by dividing by the maximum likelihood value in the training set. We also empirically show this in Fig.7 in our paper. Regarding the proof of Proposition 4.7, if $c>1$, then Proposition 4.7 will not be true because the inequality in Eq. (46) does not hold.

3. Why the gradient penalty option to satisfy the Lipschitz constraint was chosen? Why Spectral Normalization is not used?. Firstly, Spectral Normalization needs to use the power iteration method to find the eigenvalues of every weight in the model. This is more computationally expensive because it requires non-trainable vectors in the iterative process (see our Tab.2-3-4). Secondly, we did try Spectral Normalization and gradient-penalty regularization, and we found that the gradient-penalty regularization is better in terms of robustness. These observations are similar to the gap between Density-Softmax versus SNGP in terms of accuracy and latency in Tab.2-3-4 of our paper.

4. Minor Comments. Thanks for your valuable suggestions. We will fix them in our revised version. Regarding the citation in our claim about distribution shifts, please note that we cite [4] and [5] to support our claim. Regarding Proposition 4.7, we only consider overconfidence because this is a major problem of current DNNs under distribution shifts. Regarding flow used in the experiments, we use radial flow and Real NVP.

We hope we have thoroughly addressed your concerns. We are willing to answer any additional questions.

[1] Bertrand Charpentier, Oliver Borchert, Daniel Z ugner, Simon Geisler, and Stephan Gunnemann. Natural posterior network: Deep bayesian predictive uncertainty for exponential family distributions. In International Conference on Learning Representations, 2022.

[2] Bertrand Charpentier, Daniel Zugner, and Stephan Gunnemann. Posterior network: Uncertainty estimation without ood samples via density-based pseudo-counts. In Advances in Neural Information Processing Systems, 2020.

[3] Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. In Advances in Neural Information Processing Systems, 2020.

[4] Yaniv Ovadia, Emily Fertig, Jie Ren, Zachary Nado, D. Sculley, Sebastian Nowozin, Joshua Dillon, Balaji Lakshminarayanan, and Jasper Snoek. Can you trust your model's uncertainty? evaluating predictive uncertainty under dataset shift. In Advances in Neural Information Processing Systems, 2019.

[5] Matthias Minderer, Josip Djolonga, Rob Romijnders, Frances Hubis, Xiaohua Zhai, Neil Houlsby, Dustin Tran, and Mario Lucic. Revisiting the calibration of modern neural networks. In Advances in Neural Information Processing Systems, 2021.

We gratefully appreciate your prompt response to our rebuttal. We discuss the similarities and differences of our method with NatPN as follows.

1. Similarities. We agree with the reviewer that both of us have the same idea in terms of: (1) Estimating a marginal density model on feature space. (2) Combining this density model with the classifier to improve the model's uncertainty performance. (3) From the Bayesian view, both of us can be considered as a point estimate approach.

2. Differences. To easily compare, following the reviewer's notation, let us derive as follows. Since $\tilde{g}$ denotes the vector of softmax, we have the Categorial distribution from the closed-form of Nat-PN is

$p(y|x,D) = Cat\left(\frac{\alpha(x)}{\sum_i \alpha_i(x)}\right) = Cat\left(\frac{\alpha_{\text{prior}} + p(z) \tilde{g}(z)}{\sum_i \alpha_{\text{ith-prior}} + p(z) \tilde{g}(z)_{i}}\right)$

$= Cat\left(\frac{\alpha_{\text{prior}} + p(z) \frac{\exp(g(z))}{\sum_{k=1}^K \exp(g(z)_k)}}{\sum_i \alpha\_{\text{ith-prior}} + p(z) \frac{\exp(g(z)_i)}{\sum\_{k=1}^K\exp(g(z)_k)}}\right), \text{(1)}$

where $\frac{\exp(g(z)\_i)}{\sum\_{k=1}^K \exp(g(z)\_k)} \leq 1, \forall i \in [1,\cdots, K]$, $\alpha_{\text{ith-prior}}$ and $g(z)\_i$ represent the $i$-th entry in the prior vector $\alpha_{\text{prior}}$ and logit vector $g(z)$ respectively. And, the Categorial distribution of our Density-Softmax is

$

p(y|x,D) = Cat\left(\text{Softmax}\left[p(z)g(z)\right]\right) = Cat\left(\frac{\exp(p(z) g(z))}{\sum_{i=1}^K \exp(p(z) g(z)_k)}\right). \text{(2)}

$

From the equation, we have the following key differences:

• (1) From the Bayesian view, optimizing with Eq.1 is equivalent to doing Maximum a posteriori estimation (MAP) while optimizing with Eq.2 is equivalent to doing Maximum likelihood estimation (MLE).

• (2) Eq.1 shows that the $p(y|x,D)$ of NatPN is not a standard softmax output because it is not normalized by a natural exponent function. Meanwhile, Eq.2 is normalized by a natural exponent function.

3. Benefits of our approach. From the difference (1) and (2) above, we have the benefit of Density-Softmax over NatPN as follows:

• (1) Eq.2 shows that Density-Softmax can calculate $p(y|x,D)$ without the need to consider how to define $\alpha(x)$. Please note that the $\alpha(x)$ can hurt $p(y|x,D)$ performance if we pre-define a bad prior.

• (2) The natural exponential in Eq.2 helps Density-Softmax easily optimize with the cross-entropy loss by the nice derivative property of $\ln(\exp(a)) = a$, and produce a sharper distribution with a higher probability of the largest entry in the logit vector $g(z)$ and lower probabilities of the smaller entries when compared with the normalization of NatPN in Eq.1.

4. Why NatPN can have a bad performance. Given the two benefits of Density-softmax mentioned, let us give two examples to show why NatPN can potentially have a bad performance:

• (1) Let's consider we have three classes. As mentioned by the reviewer, NatPN will pre-define $\alpha_{\text{prior}} \sim \mathbb{U}$, where $\mathbb{U}$ stands for a uniform distribution. Plugin to Eq.1, we can see that if we choose $\alpha_{\text{ith-prior}} \gg p(z)$, then $\alpha_{\text{ith-prior}} \gg p(z)\frac{\exp(g(z)\_i)}{\sum_{k=1}^K \exp(g(z)_k)}$, yielding $p(y|x,D) \rightarrow \mathbb{U}$, i.e., $p(y|x,D) = (0.33,0.33,0.33)$. Therefore, NatPN's accuracy can be low and the predictive distribution is not sharp in this classification task.

• (2) Now consider a sample $x$ with the true label $y=(1,0,0)$, $z=f(x)$, $g(z) = (4,2,2)$, $p(z) = 1$, $\alpha_{\text{prior}} = (1,1,1)$, then plug into Eq.(1), NatPN prediction is $\hat{y}_1=(0.4468, 0.2766, 0.2766)$. Meanwhile, from Eq.(2), Density-Softmax prediction is $\hat{y}_2=(0.7870, 0.1065, 0.1065)$. We can see that from $\hat{y}_1$ of NatPN, the probability of the largest entry in $g(z)$ is lower, and the probability of lower entries is higher when compared to $\hat{y}_2$. This potentially leads to a more flat distribution and results in lower accuracy and higher ECE for correct predictions.

We hope we have thoroughly addressed your concerns. We are willing to answer any additional questions.

Thanks for your further questions. As suggested, we added a detailed discussion about the comparison with NatPN in Appendix B.4 in our revised version. We will add more detail following our further discussions in the next version. We would like to continue to resolve your remaining concerns as follows.

1. The selection of $\alpha_{\text{prior}} = 1$. We agree with the reviewer that if $p(z)$ is small enough, then $\arg\max$ between NatPN and Density-Softmax are the same at test-time. However, in training time, since optimizing in Eq.(1) with NatPN equivalents to do MAP, we can see that the predictive distribution of the posterior can be biased by the prior information. Therefore, if we select $\alpha_{\text{ith-prior}} = 1, \forall i \in \\{1,\cdots, K\\}$, i.e., $\alpha_{\text{prior}} \sim \mathbb{U}$, then the lower $p(z)$, the flatter of $p(y|x,D)$, and if $p(z)\ll 1$, then $p(y|x,D) \rightarrow \mathbb{U}$. We think this potentially makes $p(y|x,D)$ of NatPN under-fitting because its predictive distribution is flatter and not as sharp as Density-Softmax. Please note that Density-Softmax is trained by MLE and without $p(z)$ information in our pre-training step. In fine-tuning step, we fix $p(z)$ and only optimize with the classifier weight.

2. Benefits of $\log(\exp(a))$ in training. In this regard, let us follow the explanation in 6.2.Gradient-Based Learning, 6.2.2.Output Units, and 6.2.2.3.Softmax Units for Multinoulli Output Distributions, of the Deep Learning book [1]. As we mentioned in the different (1) in our previous response, the optimization in Eq.(2) with Density-Softmax equivalents to MLE, i.e., minimizing the negative log-likelihood, equivalently described as the cross-entropy between the training data and the model distribution as follows

\min_{\theta}\\{-\mathbb{E}\_{(x,y)\sim D} \log(p_{\theta}(y|x))\\} = \min_{\theta}\left\\{-\sum_{j=1}^N \log(\text{Softmax}(p(z)g_\theta(z)_i))\right\\},

where $\theta$ is model parameter, $N$ is the number of training data, and $i$ is the correct class for sample $j$-th. Therefore, let $a = p(z)g_\theta(z)$, i.e., $\text{Softmax}(a)\_i = \frac{\exp(a_i)}{\sum_{k=1}^K\exp(a_k)}$, then the log-likelihood can undo the $\exp$ of the softmax as

$\log(\text{Softmax}(a)\_i) = \log\left(\frac{\exp(a_i)}{\sum_{k=1}^K\exp(a_k)}\right) = \log(\exp(a_i)) - \log\left(\sum_{k=1}^K\exp(a_k)\right) = a_i - \log\left(\sum_{k=1}^K\exp(a_k)\right). \text{(3)}$

From Eq.(3), we can see that the $\exp$ in Density-Softmax roughly cancels out the $\log$ in the cross-entropy loss causing the loss to be roughly linear in $a_i$, i.e., if the correct answer already has the largest input to the softmax, then the $-a_i$ term and the $\log(\sum_{k=1}^K \exp(a_k)) \approx \max_{k \in \\{1,\cdots,K\\}} a_k = a_i$ terms will roughly cancel. This leads to a roughly constant gradient and the correct sample contributes little to the overall training cost, which will be dominated by other examples that are not yet correctly classified, allowing the model to correct itself quickly in training [1]. Therefore, a wrong saturated Density-softmax does not cause a vanishing gradient and its predictive distribution is potentially sharp in training [1].

We hope we have thoroughly addressed your concerns. We are willing to answer any additional questions.

[1] Deep Learning (Ian J. Goodfellow, Yoshua Bengio and Aaron Courville), MIT Press, 2016.

Thanks for your valuable feedback. We additionally address your concerns below.

1. Comparison with DUQ. From the tables in our general response, we can see that our method is much better than DUQ [1] in all criteria. Please note that our reported number is consistent with Tab.2-3 in SNGP [3].

2. Comparison with DDU. Firstly, according to Tab.1 in DDU [2], DDU's ECE is $0.85$ and $4.10$, while Deep Ensembles is $0.76$ and $3.32$ on CIFAR-10 and CIFAR-100. So, their reported uncertainty performance is worse than Deep Ensembles in terms of calibration error by having a higher ECE. Secondly, DDU does not consider improving accuracy as their motivation, especially the robustness under distribution shifts. Finally, we want to emphasize that the DDU method uses the post-hoc re-calibration technique to optimize temperature variables to improve the uncertainty performance, while, our Denisty-Softmax does not use any re-calibration set. We believe that the number of samples in the re-calibration set for temperature scaling is a signification number to be considered. For example, DDU needs 5000 samples in the calibration set to optimize, while the training set has 45000 samples in CIFAR-10. Therefore, we only compare the performance of DDU without temperature scaling. From the figures, we can see that our method is also better than them in all criteria.

3. How aleatoric and epistemic uncertainty is output by your model?. Similarly to DDU, we can easily disentangle aleatoric uncertainty by computing the softmax probability without multiplying with the density model by our Eq.(5) in the paper, and epistemic uncertainty by using the likelihood value of the density model.

4. Minor Comments. Thanks for your valuable suggestions. We will fix them in our revised version.

We hope we have thoroughly addressed your concerns. We are willing to answer any additional questions.

[1] Joost Van Amersfoort, Lewis Smith, Yee Whye Teh, and Yarin Gal. Uncertainty estimation using a single deep deterministic neural network. In Proceedings of the 37th International Conference on Machine Learning, 2020.

[2] Jishnu Mukhoti, Andreas Kirsch, Joost van Amersfoort, Philip HS Torr, and Yarin Gal. Deep deterministic uncertainty: A new simple baseline. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, 2023.

[3] Jeremiah Liu, Zi Lin, Shreyas Padhy, Dustin Tran, Tania Bedrax Weiss, and Balaji Lakshminarayanan. Simple and principled uncertainty estimation with deterministic deep learning via distance awareness. In Advances in Neural Information Processing Systems, 2020.

Thanks for your valuable feedback. We additionally address your concerns below.

1. Comparison with [1-3]. [1-2] are sampling-based methods which suffer from a worse efficiency performance. Additionally, according to page 15 in DEUP [1], DEUP only archives 93.89% accuracy on CIFAR-10, which is lower than our method (96%) and Deep Ensembles (96.6%). Please note that DEUP also uses the same uncertainty and robustness baseline codebase. Meanwhile, [2] is only a technical report, has not been accepted to any conference, and also does not publish source code. DUE [3] can be a sampling-free method and we add them for comparison in our revised version. From the tables in our general response, we observe that our method outperforms DUE in all criteria. It is also worth noting that DUE [3] has not been accepted to a main conference yet.

2. What if we pass the density term with the classifier outputs through a parameterized model?. Our motivation is to use the density function to reduce the over-confidence of the DNN classifier outputs under distribution shifts. If we pass the density term with the classifier outputs through an MLP model, the accuracy of the parameterized MLP model can be different from the classifier outputs, and the over-confidence issue of the classifier outputs may not be resolved. In contrast, if we multiply the density term with the classifier outputs, the model's accuracy still holds. Therefore, we can show in Prop.4.7. that our method can significantly mitigate the over-confidence issue of the classifier.

3. How Density-Softmax compare with other distillation approaches in [4-5]?. Thank you for your suggestion. We are aware that distillation approaches are trying to also speed up the inference steps. We have cited them and added more discussion about these distillation approaches in our revised version. However, in the distillation approaches, [4-5] has shown the student model still underperforms the teacher model BNNs and Deep Ensembles in terms of uncertainty and robustness. Given that we directly compare our method with BNNs and Deep Ensembles and show that we outperform BNNs and have a competitive result with Deep Ensembles, it implies our model should be better than distillation approaches (e.g., [4-5]) in terms of uncertainty quality and robustness.

4. Do the authors have some insights on how the entire algorithm can be made faster?. We are slightly unsure whether you mean faster training or testing performance. Regarding testing performance, as mentioned in Remark 3.4, our algorithm is sampling-free by requiring only a single forward pass to make an inference. Therefore, our method is faster than other sampling-based baselines. Regarding training performance, as discussed in our future work, we expect to develop new techniques to avoid computing the Jacobian matrix to improve training time performance in future work.

We hope we have thoroughly addressed your concerns. We are willing to answer any additional questions.

[1] Lahlou, S., Jain, M., Nekoei, H., Butoi, V. I., Bertin, P., Rector-Brooks, J., ... & Bengio, Y. (2022). DEUP: Direct Epistemic Uncertainty Prediction. Transactions on Machine Learning Research.

[2] Zhang, J., Dai, Y., Xiang, M., Fan, D. P., Moghadam, P., He, M., ... & Barnes, N. (2021). Dense uncertainty estimation. arXiv preprint arXiv:2110.06427.

[3] van Amersfoort, J., Smith, L., Jesson, A., Key, O., & Gal, Y. (2021). On feature collapse and deep kernel learning for single forward pass uncertainty. arXiv preprint arXiv:2102.11409.

[4] Vadera, M., Jalaian, B., & Marlin, B. (2020, August). Generalized bayesian posterior expectation distillation for deep neural networks. In Conference on Uncertainty in Artificial Intelligence (pp. 719-728). PMLR.

[5] Malinin, A., Mlodozeniec, B., & Gales, M. (2019, September). Ensemble Distribution Distillation. In International Conference on Learning Representations.

Dear reviewers,

We would like to thank you again for spending your time reading the rebuttals as well as evaluating our paper. We truly appreciate that.

As the end of the discussion period is approaching, we look forward to hearing your feedback if our answers haven’t addressed your questions. We would be happy to discuss if you still have any concerns. If our answer resolves your concerns, we kindly ask whether the score can be raised.

We did update our revised paper following your valuable feedback. We hope we have thoroughly addressed your concerns.

All the best,

Authors.