Characterizing Robust Overfitting in Adversarial Training via Cross-Class Features

Q5: There are many adversarial robust distillation methods in previous works, so it's better to compare them in details like [a].

A5: Thank you for your thoughtful suggestions. The primary distinction between our method and adversarial robust distillation methods lies in their respective purposes and settings. Our method is specifically designed to address robust overfitting, whereas robust distillation methods are focused on distilling adversarial robustness from pre-trained teacher models. Furthermore, it's noteworthy that our proposed method is entirely trained from scratch, in contrast to robust distillation methods that require a pre-trained teacher.

Q6: This paper mainly focuses on robust overfitting. How about catastrophic overfitting in Fast-AT?

A6: Thank you for offering this valuable perspective. It's essential to note that the majority of existing papers [2-4] on understanding robust overfitting do not explicitly delve into the topic of catastrophic overfitting in the context of Fast-AT, which indeed remains an open research area.

In line with your suggestion, we have extended our investigations to include Fast-AT for the CIFAR-10 dataset, employing an $\ell_\infty$-norm perturbation bound of $\epsilon=8/255$. The detailed results are presented in Appendix D.6. Notably, after catastrophic overfitting, there is a significant reduction in the usage of cross-class features. This observation aligns with our understanding, indicating that the model also tends to forget cross-class features after experiencing catastrophic overfitting.

[a] Revisiting Adversarial Robustness Distillation: Robust Soft Labels Make Student Better. ICCV 2021

[2] Understanding Robust Overfitting of Adversarial Training and Beyond. ICML 2022

[3] On the Onset of Robust Overfitting in Adversarial Training. arxiv preprint

[4] Understanding and combating robust overfitting via input loss landscape analysis and regularization. Pattern Recognition

Thank you for your valuable feedback. We have revised our manuscript based on your comments and address your concerns as follows.

Q1: In Eq. 10, the authors use label smoothing to simulate knowledge distillation. I am wondering if this simulation is reasonable. The authors only said that ''due to symmetry'', but it is difficult to understand.

A1: Thanks for pointing that out. In this context, the term 'symmetry' specifically refers to the symmetry of logits for the other two classes when taking the expectation in the loss function (equation 10). When considering data from class $y$, both the distribution of features $x_{E,i}$ and $x_{C_i}$ for the other two classes, as well as their respective weights $w_1$ and $w_2$, exhibit symmetry respectively. Consequently, after applying knowledge distillation, the expectation for logits of the other two classes in the objective loss function (equation 10) becomes identical. To simplify this process, we can employ label smoothing.

We have incorporated this discussion into our revision.

Q2: On page 6, the authors claimed that "... derive a large enough logit" and "... Since the final checkpoint makes decisions based only on these limited features, the logit for the correct class is not large enough". These two conclusions are very difficult to understand because I find it challenging to infer them from the existing experiments. Figure 4 only shows the correlation matrix which cannot reflect the logit information.

A2: Thank you for your careful reading. We have toned down such claims in the updated paper as

Since the final checkpoint makes decisions based only on these limited features, it fails to leverage comprehensive features for classification, making the model more vulnerable to adversarial attacks on these samples.

Q3: In Equation 7, why do the class-specific or cross-class features only appear in robust features?

A3: Thanks for your careful reading. We want to emphasize that our claim does not assert that features of this nature exclusively appear in robust features. As clarified earlier, our data model, inspired by [1], primarily centers on the inherent relationship between these two types of robust features, and non-robust features do not affect our main claims. To illustrate, consider a non-robust, class-specific feature $x_n$ with a Gaussian mean $\eta$ (as introduced in [1]), where $\eta$ is a small positive value. For adversarial training with perturbation bound $\epsilon>\eta$ (generally holds since $\eta$ is small), the weight $w_n$ for $x_n$ will be set to 0, as demonstrated by the proof of Theorem 1. Consequently, under adversarial training, $x_n$ does not exert an influence on the decision logic. This same conclusion extends to cross-class features as well.

[1] Robustness may be at odds with accuracy. ICLR 2019

