Venue: ICLR 2022 Conference

Reviews:

final_score: rating: 3: reject, not good enough
review_content: summary_of_the_paper: This paper proposed an approach to propagate probability distribution through neural networks. In particular, the authors use local linearization to handle nonlinearity and show its optimality in terms of total variation for ReLU. Empirically they show that the proposed method can provide calibrated confidence intervals for regression problems and improve detection of OOD data in classification problems compared to some baselines. 
main_review: The research problem of improving uncertainty estimates is very interesting and useful. However, the idea of propagating probability distribution through neural networks is not new. Specifically, the introduction section is a bit inaccurate and misleading. The authors claim that they focus on obtaining the propagated distribution in analytical form, or closed form, and that prior work [17] focuses only on Gaussian distributions (note that [17] is very early work from 1999). This is not true. For example, natural-parameter networks, or NPN [23] as cited in the paper, as well as its variants [12,26], considers exactly the same problem, and already handles distributions beyond Gaussian distributions; in fact their methods could handle any exponential-family distributions (including Gaussian ones) as well as various activation functions beyond ReLU. 

This brings me to another of my concerns, which is the lack of proper baselines. Given the similarity the authors should at least include baselines from [23, 12, 26]. That would make the paper more convincing and maybe highlight what local linearization is a better solution. Currently, it seems [27] is the only modern baseline method (and note that [27] did not establish its improvement over methods such as [23, 12, 26]).

In terms of technical novelty, since the results for linear transformations are well known, the main novelty rests on the nonlinear part, i.e., the use of local linearization. Here the authors are able to prove that in terms of TV, local linearization is optimal in the Gaussian/ReLU case and Cauchy/ReLU case. In the introduction, the author claim that they are able to be ‘more general’ than some previous methods that can only handle Gaussian distributions. It is a bit disappointing to see that the proposed version can only handle Gaussian and Cauchy distributions and it has to be ReLU activation. 

In Section 4.2, the authors mentioned that they can compute the exact probability, which is interesting. The trick is to change the softmax into a pariwise Gaussian comparison. However, it seems a lot of details are missing.

Table 4 is a bit confusing. It shows that the proposed method actually underperforms the vanilla softmax CE loss. It is also unclear what is the setting for Section 5.4. For example, what is the input variance for ‘Pairwise Gaussian (PG)’, does PG produce better uncertainty estimates, etc.

Minor:
Grammar issue: Page 4: Note that when we … we refer to Dirac’s delta distribution.



summary_of_the_review: Overall, the paper proposed a method for propagating uncertainty through neural network that could potentially be useful. However, it seems some important baselines are missing, the introduction is somehow misleading, the proposed method only works in very limited cases, i.e., Gaussian+ReLU and Cauchy+ReLU (while theoretical results on other activation such as sigmoid/tanh/other ReLU variants are not provided).
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: This paper studies the problem of propagating probability distributions through neural networks and applies the results to quantify prediction uncertainties. It proposes a local linearization method to approximate a distribution transformed by a ReLU network, as well as new loss function for learning with distribution-valued inputs. It also provides empirical results, showing that the method can quantify two kinds of uncertainty (aleatoric and epistemic) in classification and regression tasks, and training with the new loss function can improve robustness to random and adversarial perturbations.
main_review: Strengths: 
- Interesting theoretical result showing that the local linearization is in fact an optimal approximation in terms of total variation for ReLU networks
- Nice application of the method to study uncertainty quantification in neural networks

Weaknesses: 
- Only consider the case where the activation function is ReLU (going beyond ReLU, can the linearization method still give reasonable approximations, under some conditions? Would be nice to have some theoretical results in this direction)

Comments/questions:
- What is the role of depth in the uncertainty quantification and improvement in robustness to random/adversarial perturbations? 
- Can one formulate new loss function for learning with other non-Gaussian distributions (other than only Cauchy)?
- On page 3 (beginning of Section 3.1), is there a typo? There, A is a m by n matrix and b is a 1 by n vector?
summary_of_the_review: Overall, the strengths of the paper outweigh the weaknesses, although going beyond the ReLU activation (see Weaknesses above) could significantly improve the quality of the paper. Therefore, I am inclined to accept the paper.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a method to propagate uncertainties through neural networks by performing a first-order approximation at nonlinear activation functions. By injecting uncertainty into the inputs, prediction uncertainties are obtained. These prediction uncertainties are exploited at training time by using new uncertainty-aware loss functions, and at test-time to obtained prediction uncertainties, out-of-distribution detection, and adversarial robustness.
main_review: The paper presents a simple new technique to propagate distributions through neural networks. A nice feature of the method is that it can be readily used for pre-trained models. To justify the proposed method, a theoretical argument is provided that for the ReLU activation function a total variation distance is minimized. The paper is mostly clear and well written but lacks some details (see below).

