Summary
This paper explores the optimal transport distributional robustness framework in the model parameters spaces, and provides practical algorithms to solve the corresponding optimization problems. This flexible framework can be used to train a single model, ensemble models and Bayesian neural networks in a sharpness-aware fashion. The authors additionally establish that the sharpness-aware minimization serves as a special instance within their proposed framework. The experimental results demonstrate the usefulness of this framework.
Strengths
This paper considers an important question of neural networks training. It is clearly written and easy to follow. The idea of extending the form of SAM to the DRO framework is natural but very meaningful, since it provides more well-developed tools to understand and improve the generalization of training models.
Weaknesses
1. In the theory, this paper shows that SAM is equivalent to the optimization problem in the OT-MDR framework when using a particular distance metric and picking the center distribution as a Dirac delta distribution. In the experiments, the results show (1) OT-MDR performs better than SAM for both single model and ensemble models cases; (2) The loss landscape found by OT-MDR is more flatten than the one found by SAM. This probably implies the training algorithm of OT-MDR is better than directly using SGD to solve SAM. But the paper doesn’t provide much intuitive understanding or analysis of these observations.
2. Most of the experimental results don’t provide standard deviation, which makes the results less solid.
Questions
Q1. In SAM [1], their experiments show that using 2-Wasserstein distance is better than using \infty-Wasserstein distance. For the current paper, in theorem 4.3, it shows that choosing center distribution as a Dirac delta in OT-MDR with 2-Wasserstein distance is equivalent to SAM (using 2-Wasserstein distance). But based on theorem 4.2, the algorithm solves the minimization problem of OT-MDR with \infty-Wasserstein distance. And the algorithm of OT-MDR still outperforms the SAM (with 2-Wasserstein distance). Does this probably imply the performance of OT-MDR can be further improved when using algorithms that solves the minimization problem of OT-MDR with 2-Wasserstein norm? Here are a few more delicate methods [2,3] to solve the general distributional robustness learning problem.
[1] Foret, P., Kleiner, A., Mobahi, H., & Neyshabur, B. (2020). Sharpness-aware minimization for efficiently improving generalization.
[2] Wang, G., & Chizat, L. (2022). An exponentially converging particle method for the mixed nash equilibrium of continuous games.
[3] Trillos, C. G., & Trillos, N. G. (2023). On adversarial robustness and the use of Wasserstein ascent-descent dynamics to enforce it.
Q2. In figure 1, the difference of accuracy for using different numbers of particles is actually small. Since this figure doesn’t provide the information standard deviation, it might be difficult to conclude that using 2 particles gives the best results. It would be better if the author could provide different random seeds results of this experiment. It will probably provide more insight.
Q3. As mentioned in the weakness, the experimental results in table 1, 3 and 4 don't contain standard deviation. Could the author provide more results on those experiments?
Rating
6: Weak Accept: Technically solid, moderate-to-high impact paper, with no major concerns with respect to evaluation, resources, reproducibility, ethical considerations.
Confidence
3: You are fairly confident in your assessment. It is possible that you did not understand some parts of the submission or that you are unfamiliar with some pieces of related work. Math/other details were not carefully checked.
Limitations
In general, I like the idea and the presentation of this paper. The authors could probably further improve this paper through more experiments and add some analysis of the observations from comparing SAM and OT-MDR.