Q4: The results shown in Section 4.2 are somehow weak I think. For example, it is difficult to see that feature $x_{E,i}$ will dominate the whole training objective when $w_1>0$.

A4: Thank you for your valuable insights. We would like to clarify that the assumption regarding the dominance of $x_{E,i}$ is primarily supported by our empirical findings, which reveal a gradual forgetting of cross-class features $x_{C,i}$ and an increasing dominance of class-specific features $x_{E,i}$ during adversarial training (AT). These observed tendencies of these features during AT lend credibility to the reasonable supposition that feature $x_{E,i}$ will indeed dominate the training objective.

Dear Reviewer iycg, thanks for your time reviewing our paper. We have meticulously prepared a detailed response addressing your concerns and revised our paper accordingly. Could you please have a look to see if there are further questions? Your invaluable input is greatly appreciated. Thank you once again, and we hope you have a wonderful day!

Thank you for your detailed and valuable feedback. We have revised our manuscript based on your comments and addressed your concerns as follows.

Q1: Robust overfitting is specific to adversarial training. Thus, the discussion should be emphasized more on special properties of adversarial attack/examples/training. However, the idea given in this paper fails in this aspect: in regular training, it is also the case that the model first learn cross class information and then focusing on a specific model, e.g., the well-known neural collapse saying that the features of examples of one class will collapse to a direction.

A1: Thank you for reminding this viewpoint.

First, we want to emphasize that we have indeed taken into account the distinction in the explanation for robust overfitting between adversarial training and regular training. In Section 3.3 and illustrated in Figure 3, our experiments encompass different values of $\epsilon$, with a small $\epsilon$ (2/255) representing a scenario closer to regular training. Gradually increasing $\epsilon$ can be seen as a transition towards more rigorous adversarial training. Notably, the significance of cross-class feature forgetting becomes more pronounced as $\epsilon$ increases. This observation aligns with the phenomenon that robust overfitting is uniquely evident in adversarial training scenarios.

Following your suggestions, we have extended our experimental scope to include regular training on the CIFAR-10 dataset. The experimental settings mirror those outlined in Section 5, with the sole distinction being the absence of perturbations in regular training. The results are presented in Appendix D.5. Specifically, considering that regular training prioritizes natural generalization and exhibits minimal robustness, we have calculated the feature attribution vectors using clean examples. These vectors were computed for epochs $\\{50, 100, 150, 200\\}$. Notably, the results reveal a lack of clear differences between them, particularly in the latter stages (150th and 200th), where the training tends to converge. This observation is consistent with the characteristic of regular training, which typically does not exhibit overfitting.

Q2: The unclearness to the regular training also could be observed in their proposed method: stochastic weight averaging (SWA) is a standard technique to enhance generalization capability in regular training. In the current experiment, it is hard to say the advantage of the proposed WAKE comes from SWA or it comes specific property of adversarial training, then now it cannot well support the main contribution.

A2: Thank you for pointing this out. It's essential to recall that in our comparison with the baseline KDWSA [1], which combines knowledge distillation (KD) and stochastic weight averaging (SWA), the performance of WAKE is notably superior to KDSWA, not to mention using SWA alone. To provide further clarification, we have added a table below for a direct comparison. The experimental settings are identical to those in Table 1 for the CIFAR-10 dataset. Consequently, it is clear that the improvement achieved by WAKE is not solely attributed to SWA; rather, it stems from the specific properties of adversarial training.

Q3: Another weakness is that the metric used to support the conclusion is not very convincing. For example, it is believed that features in different layers capture different information, e.g., cross-class or in-class information. Simply put all features together seems too weak.

A3: Thanks for your kind suggestion. Firstly, it's important to clarify that CAS does not aggregate all features. Instead, CAS exclusively utilizes the features from the penultimate layer, just before the linear classification head. This selection is based on the understanding that features in this layer hold the highest influence on the classification logits, providing a clear reflection of the decision-making process.

Q4: There are many metric that can measure the information about class and examples learned by DNNs. Could the authors use other metric and also observe the similar phoneme as measured by CAS？