- Regarding experiments showing the total variation metric (i.e., Table 1 and several tables in the supplementary):

How is the total variation between 10^6 (discrete) samples and the (continuous) output of probability propagation computed? I have another concern about the significance of the total variation distance. In particular, it is mentioned in the introduction that f(x + \epsilon) is a random variable where the variance is used to quantify its uncertainty. How well are the prediction uncertainties between the proposed method and sampling 10^6 models in these experiments aligned? Or can the total variation distance be used to give guarantees about the predicted uncertainty in these experiments?

- Regarding activation functions:

How well does the proposed method perform with other activation functions. The experiments in the main paper are only concerned with ReLU activations and the experiments in the supplementary material are concerned with ReLU-like activation functions (leaky ReLU, GELU, ...). Does the proposed method also produce reasonable results for saturating activation functions such as tanh or the logistic sigmoid?

- Failure modes:

The paper presents mostly results where the method succeeds although there are some obvious failure modes. For instance, given a Gaussian input with either a mean very close to zero or a very large variance, the approximation made by the proposed method is not very accurate. However, there are some experiments where the injected variances are rather large (e.g., Table 1 or Figure 3) but still the method does not seem to break down. How can you explain this?

Furthermore, I see another failure mode regarding the proposed objective for classification where the sum over all classes is computed. I do not see why this objective is enforcing the model to make a correct prediction. If P(Z_c > Z_e) is close to one for all classes except for a single class where P(Z_c > Z_e) < 0.5, then a wrong prediction will be made but still a high score is assigned by the proposed loss. If, on the other side, P(Z_c > Z_e) is slightly above 0.5 for all classes, then the correct class will be predicted but only a low score is assigned by the proposed loss. This might be more severe for classification tasks with a very large number of classes (e.g., ImageNet). In my opinion, it would make more sense to define the score by maximizing "min_e P(Z_c > Z_e)". This is similar to the concept of a probabilistic margin, e.g., see [1].

- Regarding presentation:

Please emphasize somewhere in the paper that the only point where stochasticity is injected is at the inputs (i.e., no weight uncertainty, no injected noise at hidden neurons, etc). Although I (now) see that this is clearly stated in the introduction, I would really appreciate if a formal definition of the considered setup is made at the beginning of Section 3. If space is an issue, I think Figure 1 is not that important and can be moved into the supplementary material. 

- Regarding experiments:

  - How do you set the input variances in your experiments in Section 5.1? Is the input variance tuned so that the PICP is close to 95%? Is the input standard deviation in Section 5.2 always 0.1?
  - How is the prediction interval in Section 5.1 computed? Is the MPIW computed at the test data? In general, I would appreciate a clearer explanation of the used metrics and the experimental setup in Section 5.1.
  - In Section 5.2 it is stated that training is performed using both Pairwise Gaussian and Pairwise Cauchy. However, I do not see any results for training with Pairwise Cauchy.
  - What architecture is used in Section 5.2?
  - Results on Cifar-10: It would be interesting to see how well prediction uncertainties work for CNNs? For instance, a similar experiment as in Section 5.2 for out-of-distribution detection would be interesting (e.g., using samples from Cifar-100 for out-of-distribution data).
  - An evaluation of "in-between uncertainties" [2] would be interesting.


- Regarding computational complexity:

Section 5.5 states that the complexity for propagating distributions with and without covariances is linear and constant in the number of outputs, respectively. How can the complexity be linear if there are quadratically many entries in the covariance matrix? Similarly, how can the complexity be constant if we need to compute a variance entry for each output?

- Notation:

  - x used multiple times (e.g., (i) as input value of a neural network in the introduction and (ii) as parameter of the Cauchy distribution)
  - \sigma used multiple times (e.g., (i) as standard deviation of a Gaussian and (ii) as standard deviation of the input noise)

