BLOOD: Bi-level Learning Framework for Out-of-distribution Generalization

1 INTRODUCTION . Conventional machine learning algorithms are relying on the empirical risk minimization ( ERM ) method when they should learn from given data , and in many application areas , this approach has shown successful performance with high prediction accuracy . However , if a model learns spurious correlations during training , it can often fail with poor generalization performance , which is known as the out-of-distribution ( OOD ) generalization problem . Furthermore , in recent studies , it has been shown that ERM-based methods more easily learn such unstable correlations in the dataset and result in a poor generalization performance on real-world applications ( Beery et al. , 2018 ; Ilyas et al. , 2019 ; Geirhos et al. , 2018 ; de Haan et al. , 2019 ; Koh et al. , 2021 ) . To address this problem and obtain a robust de-biased model , many approaches have been proposed for the cases where biases are known beforehand or not . When biases are known as a priori , some studies applied adversarial training to remove biases from representations ( Belinkov et al. , 2019a ; b ) or re-weighting training samples ( Schuster et al. , 2019 ) , and assembling predictions of a biased model and the base model for ensemble with a product of experts ( Hinton , 2002 ; He et al. , 2019 ; Clark et al. , 2019 ; Mahabadi et al. , 2020 ) . However , these works are designed for a specific type of bias and thus require extra domain knowledge to generalize to new tasks . Moreover , without such prior knowledge for biases in the data , they are hard to be applied to practical applications . For the case of having no prior knowledge of spurious correlations , the most popular approach is leveraging the prediction result of a shallow model or weak learner while assuming them as a biased classifier . Since predictions of a biased classifier can provide useful clues for the spurious correlation it has learned , to learn from the weak models ’ mistakes , they down-weight the potentially biased examples while training a robust model ( Mahabadi et al. , 2020 ) . Although these works are more general approaches and save a lot of human efforts for finding biases in established datasets , we found that they still can not effectively remove multiple types of biases existing in data groups collected from different environments . Another type of effective OOD generalization approach is group distributionally robust optimization ( group DRO ) algorithm which alleviates model biases by minimizing the worst-case risk over a set of human-defined training groups ( Hu et al. , 2018 ; Sagawa * et al. , 2020 ) . In this method , the choice of how to group the training data allows us to introduce a prior knowledge of spurious correlations into optimization . However , finding multiple types of biases and accordingly constructing data groups are laborious processes . Therefore , a simple grouping algorithm is proposed in a recent study ( Bao et al. , 2021 ) , which splits the training dataset based on the prediction results of biased classifiers . Our approach is also aimed to create data groups that are informative for the multiple underlying biases in the training dataset so that minimizing the worst-case risk over all those data groups can provide a robust classifier . In practical settings where there is little or no prior information about the biases , most de-biasing methods , which automatically identify potential biases in the training data , can not discover all spurious correlations existing in the dataset . Furthermore , splitting training data into several static data groups can not effectively represent the effect of multiple biases existing in the dataset . In this work , first , we propose a novel strategy for discovering bias and splitting training data for a group-based de-biasing algorithm . Based on the prevailing automatic bias identifying approaches ( Utama et al. , 2020b ; Sanh et al. , 2021 ) , we train a shallow model for each group in a batch and dynamically re-group the environments according to the prediction correctness of the shallow model over all other environments in the batch . This batch-wise dynamic data re-grouping strategy allows us gradually uncover multiple unknown biases in the dataset while training a model . A shallow model tends to quickly overfit to surface form information , especially when they are trained with a small training data setting ( Utama et al. , 2020b ) ; therefore , if we re-group the samples in a batch , based on its biased prediction results , we can more effectively account for the various unknown biases in the training dataset . Furthermore , when this approach is combined with the group-based de-biasing method ( group DRO ) , we can train a more robust classifier by minimizing the worst-case risk over all interpolations of those dynamic data partitions . Second , we also propose a unified end-to-end learning framework for a stable classifier . Our framework is a bi-level learning process which extends the min-max objective of group DRO with the unknown bias discovering and grouping method . In the inner level of optimization , it discovers the biases in the environment of each batch by applying the dynamic data re-grouping method , and these re-partitioned data groups are used for the group DRO algorithm , which minimizes the worstcase risk in the outer level of optimization . We coin our novel learning framework as a Bi-level Learning framework for OOD generalization ( BLOOD ) and evaluate its OOD performance in both synthetic and real-world environments . In the empirical evaluation , our framework shows 47 % percent improvement on the Colored MNIST dataset and achieves the best results in real-world datasets ( Camelyon17-wilds , FMoW-wilds ) compared to other OOD methods . The main contributions of our work are the following : ( a ) we show that dynamically re-grouping the subset of environments , based on the predictions of the shallow model , gradually uncovers multiple types of spurious correlations existing in a dataset ; ( b ) we integrate automatic unknown bias identifying and grouping process to the group DRO by formulating a bi-level optimization objective ; ( c ) we propose a unified end-to-end learning framework which does not need prior knowledge of multiple dataset biases to obtain robust models , but automatically removes various unknown biases from out-of-distribution data . 2 METHOD . 2.1 OVERVIEW . Consider a set ofN training environments Etr = { ei } Ni=1 where each environment ei is composed of input-label pairs { ( xek , yek ) } n k=1 . Our main goal is to train a stable classifier from these environments so that it can be generalized to any new test dataset . We do not make any assumption on the biases present ( or not ) in the dataset and rely on letting the shallow model discover them during training . For this type of problem , recent de-biasing methods ( Utama et al. , 2020b ; Sanh et al. , 2021 ; Liu et al. , 2021 ; Bao et al. , 2021 ) apply a two-stage process : ( a ) train a separate biased classifier by learning from spurious correlations in the environments and use the prediction correctness of the biased classifier for identifying a biased sample ( Utama et al. , 2020b ; Liu et al. , 2021 ) or creating new environment partitions for group-based de-biasing methods ( Bao et al. , 2021 ) . ( b ) train a robust classifier by re-weighting biased samples or applying product-of-experts ( Hinton , 2002 ; Sanh et al. , 2021 ) , confidence regularization ( Utama et al. , 2020a ) to the outcome of the first stage . In our approach , we integrate these two stages into a single unified learning framework by leveraging a bi-level optimization structure . With the bi-level setting , we can identify unknown biases in the inner-level and remove the discovered biases with group DRO in the outer-level . In the innerlevel , to automatically discover unknown biases in the environments , we train a shallow model for each environment in a batch and apply an environment-specific classifier to partition all other environments in the batch , based on its prediction correctness . Since a shallow model is more prone to rely on shallow heuristics , we can obtain a biased classifier after few-shot learning in the batch , and this procedure is iteratively performed while training a model . In the outer-level , the dynamically partitioned environments , which are obtained from inner-level optimization , are provided to the group DRO algorithm and used for minimizing the worst-case risk over those partitions . In the next section , we describe each component of our bi-level learning structure with its corresponding contribution . 2.2 BI-LEVEL LEARNING FRAMEWORK FOR DE-BIASING . Our learning framework consists of two training objectives for bi-level optimization , in which the inner objective is nested within the outer objective . The inner training objective learns spurious correlations in the dataset , which are unknown , and uses them for re-grouping other environments . The outer objective learns only stable correlations by minimizing the worst-case risk over these groups . 2.2.1 INNER OBJECTIVE : LEARNING BIAS TO DE-BIAS . The inner-level optimization is aimed to automatically discover unknown spurious correlations , based on the prediction result of a biased classifier ( weak learner ) , and accordingly re-partitions the environments to provide groups to the outer-level process . Therefore , we need to train a biased model which mostly follows spurious correlations in each environment . Most of the other de-biasing methods , exploiting the prediction results of a biased classifier , use a pre-built biased model by training it with the full set or a small subset of training data . In contrast to those works , we do not train a separate biased model ; instead , we dynamically obtain a shallow model for each environment in a batch during inner-level optimization . To obtain biased models while training stable classifier , we get insight from the deep neural network ’ s tendency to exploit simple patterns in the early stage of the training , which is also observed in other researches ( Arpit et al. , 2017 ; Liu et al. , 2020 ) . Since spurious correlations are commonly characterized as simple surface patterns , we expect that models ’ rapid performance gain is attributed to their reliance on simple surface patterns ( Utama et al. , 2020a ) . In the same context , after a model is trained with only a small number of samples from an environment , we expect it to perform as a biased classifier , which achieves high accuracy mostly on the biased examples while still performing poorly on the rest of the samples from other environments . For each batch learning step in the inner optimization , a shallow model is trained and used for re-grouping as follows : Step 1 : For each sampled environment ei in a batch , temporarily train an environment-specific classifier fφi with few-shot learning . Step 2 : For all sampled environments in the batch where j 6= i , use the trained classifier fφi to partition each of them into two parts based on the prediction correctness of the classifier . Training a shallow model At the beginning of a model training , a model fθ is randomly initialized with θ . While training a classifier , we obtain a shallow model for each batch training step by applying a few gradient descent steps to the model ’ s parameters θ , so that the model can quickly overfit to Algorithm 1 BLOOD : Bi-level Learning Framework for OOD Generalization Input : Step sizes α , β , γ ; A set of training environments Etr = { ei } Ni=1 Initialize : θ and q = [ q 1 , q ⊗ 1 , · · · , q N , q ⊗ N ] 1 : while not done do 2 : for all ei ∈ Etr do 3 : φi = θ − α∇θLei ( fθ ) { Inner-loop optimization for a shallow model } 4 : ej ∼ Etr \ ei { Sample training environment ej } 5 : G i→j , G ⊗ i→j = Partition ( fφi , ej ) { Partition ej into G i→j and G ⊗ i→j } 6 : q j ← q j exp ( γLG i→j ( fφi ) ) , q ⊗ i ← q ⊗ i exp ( γLG⊗i→j ( fφi ) ) { Update group weights } 7 : end for 8 : q ← q/ ∑ i ( q i + q ⊗ i ) { Normalize group weights q } 9 : θ ← θ − β ∇θ ∑ i ∑ j q j LG i→j ( fφi|θ ) + q ⊗ j LG⊗i→j ( fφi|θ ) { Outer-loop optimization } 10 : end while surface form information while fθ is still conditioned on the previously learned θ . This temporarily trained shallow model parameter φi , which are optimized for surface patterns of ei , is updated by following stochastic gradient descent ( SGD ) step : φi = θ − α∇θLei ( fθ ) ( 1 ) where α is a step size for the inner optimization , and Lei ( fθ ) is an expected loss of fθ over the data sampled from ei . We show only a single gradient update procedure in Equation 1 for the notational simplicity . Batch-wise dynamic data re-grouping After a biased classifier fφi is trained for the environment ei , we dynamically re-partition all other training environments in the batch . The fφi is applied to all other environments ej in the batch , where j 6= i , and according to its prediction correctness , samples in each environment ej are re-grouped into two parts , a correctly predicted sample group G i→j and a incorrectly predicted sample group G⊗i→j . These dynamically partitioned groups are then used for training a stable classifier by minimizing the worst-case risk over these groups in the outer-level . As a similar group-based de-biasing approach , Predict then Interpolate ( PI ) ( Bao et al. , 2021 ) also uses a biased classifier for re-partitioning other environments . Its main difference compared to our method is that they fully train a separate model with the entire dataset and use it for partitioning all environments into static groups . Once their data groups are statically assigned , it is fixed throughout the optimization process of group DRO . Although PI has shown its theoretical and empirical effectiveness as a de-biasing method , its static grouping strategy , which relies on a fully-trained biased classifier , constrains it from discovering multiple types of spurious correlations in the environments . In contrast , our framework uses a shallow model , which is trained with a small subset of each environment in the batch . Since we are using a different biased classifier per batch , we are continuously learning a bias and producing new data partitions for each batch , representing our dynamic data re-grouping strategy . This dynamic re-grouping strategy provides various group combinations for the outer-level optimization process , so that it can effectively remove various types of biases in the environments . Furthermore , since our shallow model has overfitted to the surface patterns in a small subset of data , it is more likely to capture multiple types of bias in the environments , as shown in other research ( Utama et al. , 2020a ) . This property can be a disadvantage for a de-biasing method , which applies re-weighting to the biased examples , because it reduces the effective training data size for a model . However , in our approach , it enables us to uncover multiple types of biases in the dataset . Because , multiple types of biases can be considered while minimizing the worst-case risk over various types of group configurations .

Summary: The paper studies the problem of OOD classification: the test data and training data distribution can have different spurious feature-class dependencies. The goal is to have high test accuracy of a model trained in one environment. The paper designs a bi-level optimization problem: the lower-level is to re-group data in an environment (distribution) into two groups (correct and incorrect predictions), and the upper-level is to train a model that has different weights over the two groups to focus on the examples that have the changed spurious feature-class dependencies. Experiments on two real-world datasets and synthetic datasets show significantly improved OOD accuracy. There is no theoretical analysis of the method.