A4: We appreciate your suggestion regarding alternative metrics to quantify the usage of cross-class features, but we do not know which metric you are referring. Could you clarify which specific metric you are referring? This will assist us in addressing your suggestion more accurately.

[1] Robust overfitting may be mitigated by properly learned smoothening. ICLR 2021

We truly appreciate your valuable and detailed feedback. If you have any further questions or concerns, please let us know.

Q3: When taking average of feature vectors for a class to visualize the matrix and CAS, this cannot reflect the tendency/behavior of each feature vector. To me, a better way is to compute and visualize instance-wisely.

A3: Thank you for your thoughtful suggestion. Our choice of using vectors over classes is based on the belief that averaging over samples can yield a more comprehensive and holistic understanding of the attribution for different classes.

Following your suggestion, we additionally calculated feature attribution correlation matrices on an instance-wise basis for both $\ell_\infty$ and $\ell_2$ adversarial training. The results are presented in Appendix D.4. To elaborate, when considering classes $i$ and $j$, for each sample $x$ from class $i$, we identify its most similar sample $x'$ from class $j$. We then calculate their cosine similarity and average the results over all samples in class $i$.

In this context, $x'$ can be interpreted as the sample in class $j$ that shares the most cross-class features with $x$ among all samples in class $j$. This metric provides a meaningful way to quantify the utilization of cross-class features. We did attempt to average over all sample pairs $(x, x')$ in classes $i$ and $j$, but due to high variance among samples, each element in the correlation matrix $C$ hovered near 0 throughout all epochs in adversarial training, rendering it unable to provide meaningful information.

Consistent with the results for class-wise attribution vectors, it is still observed that there is a significant decrease in the usage of cross-class features from the best checkpoint to the last for both $\ell_\infty$ and $\ell_2$-AT. This observation further substantiates our understanding of robust overfitting.

Q4: The proposed practical method lacks novelty and does not have a strong connection to theory developed and experiments are humble

A4:

• In terms of novelty, it's crucial to highlight a key distinction of our method from existing distillation methods. Unlike approaches like KDWSA, which necessitates a pre-trained teacher model, our proposed method, WAKE, stands out in its efficiency and ease of implementation. WAKE does not rely on a pre-trained teacher model, making it a more straightforward and resource-efficient solution. Regarding your criticism on novelty, could you kindly specify any existing works that significantly overlap with our proposed WAKE?
• In terms of connection to theory, in this paper, we showed the importance of cross-class features for improving adversarial robustness, and the effectiveness of knowledge distillation to preserve such features. Thus, a natural solution is finding a better distillation teacher model for more precise cross-class feature information. This shows a clear connection between our theory and the proposed method.
• In terms of experiments, we also conducted comprehensive experiments which are shown in the Appendix. First, we delved into empirical investigations to uncover the relationship between cross-class features and robust overfitting (as detailed in Appendix D). Additionally, we conducted a series of experiments for WAKE (as detailed in Appendix F). These encompassed examinations involving $\ell_2$-norm, synergies with other adversarial training methods, and the exploration of alternative model architectures.

We truly appreciate your valuable and detailed feedback. If you have any further questions or concerns, please let us know.

Thank you for your detailed and valuable feedback. We have revised our manuscript based on your comments and addressed your concerns as follows.

Q1: The theories developed are very restrictive and under extremely unrealistic assumptions.

• Data have 6 features and each features follow a Gaussian with mean zero.
• It only considers a linear model with two parameters $w_1$ for the exclusive feature and $w_2$ for the shared feature.

A1: We appreciate your feedback, but we have to point out this is an unjustified criticism to our paper.

• Firstly, our choice of data distribution is rooted in a well-established model from [1], which separates robust and non-robust features through similar Gaussian distributions. This model has been widely adopted in subsequent studies, such as [2], for the analysis of adversarial robustness.
• Moreover, though [1] considers $d$ non-robust features, they simplified the linear function due to symmetry and they only use 2 parameters (one for robust features and another for non-robust features) in the linear model. Therefore, we believe modeling the 6 features among 3 classes is sufficient to analyze the relationship between cross-class features and adversarial robustness.
• Additionally, in the realm of adversarial robustness, many theoretical analyses, including but clearly not limited to [3,4,5], are also based on linear models. It is worth noting that linear models are sufficient to convey valuable insights, and we don't perceive the analysis within a linear model as a weakness of our paper.

