Generalist: Decoupling Natural and Robust Generalization

The goal of the adversary is to generate a malignant example $x^{\prime}$ by adding an imperceptible perturbation $\varepsilon\in\mathbb{R}^{d}$ to $x$. And the generated adversarial example $x^{\prime}$ should be in the vicinity of $x$ so that it looks visually similar to the original one. This neighbor region $\mathbb{B}_{\varepsilon}(x)$ anchored at $x$ with apothem $\varepsilon$ can be defined as $\mathbb{B}_{\varepsilon}(x)=\left\{(x^{\prime},y)\in\mathcal{D}_{2}\mid\left\|x-x^{\prime}\right\|_{\infty}\leq\varepsilon\right\}$. For adversarial training, it first generates adversarial examples and then updates the parameters over these samples. The iteration process of adversarial training can be summed up as:

where $\Pi_{\mathbb{B}\left(x,\epsilon\right)}$ is the projection operator, $\alpha$ is the step size, $\tau$ is the learning rate, and $\mathcal{R}(\cdot)$ is the loss difference of $\ell_{2}(x^{\prime},y;\boldsymbol{\theta}^{t})-\ell_{1}(x,y;\boldsymbol{\theta}^{t})$. The tradeoff factor $\beta$ balances the importance of natural and robust errors. Various adversarial training methods can be derived from Eq. 1. For instance, when $\beta=1$, it is equivalent to the vanilla PGD training [DBLP:conf/iclr/MadryMSTV18], and when $\beta=1/2$, it is transformed into the half-half loss in [DBLP:journals/corr/GoodfellowSS14]. The formulation degenerates to standard natural training as $\beta=0$. Besides, we can get the formulation in TRADES [DBLP:conf/icml/ZhangYJXGJ19] when replacing $\mathcal{R}(\cdot)$ with the KL-divergence.

Multi-Task Learning. Multi-Task Learning (MTL) is to improve performance across tasks through joint training of different models [DBLP:conf/nips/BilenV16, DBLP:conf/cvpr/LuKZCJF17, DBLP:conf/iclr/YangH17]. Consider a set of assignments containing data distribution and loss function defined as $\mathcal{A}=\{\mathcal{D},\ell\}$ with corresponding models $\left\{\mathcal{M}_{a}\right\}_{a=1}^{|\mathcal{A}|}$ parameterized by trainable tensors $\boldsymbol{\theta}_{\mathcal{M}_{a}}$. In MTL, these sets have non-trivial pairwise intersections, and are trained in a joint model to find optimal parameters $\boldsymbol{\theta}_{\mathcal{M}_{a}}^{\star}$ for each task:

where $\ell_{a}(\mathcal{D}_{a};\boldsymbol{\theta}_{\mathcal{M}_{a}})$ measures the performance of a model trained using $\boldsymbol{\theta}_{\mathcal{M}_{a}}$ on dataset $\mathcal{D}_{a}$. Our approach Generalist is directly related to MTL at first glance because both of them tend to learn a specific predictive model for different sources. However, Generalist differs significantly from MTL, i.e., multiple tasks are still learned jointly under a unified form in MTL while each assignment can be optimized by heterogeneous strategies in Generalist.

Meta-Learning. Meta-learning is to train a model that can quickly adapt to a new task. Suppose $\mathcal{A}$ is divided into non-overlapping splits $\mathcal{V}$ and $\mathcal{W}$, the model is first trained on the training sets and then guided by a small validation set on a set of tasks to make the trained model can be well adapted to new tasks:

meta-learning [DBLP:conf/icml/FinnAL17, DBLP:journals/corr/abs-1803-02999] is often designed to generalize across unseen tasks, whereas the goal of MTL is to tackle a series of known tasks. Nonetheless, our approach Generalist uses the technique of meta-learning to set good initializations for base learners to transfer knowledge between tasks.

The overall procedure of our proposed algorithm is shown in Algorithm 1, which mainly comprises two steps: optimizing parameters of the base learner $\boldsymbol{\theta}_{a}$ in its assigned data distribution $\mathcal{D}_{a}$ and distributing parameters of the global learner $\boldsymbol{\theta}_{g}$ to all base learners. Base learners and the global learner share the same architecture, i.e., $\mathcal{M}_{1}=\mathcal{M}_{2}=\cdots=\mathcal{M}_{|\mathcal{A}|}$. Since we only focus on recognizing natural examples and adversarial ones in our setting, the total number of tasks $\mathcal{W}$ is set to two.

