SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
A _ {1} ^ {e} = \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} \mathrm {R a d} (\alpha) _ {i}
$$
$$
A _ {2} ^ {e} = \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} y _ {i} ^ {e} \operatorname {R a d} (\alpha) _ {i}
$$
. Similarly, we define
$$
B _ {1} ^ {e} = \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} \operatorname {R a d} (\beta_ {e}) _ {i}
$$
$$
B _ {2} ^ {e} = \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} y _ {i} ^ {e} \mathrm {R a d} (\beta_ {e}) _ {i}.
$$
In the limit of  $n\to \infty$  , we have:
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} A _ {1} ^ {1} (t _ {1}) = - 1 / (1 + e ^ {(\gamma_ {1} + \gamma_ {2})}) (1 - \alpha) (1 - \beta_ {1}) - 1 / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) (1 - \alpha) \beta_ {1} \\ + 1 / \left(1 + e ^ {\gamma_ {2} - \gamma_ {1}}\right) \alpha \left(1 - \beta_ {1}\right) + 1 / \left(1 + e ^ {- \gamma_ {1} - \gamma_ {2}}\right) \alpha \beta_ {1}, \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} A _ {1} ^ {2} (t _ {1}) = - 1 / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) (1 - \alpha) (1 - \beta_ {2}) - 1 / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) (1 - \alpha) \beta_ {2} \\ + 1 / \left(1 + e ^ {- \gamma_ {1} + \gamma_ {2}}\right) \alpha \left(1 - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {1} - \gamma_ {2}}\right) \alpha \beta_ {2}, \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} B _ {1} ^ {1} (t _ {1}) = - 1 / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) (1 - \alpha) (1 - \beta_ {1}) + 1 / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) (1 - \alpha) \beta_ {1} \\ - 1 / \left(1 + e ^ {- \gamma_ {1} + \gamma_ {2}}\right) \alpha \left(1 - \beta_ {1}\right) + 1 / \left(1 + e ^ {- \gamma_ {1} - \gamma_ {2}}\right) \alpha \beta_ {1}, \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} B _ {1} ^ {2} (t _ {1}) = - 1 / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) (1 - \alpha) (1 - \beta_ {2}) + 1 / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) (1 - \alpha) \beta_ {2} \\ - 1 / \left(1 + e ^ {- \gamma_ {1} + \gamma_ {2}}\right) \alpha \left(1 - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {1} - \gamma_ {2}}\right) \alpha \beta_ {2}. \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} A _ {2} ^ {1} (t _ {1}) = e ^ {\gamma_ {1} + \gamma_ {2}} / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {1}) (\gamma_ {1} + \gamma_ {2}) + e ^ {\gamma_ {1} - \gamma_ {2}} / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {1} (\gamma_ {1} - \gamma_ {2}) \\ + e ^ {- \gamma_ {1} + \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} + \gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {1}) (\gamma_ {1} - \gamma_ {2}) + e ^ {- \gamma_ {1} - \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} - \gamma_ {2}}) ^ {2} \alpha \beta_ {1} (\gamma_ {1} + \gamma_ {2}), \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} A _ {2} ^ {2} (t _ {1}) = e ^ {\gamma_ {1} + \gamma_ {2}} / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {2}) (\gamma_ {1} + \gamma_ {2}) + e ^ {\gamma_ {1} - \gamma_ {2}} / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {2} (\gamma_ {1} - \gamma_ {2}) \\ + e ^ {- \gamma_ {1} + \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} + \gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {2}) (\gamma_ {1} - \gamma_ {2}) + e ^ {- \gamma_ {1} - \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} - \gamma_ {2}}) ^ {2} \alpha \beta_ {2} (\gamma_ {1} + \gamma_ {2}), \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \to \infty} B _ {2} ^ {1} (t _ {1}) = e ^ {\gamma_ {1} + \gamma_ {2}} / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {1}) (\gamma_ {1} + \gamma_ {2}) + e ^ {\gamma_ {1} - \gamma_ {2}} / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {1} (- \gamma_ {1} + \gamma_ {2}) \\ + e ^ {- \gamma_ {1} + \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} + \gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {1}) (- \gamma_ {1} + \gamma_ {2}) + e ^ {- \gamma_ {1} - \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} - \gamma_ {2}}) ^ {2} \alpha \beta_ {1} (\gamma_ {1} + \gamma_ {2}), \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} B _ {2} ^ {2} (t _ {1}) = e ^ {\gamma_ {1} + \gamma_ {2}} / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {2}) (\gamma_ {1} + \gamma_ {2}) + e ^ {\gamma_ {1} - \gamma_ {2}} / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {2} (- \gamma_ {1} + \gamma_ {2}) \\ + e ^ {- \gamma_ {1} + \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} + \gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {2}) (- \gamma_ {1} + \gamma_ {2}) + e ^ {- \gamma_ {1} - \gamma_ {2}} / (1 + e ^ {- \gamma_ {1} - \gamma_ {2}}) ^ {2} \alpha \beta_ {2} (\gamma_ {1} + \gamma_ {2}). \\ \end{array}
$$
Because  $\mathrm{Rad}(\alpha)_i$  and  $\mathrm{Rad}(\beta)_i$  are random variables, applying Hoeffding's inequality, we have with probability at least  $1 - \delta$ ,
$$
\left| A _ {1} ^ {1} \left(t _ {1}\right) - \lim  _ {n \rightarrow \infty} A _ {1} ^ {1} \left(t _ {1}\right)\right| \leq \sqrt {\frac {4 \log (1 / \delta)}{n}}.
$$
Similarly, we can apply the concentration bound to other quantities and obtain the same bound.
By the assumption that  $\gamma_{j,r,1}(t_1) = \gamma_{j,r,1}(t_1 - 1)$  and  $\gamma_{j,r,2}(t_1) = \gamma_{j,r,2}(t_1 - 1)$ , we have that  $\sum_{e}A_{1}^{e}(t_{1}) = \sum_{e}B_{1}^{e}(t_{1}) = 0$ :
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} (A _ {1} ^ {1} (t _ {1}) + A _ {1} ^ {2} (t _ {1})) = - 1 / (1 + e ^ {\gamma_ {1} + \gamma_ {2}}) (1 - \alpha) (2 - \beta_ {1} - \beta_ {2}) - 1 / (1 + e ^ {\gamma_ {1} - \gamma_ {2}}) (1 - \alpha) (\beta_ {1} + \beta_ {2}) \\ + 1 / \left(1 + e ^ {- \gamma_ {1} + \gamma_ {2}}\right) \alpha \left(2 - \beta_ {1} - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {1} - \gamma_ {2}}\right) \alpha \left(\beta_ {1} + \beta_ {2}\right) = 0 \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} \left(B _ {1} ^ {1} (t _ {1}) + B _ {1} ^ {2} (t _ {1})\right) = - 1 / \left(1 + e ^ {\gamma_ {1} + \gamma_ {2}}\right) (1 - \alpha) \left(2 - \beta_ {1} - \beta_ {2}\right) + 1 / \left(1 + e ^ {\gamma_ {1} - \gamma_ {2}}\right) (1 - \alpha) \left(\beta_ {1} + \beta_ {2}\right) \\ + 1 / \left(1 + e ^ {- \gamma_ {1} + \gamma_ {2}}\right) \alpha \left(2 - \beta_ {1} - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {1} - \gamma_ {2}}\right) \alpha \left(\beta_ {1} + \beta_ {2}\right) = 0 \\ \end{array}
$$
Solving the above equations, we have,
$$
\gamma_ {1} ^ {\infty} \left(t _ {1}\right) = \frac {1}{2} \log \left(G _ {m} G _ {b}\right) \quad \gamma_ {2} ^ {\infty} \left(t _ {1}\right) = \frac {1}{2} \log \left(G _ {m} / G _ {b}\right)
$$
where we denote  $\gamma_1^\infty(t_1) \triangleq \lim_{n\to \infty}\gamma_1(t_1)$  and  $\gamma_2^\infty(t_1) \triangleq \lim_{n\to \infty}\gamma_1(t_2)$ ,  $G_m = ((1 - A) + \sqrt{(A - 1)^2 + 4A}) / (2A)$  and  $G_b = ((1 - B) + \sqrt{(B - 1)^2 + 4B}) / (2B)$ , with  $A = \alpha(\beta_1 + \beta_2) / ((1 - \alpha)(2 - \beta_1 - \beta_2))$  and  $B = \alpha(2 - \beta_1 - \beta_2) / ((1 - \alpha)*(\beta_1 + \beta_2))$ .
By the convexity of function  $f(x) = e^{x}$ , with a constant  $C$ , we have:
$$
\left| \gamma_ {1} - \gamma_ {1} ^ {\infty} \right| <   \left| e ^ {\gamma_ {1}} - e ^ {\gamma_ {1} ^ {\infty}} \right| \leq C \left| 1 / \left(1 + e ^ {\gamma_ {1}}\right) - 1 / \left(1 + e ^ {\gamma_ {1} ^ {\infty}}\right) \right| \leq \sqrt {\frac {4 \log (1 / \delta)}{n}},
$$
$$
\left| \gamma_ {2} - \gamma_ {2} ^ {\infty} \right| <   \left| e ^ {\gamma_ {2}} - e ^ {\gamma_ {2} ^ {\infty}} \right| \leq C \left| 1 / \left(1 + e ^ {\gamma_ {2}}\right) - 1 / \left(1 + e ^ {\gamma_ {2} ^ {\infty}}\right) \right| \leq \sqrt {\frac {4 \log (1 / \delta)}{n}}.
$$
Then we know that,
$$
C _ {\mathrm {I R M v 1}} ^ {1} = \frac {1}{n _ {1}} \sum_ {i = 1} ^ {n _ {1}} \ell_ {i} ^ {\prime} \hat {y} _ {i} ^ {1} y _ {i} ^ {1} = \gamma_ {1} A _ {1} ^ {1} + \gamma_ {2} B _ {1} ^ {1}
$$
$$
C _ {\mathrm {I R M v 1}} ^ {2} = \frac {1}{n _ {2}} \sum_ {i = 1} ^ {n _ {2}} \ell_ {i} ^ {\prime} \dot {y} _ {i} ^ {2} y _ {i} ^ {2} = \gamma_ {1} A _ {1} ^ {2} + \gamma_ {2} B _ {1} ^ {2}
$$
Therefore, we have that:
$$
C _ {\mathrm {I R M v 1}} ^ {1} + C _ {\mathrm {I R M v 1}} ^ {2} = 0
$$
Then the evolution of coefficients reduces to
$$
\gamma_ {j, r, 1} (t + 1) = \gamma_ {j, r, 1} (t) - \frac {\eta}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} (1 + 2 \lambda C _ {\mathrm {I R M v l}} ^ {e} (t)) A _ {1} ^ {e} (t) - \frac {\eta \lambda}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} 2 C _ {\mathrm {I R M v l}} ^ {e} A _ {2} ^ {e} (t)
$$
$$
\gamma_ {j, r, 2} (t + 1) = \gamma_ {j, r, 2} (t) - \frac {\eta}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} (1 + 2 \lambda C _ {\mathrm {I R M v 1}} ^ {e} (t)) B _ {1} ^ {e} (t) - \frac {\eta \lambda}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} 2 C _ {\mathrm {I R M v 1}} ^ {e} B _ {2} ^ {e} (t)
$$
Taking the solution of  $\gamma_{j,r,1}(t_1),\gamma_{j,r,2}(t_1)$  and value of  $\alpha ,\beta_{1},\beta_{2}$ , we arrive at the conclusion that with a high probability at least  $1 - \delta$  and  $n > C_1\log (1 / \delta)$  with  $C_1$  being a positive constant, we have:
$$
\gamma_ {j, r, 1} (t _ {1} + 1) > \gamma_ {j, r, 1} (t _ {1}),
$$
$$
\gamma_ {j, r, 2} (t _ {1} + 1) <   \gamma_ {j, r, 2} (t _ {1}).
$$
D.5 Proof for Corollary 4.4
Corollary D.10 (Restatement of Corollary 4.4). Consider training the CNN model with the data generated from Def. 3.1, suppose that  $\psi(x) = x$ ,  $\gamma_{j,r,1}(t_1) = o(1)$ , and  $\gamma_{j,r,2}(t_1) = \Theta(1)$  at the end of ERM pre-train  $t_1$  and  $\mathcal{E}_{tr} = \{(0.25, 0.1), (0.25, 0.2)\}$ . 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, 1} (t _ {1} + 1) <   \gamma_ {j, r, 1} (t _ {1}).
$$
Proof of Corollary D.10. Recall that the feature learning update rule:
$$
\begin{array}{l} \gamma_ {j, r, 1} (t + 1) = \gamma_ {j, r, 1} (t) - \frac {\eta}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} \left(1 + 2 \lambda C _ {\mathrm {I R M v 1}} ^ {e} (t)\right) \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \operatorname {R a d} (\alpha) _ {i} \\ - \frac {\eta \lambda}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} 2 C _ {\mathrm {I R M v 1}} ^ {e} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} \cdot y _ {i} ^ {e} \operatorname {R a d} (\alpha) _ {i}, \\ \end{array}
$$
$$
\begin{array}{l} \gamma_ {j, r, 2} (t + 1) = \gamma_ {j, r, 2} (t) - \frac {\eta}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} (1 + 2 \lambda C _ {\mathrm {I R M v l}} ^ {e} (t)) \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \operatorname {R a d} (\beta_ {e}) \\ - \frac {\eta \lambda}{m} \cdot \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} 2 C _ {\mathrm {I R M v 1}} ^ {e} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} \cdot y _ {i} ^ {e} \operatorname {R a d} \left(\beta_ {e}\right) _ {i}, \\ \end{array}
$$
Taking the value of  $\gamma_{j,r,1}(t_1),\gamma_{j,r,2}(t_1)$  and, we can conclude that:
$$
\begin{array}{l} \lim _ {n \to \infty} A _ {1} ^ {1} (t _ {1}) = - 1 / (1 + e ^ {\gamma_ {2}}) (1 - \alpha) (1 - \beta_ {1}) - 1 / (1 + e ^ {- \gamma_ {2}}) (1 - \alpha) \beta_ {1} + \\ 1 / \left(1 + e ^ {\gamma_ {2}}\right) \alpha \left(1 - \beta_ {1}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \alpha \beta_ {1} \\ = 1 / \left(1 + e ^ {\gamma_ {2}}\right) \left(2 \alpha - 1\right) \left(1 - \beta_ {1}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \left(2 \alpha - 1\right) \left(\beta_ {1}\right) \\ = (2 \alpha - 1) \left[ 1 / \left(1 + e ^ {\gamma_ {2}}\right) \left(1 - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \beta_ {1} \right] \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} A _ {1} ^ {2} (t _ {1}) = 1 / (1 + e ^ {\gamma_ {2}}) (2 \alpha - 1) (1 - \beta_ {2}) + 1 / (1 + e ^ {- \gamma_ {2}}) (2 \alpha - 1) (\beta_ {2}) \\ = (2 \alpha - 1) \left[ 1 / \left(1 + e ^ {\gamma_ {2}}\right) \left(1 - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \beta_ {2} \right] \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \to \infty} B _ {1} ^ {1} (t _ {1}) = - 1 / (1 + e ^ {\gamma_ {2}}) (1 - \alpha) (1 - \beta_ {1}) + 1 / (1 + e ^ {- \gamma_ {2}}) (1 - \alpha) \beta_ {1} - \\ 1 / \left(1 + e ^ {\gamma_ {2}}\right) \alpha \left(1 - \beta_ {1}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \alpha \beta_ {1} \\ = - 1 / \left(1 + e ^ {\gamma_ {2}}\right) \left(1 - \beta_ {1}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \beta_ {1} \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} B _ {1} ^ {2} (t _ {1}) = - 1 / (1 + e ^ {\gamma_ {2}}) (1 - \alpha) (1 - \beta_ {2}) + 1 / (1 + e ^ {- \gamma_ {2}}) (1 - \alpha) \beta_ {2} - \\ 1 / \left(1 + e ^ {\gamma_ {2}}\right) \alpha \left(1 - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \alpha \beta_ {2} \\ = - 1 / \left(1 + e ^ {\gamma_ {2}}\right) \left(1 - \beta_ {2}\right) + 1 / \left(1 + e ^ {- \gamma_ {2}}\right) \beta_ {2} \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} A _ {2} ^ {1} (t _ {1}) = e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {1}) (\gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {1} (- \gamma_ {2}) \\ + e ^ {+ \gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {1}) (- \gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} \alpha \beta_ {1} (\gamma_ {2}) \\ = e ^ {\gamma_ {2}} / \left(1 + e ^ {\gamma_ {2}}\right) ^ {2} (1 - 2 \alpha) \left(1 - \beta_ {1}\right) + e ^ {- \gamma_ {2}} / \left(1 + e ^ {- \gamma_ {2}}\right) ^ {2} (2 \alpha - 1) \beta_ {1} \gamma_ {2} \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} A _ {2} ^ {2} (t _ {1}) = e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {2}) (\gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {2} (- \gamma_ {2}) \\ + e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {2}) (- \gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} \alpha \beta_ {2} (\gamma_ {2}) \\ = e ^ {\gamma_ {2}} / \left(1 + e ^ {\gamma_ {2}}\right) ^ {2} (1 - 2 \alpha) \left(1 - \beta_ {2}\right) + e ^ {- \gamma_ {2}} / \left(1 + e ^ {- \gamma_ {2}}\right) ^ {2} (2 \alpha - 1) \beta_ {2} \gamma_ {2} \\ \end{array}
$$
$$
\begin{array}{l} \lim  _ {n \rightarrow \infty} B _ {2} ^ {1} (t _ {1}) = e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {1}) (\gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {1} (\gamma_ {2}) \\ + e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {1}) (\gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} \alpha \beta_ {1} (\gamma_ {2}), \\ \end{array}
$$
$$
\begin{array}{l} \lim _ {n \rightarrow \infty} B _ {2} ^ {2} (t _ {1}) = e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} (1 - \alpha) (1 - \beta_ {2}) (\gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} (1 - \alpha) \beta_ {2} (\gamma_ {2}) \\ + e ^ {\gamma_ {2}} / (1 + e ^ {\gamma_ {2}}) ^ {2} \alpha (1 - \beta_ {2}) (\gamma_ {2}) + e ^ {- \gamma_ {2}} / (1 + e ^ {- \gamma_ {2}}) ^ {2} \alpha \beta_ {2} (\gamma_ {2}). \\ \end{array}
$$
Finally, taking the value of environment of  $(\alpha, \beta_1, \beta_2) = (0.25, 0.1, 0.2)$ , we conclude that with a high a probability at least  $1 - \delta$  and  $n > C_1 \log(1 / \delta)$  with  $C_1$  being a positive constant, we have:
$$
\gamma_ {j, r, 1} \left(t _ {1} + 1\right) <   \gamma_ {j, r, 1} \left(t _ {1}\right).
$$
E More Details about iFeAT
As mentioned in Sec. 5.2 that, when the featurizer is implemented as a deep net that have a massive amount of parameters, backpropagating through Algorithm 1 can allocate too much memory for propagating with  $2K - 1$  batches of data. It is common for many realistic benchmarks such as Camelyon17 and FMoW in wilds benchmark [39] that adopts a DenseNet [31] with 121 layers as the featurizer. To relieve the exceeding computational and memory overhead, we propose a lightweight version of FeAT, denoted as FeAT. Instead of storing all of historical subsets and classifiers, iFeAT iteratively use the augmentation and retention sets and historical classifier from only the last round. In contrast, previous rich feature learning algorithm [59, 83] incurs a high computational and memory overhead as the round grows. For example, in RxRx1, we have to reduce the batch size of Bonsai to allow the proceeding of rounds  $\geq 3$ .
We elaborate the detailed algorithmic description of iFeAT in Algorithm 2.
Algorithm 2 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  $e$  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,\ldots ,e]$  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: if  $k > 1$  then   
14: // Hence it doesnot need to maintain all historical classifiers.   
15: Update  $w_{k}\leftarrow (w_{k - 1},w_{k})$  .   
16: end if   
17: Split  $\mathcal{D}_{\mathrm{tr}}$  into groups  $\mathcal{D}_k^a,\mathcal{D}_k^r$  according to  $f_{k}$  .   
18: // Hence it doesnot need to maintain all historical subsets.   
19: Update groups  $G^{a}\gets \{\mathcal{D}_{k}^{a}\} ,G^{r}\gets \{\mathcal{D}_{k}^{r}\}$  .   
20: end for   
21: return  $f = w\circ \varphi$
F More Details about the Experiments
In this section, we provide more details and the implementation, evaluation and hyperparameter setups in complementary to the experiments in Sec. 6.
F.1 More details about COLOREDMNIST experiments
Datasets. In the controlled experiments with COLOREDMNIST, we follow the evaluation settings as previous works [4, 16, 83]. In addition to 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, we also incorporate the modified one (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, since both cases can happen at real world.
Architecture and optimization. To ensure a fair comparison, we use 4-Layer MLP with a hidden dimension of 256 as the backbone model for all methods, where we take the first 3 layers as the featurizer and the last layer as the classifier, following the common practice [25, 39]. For the optimization of the models, we use the Adam [37] optimizer with a learning rate of  $1e - 3$  and a
weight decay of  $1e - 3$ . We report the mean and standard deviation of the performances of different methods with each configuration of hyperparameters 10 times with the random seeds from 1 to 10.
Implementation of ERM-NF and OOD objectives. For the common pre-training protocol with ERM, our implementation follows the previous works [83]. Specifically, we first train the model with  $\{0,50,100,150,200,250\}$  epochs and then apply the OOD regularization of various objectives with a penalty weight of  $\{1e1,1e2,1e3,1e4,1e5\}$ . We adopt the implementations from Zhang et al. [83] for various OOD objectives, including IRMv1 [4], VREx [41], IB-IRM [1], CLOvE [76], IGA [40] and Fishr [57]. Besides, we also incorporate the state-of-the-art OOD objective proposed by Chen et al. [16] that is able to resolve both COLOREDMNIST-025 and COLOREDMNIST-01.
Evaluation of feature learning methods. For the sake of fairness in comparison, by default, we train all feature learning methods by the same number of epochs and rounds (if applicable). For the implementation Bonsai, we strictly follow the recommended setups provided by Zhang et al. [83], where we train the model with Bonsai by 2 rounds with 50 epochs for round 1, 500 epochs for round 2, and 500 epochs for the synthesize round in COLOREDMNIST-025. While in COLOREDMNIST-01, round 1 contains 150 epochs, round 2 contains 400 epochs and the synthesize round contains 500 epochs. For the implementation of FeAT, we train the model with 2 rounds of FeAT in COLOREDMNIST-025, and 3 rounds of FeAT in COLOREDMNIST-01, where each round contains 150 epochs. While for the retain penalty, we find using a fixed number of 0.01 already achieved sufficiently good performance. ERM only contains 1 round, for which we train the model with 150 epochs in COLOREDMNIST-025 as we empirically find more epochs will incur severe performance degeneration in COLOREDMNIST-025. While in COLOREDMNIST-01, we train the model with ERM by 500 epochs to match up the overall training epochs of FeAT and Bonsai. We provide a detailed distribution of the number of epochs in each round in Table 5. It can be found
Table 5: Number of epochs in each round of various feature learning algorithms.
CMNIST-025	ROUND-1	ROUND-2	ROUND-3	SYN. ROUND
ERM	150	-	-	-
BONSAI	50	150	-	500
FEAT	150	150	-	-
CMNIST-01	ROUND-1	ROUND-2	ROUND-3	SYN. ROUND
ERM	500	-	-	-
BONSAI	150	400	-	500
FEAT	150	150	150	-
that, although Bonsai costs 2 - 3 times of training epochs more than ERM and FeAT, Bonsai does not necessarily find better feature representations for OOD training, as demonstrated in Table. 1. In contrast, FeAT significantly and consistently learns richer features given both COLOREDMNIST-025 and COLOREDMNIST-01 than ERM, which shows the superiority of FeAT.
The termination check in FeAT. A key difference between FeAT and previous rich feature learning algorithms is that FeAT is able to perform the automatic termination check and learn the desired features stably. As elaborated in Sec. 5.2, FeAT can terminate automatically by inspecting the retention accuracy. To verify, we list the FeAT performances in various subsets of COLOREDMNIST-025 and COLOREDMNIST-01 at different rounds. We use a termination accuracy of  $130\%$ , which trades off the exploration (i.e., training accuracy as  $80\%$ ) and the retention (i.e., retention accuracy as  $50\%$ ) properly. As shown in Table 6, in COLOREDMNIST-025 (COLOREDMNIST-01), after FeAT learns sufficiently good features at Round 2 (3), respectively, it is not necessary to proceed with Round 3 (4) as it will destroy the already learned features and lead to degenerated retention performance (i.e., the sum of training and retention accuracies is worse than  $130\%$ .
<sup>5</sup>https://github.com/TjuJianyu/RFC
Table 6: Performances in various sets 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
COLOREDMNIST-01	ROUND-1	ROUND-2	ROUND-3	ROUND-4
TRAINING ACC.	88.63±0.15	74.25±1.23	86.07±0.36	77.29±0.24
RETENTION ACC.	-	85.91±1.78	48.05±1.39	29.09±1.15
OOD ACC.	73.50±0.41	17.32±2.69	85.40±0.54	12.48±2.85
F.2 More details about WILDS experiments
In this section, we provide more details about the WILDS datasets used in the experiments as well as the evaluation setups.
F.2.1 Dataset description.
To evaluate the feature learning performance given data from realistic scenarios, we select 6 challenging datasets from WILDS [39] benchmark. The datasets contain various realistic distribution shifts, ranging from domain distribution shifts, subpopulation shifts and the their mixed. A summary of the basic information and statistics of the selected WILDS datasets can be found in Table. 7, Table. 8, respectively. In the following, we will give a brief introduction to each of the datasets. More details can be found in the WILDS paper [39].