Q2: The theoretical results obtained are humble and do not really reflect the empirical finding. Moreover, the statement of Theorem 2 is ambiguous. It is not clear what it does mean by "a larger w2 increases the possibility of the model distinguishing the adversarial examples from any other given class".

A2: We appreciate your feedback, but the theoretical results are well aligned with the empirical findings.

• Theorem 1 explains why larger $\epsilon$ in adversarial training forgets more cross-class features by showing cross-class features are more sensitive to robust loss in adversarial training, which is consistent with Figure 3.
• Theorem 2 explains why forgetting cross-class features leads to a decrease of robustness by showing larger weight of cross-class features $w_2$ can bring better robustness, which is consistent with the overall finding of this paper.
• Theorem 3 and Corollary 1 show how knowledge distillation mitigates robust overfitting which is consistent with Figure 4(b)(c).

As for theorem 2, following your feedback, we have refined the claim in Theorem 2 to make it clearer in our revised paper, and we copied it below:

For any class $y$, consider weights $w_1 > 0$, $w_2 \in [0, w_1]$, and $\epsilon \in (0, \frac{\mu}{2})$. When sampling $x$ from the distribution of class $y$, increasing the value of $w_2$ enhances the possibility of the model assigning a higher logit to class $y$ than to any other classes $y'\ne y$ under adversarial attack. In other words, the probability $\Pr_{x\sim\mathcal D_y}[f_w(x+\delta))\_{y}>f_w(x+\delta)\_{y'},\forall \delta:\|\delta\|_\infty\le\epsilon]$ monotonically increases with $w_2$ within the range $[0, w_1]$.

[1] Robustness may be at odds with accuracy. ICLR 2019

[2] Adversarial Examples Are Not Bugs, They Are Features. NeurIPS 2019

[3] Theoretically Principled Trade-off between Robustness and Accuracy. ICML 2019

[4] To be Robust or to be Fair: Towards Fairness in Adversarial Training. ICML 2020

[5] On the Tradeoff Between Robustness and Fairness. NeurIPS 2023

Dear Reviewer NSUP, thanks for your time reviewing our paper. We have meticulously prepared a detailed response addressing your concerns and revised our paper accordingly. Could you please have a look to see if there are further questions? Your invaluable input is greatly appreciated. Thank you once again, and we hope you have a wonderful day!

Q15: As pointed out in point 2 under Weaknesses: Is the weight averaged model created on the fly during adversarial training, similar to the self ensemble adversarial training (SEAT) method of (Wang & Wang, 2022)? The piecewise linear scheduling should be described precisely. More clarity is needed on these aspects, and an algorithm block would also be helpful (appendix can be used if space is limited).

A15: Thank you for bringing this to our attention. The weight-averaged model, denoted as $\bar\theta$, is dynamically constructed on-the-fly and continues until the warm-up procedure concludes.

During the warm-up stage of the weight-averaged model, the piecewise linear schedule for $\lambda$ is defined as follows: $\lambda$ starts at 0 and linearly increases to 1 until a certain epoch. Subsequently, $\lambda$ remains constant at 1 until the end of the training process.

In response to your suggestion, we have incorporated a detailed algorithm outlining our proposed method in Appendix G.1.

Q16: In Eqn (11), please clarify that the max over $\delta$ applies to both the loss function terms.

A16: Thanks for your careful reading. We have corrected this in our revision.

Q17: What is the temperature parameter $T$ referred to in section 5.2 under settings?

A17: Thanks for your kind consideration, $T$ is the temperature for the knowledge distillation function. We have added more details on knowledge distillation in Appendix G.2, as discussed in our response in A2.

Q18. Finally, the number of references seems to be small given the breadth of literature in this area.

A18. Thanks for the kind advice. We have added additional related work in Appendix H.