Given a global data distribution $\mathcal{D}$ for the tradeoff problem, as denoted in Section 2, $\mathcal{D}_{1},\mathcal{D}_{2}$ are subject to the distribution of training data $\mathcal{D}_{\mathcal{W}}$. And natural images $(x,y)\sim\mathcal{D}_{1}$ while adversarial examples $(x^{\prime},y)\sim\mathcal{D}_{2}$ generated by Eq. 1. So the training process of base learners is to solve the inner minimization of Eq. 3 over different distributions in a distributed manner:

Specifically, during the process, base learners $f_{\boldsymbol{\theta}_{n}}$ and $f_{\boldsymbol{\theta}_{r}}$ are assigned different subproblems that only requires accessing their own data distribution, respectively. Note that two base learners work in a complementary manner, meaning the update of parameters is independent among base learners and the global learner always collects parameters of both base learners. So the subproblem for each base learner is defined as:

where the task-aware optimizer $\mathcal{Z}_{\mathcal{W}}^{T}(\cdot,\cdot)$ search the optimal parameter states $\boldsymbol{\theta}_{\mathcal{W}}^{\star}$ over the subproblem $\mathcal{W}$ in $T$ rounds. Loss functions can also be task-specific and applied to each base learner separately. It is natural to consider minimizing the 0-1 loss in the natural and robust errors, however, solving the optimization problem is NP-hard thus computationally intractable. In practice, we select cross-entropy as the surrogate loss for both $\ell_{1}$ and $\ell_{2}$ since it is simple but good enough.

During the initial training periods, base learners are less instrumental since they are not adequately learned. Directly initializing parameters of base learners may mislead the training procedure and further accumulate bias when mixing them. Therefore, we set aside $t^{\prime}$ epochs from the beginning for fully training base learners and just aggregates states on the searching trajectory of base learners through optimization by exponential moving average (EMA), computed as: $\boldsymbol{\theta}_{g}\leftarrow\alpha^{\prime}\boldsymbol{\theta}_{g}+(1-\alpha^{\prime})(\gamma\boldsymbol{\theta}_{r}+(1-\gamma)\boldsymbol{\theta}_{t})$, where $\alpha^{\prime}$ is the exponential decay rates for EMA and $\gamma$ is the mixing ratio for base learners. They then learn an initialization from parameters of the global learner every $c$ epochs when each base learner is well trained in its field. Thus, the optimization of each base learner for every interlude can be expressed in Eq. 6:

Note that $\boldsymbol{\theta}_{g}$ contains both $\boldsymbol{\theta}_{n}$ and $\boldsymbol{\theta}_{r}$, meaning there always exists a term updated by gradient information of distribution different from the current subproblem. This mechanism enables fast learning within a given assignment and improves generalization, and the acceleration is applicable to the given assignment for its corresponding base learner only (proof in Appendix B.1).

With all discussed above, the learning progress of Generalist can be constructed by decending the gradient of $\boldsymbol{\theta}_{r},\boldsymbol{\theta}_{n}$ and mixing both of them. The calculating steps in Algorithm 1 can be summarized in Eq. 7.

where $\mathcal{B}(t,t^{\prime},c)$ is a Boolean function that returns one only when both $t\geq t^{\prime}$ and $t\ \operatorname{mod}c==0$, otherwise it returns zero. $\mathcal{Z}_{n}$ and $\mathcal{Z}_{r}$ are optimizers for natural training and adversarial training assignments.

In this part, we theoretically analyze how base learners help global learner in Generalist. For brevity, we omit the expectation notation over samples from each distribution without losing generalization.

The last term obtains the oracle state $\boldsymbol{\theta}_{a}^{\star}$, theoretically optimal parameters for each task $a$. $\mathbf{R}_{T}$ is the sum of the difference between the parameters of each base learner and the theoretically optimal parameters for each task. Based on the definition, we can give the following upper bound on the expected error of classifier trained by Generalist with respect to $\mathbf{R}_{T}$ as:

So the above inequality indicates that any strategy beneficial to reducing the error of each task that makes $\mathbf{R}_{T}$ smaller will decrease the error bound of the global learner. Considering Generalist divides the tradeoff problem into two independent tasks, Theorem 1 guarantees the upper bound of the risks given by the global learner trained by Generalist will get lower once the error for each task becomes lower. In practice, we can apply customized learning rate strategies, optimizers, and weight averaging to guarantee the error reduction of each base learner.

To comprehensively manifest the power of our Generalist method, we present the results of both ResNet-18 and WRN-32-10 on CIFAR-10 in Table 1.

In Table 1(a), Generalist consistently improves standard test error relative to models trained by several robust methods, while maintaining adversarial robustness at the same level. More specifically, Generalist achieves the second highest standard accuracy of 89.09% (only lower than 93.04% obtained by natural training (NT)), while meantime robust accuracy against AA is 46.07%, hanging on to 48.2% from TRADES. If we force TRADES to meet the same level of clean accuracy as Generalist (89%), the robustness of TRADES against APGD will drop to 30% (see TRADES in Appendix A.4), which is significantly worse than Generalist. That means it is hard to obtain acceptable robustness but maintain clean accuracy above 89% in the joint training framework even if it is equipped with an advanced loss function, while the improvement of Generalist is notable since we only use the naive cross-entropy loss. Contrary to FAT managing the tradeoff through adaptively decreasing the step size of PGD, which still hurts robustness a lot, Generalist is the only method with clean accuracy above 89% and robust accuracy against AA above 46%. We should emphasize the final obtained model of Generalist is the same size as other trained models are. For the training time, Generalist does perform both NT and naive AT but the cost of NT is negligible, so the overhead of Generalist is smaller than TRADES, and whatever serial and parallel versions of Generalist are even faster than TRADES (see Appendix A.4).

Things become more obvious when it comes to WRN-32-10. In Table 1(b), the gap between test natural accuracy of Generalist and NT is reduced to 2.27%, a relative decrease of 3.65% in standard test error as compared to the second highest natural accuracy (except NT) achieved by FAT. It is also remarkable that the boost of accuracy does not hurt the robustness of Generalist, instead, Generalist even outperforms TRADES across multiple types of adversarial attacks. In particular, we find that Generalist has a standard test error of 6.7% while TRADES with $\lambda=6$ has a standard test error of 14.89% only. And the improved robustness of Generalist among PGD20/100, MIM, CW, $\operatorname{FAT}_{t}$ and Square is conspicuous. Besides, the best performance on AA, which is an ensemble of different attacks and the most powerful adaptive adversarial attack so far, demonstrates the reliability of Generalist. Likewise, only Generalist attains robust accuracy of AA higher than 52% along with clean accuracy higher than 90%. It should be emphasized that these features confirm the practicability of Generalist. In short, Generalist has consistently improved robustness without loss of natural accuracy. More results on benchmark datasets of MNIST, SVHN, and CIFAR-100 are in Appendix A.2 and A.3.

We run a number of ablations to analyze the Generalist framework in this part. As illustrated in Algorithm 1, two factors control the tradeoff between accuracy and robustness of the global learner: frequency of communication $c$ and mixing ratio $\gamma$. Here, we investigate how these parameters affect performance. If not specified otherwise, the experiments are conducted on CIFAR-10 using ResNet-18.

As stressed above, one of the major advantages of Generalist in comparison with the standard joint training framework is that each base learner enables to customize the corresponding strategy for their own tasks freely rather than using the same strategy for all tasks. In this part, we investigate whether Generalist performs better when cooperating with diverse techniques.

