$\displaystyle=\sum_{j=\pm 1}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{d{\mathbf{w}}_{j,r}(t)}{dt}\right\rangle$
$\displaystyle=\sum_{e^{\prime}}2\lambda
C^{e^{\prime}}_{\text{IRMv1}}({\mathbf{W}},t)\sum_{j}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{\partial
C^{e^{\prime}}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\right\rangle$
$\displaystyle\quad+\sum_{j=\pm 1}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{\partial
L({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\right\rangle$
$\displaystyle=2\lambda\sum_{e^{\prime}}C^{e^{\prime}}_{\text{IRMv1}}({\mathbf{W}},t)\cdot{\mathbf{H}}_{e,e^{\prime}}(t)+\mathbf{g}_{e}(t),$

where we define
${\mathbf{H}}_{e,e^{\prime}}(t)=\sum_{j}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{\partial
C^{e^{\prime}}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\right\rangle$
and ${\mathbf{g}}_{e}(t)=\sum_{j=\pm
1}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{\partial
L({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\right\rangle$. Thus
${\mathbf{H}}(t)$ is an
$|{{\mathcal{E}}_{\text{tr}}}|\times|{{\mathcal{E}}_{\text{tr}}}|$ matrix. We
can write the dynamics of
${\mathbf{c}}(t)=\left[C^{1}_{\text{IRMv1}}({\mathbf{W}},t),C^{2}_{\text{IRMv1}}({\mathbf{W}},t),\cdots,C^{|{{\mathcal{E}}_{\text{tr}}}|}_{\text{IRMv1}}({\mathbf{W}},t)\right]$
in a compact way:

$\displaystyle\frac{d{\mathbf{c}}(t)}{dt}=2\lambda\cdot{\mathbf{H}}(t){\mathbf{c}}(t)+{\mathbf{g}}(t).$
(22)

Our next step is to show ${\mathbf{H}}(t)$ is stable during training. To
proceed with the analysis, we write down the expression for $\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\in\mathbb{R}^{d}$:

$\displaystyle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}}(t))}{\partial{\mathbf{w}}_{j,r}(t)}$
$\displaystyle=\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
j{\mathbf{v}}_{i}^{e}+\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
jy_{i}^{e}\bm{\xi}_{i}$
$\displaystyle\quad+\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
jy_{i}^{e}{\mathbf{v}}_{i}^{e}+\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
j\bm{\xi}_{i}.$

When we consider non-linear activation function ${\psi}(x)$, the entry of
matrix ${\mathbf{H}}(t)$ can be computed as follows:

$\displaystyle{\mathbf{H}}_{e,e^{\prime}}(t)=\sum_{j}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{\partial
C^{e^{\prime}}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\right\rangle$
$\displaystyle=\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\bigg{[}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi^{\prime}}j{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\ell^{\prime}_{i^{\prime}}(t){\psi^{\prime}}j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}+\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)jy_{i}^{e}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)jy_{i^{\prime}}^{e^{\prime}}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\bigg{]}$
$\displaystyle+\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\bigg{[}\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)jy_{i}^{e}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime}_{i^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}+\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime}_{i}(t)j{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\bigg{]}$
$\displaystyle+\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\bigg{[}\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime}_{i}(t)jy^{e}_{i}\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime}_{i^{\prime}}(t)jy^{e^{\prime}}_{i^{\prime}}\bm{\xi}_{i^{\prime}}^{e^{\prime}}+\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)j\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)j\bm{\xi}_{i^{\prime}}^{e^{\prime}}\bigg{]}$
$\displaystyle+\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\bigg{[}\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}^{e}_{i}j\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime}_{i^{\prime}}(t)jy^{e^{\prime}}_{i^{\prime}}\bm{\xi}_{i^{\prime}}^{e^{\prime}}+\sum_{i=1}^{n_{e}}{\psi^{\prime}}y^{e}_{i}\ell^{\prime}_{i}(t)j\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)j\bm{\xi}_{i^{\prime}}^{e}\hat{y}^{e^{\prime}}_{i^{\prime}}(t)\bigg{]}$
$\displaystyle\triangleq{\mathbf{H}}^{1}_{e,e^{\prime}}(t)+{\mathbf{H}}^{2}_{e,e^{\prime}}(t)+{\mathbf{H}}^{3}_{e,e^{\prime}}(t)+{\mathbf{H}}^{4}_{e,e^{\prime}}(t)+{\mathbf{H}}^{5}_{e,e^{\prime}}(t)+{\mathbf{H}}^{6}_{e,e^{\prime}}(t)+{\mathbf{H}}^{7}_{e,e^{\prime}}(t)+{\mathbf{H}}^{8}_{e,e^{\prime}}(t).$

The matrix $\mathbf{H}$ is composed of eight elements. In addition, we define

$\displaystyle{\mathbf{H}}^{1,\infty}_{e,e^{\prime}}$
$\displaystyle=\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\left[\sum_{i=1}^{n_{e}}-\frac{1}{2}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(0),\mathbf{v}^{e}_{i}\rangle)j{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}-\frac{1}{2}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(0),\mathbf{v}^{e^{\prime}}_{i^{\prime}}\rangle)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right]$
$\displaystyle=\frac{1}{2mn_{e}n_{e^{\prime}}}\sum_{i=1}^{n_{e}}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(0),\mathbf{v}^{e}_{i}\rangle){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(0),\mathbf{v}^{e^{\prime}}_{i^{\prime}}\rangle){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}.$

Then we can show that:

$\displaystyle\left|{\mathbf{H}}^{1}_{e,e^{\prime}}(t)-{\mathbf{H}}^{1,\infty}_{e,e^{\prime}}\right|$
$\displaystyle=$
$\displaystyle\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(t)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(t)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}\frac{1}{2}{\psi}^{\prime}(0){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\frac{1}{2}{\psi}^{\prime}(0){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle=$
$\displaystyle\frac{2}{mn_{e}n_{e^{\prime}}}\bigg{|}\sum_{i=1}^{n_{e}}{\psi}^{\prime}(t)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(t)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}$
$\displaystyle\quad+\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}\frac{1}{2}{\psi}^{\prime}(0){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\frac{1}{2}{\psi}^{\prime}(0){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\bigg{|}$
$\displaystyle\leq$
$\displaystyle\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(t)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(t)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(t)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\quad+\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(t)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\quad+\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\quad+\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}-\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\frac{1}{2}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq$
$\displaystyle\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}({\psi}^{\prime}(t)-{\psi}^{\prime}(0))\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\ell^{\prime}_{i^{\prime}}(t){\psi}^{\prime}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|+\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(t)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}({\psi}^{\prime}(t)-{\psi}^{\prime}(0))\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle+\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\left(\ell^{\prime}_{i^{\prime}}(t)+\frac{1}{2}\right){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|+\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\left(\ell^{\prime}_{i}(t)+\frac{1}{2}\right){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\triangleq$ $\displaystyle I_{1}+I_{2}+I_{3}+I_{4}$

where we calculate each item as follows:

$\displaystyle I_{1}$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}({\psi}^{\prime}(t)-{\psi}^{\prime}(0))\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\ell^{\prime}_{i^{\prime}}(t){\psi}^{\prime}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\overset{(a)}{\leq}\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\beta\|\mathbf{w}_{j,r}(t)-\mathbf{w}_{j,r}(0)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\ell^{\prime}_{i^{\prime}}(t){\psi}^{\prime}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\overset{(b)}{\leq}\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)-\mathbf{w}_{j,r}(0)\|_{2}\|{\mathbf{v}}^{e}_{i}\|^{2}_{2}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\|^{2}_{2}$
$\displaystyle\overset{(c)}{\leq}\frac{32\beta^{2}R(R+\frac{3}{2}\sigma_{0}d)}{m},$

