SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Theorem 4.2. Consider training a CNN model (1) with data model (3.1), define  $\mathbf{c}(t) \triangleq \left[C_{IRMvI}^{1}(\mathbf{W},t), C_{IRMvI}^{2}(\mathbf{W},t), \dots, C_{IRMvI}^{|\mathcal{E}_{r}|}(\mathbf{W},t)\right]$ , and  $\lambda_0 = \lambda_{\min}(\mathbf{H}^\infty)$ , where  $\mathbf{H}_{e,e'}^{\infty} \triangleq \frac{1}{2mn_{e}n_{e'}} \sum_{i=1}^{n_e} \psi'(\langle \mathbf{w}_{j,r}(0), \mathbf{x}_{1,i}' \rangle) \mathbf{x}_{1,i}' \sum_{i'=1}^{n_{e'}} \psi'(\langle \mathbf{w}_{j,r}(0), \mathbf{x}_{1,i'}^{e'} \rangle) \mathbf{x}_{1,i'}^{e'}$ . Suppose that dimension  $d = \Omega(\log(m / \delta))$ , network width  $m = \Omega(1 / \delta)$ , regularization factor  $\lambda \geq 1 / (\sigma_0 |\mathcal{E}_{tr}|^{3/2})$ , noise variance  $\sigma_p = O(d^{-2})$ , weight initial scale  $\sigma_0 = O(\frac{|\mathcal{E}_{tr}|^{7/2} \beta^3 L}{d^{1/2} m^2 \lambda_0^2 \log(1 / \epsilon)})$ , then with probability at least  $1 - \delta$ , after training time  $T = \Omega\left(\frac{\log(1 / \epsilon)}{\eta \lambda \lambda_0}\right)$ , we have  $\| \mathbf{c}(T) \|_2 \leq \epsilon$ ,  $\gamma_{j,r}^{inv}(T) = o(1)$ ,  $\gamma_{j,r}^{spu}(T) = o(1)$ .
The proof is given in Appendix D.3. We highlight that Theorem 4.2 allows any number of training environments, which indicates a fundamental limitation of pure IRMv1 training. Intuitively, Theorem 4.2 implies that, when a heavy regularization of IRMv1 is applied, the model will not learn any features, corresponding to Figure 2(d). Instead, IRMv1 suppresses any feature learning, even at the beginning of the training. Then, what would happen when given a properly pre-trained network?
After ERM pre-training, according to Theorem 4.1, we have  $|\langle \mathbf{w}_{j,r},\mathbf{v}_1\rangle | = \Omega (1),|\langle \mathbf{w}_{j,r},\mathbf{v}_2\rangle | = \Omega (1)$ , and  $|\langle \mathbf{w}_{j,r},\pmb {\xi}\rangle | = O(\sigma_0\sigma_p\sqrt{d})$ . Then, we have the following hold.
Proposition 4.3. Given the same setting as Theorem 4.2, suppose that  $\psi(x) = x$ ,  $\gamma_{j,r}^{inv}(t_1) = \gamma_{j,r}^{inv}(t_1 - 1)$ , and  $\gamma_{j,r}^{spu}(t_1) = \gamma_{j,r}^{spu}(t_1 - 1)$  at the end of ERM pre-train  $t_1$ ,  $\delta > 0$ , and  $n > C\log(1/\delta)$ , with  $C$  being a positive constant, then with a high probability at least  $1 - \delta$ , we have  $\sum_{e} C_{IRMvl}^{e}(t_1) = 0$ ,  $\gamma_{j,r}^{inv}(t_1 + 1) > \gamma_{j,r}^{inv}(t_1)$ , and  $\gamma_{j,r}^{spu}(t_1 + 1) < \gamma_{j,r}^{spu}(t_1)$ .
The proof is given in Appendix D.4, which takes converged feature learning terms from Theorem 4.1 as the inputs. Proposition 4.3 demonstrates that with sufficient ERM pre-training, IRMv1 can enhance the learning of invariant features while suppressing the learning of spurious features, which is verified in Figure 2(b) and 2(a). Thus, when given the initialization with better learned invariant features, i.e., longer ERM pre-training epochs, IRMv1 improves the invariant feature better. Proposition 4.3 explains why the final OOD performance highly depends on the ERM pre-training [16, 83].
4.3 Limitations of ERM Feature Learning
Combining results from both Sec. 4.1 and Sec. 4.2, we know that the invariant features will be learned during ERM pre-training and discovered during OOD training. However, given poorly learned invariant features, can IRMv1 still improve it? In practice, there often exist some invariant features that are not properly learned by ERM. For example, in our data model Def. 3.1 when the invariant correlation is much weaker than the spurious correlation, given a limited number of training steps, the spurious feature learning can dominate the invariant feature learning. Besides, when considering other factors such as the simplicity bias of ERM [67] or the inductive biases of the network architecture [28], it is more likely that there exist invariant features that are not properly learned [60]. Then we have:
Corollary 4.4. Consider training the CNN with the data generated from Def. 3.1, suppose that  $\psi(x) = x$ ,  $\gamma_{j,r}^{inv}(t_1) = o(1)$ , and  $\gamma_{j,r}^{spu}(t_1) = \Theta(1)$  at the end of ERM pre-training  $t_1$ . Suppose that  $\delta > 0$ , and  $n > C\log(1/\delta)$ , with  $C$  being a positive constant, then with a high probability at least  $1 - \delta$ , we have  $\gamma_{j,r}^{inv}(t_1 + 1) < \gamma_{j,r}^{inv}(t_1)$ .
Corollary 4.4 shows that IRMv1 requires sufficiently well-learned features for OOD generalization. It is also consistent with the experimental results in Fig. 2(b), 2(c), and Fig. 1, where all the OOD objectives only achieve a performance comparable to random guesses.
5 Feature Augmented Training
5.1 Rich Features for OOD Generalization
The results in Sec. 4 imply the necessity of learning all potentially useful features during the pretraining stage for OOD generalization. Otherwise, the OOD training is less likely to enhance the poorly learned features. It also explains the success of learning diverse and rich features by weight averaging [5, 59] and rich feature construction (or Bonsai) [83], and other approaches [58, 81].
Despite the empirical success, however, the learning of rich features in both Bonsai and weight averaging is unstable and expensive. On the one hand, they may discard previously learned useful features or fail to explore all the desired features as it is hard to evaluate the quality of the intermediate learned features. On the other hand, they also need multiple initializations and training of the whole networks with different random seeds to encourage the diversity of feature learning, which brings more instability and computational overhead, especially when applied to large and deep networks.
5.2 The FeAT Algorithm
To overcome the limitations of previous rich feature learning algorithms, we propose Feature Augmented Training (FeAT), that directly augment the feature learning in an iterative manner.
Intuitively, the potentially useful features presented in the training data are features that have non-trivial correlations with labels, or using the respective feature to predict the labels is able to achieve a non-trivial training performance. Moreover, the invariance principle assumes that the training data comes from different environments [4], which implies that each set of features can only dominate the correlations with labels in a subset of data. Therefore, it is possible to differentiate the distinct sets of useful features entangled in the training data into different subsets, where ERM can effectively learn the dominant features presented in the corresponding subset as shown in Theorem 4.1.
The intuition naturally motivates an iterative rich feature learning algorithm, i.e., FeAT, that identifies the subsets containing distinct features and ex
Algorithm 1 FeAT: Feature Augmented Training
1: Input: Training data  $\mathcal{D}_{\mathrm{tr}}$ ; the maximum augmentation rounds  $K$ ; predictor  $f \coloneqq w \circ \varphi$ ; length of inner training epochs  $t$ ; termination threshold  $p$ ;
2: Initialize groups  $G^{a} \gets \mathcal{D}_{\mathrm{tr}}, G^{r} \gets \{\}$ ;
3: for  $k \in [1, \dots, K]$  do
4: Randomly initialize  $w_{k}$
5: for  $j\in [1,\dots ,t]$  do
6: Obtain  $\ell_{\mathrm{FeAT}}$  with  $G$  via Eq. 7;
7: Update  $w_{k},\varphi$  with  $\ell_{\mathrm{FeAT}}$
8: end for
9: // Early Stop if  $f_{k} = w_{k} \circ \varphi$  fails to find new features.
10: if Training accuracy of  $f_{k}$  is smaller than  $p$  then
11: Set  $K = k - 1$  and terminate the loop;
12: end if
13: Split  $\mathcal{D}_{\mathrm{tr}}$  into groups  $\mathcal{D}_k^a, \mathcal{D}_k^r$  according to whether  $f_k$  classifies the examples in  $\mathcal{D}_{\mathrm{tr}}$  correctly or not;
14: Update groups  $G^{a} \gets G^{a} \cup \{\mathcal{D}_{k}^{a}\}$ ,  $G^{r} \gets G^{r} \cup \{\mathcal{D}_{k}^{r}\}$ ;
15: end for
16: Synthesize the final classifier  $w \gets \frac{1}{K} \sum_{i=1}^{K} w_i$ ;
17: return  $f = w \circ \varphi$
plies to learn new features in multiple rounds. The details of FeAT are given in Algorithm 1, where we are given a randomly initialized or pre-trained model  $f = w \circ \varphi$  that consists of a featurizer  $\varphi$  and a classifier  $w$ . In round  $k$ , FeAT first identifies the subset that contains the already learned features by collecting the samples where  $f$  yields the correct prediction, denoted as  $G_k^r$ , and the subset of samples that contains the features that have not been learned, denoted as  $G_k^a$ .
At the  $k$ -th round, given the grouped subsets  $G = \{G^r, G^a\}$  with  $2k - 1$  groups, where  $G^a = \{\mathcal{D}_i^a\}_{i=0}^{k-1}$  is the grouped sets for new feature augmentation, and  $G^r = \{\mathcal{D}_i^r\}_{i=1}^{k-1}$  is the grouped sets for already learned feature retention (notice that  $\mathcal{D}_0^r$  is the empty set), where  $\mathcal{D}_i^a$  and  $\mathcal{D}_i^r$  are the corresponding augmentation and retention set elicited at  $i$ -th round. FeAT performs distributionally robust optimization (DRO) [50, 83] on  $G^a$  to explore new features that have not been learned in previous rounds. Meanwhile, FeAT also needs to retain the already learned features by minimizing the empirical risk at  $G^r$ , for which we store and use the historical classifiers  $w_i$  with the current featurizer to evaluate the feature retention degree. Then, the FeAT objective at round  $k$  is
$$
\ell_ {\text {F e A T}} = \max  _ {\mathcal {D} _ {i} ^ {a} \in G ^ {a}} \ell_ {\mathcal {D} _ {i} ^ {a}} \left(w _ {k} \circ \varphi\right) + \lambda \cdot \sum_ {\mathcal {D} _ {i} ^ {r} \in G ^ {r}} \ell_ {\mathcal {D} _ {i} ^ {r}} \left(w _ {i} \circ \varphi\right), \tag {7}
$$
Table 1: OOD performance on COLOREDMNIST datasets initialized with different representations.
COLOREDMNIST-025	COLOREDMNIST-01
ERM-NF	ERM	BONSAI	FEAT	ERM-NF	ERM	BONSAI	FEAT
ERM	17.14 (±0.73)	12.40 (±0.32)	11.21 (±0.49)	17.27 (±2.55)	73.06 (±0.71)	73.75 (±0.49)	70.95 (±0.93)	76.05 (±1.45)
IRMV1	67.29 (±0.99)	59.81 (±4.46)	70.28 (±0.72)	70.57 (±0.68)	76.89 (±3.25)	73.84 (±0.56)	76.71 (±4.10)	82.33 (±1.77)
V-REX	68.62 (±0.73)	65.96 (±1.29)	70.31 (±0.66)	70.82 (±0.59)	83.52 (±2.52)	81.20 (±3.27)	82.61 (±1.76)	84.70 (±0.69)
IRMX	67.00 (±1.95)	64.05 (±0.88)	70.46 (±0.42)	70.78 (±0.61)	81.61 (±1.98)	75.97 (±0.88)	80.28 (±1.62)	84.34 (±0.97)
IB-IRM	56.09 (±2.04)	59.81 (±4.46)	70.28 (±0.72)	70.57 (±0.68)	75.81 (±0.63)	73.84 (±0.56)	76.71 (±4.10)	82.33 (±1.77)
CLOVE	58.67 (±7.69)	65.78 (±0.00)	65.57 (±3.02)	65.78 (±2.68)	75.66 (±10.6)	74.73 (±0.36)	72.73 (±1.18)	75.12 (±1.08)
IGA	51.22 (±3.67)	62.43 (±3.06)	70.17 (±0.89)	67.11 (±3.40)	74.20 (±2.45)	73.74 (±0.48)	74.72 (±3.60)	83.46 (±2.17)
FISHR	69.38 (±0.39)	67.74 (±0.90)	68.75 (±1.10)	70.56 (±0.97)	77.29 (±1.61)	82.23 (±1.35)	84.19 (±0.66)	84.26 (±0.93)
ORACLE	71.97 (±0.34)	86.55 (±0.27)
where  $\ell_{\mathcal{D}_i}(w\circ \varphi)$  refers to the empirical risk of  $w\circ \varphi$  evaluated at the subset  $\mathcal{D}_i$ , and  $\{w_i|1\leq i\leq k - 1\}$  are the historical classifiers trained in round  $i$ . The final classifier is the average of all historical classifiers as they already capitalize all the learned features in each round.
Relations with previous rich feature learning algorithms. Compared with previous rich feature learning algorithms, FeAT directly trades off the exploration of new features and the retention of the already learned features. Although Bonsai also adopts DRO to explore new features, the isolation of new feature exploration and already learned feature synthesis makes the feature learning in Bonsai more unstable. In other words, Bonsai can not evaluate the intermediate feature learning results due to the indirect feature exploration and synthesis. Consequently, Bonsai can not control when to stop the new feature exploration, and thus may fail to explore all of the desired features or discard important features. Besides, the multiple re-initializations and re-training of the whole network in Bonsai could also lead to suboptimal performance and meanwhile require more computational overhead.
Practical implementations. Algorithm 1 requires to store  $2K - 1$  subsets and a larger memory cost in training the network, which may cause additional storage burden when  $\varphi$  contains a massive amount of parameters [39]. Hence, we propose a lightweight variant of FeAT (denoted as iFeAT) which only retains the latest subsets and historical classifiers  $(\mathcal{D}_{k - 1}^{a},\mathcal{D}_{k - 1}^{r},w_{k - 1}$  at the  $k$ -th round). Throughout the whole experiment, we will use iFeAT and find that iFeAT already achieves state-of-the-art. More details are given in Appendix E.
As iFeAT stores only the latest augmentation and retention subsets, inspecting the training performance for termination check (line 10 of Algorithm 1) may not be suitable. However, one can still inspect the performance in either an OOD validation set to check the quality of the intermediate feature representations, or the retention set to check whether learning new features leads to a severe contradiction of the already learned features (FeAT should terminate if so).
Compared to ERM, the additional computational and memory overhead introduced in FeAT mainly lie in the FeAT training and partitioning. At each training step, FeAT needs  $(k - 1)$  additional forward and backward propagation, the same as Bonsai, while FeAT only needs  $\min(1, k - 1)$  additional propagation. Besides, Bonsai additionally require another round of training with  $(K - 1)$  additional propagation, given  $K$  total rounds. More details are given in Appendix F.4.
6 Empirical Study
We conduct extensive experiments on COLOREDMNIST [16] and WILDS [39] to verify the effectiveness of FeAT in learning richer features than ERM and the state-of-the-art algorithm Bonsai [83].
Proof-of-concept study on COLOREDMNIST. We first conduct a proof-of-concept study using COLOREDMNIST [16] and examine the feature learning performance of FeAT under various conditions. We consider both the original COLOREDMNIST with  $\mathcal{E}_{\mathrm{tr}} = \{(0.25, 0.1), (0.25, 0.2)\}$  (denoted as COLOREDMNIST-025), where spurious features are better correlated with labels, and the modified COLOREDMNIST (denoted as COLOREDMNIST-01) with  $\mathcal{E}_{\mathrm{tr}} = \{(0.1, 0.2), (0.1, 0.25)\}$ , where invariant features are better correlated with labels. We compare the OOD performance of the features learned by FeAT, with that of ERM and the state-of-the-art rich feature learning algorithm Bonsai [83]. Based on the features learned by ERM, Bonsai, and FeAT, we adopt various state-of-the-art OOD objectives including IRMv1 [4], VREx [41], IRMX [16], IB-IRM [1], CLOvE [76], IGA [40] and Fishr [57] for OOD training, in order to evaluate the practical quality of the learned features. The feature representations are frozen once initialized for the OOD training as fine-tuning the featurizer can distort the pre-trained features [43]. We also compare FeAT with the common training approach that uses unfrozen ERM features, denoted as ERM-NF. For Bonsai, we trained 2
Table 2: OOD generalization performances on WILDS benchmark.
INIT.	METHOD	CAMELYON17	CIVILCOMMENTS	FMOW	IWILDCAM	AMAZON	RxRX1
Avg. acc. (%)	Worst acc. (%)	Worst acc. (%)	Macro F1	10-th per. acc. (%)	Avg. acc. (%)
ERM	DFR†	95.14 (±1.96)	77.34 (±0.50)	41.96 (±1.90)	23.15 (±0.24)	48.00 (±0.00)	-
ERM	DFR-s†	-	82.24 (±0.13)	56.17 (±0.62)	52.44 (±0.34)	-	-
Bonsai	DFR†	95.17 (±0.18)	77.07 (±0.85)	43.26 (±0.82)	21.36 (±0.41)	46.67 (±0.00)	-
Bonsai	DFR-s†	-	81.26 (±1.86)	58.58 (±1.17)	50.85 (±0.18)	-	-
FeAT	DFR†	95.28 (±0.19)	77.34 (±0.59)	43.54 (±1.26)	23.54 (±0.52)	49.33 (±0.00)	-
FeAT	DFR-s†	-	79.56 (±0.38)	57.69 (±0.78)	52.31 (±0.38)	-	-
ERM	ERM	74.30 (±5.96)	55.53 (±1.78)	33.58 (±1.02)	28.22 (±0.78)	51.11 (±0.63)	30.21 (±0.09)
ERM	GroupDRO	76.09 (±6.46)	69.50 (±0.15)	33.03 (±0.52)	28.51 (±0.58)	52.00 (±0.00)	29.99 (±0.13)
ERM	IRMv1	75.68 (±7.41)	68.84 (±0.95)	33.45 (±1.07)	28.76 (±0.45)	52.00 (±0.00)	30.10 (±0.05)
ERM	V-REx	71.60 (±7.88)	69.03 (±1.08)	33.06 (±0.46)	28.82 (±0.47)	52.44 (±0.63)	29.88 (±0.35)
ERM	IRMX	73.49 (±9.33)	68.91 (±1.19)	33.13 (±0.86)	28.82 (±0.47)	52.00 (±0.00)	30.10 (±0.05)
Bonsai	ERM	73.98 (±5.30)	63.34 (±3.49)	31.91 (±0.51)	28.27 (±1.05)	48.58 (±0.56)	24.22 (±0.44)
Bonsai	GroupDRO	72.82 (±5.37)	70.23 (±1.33)	33.12 (±1.20)	27.16 (±1.18)	42.67 (±1.09)	22.95 (±0.46)
Bonsai	IRMv1	73.59 (±6.16)	68.39 (±2.01)	32.51 (±1.23)	27.60 (±1.57)	47.11 (±0.63)	23.35 (±0.43)
Bonsai	V-REx	76.39 (±5.32)	68.67 (±1.29)	33.17 (±1.26)	25.81 (±0.42)	48.00 (±0.00)	23.34 (±0.42)
Bonsai	IRMX	64.77 (±10.1)	69.56 (±0.95)	32.63 (±0.75)	27.62 (±0.66)	46.67 (±0.00)	23.34 (±0.40)
FeAT	ERM	77.80 (±2.48)	68.11 (±2.27)	33.13 (±0.78)	28.47 (±0.67)	52.89 (±0.63)	30.66 (±0.42)
FeAT	GroupDRO	80.41 (±3.30)	71.29 (±0.46)	33.55 (±1.67)	28.38 (±1.32)	52.58 (±0.56)	29.99 (±0.11)
FeAT	IRMv1	77.97 (±3.09)	70.33 (±1.14)	34.04 (±0.70)	29.66 (±1.52)	52.89 (±0.63)	29.99 (±0.19)
FeAT	V-REx	75.12 (±6.55)	70.97 (±1.06)	34.00 (±0.71)	29.48 (±1.94)	52.89 (±0.63)	30.57 (±0.53)
FeAT	IRMX	76.91 (±6.76)	71.18 (±1.10)	33.99 (±0.73)	29.04 (±2.96)	52.89 (±0.63)	29.92 (±0.16)
$\dagger$ DFR/DFR-s use an additional OOD dataset to evaluate invariant and spurious feature learning, respectively.
rounds following Zhang et al. [83], while for FeAT the automatic termination stopped at round 2 in COLOREDMNIST-025 and round 3 in COLOREDMNIST-01. For ERM, we pre-trained the model with the same number of overall epochs as FeAT in COLOREDMNIST-01, while early stopping at the number of epochs of 1 round in COLOREDMNIST-025 to prevent over-fitting. All methods adopted the same backbone and the same training protocol following previous works [16, 83]. More details are given in Appendix F.1.
The results are reported in Table 1. It can be found that ERM will learn insufficiently good features under both stronger spurious correlations and invariant correlations, confirming our discussion in Sec. 4.3. Besides, Bonsai learns richer features in COLOREDMNIST-025 and boosts OOD performance, but Bonsai sometimes leads to suboptimal performances in COLOREDMNIST-01, which could be caused by the unstable feature learning in Bonsai. In contrast, FeAT consistently improves the OOD performance of all OOD objectives for all the COLOREDMNIST datasets, demonstrating the advances of direct feature learning control in FeAT than Bonsai and ERM.
Experiments on real-world benchmarks. We also compare FeAT with ERM and Bonsai in 6 real-world OOD generalization datasets curated by Koh et al. [39] that contain complicated features and distribution shifts. The learned features are evaluated with several representative state-of-the-art OOD objectives in WILDS, including GroupDro [64], IRMv1 [4], VREx [41] as well as IRMX [16]. By default, we train ERM, Bonsai and FeAT the same number of steps, and kept the rounds of Bonsai and FeAT the same (though Bonsai still requires one more round for feature synthesis). The only exception is in RxRx1 where both Bonsai and FeAT required more steps than ERM to converge. We use the same evaluation protocol following the practice in the literature [16, 39, 69, 83] to ensure a fair comparison. More details are given in Appendix F.2.
In addition to OOD objectives, we evaluate the learned features with Deep Feature Reweighting (DFR) [38]. DFR uses an additional OOD validation set where the spurious correlation does not hold, to perform logistic regression based on the learned features. Intuitively, DFR can serve as a proper measure for the quality of learned invariant features [33]. When the original dataset does not provide a proper OOD validation set, e.g., CAMELYON17, we use an alternative implementation based on a random split of the training and test data to perform the invariant feature quality measure [63]. Similarly, we also report DFR-s by regression with the environment labels (when available) to evaluate the spurious feature learning of different methods. More details are given in Appendix F.2.2.
The results are presented in Table 2. Similarly, when the tasks grow more challenging and neural architectures become more complicated, the ERM learned features can have a lower quality as discussed Sec. 4.3. For example, ERM can not sufficiently learn all useful features in FMoW, while ERM can learn more spurious correlations in CivilComments. Moreover, it can also be observed the instability of Bonsai in learning richer features that Bonsai even underperforms ERM in rich feature learning and OOD generalization in multiple datasets. In con
(a) ERM 300 epochs
(b) FeAT 2 rounds
Figure 3: GradCAM visualization on COLOREDMNIST-025, where the shortcuts are now concentrated to a colored path at the up left. Three visualizations are drawn for each sample: the original figure, the gray-colored gradcam, and the gradcam. It can be found that ERM can not properly capture the desired features while FeAT can stably capture the desired features.
trast, FeAT consistently achieves the best invariant feature learning performance across various challenging realistic datasets. Meanwhile, compared to ERM and Bonsai, FeAT also reduces over-fitting to the spurious feature learning led by spurious correlations. As a result, FeAT achieves consistent improvements when the learned features are applied to various OOD objectives.
The termination check in FeAT. As elaborated in Sec. 5.2, a key difference between FeAT and previous rich feature learning algorithms such as Bonsai is that FeAT is able to access the intermediate feature representa
Table 3: Performances at different FeAT rounds.
COLOREDMNIST-025	ROUND-1	ROUND-2	ROUND-3
TRAINING ACC.	85.08±0.14	71.87±0.96	84.93±1.26
RETENTION ACC.	-	88.11±4.28	43.82±0.59
OOD ACC.	11.08±0.30	70.64±0.62	10.07±0.26
tions and thus can perform the automatic termination check and learn the desired features stably. To verify, we list the FeAT performances in various subsets of COLOREDMNIST-025 at different rounds in Table 3. By inspecting the retention accuracy, after FeAT learns sufficiently good features at Round 2, it is not necessary to proceed with Round 3 as it will destroy the already learned features and lead to degenerated retention and OOD performance. More details and results are given in Appendix F.1.
Computational analysis. We also analyze the computational and memory overhead of different methods, for which the details are given in Appendix F.4. Compared to ERM and Bonsai, iFeAT achieves the best performance without introducing too much additional overhead.
Feature learning analysis. We visualize the feature learning of ERM and FeAT on ColoredMNIST-025. As shown in Fig. 3, ERM can learn both invariant and spurious features to predict the label, aligned with our theory. However, ERM focuses more on spurious features and even forgets certain features with longer training epochs, which could be due to multiple reasons such as the simplicity biases of ERM. Hence predictions based on ERM learned features fail to generalize to OOD examples. In contrast, FeAT effectively captures the meaningful features for all samples and generalizes to OOD examples well. More analysis including results on WILDS benchmark can be found in Appendix F.5.
In this paper, we conducted a theoretical investigation of the invariant and spurious feature learning of ERM and OOD objectives. We found that ERM learns both invariant and spurious features when OOD objectives rarely learn new features. Thus, the features learned in the ERM pre-training can greatly influence the final OOD performance. Having learned the limitations of ERM pre-training, we proposed FeAT to learn all potentially useful features. Our extensive experimental results verify that FeAT significantly boosts the OOD performance when used for OOD training.
We thank the reviewers for their valuable comments. This work was supported by the RIKEN Incentive Research Project 100847-202301062011, by CUHK direct grant 4055146. BH was supported by the NSFC Young Scientists Fund No. 62006202, NSFC General Program No. 62376235, Guangdong Basic and Applied Basic Research Foundation No. 2022A1515011652, HKBU Faculty Niche Research Areas No. RC-FNRA-IG/22-23/SCI/04, and Tencent AI Lab Rhino-Bird Gift Fund.