Minor comments

W3: Experiments in the main paper are limited, and only two baselines are used for comparison. However, there are more experiments and other baselines like TRADES in the appendix.

A: Thank you for recognizing that more comprehensive experiments are included in the appendix, where experiments over $\ell_2$-norm AT, more model architectures, and combinations with TRADES are included. As the proposed knowledge distillation method is designed to mitigate robust overfitting, to the best of our knowledge, KDSWA is the only existing knowledge distillation method for the same purpose. Thus, we only consider vanilla training and KDSWA as baselines.

W4: Code has not been made available.

A: We commit that we will release our code upon publication.

We truly appreciate your valuable and detailed feedback. If you have any further questions or concerns, please let us know.

Q9: In Eqn (9), it would be better to say the outer expectation is over the class prior distribution since $\mathrm{E}_i[\cdot]$ is pretty vague.

A9: Thanks for pointing this out. Following your suggestion, we have modified $\mathbb E_i$ as $\mathbb E_{y_i\sim\{y_1,y_2,y_3\}}$ in both equations (9) and (10) to clarify the class prior for the expectation.

Q10: In Eqn (9), it is not clear how the cross-entropy loss simplifies to the term $\max_{j \neq i} f_w(x + \delta)_j - f_w(x + \delta)_i$.

A10: Thanks for your careful consideration, but we didn't claim that we apply cross-entropy loss in this framework, which may be hard to analyze. Instead, we use this simple loss function: $\max_{j\ne i} [f_w(x+\delta)\_j]-f_w(x+\delta)\_{y_i}$, where $y_i$ is the index of the true class. Both cross-entropy and this loss function encourage the model to increase the logit margin between the correct class $f_w(x+\delta)_{y_i}$ and any other class $f_w(x+\delta)_j\quad (j\ne y_i)$.

Q11: What is the significance of introducing the regularization term in this objective?

A11: Thanks for your careful reading, we explain the significance of this regularization term as follows. Adding this regularization term prevents both $w_1$ and $w_2$ from going infinitely large. For example, by multiplying a positive term $k>0$ to both $w_1$ and $w_2$, the weights become $kw_1$ and $kw_2$, and these parameters can also cause the loss function to be multiplied by $k$ due to the linearity of $\mathcal L$, but this will not change the classification results due to the linearity of $f_w$. Therefore, for any $w_1$ and $w_2$ that imply a negative loss function $\mathcal L(f_{w_1,w_2})$, multiplying $k>1$ makes $\mathcal L(f_{kw_1,kw_2})<\mathcal L(f_{w_1,w_2})$. Thus, there is no optimal solution for $\min \mathcal L(f_w)$ without a regularization term.

Q12: In Section 4.2, 2nd line: should it be abandon $x_C$ rather than abandon $x_E$?

A12: Thank you for your careful reading. It is exactly abandon $x_C$ rather than abandon $x_E$, and we have fixed this typo in our revision.

Q13: The statements of theorem 2 and theorem 3 are vague. They could be stated more formally.

A13: Thank you for your valuable comment, we have clarified the possibly confusing part in the statements in our revised paper and explained them as follows.

• We have revised the claim in Theorem 2 more clearly as

  For any class $y$, consider weights $w_1 > 0$, $w_2 \in [0, w_1]$, and $\epsilon \in (0, \frac{\mu}{2})$. When sampling $x$ from the distribution of class $y$, increasing the value of $w_2$ enhances the possibility of the model assigning a higher logit to class $y$ than to any other classes $y'\ne y$ under adversarial attack. In other words, the probability $\Pr_{x\sim\mathcal D_y}[f_w(x+\delta))\_{y}>f_w(x+\delta)\_{y'},\forall \delta:\|\delta\|_\infty\le\epsilon]$ monotonically increases with $w_2$ within the range $[0, w_1]$.

• For $\epsilon_0$ in Theorem 3, it is exactly the $\epsilon_0$ derived in Theorem 1. We have added this explanation in our revision.