where we have used $R\triangleq\|\mathbf{w}_{j,r}(t)\|_{2}$. Besides,
inequality (a) results from applying the smoothness property of the activation
function and Cauchy–Schwarz inequality; inequality (b) is by smoothness
property of the activation function and Cauchy–Schwarz inequality. Besides, we
have used $|\ell^{e}_{i}|\leq 1$ for all $i\in n_{e}$ and
$e\in{{\mathcal{E}}_{\text{all}}}$; inequality (c) is by the fact that
$\|\mathbf{v}^{e}_{i}\|_{2}\leq 2$ for all $i\in n_{e}$ and
$e\in{{\mathcal{E}}_{\text{all}}}$ and Lemma D.8.

Similarly, we calculate the upper bound for $I_{2}$ as follows:

$\displaystyle I_{2}$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(t)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}({\psi}^{\prime}(t)-{\psi}^{\prime}(0))\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{32\beta^{2}R(R+\frac{3}{2}\sigma_{0}d)}{m}.$

Next, we give the upper bound of $I_{3}$:

$\displaystyle I_{3}$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\left(\ell^{\prime}_{i^{\prime}}(t)+\frac{1}{2}\right){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\beta\|\mathbf{w}_{j,r}(0)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\beta\|\mathbf{w}_{j,r}(0)\|_{2}\|{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\|_{2}\left(\ell^{\prime}_{i^{\prime}}(t)+\frac{1}{2}\right){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{64\beta^{2}LR(\frac{3}{2}\sigma_{0}d)^{2}}{m},$

where we have used $\gamma$ which is defined as follows:

$\displaystyle|\hat{y}^{e}_{i}(t)|$
$\displaystyle=\left|\frac{1}{m}\sum_{j}\sum_{r=1}^{m}\left[{\psi}({\mathbf{w}}_{j,r}^{\top}(t){\mathbf{x}}_{1})+{\psi}({\mathbf{w}}_{j,r}^{\top}(t){\mathbf{x}}_{2})\right]\right|$
$\displaystyle\overset{(a)}{\leq}2LR,$

where inequality (a) is by the Lipschitz property of non-linear activation
function and we have used the bound for $\ell^{\prime}_{i}(t)+\frac{1}{2}$:

$\displaystyle\left|\ell^{\prime}_{i}(t)+\frac{1}{2}\right|$
$\displaystyle=\left|-\frac{\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t))}{1+\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t))}+\frac{1}{2}\right|$
$\displaystyle=\left|\frac{1}{2}-\frac{1}{1+\exp(y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t))}\right|$
$\displaystyle\leq\max\left\\{\left|\frac{1}{2}-\frac{1}{1+\exp(2LR)}\right|,\left|\frac{1}{2}-\frac{1}{1+\exp(-2LR)}\right|\right\\}$
$\displaystyle\leq\max\left\\{\left|\frac{1}{2}-\frac{1}{2+\frac{7}{4}2LR}\right|,\left|\frac{1}{2}-\frac{1}{2-2LR}\right|\right\\}=\Theta(LR).$

and we provide the bound of $\ell^{\prime\prime}_{i}(t)-\frac{1}{4}$:

$\displaystyle\left|\ell^{\prime\prime}_{i}(t)-\frac{1}{4}\right|$
$\displaystyle=\left|\frac{\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t))}{(1+\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t)))^{2}}-\frac{1}{4}\right|$
$\displaystyle=\left|\frac{1}{\exp(y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t))+2+\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i},t))}-\frac{1}{4}\right|$
$\displaystyle\leq\left|\frac{1}{4}-\frac{1}{2+2\exp((2LR)^{2}/2)}\right|=\Theta((LR)^{2}).$

Similarly, we give the upper bound of $I_{4}$:

$\displaystyle I_{4}$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}{\psi}^{\prime}(0)\left(\ell^{\prime}_{i}(t)+\frac{1}{2}\right){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi}^{\prime}(0)\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\beta\|\mathbf{w}_{j,r}(0)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}(\ell^{\prime}_{i}(t)+\frac{1}{2}){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\beta\|\mathbf{w}_{j,r}(0)\|_{2}\|{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\|_{2}\frac{1}{2}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{64\beta^{2}LR\gamma(\frac{3}{2}\sigma_{0}d)^{2}}{m}.$

Together, we obtain the upper bound for
$\left|{\mathbf{H}}^{1}_{e,e^{\prime}}(t)-{\mathbf{H}}^{1,\infty}_{e,e^{\prime}}\right|$:

$\displaystyle\left|{\mathbf{H}}^{1}_{e,e^{\prime}}(t)-{\mathbf{H}}^{1,\infty}_{e,e^{\prime}}\right|\leq\frac{64\beta^{2}R(R+\frac{3}{2}\sigma_{0}d)}{m}+\frac{128\beta^{2}LR(\frac{3}{2}\sigma_{0}d)^{2}}{m}.$

Then we calculate the upper bound for the residual terms:

$\displaystyle\left|{\mathbf{H}}^{2}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)jy_{i}^{e}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)jy_{i^{\prime}}^{e^{\prime}}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\psi^{\prime}(t)\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)jy_{i}^{e}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\psi^{\prime}(t)\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)jy_{i^{\prime}}^{e^{\prime}}{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\overset{(a)}{\leq}\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\overset{(b)}{\leq}\frac{128\beta^{2}L^{2}R^{4}}{m},$

where inequality (a) is by the smoothness property of the activation function
and Cauchy–Schwarz inequality, and inequality (b) is by triangle inequality
and the fact that $\|\mathbf{v}^{e}_{i}\|_{2}\leq 2$ for all $i\in n_{e}$ and
$e\in{{\mathcal{E}}_{\text{all}}}$, and $|\ell^{\prime\prime}_{i}|\leq 1$ for
all $i\in[n]$ and Lemma D.8. Similarly, we further provide the upper bound of
residual terms:

$\displaystyle\left|{\mathbf{H}}^{3}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)jy_{i}^{e}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime}_{i^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\psi^{\prime}(t)\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)jy_{i}^{e}{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\psi^{\prime}(t)\ell^{\prime}_{i^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{64\beta^{2}LR^{3}}{m}.$

Similarly, we further have that:

$\displaystyle\left|{\mathbf{H}}^{4}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime}_{i}(t)j{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle=\frac{2}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\psi^{\prime}(t)\ell^{\prime}_{i}(t)j{\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\psi^{\prime}(t)\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)j{\mathbf{v}}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{64\beta^{2}LR^{3}}{m}.$

Keep going on, we provide the computation results further:

$\displaystyle\left|{\mathbf{H}}^{5}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime}_{i}(t)jy^{e}_{i}\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime}_{i^{\prime}}(t)jy^{e^{\prime}}_{i^{\prime}}\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\overset{(a)}{\leq}\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}_{i}\|_{2}\ell^{\prime}_{i}(t)\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime}_{i^{\prime}}(t)\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\overset{(b)}{\leq}\frac{2\beta^{2}R^{2}\sigma^{2}_{p}d}{m},$

where inequality (a) is by smoothness property of non-linear activation
function and Cauchy inequality, inequality (b) is by Lemma D.6 and Lemma D.8.
Next, we calculate the

$\displaystyle\left|{\mathbf{H}}^{6}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}_{i}^{e}(t)j\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)j\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}_{i}\|_{2}\ell^{\prime\prime}_{i}(t)\hat{y}^{e}_{i}\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}_{i^{\prime}}^{e^{\prime}}(t)\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{8\beta^{2}dL^{4}R^{2}}{m}.$