- References:
  - [1] Roth et al., Hybrid Generative-Discriminative Training of Gaussian Mixture Models, Pattern Recognition Letters
  - [2] Foong et al., 'In-Between' Uncertainty in Bayesian Neural Networks, ICML 2019 Workshop on Uncertainty and Robustness in Deep Learning
summary_of_the_review: The proposed method is a simple and effective method to propagate distributions through a neural network. I am fine with the experimental results. However, I am not voting for accept right now since in my opinion there are some important details missing (as outlined in my review above) that are required to reproduce results and that would improve the quality of the paper. I am also missing a discussion of failure modes.

--- POST REBUTTAL ---

I have updated my rating from 5 to 6.
review_title: null

final_score: rating: 6: marginally above the acceptance threshold
review_content: summary_of_the_paper: The paper proposes a simple method that propagates the uncertainty in the data through a neural network with ReLU non-linearity. The proposed method simply transforms the mean and covariance of the input Gaussian/Cauchy distribution through linear layers, and applies a local linear approximation for the ReLU activation. The proposed transformation is shown to be optimal under the total variation criterion. The paper applies the method to train neural networks with regression or classification targets, and show that the proposed method can well estimate the uncertainty of the network's output. Moreover, they also show some improvements in the robustness of the trained networks.
main_review: Strength:
1. The proposed method is very simple, and has theoretical justifications with respect to total variation. The method is shown to be applied to relatively larger neural network structures where some previous methods are intractable.

2. The improvements on the networks' robustness are interesting. This may bring more focus about studying the robustness and uncertainty of a network together.

Weakness:
1. The proposed method is only shown to have theoretical justifications for ReLU non-linearity. I think to extend it to more general non-linearities or give some more empirical analysis about the approximation, at least for the class of piece-wise linear functions, is important for this work. Also, it is not clear how to apply the method on max-pooling layers.

2. Even though the method is simple, and efficient to some extend, the running time is still quite significant for large Resnet structures (with covariance). This limits the applicability of the proposed method to more realistic settings and datasets.

3. The experimental settings are quite limited. Especially for the out-of-distribution data prediction, I think it's worth investigating more datasets, and different ways of introducing OOD data: for example, separating a dataset into different classes, or creating certain artifacts on a portion of the data.

4. Given the network structure that the method can apply, the network is a piecewise linear model overall. Particularly, the model is a linear model for every data point. What will be the performance of the method if we propagate the distribution through every local linear model?
summary_of_the_review: The main strength of the paper is that the proposed method is simple. The main weakness is that the method is only applicable to limited kinds of layers of neural networks, and the paper lacks more comprehensive experiments to demonstrate the method's performance.
review_title: null

Meta-review: The paper discusses propagating input uncertainty through non-linear layers by a simple local linearization approach. This is a straightforward idea and the authors explain how this is an optimal approximation of the propagated distribution for Total Variation and reLU non-linearity (for a single layer). This is an interesting (if quite limited) theoretical result. What is not clearly stated is that this result only holds for a single layer. It does not mean that the local approach is the best way to approximate (in the total variation sense) a distribution passed through multiple reLU layers.

By repeating the procedure, the authors are able to define closed form objective functions for noise-robust training of deep networks.

The reviewers found this an interesting paper and there was a good effort by the authors to improve the results. However, technical innovation is modest and reviewer doubts still remain. For that reason, clarity of presentation is critically important. The overall numerical score isn't convincing and with one reviewer remaining very unconvinced.

I agree with the reviewers that the technical contribution is quite limited and I would argue is not particularly well explained. For example, a simple alternative would be to use a "global" linearization in which one can consider the network function $f$ as a whole, and then simply linearize this (rather than linearizing each layer). Indeed, the way that the paper is written, this would be a natural interpretation since $f$ is defined in the introduction as the network function, but is used later differently (eg section 3.2) as the transfer function. The approach is to recursively compute a new mean and covariance for each layer, propagating these through the network (similar to moment matching approaches). It would have helped if the authors had made pseudocode for their approach. It would be natural and interesting to compare to the simple global linearization approach (which is computationally faster).

The presentation of results and experiments could be improved. For example, in figure 3, it is not clear (nor is it explained in the text) what the definition of "robust accuracy" is.  

Given the modest technical innovation, I also feel that the clarity of presentation isn't yet at the level that would merit acceptance. The paper would also benefit from some deeper insight into why the approach might perform better than other approaches (such as local moment matching) at the network level (rather than a single layer).
Decision: Reject