Q14: In the subsection Knowledge distillation preserves cross-class features, could the authors clarify how label smoothing applies here due to symmetry? And what does symmetry refer to here?

A14: Thank you for your careful reading. In this context, the term 'symmetry' specifically refers to the symmetry of logits for the other two classes when taking the expectation in the loss function (equation 10). When considering data from class $y$, both the distribution of features $x_{E,i}$ and $x_{C_i}$ for the other two classes, as well as their respective weights $w_1$ and $w_2$, exhibit symmetry respectively. Consequently, after applying knowledge distillation, the expectation for logits of the other two classes in the objective loss function (equation 10) becomes identical. To simplify this process, we can employ label smoothing.

We have incorporated this discussion into our revision.

Thank you for your detailed and valuable feedback. We have revised our manuscript based on your comments and address your concerns as follows.

Q1: In Eqn (3), the max over the perturbation $\delta$ should include all three loss terms. It seems like the max is only over the cross-entropy loss term. Please make it clear by moving the $(1-\lambda_1-\lambda_2)$ term inside the max and using parentheses around the three loss terms.

A1: Thanks for your careful reading. We have fixed this typo in Equation (3) in our manuscript.

Q2: The terms $\ell_{CE}$ and $\mathcal{KD}$ in Eqn (3) should be defined as the cross entropy loss and the Kullback-Leibler divergence respectively.

A2: Thank you for your valuable comment. We'd like to clarify that $\mathcal {KD}$ is not the Kullback-Leibler divergence; rather, it represents the knowledge distillation function. This function initially computes the distilled probability prediction for both the student and teacher models by dividing the temperature $T$. Subsequently, it calculates the KL-divergence between these predictions:

$\mathcal{KD}(f(\boldsymbol\theta_{\text{student}};x), f(\boldsymbol\theta_{\text{teacher}};x)) = \text{KL}[\text{softmax}(\frac{f(\boldsymbol\theta_{\text{student}};x)}{T}), \text{softmax}(\frac{f(\boldsymbol\theta_{\text{teacher}};x)}{T}$)]

where $\text{KL}(\cdot,\cdot)$ is the Kullback-Leibler divergence and $T$ is the distillation temperature. We have clarified this in page 3 in our revision and added the details in Appendix G.2.

Q3: Nit: the $\times$ operator is not needed in Eqn (4).

A3: Thanks for pointing that out. We have removed that $\times$ operator in Eqn (4) to improve readability.

Q4: First paragraph of section 3.1: Notations for $f(\cdot), W$ and $f(x)_j$.

A4: Thank you for providing such detailed feedback. Based on your suggestions, we have addressed the typos in Section 3.1 on page 4 to enhance the overall readability of the manuscript.

Q5: Referring to the subsection Knowledge distillation mitigates robust overfitting, the following statement is not clear: “. . . using knowledge distillation can preserve the weight of these features by smoothing the labels of the overlapping classes”.

A5: Thanks for your careful reading. In order to provide further clarity regarding the effectiveness of knowledge distillation for preserving cross-class features, we have refined the claim in our revised paper as follows:

In the process of AT with knowledge distillation, the teacher model adeptly captures the cross-class features present in the training data, and provides more precise labels by considering both class-specific and cross-class features. This stands in contrast to vanilla AT with one-hot labels, which primarily emphasizes class-specific features and may inadvertently suppress cross-class features in the model weights. The incorporation of cross-class features, backed by both our empirical findings and theoretical insights highlighting their significance for enhanced robustness, enables knowledge distillation to effectively mitigate robust overfitting by preserving these crucial features.

Q6: Use of $i$ to index the class is confusing. Instead, $y$ could be used to denote the class and or to index the features.

A6: Thanks for your kind consideration. Following your suggestion, we have replaced the index $i$ for classes with $y_i$ to differentiate them from the index of features in the revised paper.

Q7: It is mentioned that the data distribution $\mathcal{D}\_i = \\{ x\_{E_j}, x\_{C\_j} | 1 \leq j \leq 3 \\} \in \mathrm{R}^6$. How can the data distribution be a set of features which lives in $R^6$? I can follow what is implied, but it is not formally correct.