Similarly, the next $\mathbf{H}$ term can be calculated as follows:

$\displaystyle\left|{\mathbf{H}}^{7}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}\ell^{\prime\prime}_{i}(t)\hat{y}^{e}_{i}j\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime}_{i^{\prime}}(t)jy^{e^{\prime}}_{i^{\prime}}\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}_{i}\|_{2}\ell^{\prime\prime}_{i}(t)\hat{y}^{e}_{i}(t)\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime}_{i^{\prime}}(t)\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{4\beta^{2}\sigma^{2}_{p}dLR^{3}}{m}.$

Finally, we have the upper for the last term:

$\displaystyle\left|{\mathbf{H}}^{8}_{e,e^{\prime}}(t)\right|$
$\displaystyle=\left|\sum_{j}\sum_{r=1}^{m}\left(\frac{1}{n_{e}m}\right)\left(\frac{1}{n_{e^{\prime}}m}\right)\sum_{i=1}^{n_{e}}{\psi^{\prime}}y^{e}_{i}\ell^{\prime}_{i}(t)j\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}{\psi^{\prime}}\ell^{\prime\prime}_{i^{\prime}}(t)j\bm{\xi}_{i^{\prime}}^{e}\hat{y}^{e^{\prime}}_{i^{\prime}}(t)\right|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e^{\prime}}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}_{i}\|_{2}\ell^{\prime}_{i}(t)\bm{\xi}^{e\top}_{i}\sum_{i^{\prime}=1}^{n_{e^{\prime}}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}^{e^{\prime}}_{i^{\prime}}\|_{2}\ell^{\prime\prime}_{i^{\prime}}(t)\hat{y}^{e^{\prime}}_{i^{\prime}}(t)\bm{\xi}_{i^{\prime}}^{e^{\prime}}\right|$
$\displaystyle\leq\frac{4\beta^{2}\sigma^{2}_{p}dLR^{3}}{m}.$

To summarize, we have that,

$\displaystyle\left|{\mathbf{H}}_{e,e^{\prime}}(t)-{\mathbf{H}}^{\infty}_{e,e^{\prime}}\right|$
$\displaystyle\leq\frac{32\beta^{2}R(R+\frac{3}{2}\sigma_{0}d)}{m}+\frac{128\beta^{2}(R+\frac{3}{2}\sigma^{2}_{0}d)^{2}L^{2}R^{2}}{m}+\frac{128\beta^{2}LR^{3}}{m}$
$\displaystyle\quad+\frac{2\beta^{2}R^{2}\sigma^{2}_{p}d}{m}+\frac{8\beta^{2}\sigma^{2}_{p}dL^{2}R^{4}}{m}+\frac{4\beta^{2}\sigma^{2}_{p}dLR^{3}}{m}$
$\displaystyle\leq O\left(\frac{\beta^{2}LR}{m}\right).$

where we have used $\sigma_{p}=O(d^{-2})$, $R=o(1)$, and
$\sigma_{0}=O(\sqrt{R/d})$. Furthermore, we show that the perturbation term in
Equation (22) is bounded during training. In particular, we show the complete
expression:

$\displaystyle{\mathbf{g}}_{e}(t)$ $\displaystyle=\sum_{j=\pm
1}\sum_{r=1}^{m}\left\langle\frac{\partial
C^{e}_{\text{IRMv1}}({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)},\frac{\partial
L({\mathbf{W}},t)}{\partial{\mathbf{w}}_{j,r}(t)}\right\rangle$
$\displaystyle=\sum_{j=\pm
1}\sum_{r=1}^{m}\bigg{[}\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
j{\mathbf{v}}_{i}^{e}\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
j{\mathbf{v}}_{i}^{e}$
$\displaystyle\qquad+\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
jy_{i}^{e}{\mathbf{v}}_{i}^{e}\frac{1}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),y_{i}^{e}{\mathbf{v}}_{i}^{e}\rangle)\cdot
j{\mathbf{v}}_{i}^{e}$
$\displaystyle\qquad+\frac{\eta}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
jy_{i}^{e}\bm{\xi}_{i}\frac{\eta}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
jy_{i}^{e}\bm{\xi}_{i}$
$\displaystyle\qquad+\frac{\eta}{n_{e}m}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t){\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
jy_{i}^{e}\bm{\xi}_{i}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}{\psi}^{\prime}(\langle{\mathbf{w}}_{j,r}(t),\bm{\xi}_{i}\rangle)\cdot
j\bm{\xi}_{i}\bigg{]}$ $\displaystyle\triangleq I_{1}+I_{2}+I_{3}+I_{4}.$

Similar to the computation process for matrix $\mathbf{H}$, we adopt a divide
and conquer manner:

$\displaystyle|I_{1}|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i}^{e}\right|$
$\displaystyle\leq\frac{32\beta^{2}(R+\frac{3}{2}\sigma^{2}_{0}d)^{2}}{m}.$

The techniques used are the same when deriving upper bound for matrix
$\mathbf{H}$. Next, we have

$\displaystyle|I_{2}|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime\prime}_{i}(t)\hat{y}^{e}_{i}(t){\mathbf{v}}^{e\top}_{i}\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|{\mathbf{v}}^{e}_{i}\|_{2}\ell^{\prime}_{i^{\prime}}(t){\mathbf{v}}_{i}^{e}\right|$
$\displaystyle\leq\frac{64\beta^{2}R^{2}LR}{m}.$

The last second term can be calculated as follows:

$\displaystyle|I_{3}|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}_{i}\|_{2}\ell^{\prime}_{i}(t)\bm{\xi}^{e\top}_{i}\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}^{e}_{i}\|_{2}\ell^{\prime}_{i}(t)\bm{\xi}_{i}^{e}\right|$
$\displaystyle\leq\frac{2\beta^{2}R^{2}\sigma^{2}_{p}d}{m},$

Finally, we show the upper bound of last term:

$\displaystyle|I_{4}|$
$\displaystyle\leq\frac{2\beta^{2}}{mn_{e}n_{e}}\left|\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}_{i}\|_{2}\ell^{\prime\prime}_{i}(t)\hat{y}^{e}_{i}(t)\bm{\xi}^{e\top}_{i}\sum_{i=1}^{n_{e}}\|\mathbf{w}_{j,r}(t)\|_{2}\|\bm{\xi}^{e}_{i}\|_{2}\ell^{\prime}_{i^{\prime}}(t)\bm{\xi}_{i}^{e}\right|$
$\displaystyle\leq\frac{4\beta^{2}\sigma^{2}_{p}dLR^{3}}{m}.$

In a summary, we have the following inequality:

$\displaystyle\left|\mathbf{g}_{e}(t)\right|$
$\displaystyle\leq\frac{32\beta^{2}R^{2}}{m}+\frac{64\beta^{2}LR^{3}}{m}+\frac{2\beta^{2}R^{2}\sigma^{2}_{p}d}{m}+\frac{4\beta^{2}R^{3}\sigma^{2}_{p}dL}{m}$
$\displaystyle\leq O\left(\frac{\beta^{2}LR^{2}}{m}\right),$

where we have used $\sigma_{p}=O(d^{-2})$ and $R=o(1)$. With all the bounds at
hand, we are ready to have the dynamics for $\|{\mathbf{c}}(t)\|^{2}_{2}$

$\displaystyle\frac{d\|{\mathbf{c}}(t)\|^{2}_{2}}{dt}$
$\displaystyle=-2\lambda\mathbf{c}^{\top}(t){\mathbf{H}}(t)\mathbf{c}(t)-\mathbf{c}(t)\mathbf{g}(t)\leq-\lambda_{0}\lambda\|\mathbf{c}(t)\|^{2}_{2},$
(23)

