Paper: Density-Softmax: Efficient Test-time Model for Uncertainty Estimation and Robustness under Distribution Shifts (ICLR-2024)

Reviewer comment (Theory):
It is also unclear whether the proposed training results in a 1-Lipschitz model or feature extractor. Note that training a model with a gradient norm penalty is insufficient to train a 1-Lipschitz model, as mere regularization does not imply that the constraint is satisfied. To verify whether models achieve the 1-Lipschitz guarantee, the method either needs to use provable methods such as [1] or computationally verify the existence of a small Lipschitz constant using [2].

Author response:
**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.

Reviewer comment (Novelty):
One alternate and simple way to incorporate knowledge of the input distribution is to estimate the joint probability $p(y, x) = p(y \mid x) p(x)$ [3]. This can be done by simply multiplying the softmax probabilities with the density estimates from the normalizing flow in the case of this paper. Note that this estimate satisfies all the asymptotic properties that density softmax does. Given the conceptual simplicity of this approach, I am wondering whether the authors can comment on the similarities/differences between their approach and joint probability estimation.

Author response:
**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.

Reviewer comment (Theory):
Apart from the fact that the corollary builds on top of this possibly flawed theorem, it also makes the assumption that asymptotic conclusions derived for certain points $d(z_{iid},Z) \rightarrow 0$ can be applied to **every** IID point. This is simply incorrect. There can exist IID points that simply haven’t been sampled from the distribution yet and thus do not have near zero distances from the train set points.

Author response:
**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.

Reviewer comment (Novelty):
My major concern is the novelty and overall contribution to the topic. The same idea of using a density model (e.g. Normalizing flow) to scale (by multiplication) some predicted parameters was used in NatPN [1]. I also wonder why, taking into account this similarity, the paper was not even discussed. It indeed was cited several times, but it was not discussed what is the gist of the method, and it was not compared. The authors of the paper under review compared only with PostNet [2]. But there is a difference in parametrization between NatPN [1] and PostNet [2]. In [2] they learned a flow per class (which appeared to be redundant) and in the follow-up (and generalized) paper they introduced another parametrization, which needs only one flow (or another density model) overall. This explains why Latency for Posterior Net [2] scales with the number of classes (which should not be the case for NatPN [1]). Moreover, in the problem of [1], they want to predict parameters of Dirichlet distribution, which is helpful as it allows them to distinguish between high aleatoric/epistemic uncertainties naturally. The parametrization of these predicted parameters of Dirichlet was in the form of the product of the density model and "update" to Dirichlet parameters. In the paper under review, authors do the same, but directly scale logits. So it seems to me that it is pretty same idea, the same computational overhead (train one density model). The only difference is that in [1] they did the trick for the "second-order" distribution (which is Dirichlet), but here for the "first-order" distribution (Categorical), which does not have apparent benefits.

Author response:
**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](https://github.com/google/uncertainty-baselines/tree/main/baselines/cifar). 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.

Reviewer comment (Theory):
Normalization of the p(z; \alpha) to be in [0, 1] is pretty strange and unnatural. First, why is it even an issue? How is it possible to have extremely large p(z;\alpha), given a well trained Lipschitz feature extractor? Since we train the density model on top of the fixed extracted features (and well separated), it is surprising it is a problem. Second, in the proof of Proposition 4.7, the fact that p(z; \alpha) is in [0, 1] is used. I wonder if it will be possible to prove it given p(z; \alpha) is still restricted, but say in [0, C], where C > 1.

Author response:
**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.

Reviewer comment (Novelty):
Using the Lipschitz network is also not something new. I wonder why this option to satisfy the Lipschitz constraint was chosen (gradient penalty)? In the Future Work there is a phrase "We plan to .... developing new techniques to avoid computing Jacobian matrix at training-time..." — why iterative spectral normalization method (based on power method) [3] is not used for this purpose? It was applied in SNGP for example, in a paper cited by authors.

Author response:
**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.

Reviewer comment (Theory):
In the introduction: "However, it often struggles with over-confidence and overfitting (Guo et al., 2017). This poor performance **usually occurs when the test data is far and does not come from the same distribution as the training set**" — I am a bit skeptical about it, as in [4] they observed this issue without shift when P_train(x, y) != P_test(x, y).

Author response:
**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.

Reviewer comment (Theory):
In Proposition 4.7, why only an overconfident case is considered?

Author response:
**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.

Reviewer comment (Experiments):
I believe the largest weakness in this paper are the baselines; while there are many baselines covering recent methods from the state of the art, some recent single network methods are missing and these are direct competitors of Density-Softmax: Van Amersfoort, Joost, et al. "Uncertainty estimation using a single deep deterministic neural network." International conference on machine learning. PMLR, 2020. This paper defines DUQ which uses a RBF layer instead of softmax, also being feature space distance aware. Mukhoti, Jishnu, et al. "Deep Deterministic Uncertainty: A New Simple Baseline." Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition. 2023. This paper defines DDU which uses 1-Lipschitz networks and feature space density as a signal for epistemic uncertainty, while softmax outputs are used for aleatoric uncertainty. The proposed method Density-Softmax is a direct competitor of DDU. To me it is very clear that at least DUQ and DDU are clear baselines, especially DDU, as to me the only difference between Density-Softmax and DDU is the use of normalizing flows for density estimation and combination of feature density into the softmax activation in the last layer.

Author response:
**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.

Reviewer comment (Experiments):
A notable point is that DDU outperforms most of the baselines used in this paper (like SNGP and ERM) and is very competitive with ensembles, producing better uncertainty than 5 ensembles in the out-of-distribution setting, and having slightly lower calibration error. This supports my proposal that Density-Softmax should be compared to DDU. DDU has been part of the public literature for around a year.

Author response:
**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.

Reviewer comment (Theory):
One important concept that seems to be missing from this paper is aleatoric and epistemic uncertainty. Feature density estimates epistemic uncertainty (shown and proven in the DDU paper), while softmax probabilities model aleatoric uncertainty. Since the proposed model Density-Softmax modifies softmax probabilities, these now are more closely estimating epistemic uncertainty, and then, what component of the model outputs aleatoric uncertainty? This could be a possible disadvantage of this model.

Author response:
**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.

Reviewer comment (Presentation):
The links to each abbreviation, highlighted in red, are a bit too much distraction. I suggest that for abbreviations, the authors turn off the links highlighted in red, the reader will thank you. I would even suggest to replace link boxes with colored links using hyperref options.

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

Reviewer comment (Presentation):
In Figure 1, each plot contains the x and y axis labels, these do not really add any new information to the figure and I suggest to remove them, for cleanliness.

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

Reviewer comment (Novelty):
The main weakness of this paper is that it doesn't look into other recent methods proposed - including those that involve direct uncertainty estimation [1-3], and distillation of BNNs/ensembles [4, 5].

Author response:
**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](https://github.com/google/uncertainty-baselines/tree/main/baselines/cifar). 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.

Reviewer comment (Experiments):
Rather than just multiplying the density term with the classifier outputs, what would be the impact of passing them through a parameterized model (such as MLP)? That would add more flexibility to the overall approach and improve performance, I believe.

Author response:
**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.

Reviewer comment (Evaluation):
Do the authors have some insights on how the entire algorithm can be made faster?

Author response:
**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.