DAIR: Data Augmented Invariant Regularization

1 INTRODUCTION . Deep neural networks are widely used in various applications ranging from computer vision to language processing . While deep learning has surpassed human-level performance in numerous tasks , neural networks are extremely vulnerable to overfitting to spurious correlations and therefore fail to generalize even under slight perturbations of the test distribution ( Arjovsky et al. , 2019 ) . This observation motivated the research community to tackle the problem of domain generalization ( see ( Ribeiro et al. , 2020 ) for a detailed literature review ) . Recent benchmark datasets , such as Rotated MNIST ( Arjovsky et al. , 2019 ) , Colored MNIST ( Arjovsky et al. , 2019 ) , PACS ( Li et al. , 2017 ) , VLCS ( Fang et al. , 2013 ) , Office-Home ( Venkateswara et al. , 2017 ) , Terra Incognita ( Beery et al. , 2018 ) and DomainNet ( Peng et al. , 2019 ) , have shown difficulties for the generalization of deep neural network models under distribution shifts , and have sparked invention of many new algorithmic frameworks to address domain generalization . A standard approach for improving out-of-distribution performance is to guarantee that learned models are invariant to certain transformations . For example , trained models for computer vision should generally be invariant to rotations , changes in color , or background . Geometric deep learning bakes such invariances into the neural network architecture . For example , convolutional layers ( Lecun et al. , 1998 ) are fundamentally preserving translations . There are other specifically designed networks to maintain invariances : Zaheer et al . ( 2017 ) studied the problem of designing models for machine learning tasks defined on sets and characterized the permutation invariant functions . Bloem-Reddy & Teh ( 2020 ) obtained generative functional representations of probability distributions that are invariant under the action of a compact group . Finzi et al . ( 2021 ) provided an algorithm for solving for the equivariant layers of matrix groups . Data augmentation promotes invariances in models by curating synthetic examples that exhibit the desired invariances . Tensmeyer & Martinez ( 2016 ) showed simple image transformations affect the CNN representations . Mixup ( Zhang et al. , 2017 ) , CutMix ( Yun et al. , 2019 ) and Cutout ( DeVries & Taylor , 2017 ) showed linear combination and random blocking features improves generalization of state-of-the-art neural network architectures . Volpi et al . ( 2018 ) ; Zhou et al . ( 2020 ) showed data augmentation with adversarial images could make the label classifier more robust to unknown domain shifts . Cubuk et al . ( 2018 ) ; Lim et al . ( 2019 ) introduced a procedure which automatically searches for improved data augmentation policies . Zhou et al . ( 2020 ) showed data augmentation with adversarial images could make the label classifier more robust to unknown domain shifts . Nam et al . ( 2021 ) improved domain generalization by reducing the intrinsic style bias of CNNs through training a separate network for randomizing the style of images and generating augmented data during training . Consistency regularization can be further applied on top of data augmentation to enhance invariance by enforcing similarities on the model . Engstrom et al . ( 2018 ) ; Kannan et al . ( 2018 ) ; Zhang et al . ( 2019 ) utilized consistency regularization to train robust neural networks against adversarial attacks . This has been applied to unsupervised learning ( Sinha & Dieng , 2021 ) , self-supervised learning ( Chen et al. , 2020 ; von Kügelgen et al. , 2021 ) , and semi-supervised learning to exploit unlabeled data ( Bachman et al. , 2014 ; Laine & Aila , 2016 ; Sohn et al. , 2020 ; Xie et al. , 2020 ) . Besides the directions mentioned above , researchers have proposed numerous algorithmic solutions to impose invariance and improve domain generalization such as DANN ( Ganin et al. , 2016 ) , IRM ( Ghifary et al. , 2015 ) , DRO ( Sagawa et al. , 2019 ) , MLDG ( Li et al. , 2018a ) , CORAL ( Sun & Saenko , 2016 ) , MMD ( Li et al. , 2018b ) and CDANN ( Li et al. , 2018c ) and REx ( Krueger et al. , 2021 ) . The approaches listed above are more complex than simple training mechanisms such as empirical risk minimization ( ERM ) and hence they can not be readily applied to involved tasks with non-trivial model architectures . For example , in generative language models imposing a constraint on the intermediate data representations is non-trivial , which is required by CORAL ( Sun & Saenko , 2016 ) . Recently , Gulrajani & Lopez-Paz ( 2020 ) demonstrated that ERM may even outperform many such complex methods in real-world scenarios , while ERM itself is known to generalizes poorly to distribution shift . For example , in learning neural dialog models , Qian et al . ( 2021 ) showed up to 29 % performance drop due to the memorization of named entities . Ribeiro et al . ( 2020 ) showed that both commercial and state-of-art language models fail on up to 76.4 % of the generalization tests . In this paper , we propose a consistency regularization technique , called data augmented invariant regularization ( DAIR ) . DAIR is applicable when data augmentation results in pairs of data samples expecting consistent performance , it specifically penalizes the inconsistency of loss on augmented samples with respect to the original ones . This is in contrast to many feature consistency regularizers that apply on an intermediate embedding space . As a result , DAIR only requires marginal additional cost on top of data augmentation , and is simple and broadly applicable to a wide host of supervised and unsupervised learning tasks , including generative models . We introduce the DAIR formulation , motivate it , and theoretically prove some of its properties in Section 2 . We empirically evaluate DAIR on a variety of problem setups ranging from defense against adversarial attacks to domain generalization in the presence of environment shift in Section 3 , where our experimental results show that DAIR is competitive with or even outperforms state-of-the-art algorithms specifically designed for imposing invariance in these problems . 2 DAIR : DATA AUGMENTED INVARIANT REGULARIZATION . For a data sample z = ( x , y ) , let ℓ ( z ; θ ) be its parametric loss function , where θ is the set of model parameters ( e.g. , network weights ) . The popular Empirical Risk Minimization ( ERM ) framework trains the model by minimizing the expected value of the following loss over the training data : fERM ( z ; θ ) = ℓ ( z ; θ ) . ( ERM ) We assume that we have access to a ( potentially randomized ) data augmenter function A ( · ) . Examples for A include ( random ) rotation , change of background , or change of entity names . Such augmenters aim at capturing the transformations against which we wish to be invariant to . Given a sample z , let z̃ = ( x̃ , ỹ ) = A ( z ) denote an augmented sample . Previous work has used both original and augmented examples during training , which leads to the following standard objective function , called Data Augmented Empirical Risk Minimization ( DA-ERM ) : fDA-ERM ( z , z̃ ; θ ) = 1 2 ℓ ( z ; θ ) + 1 2 ℓ ( z̃ ; θ ) . ( DA-ERM ) While DA-ERM has been successful in many applications , one natural question is whether we can further improve upon it using the knowledge that the performance on augmented samples should be consistent with the original ones . Consistency regularization further penalizes DA-ERM for any such inconsistency at the feature/loss level : fConsistency , D , λ ( z , z̃ ; θ ) = fDA-ERM ( z , z̃ ; θ ) + λD ( z , z̃ ; θ ) , where D ( z , z̃ ; θ ) is a proper divergence between the original sample representation and the augmented sample representation , and where the goal of the regularizer applied at some intermediate feature space is to maintain the performance of the model on z and z̃ consistent . In this paper , we focus on a specific type of such regularization , called data augmented invariant regularization ( DAIR ) : fDAIR , R , λ ( z , z̃ ; θ ) = fDA-ERM ( z , z̃ ; θ ) + λD ( z , z̃ ; θ ) = 1 2 ℓ ( z ; θ ) + 1 2 ℓ ( z̃ ; θ ) + λR ( ℓ ( z ; θ ) , ℓ ( z̃ ; θ ) ) , ( DAIR ) where the regularization is directly applied to the loss . The idea behind DAIR is to simply promote ℓ ( z ; θ ) ≈ ℓ ( z̃ ; θ ) , and ignore the features or even the rest of the possible outcomes of y and simply focus on the current sample ’ s loss . Hence , DAIR is a relatively weak form of consistency regularization only enforcing an original sample and an augmented one to be equally likely under the learned model ( assuming loss is a log-likelihood function ) . This weaker form of consistency is suitable for problems where feature consistency may not be conceptually meaningful . For instance , in language modeling when a pair of sentences differ in their corresponding named entities , it is not clear why we should enforce their embeddings to be similar , however , loss consistency is still meaningful promoting the probability of label given input to be the same on the original and the augmented samples . We remark that DAIR requires pairing information between original and augmented samples , which may not always be available ( e.g. , DomainBed ( Gulrajani & Lopez-Paz , 2020 ) ) . However , we show that this simple approach is still broadly applicable to various real-world problems regardless of model architecture , and is indeed competitive with state-of-the-art methods for imposing invariance . As it turns out , we are particularly interested in a particular form of the DAIR regularizer : Rsq ( ℓ ( z ; θ ) , ℓ ( z̃ ; θ ) ) : = ( √ ℓ ( z ; θ ) − √ ℓ ( z̃ ; θ ) ) 2 , ( SQ Regularizer ) and we call this variant DAIR-SQ . Note that Rsq has the same scale as the loss function ℓ , making it easier to tune λ. Empirically we observe that the optimal λ for all the experiments mentioned later in the paper falls in [ 0.2 , 100 ] , across various tasks ( from regression to sequence-to-sequence generative modeling ) . Further justification on DAIR-SQ will be provided through the rest of this section . Finally , in most ( real-world ) applications performance is measured through 0-1 metrics other than the loss function . For example , we are usually concerned with accuracy in image classification while we optimize cross-entropy loss . Let F ( z ; θ ) ∈ { 0 , 1 } denote a 0-1 evaluation performance metric of interest , e.g. , accuracy . Given the sample z ( or z̃ ) , the model performance is captured by F ( z ; θ ) ( or F ( z̃ ; θ ) ) . For any z such that F ( z ; θ ) = 1 , we define the corresponding consistency metric as : CM ( z , z̃ ; θ ) = I { F ( z̃ ; θ ) = 1 | F ( z ; θ ) = 1 } . ( Consistency Metric ) Notice that similarly to the original performance metric , which is only used for model evaluation , we use the consistency metric at evaluation time only .

Summary: The authors propose a conceptually simple regularizer which forces consistency of loss between the true samples and their corresponding augmentations, termed DAIR-SQ. The authors demonstrate the qualitative characteristics of DAIR-SQ with a clean toy example, and a supporting theoretical result of asymptotic consistency on the same. For real-world datasets, the authors demonstrate the robustness induced by DAIR-SQ under domain shift, label noise, and adversarial attacks.