which requires that
$\|{\mathbf{H}}(t)-{\mathbf{H}}^{\infty}\|_{2}\leq\lambda_{0}$. This leads to
the following inequality:

$\displaystyle\|{\mathbf{H}}(t)-{\mathbf{H}}^{\infty}\|_{2}$
$\displaystyle\leq\|{\mathbf{H}}(t)-{\mathbf{H}}^{\infty}\|_{F}\leq\sum_{i,j}|{\mathbf{H}}_{ij}(t)-{\mathbf{H}}_{ij}^{\infty}|$
$\displaystyle\leq\frac{|\mathcal{E}_{tr}|^{2}\beta^{2}LR}{m}\leq\lambda_{0}.$

which leads to the conclusion for $R$ as follows:

$\displaystyle R\leq\frac{\lambda_{0}m}{|\mathcal{E}_{tr}|^{2}\beta^{2}L}.$
(24)

Besides, we have the inequality that

$\displaystyle\|\mathbf{g}\|_{2}\leq\frac{\sqrt{|\mathcal{E}_{tr}|}\beta^{2}LR}{m}\leq\lambda\lambda_{0}\|\mathbf{c}(0)\|_{2}.$
(25)

Combined with Equation (24), we obtain the condition for $\lambda$ as follows:

$\displaystyle\lambda\geq 1/(\sigma_{0}\sqrt{|\mathcal{E}_{tr}|}^{3}).$ (26)

By inequality (23), taking the convergence time
$T=\Omega\left(\frac{\log(\sigma_{0}/\epsilon)}{\eta\lambda\lambda_{0}}\right)$
we have that:

$\displaystyle\|\mathbf{c}(T)\|_{2}\leq\epsilon.$

According to the gradient descent for IRMV1 objective function, the evolution
of coefficients can be expressed as:

$\displaystyle\gamma_{j,r}^{inv}(t+1)$
$\displaystyle=\gamma^{inv}_{j,r}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t)\psi_{i}^{\prime}(t)\textrm{Rad}(\alpha)_{i}$
$\displaystyle\ \
-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\psi_{i}^{\prime}(t)\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\alpha)_{i},$ $\displaystyle\gamma_{j,r}^{spu}(t+1)$
$\displaystyle=\gamma^{spu}_{j,r}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t)\psi_{i}^{\prime}(t)\textrm{Rad}(\beta_{e})$
$\displaystyle\ \
-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\psi_{i}^{\prime}(t)\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\beta_{e})_{i}.$

Then we have,

$\displaystyle|\gamma_{j,r}^{inv}(t+1)|$
$\displaystyle\leq|\gamma_{j,r}^{inv}(t)|+\left|\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\psi_{i}^{\prime}(t)\ell^{\prime}_{i}(t)\textrm{Rad}(\alpha)_{i}\right|$
$\displaystyle+\left|\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\psi_{i}^{\prime}(t)\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\alpha)_{i}\right|$
$\displaystyle\leq|\gamma^{inv}_{j,r}(t)|+C\frac{\eta\sqrt{|\mathcal{E}_{tr}|}\lambda\beta
R^{2}L}{m}\|{\mathbf{c}}(t)\|_{2}.$

Similarly, we have,

$\displaystyle|\gamma_{j,r}^{spu}(t+1)|$
$\displaystyle\leq|\gamma^{spu}_{j,r,2}(t)|+C\frac{\eta\sqrt{|\mathcal{E}_{tr}|}\lambda\beta
R^{2}L}{m}\|{\mathbf{c}}(t)\|_{2}.$

At the time step $T$, the feature learning satisfies that:

$\displaystyle\gamma^{inv}_{j,r}(T)\leq
C\frac{\eta\sqrt{|\mathcal{E}_{tr}|}\lambda\beta
R^{2}LT}{m}\|\mathbf{c}(0)\|_{2};\quad\gamma^{spu}_{j,r}(T)\leq
C\frac{\eta\sqrt{|\mathcal{E}_{tr}|}\lambda\beta
R^{2}LT}{m}\|\mathbf{c}(0)\|_{2}.$

To make sure that $\gamma^{inv}_{j,r}(T)=o(1)$ and
$\gamma^{spu}_{j,r}(T)=o(1)$, we need the following condition:

$\displaystyle C\frac{\eta\sqrt{|\mathcal{E}_{tr}|}\lambda\beta
R^{2}LT}{m}\|\mathbf{c}(0)\|_{2}\leq d^{-\frac{1}{2}},$ (27)

combined with inequality (24) and inequality (26), we have:

$\displaystyle\sigma_{0}\leq\frac{|\mathcal{E}_{tr}|^{7/2}\beta^{3}L}{d^{1/2}m^{2}\lambda^{2}_{0}\log(1/\epsilon)}.$

∎

###  D.4 Proof for Proposition 4.3

######  Proposition D.9 (Restatement of Proposition 4.3).

Consider training the CNN model with the same data as Theorem 4.1, suppose
that $\psi(x)=x$, $\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)$ 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

  * •

$\sum_{e}C^{e}_{\text{IRMv1}}(t_{1})=0$.

  * •

$\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})$.

###### Proof of Proposition D.9.

According to the gradient descent for IRMV1 objective function, the evolution
of coefficients can be expressed as:

$\displaystyle\gamma_{j,r,1}(t+1)$
$\displaystyle=\gamma_{j,r,1}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t)\textrm{Rad}(\alpha)_{i}$
$\displaystyle\ \
-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\alpha)_{i},$ $\displaystyle\gamma_{j,r,2}(t+1)$
$\displaystyle=\gamma_{j,r,2}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t)\textrm{Rad}(\beta_{e})$
$\displaystyle\ \
-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\beta_{e})_{i},$

where $\ell^{\prime\prime}(y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}^{e}_{i}))=\frac{\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i}))}{(1+\exp(-y_{i}^{e}\cdot
f({\mathbf{W}},{\mathbf{x}}_{i})))^{2}}$.

To simplify the notation, we further define

$A_{1}^{e}=\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}\mathrm{Rad}(\alpha)_{i}$

and

$A_{2}^{e}=\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}y_{i}^{e}\mathrm{Rad}(\alpha)_{i}$

. Similarly, we define

$B_{1}^{e}=\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}\mathrm{Rad}(\beta_{e})_{i}$

and

$B_{2}^{e}=\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}y_{i}^{e}\mathrm{Rad}(\beta_{e})_{i}.$

In the limit of $n\rightarrow\infty$, we have:

$\displaystyle\lim_{n\rightarrow\infty}A^{1}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{(\gamma_{1}+\gamma_{2})})(1-\alpha)(1-\beta_{1})-1/(1+e^{\gamma_{1}-\gamma_{2}})(1-\alpha)\beta_{1}$
$\displaystyle\quad+1/(1+e^{\gamma_{2}-\gamma_{1}})\alpha(1-\beta_{1})+1/(1+e^{-\gamma_{1}-\gamma_{2}})\alpha\beta_{1},$
$\displaystyle\lim_{n\rightarrow\infty}A^{2}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{\gamma_{1}+\gamma_{2}})(1-\alpha)(1-\beta_{2})-1/(1+e^{\gamma_{1}-\gamma_{2}})(1-\alpha)\beta_{2}$
$\displaystyle\quad+1/(1+e^{-\gamma_{1}+\gamma_{2}})\alpha(1-\beta_{2})+1/(1+e^{-\gamma_{1}-\gamma_{2}})\alpha\beta_{2},$
$\displaystyle\lim_{n\rightarrow\infty}B^{1}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{\gamma_{1}+\gamma_{2}})(1-\alpha)(1-\beta_{1})+1/(1+e^{\gamma_{1}-\gamma_{2}})(1-\alpha)\beta_{1}$
$\displaystyle\quad-1/(1+e^{-\gamma_{1}+\gamma_{2}})\alpha(1-\beta_{1})+1/(1+e^{-\gamma_{1}-\gamma_{2}})\alpha\beta_{1},$
$\displaystyle\lim_{n\rightarrow\infty}B^{2}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{\gamma_{1}+\gamma_{2}})(1-\alpha)(1-\beta_{2})+1/(1+e^{\gamma_{1}-\gamma_{2}})(1-\alpha)\beta_{2}$
$\displaystyle\quad-1/(1+e^{-\gamma_{1}+\gamma_{2}})\alpha(1-\beta_{2})+1/(1+e^{-\gamma_{1}-\gamma_{2}})\alpha\beta_{2}.$

