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: In this paper, the authors proposed a regularizer (DAIR) that penalize the difference between the model performance on original data and augmented data. The authors argued that the new regularizer imposes invariance and effectively prevented the model from overfitting to spurious features. By applying the regularizer on various types of machine learning tasks, the authors analyzed the experiments results and explained the empirical novelty.