ARCH : Efficient Adversarial Regularized Training with Caching (2021)

Abstract: Adversarial regularization can improve model generalization in many natural language processing tasks. However, conventional approaches are computationally expensive since they need to generate a perturbation for each sample in each epoch. We propose a new adversarial regularization method ARCH (adversarial regularization with caching), where perturbations are generated and cached once every several epochs. As caching all the perturbations imposes memory usage concerns, we adopt a K-nearest neighbors-based strategy to tackle this issue. The strategy only requires caching a small amount of perturbations, without introducing additional training time. We evaluate our proposed method on a set of neural machine translation and natural language understanding tasks. We observe that ARCH significantly eases the computational burden (saves up to 70% of computational time in comparison with conventional approaches). More surprisingly, by reducing the variance of stochastic gradients, ARCH produces a notably better (in most of the tasks) or comparable model generalization. Our code is publicly available.

## introduction adversarial regularization @xcite can improve model generalization in many natural language processing tasks, such as neural machine translation @xcite , natural language understanding @xcite , language modeling @xcite , and reading comprehension @xcite . even though the method has demonstrated its power in many scenarios, its computational efficiency remains unsatisfactory. conventional adversarial regularization @xcite methods involve a min-max optimization problem. specifically, a perturbation is * corresponding author. generated for each sample by solving a maximization problem, and the model parameters are subsequently updated through a minimization problem, subject to the generated perturbations. a popular algorithm @xcite for such optimization is to alternate between several projected gradient descent steps (pgd, for the maximization) and a gradient descent step (for the minimization). there are two drawbacks with the alternating gradient descent/ascent method. first, the procedure requires significant computational efforts. suppose we run pgd for s steps, then we introduce extra s forward passes and extra s backward passes in each iteration. as such, training with adversarial regularization is significantly slower than standard training. second, optimizing the min-max problem is hard. this is because the perturbations are model and data dependent, and thus, variance of them is large. that is, the model needs to adapt to drastically different "noisy data" (i.e., clean data with perturbations), such that the stochastic gradients vary significantly during training. such large variance imposes optimization challenges. we propose arch (adversarial regularization with caching) that alleviates the aforementioned issues by reusing perturbations. recall that in conventional adversarial regularization methods, a different perturbation is generated for each sample in each epoch. in contrast to this, we propose to generate perturbations less frequently. for example, for a given sample, we can generate a new perturbation every 20 epochs, and the sample's perturbation remains unchanged in other epochs. we call this method "caching". the method has two advantages. first, it alleviates the computational burden. by reusing the perturbations, we avoid the extra forward and backward passes caused by pgd for most of the iterations. second, caching stabilizes the stochastic gradients. notice that in our method, the model is optimized with respect to the same noisy data for multiple times, instead of only one. in this way, variance of the stochastic gradients is reduced. one caveat of the caching method is its memory overhead. this is because a sample's perturbation is significantly larger than itself (the perturbation has an extra embedding dimension). we propose a k-nearest neighbors-based approach to tackle this problem. specifically, instead of caching perturbations for all the samples, we only cache a small proportion of them. each uncached perturbation can then be constructed using the cached ones in its neighborhood. such a construction procedure can be executed in parallel with model training. therefore, training time will not be prolonged because of this memory saving strategy. we use a moving average approach to boost model generalization. specifically, when generating a new perturbation, we integrate information from both the current model and the current perturbation. this is different from conventional approaches, where the new perturbation only depends on the current model. the moving average approach has a smoothing effect that boosts model generalization, as demonstrated both theoretically and empirically by previous works @xcite @xcite . arguably, the perturbations introduced by our method may not constitute strong adversarial attacks, because of the "staleness" caused by infrequent updates. however, we highlight that the focus of this work is model generalization over clean data, instead of adversarial robustness (ability to defend attacks). as we will demonstrate in the experiments, the "weak" perturbations show notable improvement of model generalization. and somewhat surprisingly, arch also exhibits on par or even better robustness comparing with conventional approaches. we conduct extensive experiments on neural machine translation (nmt) and natural language understanding (nlu) tasks. in comparison with conventional adversarial regularization approaches, arch can save up to 70% computational time. moreover, in nmt tasks, our method improves about 0.5 blue over baseline methods on seven datasets. arch also achieves 0.7 average score improvement on the glue @xcite development set over existing methods. we summarize our contributions as follows: (1) we propose a caching method that needs drastically less computational efforts. the method can also improve model generalization by reducing variance of stochastic gradient. (2) we propose a memory saving strategy to efficiently implement the caching method. (3) extensive experiments on neural machine translation and natural language understanding demonstrate the efficiency and effectiveness of the proposed method. fine-tuning pre-trained language models @xcite @xcite @xcite ) is a stateof-the-art method for natural language understanding tasks such as the glue @xcite benchmark. adversarial regularization is also incorporated into the fine-tuning approach. for example, @xcite al. (2020a) combines adversarial pretraining and fine-tuning, zhu et al. (2020); jiang et al. (2020) adopt trust region-based methods, and aghajanyan et al. (2020) aims for a more efficient computation. adversarial training was originally proposed for computer vision tasks @xcite @xcite , where the goal is to train robust classifiers. such methods synthesize adversarial samples, such that the classifier is trained to be robust against them. this strategy is also effective for tasks beyond computer vision, such as in reinforcement learning @xcite . various algorithms are proposed to craft the adversarial samples, e.g., learning-to-learn @xcite and stackelberg adversarial training @xcite . moreover, adversarial training is also well-studied theoretically @xcite . in natural language processing, the goal is no longer adversarial robustness, but instead we use adver-sarial regularization to boost model generalization. note that adversarial training and adversarial regularization are different concepts. the former focuses on defending against adversarial attacks, and the latter focuses on encouraging smooth model predictions @xcite . these two goals are usually treated as mutually exclusive @xcite . ## method generating perturbations for natural language inputs faces the difficulty of discreteness, i.e., words are defined in a discrete space. a common approach to tackle this is to work on the continuous embedding space @xcite . denote f (x, θ) a neural network parameterized by θ, where x is the input embedding. further denote y the ground-truth corresponding to x. for example, in classification tasks, x is the sentence embedding, and y is its label. in sequenceto-sequence learning, x is the source sentence embedding, and y is the target sentence. in both of these cases, the model is trained by minimizing the empirical risk over the training data, i.e., here {(x i , y i )} n @xmath0, and ℓ is a taskspecific loss, e.g., cross-entropy loss for classification and mean-squared error for regression. ## experiments in all the experiments, we use pytorch 1 @xcite as the backend. all the experiments are conducted on nvidia v100 gpus. ## conclusion we propose a new caching method to speedup the training of neural models with adversarial regularization. by reusing the generated perturbations, our proposed method significantly amortizes the computational cost of the backward passes at each iteration. our thorough experiments show that the proposed method not only improves the computational efficiency, but also reduces the variance of the stochastic gradients, which leads to better model generalization.