and,

$\displaystyle\lim_{n\rightarrow\infty}A^{1}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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}),$
$\displaystyle\lim_{n\rightarrow\infty}A^{2}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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}),$
$\displaystyle\lim_{n\rightarrow\infty}B^{1}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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}),$
$\displaystyle\lim_{n\rightarrow\infty}B^{2}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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}).$

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$,

$\displaystyle\left|A^{1}_{1}(t_{1})-\lim_{n\rightarrow\infty}A^{1}_{1}(t_{1})\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^{e}_{1}(t_{1})=\sum_{e}B^{e}_{1}(t_{1})=0$:

$\displaystyle\lim_{n\rightarrow\infty}(A_{1}^{1}(t_{1})+A_{1}^{2}(t_{1}))$
$\displaystyle=-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})$
$\displaystyle\quad+1/(1+e^{-\gamma_{1}+\gamma_{2}})\alpha(2-\beta_{1}-\beta_{2})+1/(1+e^{-\gamma_{1}-\gamma_{2}})\alpha(\beta_{1}+\beta_{2})=0$
$\displaystyle\lim_{n\rightarrow\infty}(B_{1}^{1}(t_{1})+B_{1}^{2}(t_{1}))$
$\displaystyle=-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})$
$\displaystyle\quad+1/(1+e^{-\gamma_{1}+\gamma_{2}})\alpha(2-\beta_{1}-\beta_{2})+1/(1+e^{-\gamma_{1}-\gamma_{2}})\alpha(\beta_{1}+\beta_{2})=0$

Solving the above equations, we have,

$\displaystyle\gamma^{\infty}_{1}(t_{1})=\frac{1}{2}\log(G_{m}G_{b})\quad\gamma^{\infty}_{2}(t_{1})=\frac{1}{2}\log(G_{m}/G_{b})$

where we denote
$\gamma^{\infty}_{1}(t_{1})\triangleq\lim_{n\rightarrow\infty}\gamma_{1}(t_{1})$
and
$\gamma^{\infty}_{2}(t_{1})\triangleq\lim_{n\rightarrow\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:

$\displaystyle\left|\gamma_{1}-\gamma^{\infty}_{1}\right|<\left|e^{\gamma_{1}}-e^{\gamma^{\infty}_{1}}\right|\leq
C\left|1/(1+e^{\gamma_{1}})-1/(1+e^{\gamma^{\infty}_{1}})\right|\leq\sqrt{\frac{4\log(1/\delta)}{n}},$
$\displaystyle\left|\gamma_{2}-\gamma^{\infty}_{2}\right|<\left|e^{\gamma_{2}}-e^{\gamma^{\infty}_{2}}\right|\leq
C\left|1/(1+e^{\gamma_{2}})-1/(1+e^{\gamma^{\infty}_{2}})\right|\leq\sqrt{\frac{4\log(1/\delta)}{n}}.$

Then we know that,

$\displaystyle
C_{\text{IRMv1}}^{1}=\frac{1}{n_{1}}\sum_{i=1}^{n_{1}}\ell^{\prime}_{i}\hat{y}_{i}^{1}y_{i}^{1}=\gamma_{1}A_{1}^{1}+\gamma_{2}B^{1}_{1}$
$\displaystyle
C_{\text{IRMv1}}^{2}=\frac{1}{n_{2}}\sum_{i=1}^{n_{2}}\ell^{\prime}_{i}\hat{y}_{i}^{2}y_{i}^{2}=\gamma_{1}A_{1}^{2}+\gamma_{2}B^{2}_{1}$

Therefore, we have that:

$\displaystyle C_{\text{IRMv1}}^{1}+C_{\text{IRMv1}}^{2}=0$

Then the evolution of coefficients reduces to

$\displaystyle\gamma_{j,r,1}(t+1)$
$\displaystyle=\gamma_{j,r,1}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))A^{e}_{1}(t)-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}A^{e}_{2}(t)$
$\displaystyle\gamma_{j,r,2}(t+1)$
$\displaystyle=\gamma_{j,r,2}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))B^{e}_{1}(t)-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}B^{e}_{2}(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 a probability at least $1-\delta$ and $n>C_{1}\log(1/\delta)$ with
$C_{1}$ being a positive constant, we have:

$\displaystyle\gamma_{j,r,1}(t_{1}+1)>\gamma_{j,r,1}(t_{1}),$
$\displaystyle\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

$\displaystyle\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:

$\displaystyle\gamma_{j,r,1}(t+1)$
$\displaystyle=\gamma_{j,r,1}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t)\textrm{Rad}(\alpha)_{i}$
$\displaystyle\ \
-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\alpha)_{i},$ $\displaystyle\gamma_{j,r,2}(t+1)$
$\displaystyle=\gamma_{j,r,2}(t)-\frac{\eta}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}(1+2\lambda
C_{\text{IRMv1}}^{e}(t))\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime}_{i}(t)\textrm{Rad}(\beta_{e})$
$\displaystyle\ \
-\frac{\eta\lambda}{m}\cdot\sum_{e\in\mathcal{E}_{\mathrm{tr}}}2C_{\text{IRMv1}}^{e}\frac{1}{n_{e}}\sum_{i=1}^{n_{e}}\ell^{\prime\prime}_{i}\hat{y}_{i}^{e}\cdot
y_{i}^{e}\textrm{Rad}(\beta_{e})_{i},$

Taking the value of $\gamma_{j,r,1}(t_{1})$, $\gamma_{j,r,2}(t_{1})$ and, we
can conclude that:

$\displaystyle\lim_{n\rightarrow\infty}A^{1}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{\gamma_{2}})(1-\alpha)(1-\beta_{1})-1/(1+e^{-\gamma_{2}})(1-\alpha)\beta_{1}+$
$\displaystyle\ \
1/(1+e^{\gamma_{2}})\alpha(1-\beta_{1})+1/(1+e^{-\gamma_{2}})\alpha\beta_{1}$
$\displaystyle=1/(1+e^{\gamma_{2}})(2\alpha-1)(1-\beta_{1})+1/(1+e^{-\gamma_{2}})(2\alpha-1)(\beta_{1})$
$\displaystyle=(2\alpha-1)[1/(1+e^{\gamma_{2}})(1-\beta_{2})+1/(1+e^{-\gamma_{2}})\beta_{1})]$
$\displaystyle\lim_{n\rightarrow\infty}A^{2}_{1}(t_{1})$
$\displaystyle=1/(1+e^{\gamma_{2}})(2\alpha-1)(1-\beta_{2})+1/(1+e^{-\gamma_{2}})(2\alpha-1)(\beta_{2})$
$\displaystyle=(2\alpha-1)[1/(1+e^{\gamma_{2}})(1-\beta_{2})+1/(1+e^{-\gamma_{2}})\beta_{2})]$
$\displaystyle\lim_{n\rightarrow\infty}B^{1}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{\gamma_{2}})(1-\alpha)(1-\beta_{1})+1/(1+e^{-\gamma_{2}})(1-\alpha)\beta_{1}-$
$\displaystyle\ \
1/(1+e^{\gamma_{2}})\alpha(1-\beta_{1})+1/(1+e^{-\gamma_{2}})\alpha\beta_{1}$
$\displaystyle=-1/(1+e^{\gamma_{2}})(1-\beta_{1})+1/(1+e^{-\gamma_{2}})\beta_{1}$
$\displaystyle\lim_{n\rightarrow\infty}B^{2}_{1}(t_{1})$
$\displaystyle=-1/(1+e^{\gamma_{2}})(1-\alpha)(1-\beta_{2})+1/(1+e^{-\gamma_{2}})(1-\alpha)\beta_{2}-$
$\displaystyle\ \
1/(1+e^{\gamma_{2}})\alpha(1-\beta_{2})+1/(1+e^{-\gamma_{2}})\alpha\beta_{2}$
$\displaystyle=-1/(1+e^{\gamma_{2}})(1-\beta_{2})+1/(1+e^{-\gamma_{2}})\beta_{2}$

On the other hand,

$\displaystyle\lim_{n\rightarrow\infty}A^{1}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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})$
$\displaystyle=e^{\gamma_{2}}/(1+e^{\gamma_{2}})^{2}(1-2\alpha)(1-\beta_{1})+e^{-\gamma_{2}}/(1+e^{-\gamma_{2}})^{2}(2\alpha-1)\beta_{1}\gamma_{2}$
$\displaystyle\lim_{n\rightarrow\infty}A^{2}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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})$
$\displaystyle=e^{\gamma_{2}}/(1+e^{\gamma_{2}})^{2}(1-2\alpha)(1-\beta_{2})+e^{-\gamma_{2}}/(1+e^{-\gamma_{2}})^{2}(2\alpha-1)\beta_{2}\gamma_{2}$
$\displaystyle\lim_{n\rightarrow\infty}B^{1}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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}),$
$\displaystyle\lim_{n\rightarrow\infty}B^{2}_{2}(t_{1})$
$\displaystyle=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})$
$\displaystyle\quad+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}).$

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:

$\displaystyle\gamma_{j,r,1}(t_{1}+1)<\gamma_{j,r,1}(t_{1}).$

∎

##  Appendix 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 [83, 59] 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:=w\circ\varphi$; length of inner training epochs $e$;
termination threshold $p$;

2: Initialize groups
$G^{a}\leftarrow{\mathcal{D_{\mathrm{tr}}}},G^{r}\leftarrow\\{\\}$;

3: for $k\in[1,\ldots,K]$ do

4: Randomly initialize $w_{k}$;

5: for $j\in[1,\ldots,e]$ do

6: Obtain $\ell_{\text{FeAT}}$ with $G$ via Eq. 7;