Weight Averaging. Recent works [DBLP:journals/corr/abs-2103-01946, DBLP:conf/uai/IzmailovPGVW18, wang2022selfensemble] have shown that weight averaging (WA) greatly improves both natural and robust generalization. The average parameters of all history model snapshot through the training process to build an ensemble model on the fly. However, such technique cannot benefit both accuracy and robustness in the joint training framework. Therefore, we introduce WA into base learners separately. Results are shown in Figure 4 (a). We employ WA in either NT (NT_only) or AT (AT_only) or both of them (NT+AT). Overall, the results confirm that the performance of the global learner can be further improved after both base learners exploit WA. But unfortunately, an obvious tradeoff happens if only one of the base learner is equipped with WA. For instance, the standard test accuracy of NT_only continues to increase at the expense of the drop in the ability to defend attacks. A likely reason is that WA implicitly controls the learning speed of base learners. Indeed, the base learner with WA becomes an expert much faster than the one without WA in its sub-task, meaning the fast one is not in accordance with the slow one. This result is important because it not only illustrates the potential of Generalist comes from its base learners but also identifies a key challenge of tradeoff for future improvement.

Different Optimizers. We also investigate the effect of optimizers designed for different tasks. We choose AT ($\beta=1$) using SGD with momentum and Adam for piecewise learning rate schedule optimized by joint training as the baseline. The initial learning rate for Adam is 0.0001. We alternately apply these two optimizers in each subproblem. The comparison of the results is shown in Figure 4 (b). We can see that the gap of robust accuracy between models adversarially trained by Adam and the ones trained by SGD is significant. All three schemes equipped with Adam, namely NT (Adam)+AT (Adam), NT (SGD)+AT (Adam), and Baseline (Adam), perform worse than the ones using SGD when evaluated by adversarial attacks. But on the other hand, by comparing the results of Baseline (Adam) and NT (Adam)+AT (SGD), it confirms a proper optimization scheme with respect to data distribution can effectively benefit the corresponding performance without overlooking the other. That not only demonstrates the necessity of Generalist to decouple task-aware assignments from joint training but also indicates using Adam may not be the principal reason for robustness drop. It is just ill-suited for the outer and inner optimization in AT. Besides, though the best results still come from using SGD, the learning rate for different tasks can be customized which is not feasible in the joint framework, as shown in Appendix A.5.

Considering the proposed method achieves impressive clean accuracy without a harsh drop in robustness, it is naturally to ask what improvements Generalist has secured in comparison with robust methods in detail. Thus, we further investigate the predictions that robust classifiers are prone to make. As shown in Figure 5, we provide two perspectives to analyze the differences that classifiers trained by different AT methods.

To broadly study the case, we perform experiments on NT, TRADES with different $\lambda$, FAT and Generalist, then plot the distribution of the correct predictions of all methods for each class in Figure 5(a). As evident at first glance, we note that animals are more frequently misclassified, especially cats/dogs in the natural scenario and cats/deers in the adversarial scenario. In addition, the classifier trained by standard natural training does not always outperform the ones adversarially trained. Actually, they are equally skilled at most categories and the outcome is decided by specific categories (e.g. birds, cats and dogs). Generalist keeps pace with NT in the natural task, and meanwhile promotes the higher improvements in difficult items (e.g. cats and deers) against AA attack.

In Figure 5(b), we display specific samples in the testing dataset that are misclassified by robust classifiers (TRADES and FAT) but recognized by our proposed method, including both natural examples (the first two rows) and adversarial examples (the last two rows). Here, images shown in the first row are easy ones where the foreground objects stand out from the clear backgrounds, while hard samples are referred to those having confused objects with messy backgrounds. It is worth noting that TRADES delivers poor performances not only on hard examples with complex backgrounds or obscured objects but also on simple ones. For example, each image in the first row is typically plain and regular, however, TRADES fails in categorizing them into the right class. A plausible explanation for the issue is that TRADES lacks in a set of support measures specially devised for the natural classification task unlike Generalist does, highlighting design differentiation for sub-tasks is necessary.

Another interesting finding is that though both TRADES and FAT can build a robust classifier, they still rely on spurious background information and thus are easily deceived when encountering images with similar backgrounds but different objects. This phenomenon can be verified from the misclassification of the fourth and fifth images in the first row (taking white/blue backgrounds as evidence), and the fifth image in the fourth row (confused by the green background). But Generalist has the ability to sift the invariant feature of the foreground object while ignoring the background information spuriously correlated with the categories in both natural and adversarial settings. On the whole, Generalist demonstrates its strength to differentiate difficult samples close to the decision boundary and its potential to learn a background-invariant classifier.