A7: Thanks for your careful checking. In Section 4.1, page 7 of our revised paper, we redefined the features for each sample as living in $\mathbb R^6$, and defined the distribution for the features independently in equation (7).

Q8: In Eqn (7), the probability distributions are actually class-conditional rather than marginal. That is, for the exclusive features: $x_{Ej} \|\ y = i \sim \mathcal{N}(\mu, \sigma^2)$ when $j=i$ and 0 when $j\ne i$. Similarly, for the cross-class features: $x_{Cj} \|\ y = i \sim \mathcal{N}(\mu, \sigma^2)$ when $j\ne i$ and 0 when $j=i$. Specifying it this way makes it clear why these features are defined as exclusive and cross-class.

A8: Thank you for your kind suggestion. We have revised the definition in equation (7) as per your suggestion to make it clearer.

Dear Reviewer fn47, thanks for your time reviewing our paper. We have meticulously prepared a detailed response addressing your concerns and revised our paper accordingly. Could you please have a look to see if there are further questions? Your invaluable input is greatly appreciated. Thank you once again, and we hope you have a wonderful day!

In the earlier version, the class label was denoted by $y = i$, which I found to be more clear than the current version (which denotes the label by $y_i$). For instance, the notation $x_{E_i}$ is clearer than $x_{E_{y_i}}$, and $f_w(x)\_i$ is clearer than $f_w(x)\_{y_i}$.

As suggested in my original review, Eqn 7 is perhaps more clear when expressed as follows. The conditional distribution $\mathcal{D}\_i$ for class $i$ is defined as:
x_{E_j} \~|\~ y = i \~\sim \begin{cases} \mathcal{N}(\mu, \sigma^2) \~\~&\text{ if } j = i \\\\ 0 \~\text{ w.p. } 1 \~\~&\text{ if } j \neq i \end{cases}

x_{C_j} \~|\~ y = i \~\sim \begin{cases} \mathcal{N}(\mu, \sigma^2) \~\~&\text{ if } j \neq i \\\\ 0 \~\text{ w.p. } 1 \~\~&\text{ if } j = i \end{cases}
where $j \in \\{1, 2, 3\\}$.

Regarding Equation 9

My suggestion in the original review is a simple change, to indicate that the outer expectation is over the prior distribution on class labels $p_Y$, i.e.
$\mathcal{L}(f_w) = \mathbb{E}\_{y \sim p_Y}[ \mathbb{E}\_{x \sim \mathcal{D}_y}[ \cdots ] ] + \frac{\lambda}{2} \\|w\\|^2$

That's all the changes I want to suggest. Glad to know that you found my feedback useful. Good luck!

Regarding Equations 3, 11, and 10, I meant changing them as follows because even in the revised version, it is not clear that the inner max over $\delta$ applies to all three terms.

Equation 3
\min_{\theta} \mathbb{E}_{(x,y) \sim \mathcal{D}\_{train}} \left[ \max\limits\_{\\|\delta\\|_p \le \epsilon} \tilde{\ell}(\theta \~;\~ \theta_1, \theta_2, x + \delta, y) \right] , where
\tilde{\ell}(\theta \~;\~ \theta_1, \theta_2, x + \delta, y) \~=\~ (1 - \lambda_1 - \lambda_2) \ell\_{CE}(f(\theta, x + \delta), y) \~+\~ \sum\limits\_{i=1}^2 \lambda_i \\, \mathcal{KD}( f(\theta, x + \delta), f(\theta_i, x + \delta) )

Equation 11
\max\limits\_{\\|\delta\\|_p \le \epsilon} \tilde{\ell}(\theta \~;\~ \bar{\theta}, x + \delta, y), where
\tilde{\ell}(\theta \~;\~ \bar{\theta}, x + \delta, y) = (1 - \lambda) \ell\_{CE}( f(\theta, x + \delta), y ) \~+\~ \lambda \\,\mathcal{KD}( f(\theta, x + \delta), f(\bar{\theta}, x + \delta) )

Similar suggestion for Equation 10 as well.