7: Update $w_{k},\varphi$ with $\ell_{\text{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 ${\cal D}_{k}^{a},{\cal
D}_{k}^{r}$ according to $f_{k}$;

18: // Hence it doesnot need to maintain all historical subsets.

19: Update groups $G^{a}\leftarrow\\{{\cal
D}_{k}^{a}\\},G^{r}\leftarrow\\{{\cal D}_{k}^{r}\\}$;

20: end for

21: return $f=w\circ\varphi$;

##  Appendix 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, 83, 16]. In addition to the original
ColoredMNIST with ${{\mathcal{E}}_{\text{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}}_{\text{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], 555https://github.com/TjuJianyu/RFC 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.

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 | -

It can be found 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\%$.

Table 6: Performances in various sets at different FeAT rounds. ColoredMNIST-025 | Round-1 | Round-2 | Round-3 |
---|---|---|---|---
Training Acc. |  85.08$\pm$ 0.14 |  71.87$\pm$ 0.96 |  84.93$\pm$ 1.26 |
Retention Acc. | - |  88.11$\pm$ 4.28 |  43.82$\pm$ 0.59 |
OOD Acc. |  11.08$\pm$ 0.30 |  70.64$\pm$ 0.62 |  10.07$\pm$ 0.26 |
ColoredMNIST-01 | Round-1 | Round-2 | Round-3 | Round-4
Training Acc. |  88.63$\pm$ 0.15 |  74.25$\pm$ 1.23 |  86.07$\pm$ 0.36 |  77.29$\pm$ 0.24
Retention Acc. | - |  85.91$\pm$ 1.78 |  48.05$\pm$ 1.39 |  29.09$\pm$ 1.15
OOD Acc. |  73.50$\pm$ 0.41 |  17.32$\pm$ 2.69 |  85.40$\pm$ 0.54 |  12.48$\pm$ 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].

Table 7: A summary of datasets information from Wilds.

Dataset | Data ($x$) | Class information | Domains | Metric | Architecture
---|---|---|---|---|---
Amazon | Product reviews | Star ratings (5 classes) | 7,676 reviewers | 10-eth percentile acc. | DistillBERT
Camelyon17 | Tissue slides | Tumor (2 classes) | 5 hospitals | Avg. acc. | DenseNet-121
CivilComments | Online comments | Toxicity (2 classes) | 8 demographic groups | Wr. group acc. | DistillBERT
FMoW | Satellite images | Land use (62 classes) | 16 years x 5 regions | Wr. group acc. | DenseNet-121
iWildCam | Photos | Animal species (186 classes) | 324 locations | Macro F1 | ResNet-50
RxRx1 | Cell images | Genetic treatments (1,139 classes) | 51 experimental batches | Avg. acc | ResNet-50

Table 8: A summary of datasets statistics from Wilds.

Dataset |  $\\#$ Examples |  |  $\\#$ Domains
---|---|---|---
train | val | test |  | train | val | test
Amazon | 1,000,124 | 100,050 | 100,050 |  | 5,008 | 1,334 | 1,334
Camelyon17 | 302,436 | 34,904 | 85,054 |  | 3 | 1 | 1
CivilComments | 269,038 | 45,180 | 133,782 |  | - | - | -
FMoW | 76,863 | 19,915 | 22,108 |  | 11 | 3 | 2
iWildCam | 129,809 | 14,961 | 42,791 |  | 243 | 32 | 48
RxRx1 | 40,612 | 9,854 | 34,432 |  | 33 | 4 | 14

Amazon. We follow the Wilds splits and data processing pipeline for the Amazon
dataset [51]. It provides $1.4$ million comments collected from $7,676$ Amazon
customers. The task is to predict the score (1-5 stars) for each review. The
domains $d$ are defined according to the reviewer/customer who wrote the
product reviews. The evaluation metric used for the task is $10$th percentile
of per-user accuracies in the OOD test sets, and the backbone model is a
DistilBert [66], following the Wilds protocol [39].

Camelyon17. We follow the Wilds splits and data processing pipeline for the
Camelyon17 dataset [6]. It provides $450,000$ lymph-node scans from $5$
hospitals. The task in Camelyon17 is to take the input of $96\times 96$
medical images to predict whether there exists a tumor tissue in the image.
The domains $d$ refers to the index of the hospital where the image was taken.
The training data are sampled from the first $3$ hospitals where the OOD
validation and test data are sampled from the $4$-th and $5$-th hospital,
respectively. We will use the average accuracy as the evaluation metric and a
DenseNet-121 [31] as the backbone for the featurizer.

CivilComments. We follow the Wilds splits and data processing pipeline for the
CivilComments dataset [9]. It provides $450,000$ comments collected from
online articles. The task is to classify whether an online comment text is
toxic or non-toxic. The domains $d$ are defined according to the demographic
features, including male, female, LGBTQ, Christian, Muslim, other religions,
Black, and White. CivilComments is used to study the subpopulation shifts,
here we will use the worst group/domain accuracy as the evaluation metric. As
for the backbone of the featurizer, we will use a DistillBert [66] following
Wilds [39].

FMoW. We follow the Wilds splits and data processing pipeline for the FMoW
dataset [18]. It provides satellite images from $16$ years and $5$ regions.
The task in FMoW is to classify the images into $62$ classes of building or
land use categories. The domain is split according to the year that the
satellite image was collected, as well as the regions in the image which could
be Africa, America, Asia, Europe or Oceania. Distribution shifts could happen
across different years and regions. The training data contains data collected
before $2013$, while the validation data contains images collected within
$2013$ to $2015$, and the test data contains images collected after $2015$.
The evaluation metric for FMoW is the worst region accuracy and the backbone
model for the featurizer is a DenseNet-121 [31].

iWildCam. We follow the Wilds splits and data processing pipeline for the
iWildCam dataset [8]. It is consist of $203,029$ heat or motion-activated
photos of animal specifies from 323 different camera traps across different
countries around the world. The task of iWildCam is to classify the
corresponding animal specifies in the photos. The domains is split according
to the locations of the camera traps which could introduce the distribution
shifts. We will use the Macro F1 as the evaluation metric and a ResNet-50 [27]
as the backbone for the featurizer.

RxRx1. We follow the Wilds splits and data processing pipeline for the RxRx1
dataset [72]. The input is an image of cells taken by fluorescent microscopy.
The cells can be genetically perturbed by siRNA and the task of RxRx1 is to
predict the class of the corresponding siRNA that have treated the cells.
There exists $1,139$ genetic treatments and the domain shifts are introduced
by the experimental batches. We will use the average accuracy of the OOD
experimental batches as the evaluation metric and a ResNet-50 [27] as the
backbone for the featurizer.

####  F.2.2 Training and evaluation details.

We follow previous works to implement and evaluate different methods used in
our experiments [39]. The information of the referred paper and code is listed
as in Table. 9.

Table 9: The information of the referred paper and code.

Paper | Commit | Code
---|---|---
Wilds [39] | v2.0.0 | https://wilds.stanford.edu/
Fish [69] | 333efa24572d99da0a4107ab9cc4af93a915d2a9 | https://github.com/YugeTen/fish
Bonsai [83] | 33b9ecad0ce8b3462793a2da7a9348d053c06ce0 | https://github.com/TjuJianyu/RFC
DFR [38, 33] | 6d098440c697a1175de6a24d7a46ddf91786804c | https://github.com/izmailovpavel/spurious_feature_learning

The general hyperparemter setting inherit from the referred codes and papers,
and are as listed in Table 10. We use the same backbone models to implement
the featurizer [27, 31, 66]. By default, we repeat the experiments by $3$ runs
with the random seeds of $0,1,2$. While for Camelyon17, we follow the official
guide to repeat $10$ times with the random seeds from $0$ to $9$.

Table 10: General hyperparameter settings for the experiments on Wilds.

Dataset | Amazon | Camelyon17 | CivilComments | FMoW | iWildCam | RxRx1
---|---|---|---|---|---|---
Num. of seeds | 3 | 10 | 3 | 3 | 3 | 3
Learning rate | 2e-6 | 1e-4 | 1e-5 | 1e-4 | 1e-4 | 1e-3
Weight decay | 0 | 0 | 0.01 | 0 | 0 | 1e-5
Scheduler | n/a | n/a | n/a | n/a | n/a | Cosine Warmup
Batch size | 64 | 32 | 16 | 32 | 16 | 72
Architecture | DistilBert | DenseNet121 | DistilBert | DenseNet121 | ResNet50 | ResNet50
Optimizer | Adam | SGD | Adam | Adam | Adam | Adam
Domains in minibatch | 5 | 3 | 5 | 5 | 10 | 10
Group by | Countries | Hospitals |  Demographics$\times$ toxicity |  Times $\times$ regions | Trap locations | Experimental batches
Training epochs | 200 | 10 | 5 | 12 | 9 | 90

##### OOD objective implementations.

We choose $4$ representative OOD objectives to evaluate the quality of learned
features, including GroupDRO [64], IRMv1 [4], VREx [41] and IRMX [16]. We
implement the OOD objectives based on the code provided by Shi et al. [69].
For each OOD objective, by default, we follow the Wilds practice to sweep the
penalty weights from the range of $\\{1e-2,1e-1,1,1e1,1e2\\}$, and perform the
model and hyperparameter selection via the performance in the provided OOD
validation set of each dataset. Due to the overwhelming computational overhead
required by large datasets and resource constraints, we tune the penalty
weight in iWildCam according to the performance with seed $0$, which we
empirically find yields similar results as full seed tunning. Besides in
Amazon, we adopt the penalty weights tuned from CivilComments since the two
datasets share a relatively high similarity, which we empirically find yields
similar results as full seed tunning, too. On the other hand, it raises more
challenges for feature learning algorithms in iWildCam and Amazon.

##### Deep Feature Reweighting (DFR) implementations.

For the implementation of DFR [38, 33], we use the code provided in Izmailov
et al. [33]. By default, DFR considers the OOD validation as an unbiased
dataset and adopts the OOD validation set to learn a new classifier based on
the frozen features from the pre-trained featurizer. We follow the same
implementation and evaluation protocol when evaluating feature learning
quality in FMoW and CivilComments. However, since Camelyon17 does not have the
desired OOD validation set, we follow the “cheating” protocol as in Rosenfeld
et al. [63] to perform the logistic regression based the train and test sets.
Note that when “cheating”, the model is not able to access the whole test
sets. Instead, the logistic regression is conducted on a random split of the
concatenated train and test data. Moreover, for Amazon and iWildCam, we find
the original implementation fails to converge possibly due to the complexity
of the task, and the relatively poor feature learning quality. Hence we
implement a new logistic regression based on PyTorch [54] optimized with SGD,
and perform DFR using “cheating” protocol based on the OOD validation set and
test set. Besides, we find neither the two aforementioned implementations or
dataset choices can lead to DFR convergence in RxRx, which we will leave for
future investigations.

##### Feature learning algorithm implementations.

We implement all the feature learning methods based on the Fish code
framework. For the fairness of comparison, we set all the methods to train the
same number of steps or rounds (if applicable) in Wilds datasets. The only
exception is in RxRx1, where both Bonsai and FeAT require more steps to
converge, since the initialized featurizer has a relatively large distance
from the desired featurizer in the task. We did not train the model for much
too long epochs as Izmailov et al. [33] find that it only requires $2-5$
epochs for deep nets to learn high-quality invariant features. The final model
is selected based on the OOD validation accuracy during the training. Besides,
we tune the retain penalty in FeAT by searching over
$\\{1e-2,1e-1,0.5,1,2,10\\}$, and finalize the retain penalty according to the
OOD validation performance. We list the detailed training steps and rounds
setups, as well as the used retain penalty in FeAT in Table 11.

Table 11: Hyperparameter setups of feature learning algorithms for the
experiments on Wilds.

Dataset | Amazon | Camelyon17 | CivilComments | FMoW | iWildCam | RxRx1
---|---|---|---|---|---|---
Overall steps | 31,000 | 10,000 | 50,445 | 9,600 | 48,000 | 20,000
Approx. epochs | 4 | 10 | 3 | 4 | 10 | 10
Num. of rounds | 3 | 2 | 3 | 2 | 2 | 10
Steps per round | 10,334 | 5,000 | 16,815 | 4,800 | 10 | 10
FeAT Retain penalty | 2.0 | 1e-2 | 1e-2 | 1.0 | 0.5 | 10

For ERM, we train the model simply by the overall number of steps, except for
RxRx1 where we train the model by $15,000$ steps following previous setups
[69]. Bonsai and FeAT directly adopt the setting listed in the Table 11.
Besides, Bonsai will adopt one additional round for synthesizing the pre-
trained models from different rounds. Although Zhang et al. [83] requires
Bonsai to train the two rounds for synthesizing the learned features, we
empirically find additional training steps in synthesizing will incur
overfitting and worse performance. Moreover, as Bonsai requires propagating
$2K-1$ batches of the data that may exceed the memory limits, we use a smaller
batch size when training Bonsai in iWildCam ($8$) and RxRx1 ($56$).

###  F.3 Software and hardware

We implement our methods with PyTorch [54]. For the software and hardware
configurations, we ensure the consistent environments for each datasets. We
run all the experiments on Linux servers with NVIDIA V100 graphics cards with
CUDA 10.2.

###  F.4 Computational analysis

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 requires another round of training with $(K-1)$
additional propagation, given $K$ total rounds.

We calculated the computational overhead:

Table 12: Training and memory overhead of different algorithms. | Camelyon17 |  | CivilComments |
---|---|---|---|---
| Training time | Memory (%) | Training time | Memory (%)
ERM |  56.21$\pm$8.29 mins |  22.56$\pm$0.00 |  24.22$\pm$0.33 hrs |  36.46$\pm$0.00
Bonsai |  214.55$\pm$1.13 mins |  51.75$\pm$0.01 |  58.47$\pm$0.91 hrs |  64.43$\pm$0.31
FeAT |  101.14$\pm$12.79 mins |  51.92$\pm$0.04 |  28.19$\pm$1.15 hrs |  56.21$\pm$0.48

The results aligned with our discussion. Bonsai requires much more time for
the additional synthetic round and much more memory when there are 3 or more
rounds. In contrast, FeAT achieves the best performance without introducing
too much additional computational overhead.

###  F.5 Feature learning analysis

We first visualize the feature learning of ERM and FeAT on ColoredMNIST-025,
as shown in Fig. 4 It can be found that 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.

Table 13: Labels and predictions for the visualized samples.

| Label | ERM | Bonsai | FeAT |  | Label | ERM | Bonsai | FeAT
---|---|---|---|---|---|---|---|---|---
Camelyon17 | 1 | 1 | 0 | 1 | iWildCam | 113 | 68 | 0 | 113
1 | 1 | 0 | 1 | 113 | 0 | 0 | 113
1 | 1 | 0 | 1 | 36 | 36 | 36 | 36
1 | 0 | 0 | 0 | 36 | 36 | 36 | 36
FMoW | 40 | 40 | 40 | 40 | RxRx1 | 1138 | 812 | 812 | 812
40 | 40 | 40 | 40 | 1138 | 1133 | 1125 | 1133
40 | 2 | 29 | 29 | 35 | 43 | 1119 | 143
40 | 40 | 40 | 40 | 35 | 35 | 1054 | 35
CivilComments | toxic | toxic | toxic | toxic | Amazon | 2 | 3 | 3 | 2
toxic | toxic | toxic | toxic | 5 | 5 | 5 | 5
toxic | toxic | toxic | toxic | 3 | 4 | 4 | 4
nontoxic | nontoxic | nontoxic | nontoxic | 5 | 5 | 5 | 5

We also visualize the saliency maps of ERM, Bonsai, and FeAT on all real-world
datasets used in our work with https://github.com/pytorch/captum. The
visualizations are shown as in Fig. 5 to Fig. 11, for which the labels and the
predictions of different algorithms are given in Table. 13. It can be found
that, across various tasks and data modalities, FeAT effectively learns more
meaningful and diverse features than ERM and Bonsai, which serve as strong
evidence for the consistent superiority of FeAT in OOD generalization.

(a) ERM 150 epochs (b) ERM 300 epochs

(c) ERM 450 epochs (d) FeAT 2 rounds

Figure 4: 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 or even forget certain features with longer training epochs. FeAT can
stably capture the desired features. (a) CivilComments Figure 5: Saliency map
of feature learning on Wilds CivilComments benchmark. The green-colored tokens
are the learned features that contributed most to the target class, while the
red-colored tokens contributed to the other classes. It can be found that FeAT
is able to learn more meaningful and diverse features than ERM and Bonsai.

(a) Amazon-ERM (b) Amazon-Bonsai

Figure 6: Saliency map of feature learning on Wilds Amazon benchmark. The
green-colored tokens are the learned features that contributed most to the
target class, while the red-colored tokens contributed to the other classes.
It can be found that FeAT is able to learn more meaningful and diverse
features than ERM and Bonsai. (a) Amazon-FeAT Figure 7: Saliency map of
feature learning on Wilds Amazon benchmark (part 2). The green-colored tokens
are the learned features that contributed most to the target class, while the
red-colored tokens contributed to the other classes. It can be found that FeAT
is able to learn more meaningful and diverse features than ERM and Bonsai.

(a) Camelyon17-ERM (b) Camelyon17-Bonsai (c) Camelyon17-FeAT

(d) Camelyon17-ERM (e) Camelyon17-Bonsai (f) Camelyon17-FeAT

(g) Camelyon17-ERM (h) Camelyon17-Bonsai (i) Camelyon17-FeAT

(j) Camelyon17-ERM (k) Camelyon17-Bonsai (l) Camelyon17-FeAT

Figure 8: Saliency map of feature learning on Camelyon17 benchmark. The blue
dots are the salient features. A deeper blue color denotes more salient
features. It can be found that FeAT is able to learn more meaningful and
diverse features than ERM and Bonsai.

(a) FMoW-ERM (b) FMoW-Bonsai (c) FMoW-FeAT

(d) FMoW-ERM (e) FMoW-Bonsai (f) FMoW-FeAT

(g) FMoW-ERM (h) FMoW-Bonsai (i) FMoW-FeAT

(j) FMoW-ERM (k) FMoW-Bonsai (l) FMoW-FeAT

Figure 9: Saliency map of feature learning on FMoW benchmark. The blue dots
are the salient features. A deeper blue color denotes more salient features.
It can be found that FeAT is able to learn more meaningful and diverse
features than ERM and Bonsai.

(a) iWildCam-ERM (b) iWildCam-Bonsai (c) iWildCam-FeAT

(d) iWildCam-ERM (e) iWildCam-Bonsai (f) iWildCam-FeAT

(g) iWildCam-ERM (h) iWildCam-Bonsai (i) iWildCam-FeAT

(j) iWildCam-ERM (k) iWildCam-Bonsai (l) iWildCam-FeAT

Figure 10: Saliency map of feature learning on iWildCam benchmark. The blue
dots are the salient features. A deeper blue color denotes more salient
features. It can be found that FeAT is able to learn more meaningful and
diverse features than ERM and Bonsai.

(a) RxRx1-ERM (b) RxRx1-Bonsai (c) RxRx1-FeAT

(d) RxRx1-ERM (e) RxRx1-Bonsai (f) RxRx1-FeAT

(g) RxRx1-ERM (h) RxRx1-Bonsai (i) RxRx1-FeAT

(j) RxRx1-ERM (k) RxRx1-Bonsai (l) RxRx1-FeAT

Figure 11: Saliency map of feature learning on RxRx1 benchmark. The blue dots
are the salient features. A deeper blue color denotes more salient features.
It can be found that FeAT is able to learn more meaningful and diverse
features than ERM and Bonsai.