SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left| \Xi_ {j, r, i ^ {\prime}} ^ {t + 1, e ^ {\prime}} \right| \leq \left| \Xi_ {j, r, i ^ {\prime}} ^ {t, e ^ {\prime}} \right| + \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \frac {1}{1 + \exp \{y _ {i} ^ {e} \hat {y} _ {i} ^ {e} \}} \cdot | \langle \pmb {\xi} _ {i} ^ {e}, \pmb {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle | \\ \leq \left| \Xi_ {j, r, i ^ {\prime}} ^ {t, e ^ {\prime}} \right| + \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} | \langle \pmb {\xi} _ {i} ^ {e}, \pmb {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle | \\ = \left| \Xi_ {j, r, i ^ {\prime}} ^ {0, e ^ {\prime}} \right| + (t + 1) \cdot \frac {\eta}{m} \sum_ {e \in \mathcal {E} _ {t r}} \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} | \langle \pmb {\xi} _ {i} ^ {e}, \pmb {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle | \\ = | \langle \boldsymbol {\xi} _ {r}, \boldsymbol {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle | + (t + 1) \cdot \left(\frac {\eta}{m n _ {e ^ {\prime}}} \| \boldsymbol {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \| _ {2} ^ {2} + \sum_ {(i, e) \neq (i ^ {\prime}, e ^ {\prime})} \frac {\eta}{m n _ {e}} | \langle \boldsymbol {\xi} _ {i} ^ {e}, \boldsymbol {\xi} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle |\right) \\ \leq 2 \sigma_ {p} \sigma_ {0} \sqrt {(d - 2) \log \frac {1 6 n m}{\rho}} \\ + \frac {T \eta \sigma_ {p} ^ {2}}{m \underline {{n}}} \left((d - 2) + 1 6 \sqrt {2 (d - 2) \log \frac {8 n}{\rho}} + 2 n \sqrt {(d - 2) \log \frac {8 n ^ {2}}{\rho}}\right). \\ \end{array}
$$
Then, we can bound  $|\mathbb{Q}_i^e |$  as
$$
\begin{array}{l} \left| \mathbb {Q} _ {i} ^ {e} \right| \leq 2 \cdot \left| \frac {1}{m} \sum_ {r = 1} ^ {m} \xi_ {r} \right| + 2 \cdot \left| \frac {1}{m} \sum_ {r = 1} ^ {m} \xi_ {r} \right| + \frac {2}{m} \sum_ {r = 1} ^ {m} \left| \Xi_ {j, r, i} ^ {t, e} \right| \\ \leq 4 \sigma_ {0} \sqrt {\frac {2}{m} \log \frac {3 2 m}{\rho}} + 4 \sigma_ {p} \sigma_ {0} \sqrt {(d - 2) \log \frac {1 6 n m}{\rho}} \\ + \frac {2 T \eta \sigma_ {p} ^ {2}}{m \underline {{n}}} \left((d - 2) + 1 6 \sqrt {2 (d - 2) \log \frac {8 n}{\rho}} + 2 n \sqrt {(d - 2) \log \frac {8 n ^ {2}}{\rho}}\right). \\ \end{array}
$$
Thus, with sufficient small  $\sigma_0,\sigma_p$  ,i.e.,
we ensured that  $|\mathbb{Q}_i^e | = O(\underline{n}^{-1})$
D.2.2 Proof of Lemma D.3
For  $e \in \mathcal{E}_{tr}$ , using Hoeffding's inequality, it holds that
$$
\Pr \left\{\left| \frac {1}{n _ {e}} \sum_ {i = 1} ^ {n _ {e}} \mathbf {1} _ {\left\{\operatorname {R a d} (\alpha) _ {i} = + 1, \operatorname {R a d} (\beta_ {e}) _ {i} = + 1 \right\}} - (1 - \alpha) (1 - \beta_ {e}) \right| \geq a \right\} \leq 2 \exp \left\{- 2 a ^ {2} n _ {e} \right\}.
$$
Considering two environments, using a union bound, we can conclude that
$$
\Pr \left\{\left| \overline {{C}} _ {+ 1 + 1} - (1 - \alpha) (2 - \beta_ {1} - \beta_ {2}) \right| \leq a \right\} \geq 1 - 4 \exp \left\{- \frac {a ^ {2} n}{2} \right\}.
$$
Thus, for  $\rho > 0$ , with probability at least  $1 - \frac{\rho}{4}$ , we can conclude that
$$
\left| \overline {{C}} _ {+ 1 + 1} - (1 - \alpha) (2 - \beta_ {1} - \beta_ {2}) \right| \leq \sqrt {\frac {2 \log (1 6 / \rho)}{\underline {{n}}}}.
$$
Using the above arguments for other constants  $\overline{C}_{+1 - 1}$ ,  $\overline{C}_{-1 + 1}$  and  $\overline{C}_{-1 - 1}$ , and applying a union bound, we obtain the desired results.
D.2.3 ERM Feature Learning with Non-Linear Activation Functions
It was numerically observed that in the early stage of (stochastic) GD training, the learning dynamics of neural networks can be mimicked by training a simple linear model [35]. Hu et al. [30] rigorously proved this phenomenon for training two-layer neural network with  $\ell_2$  loss function in the Neural Tangent Kernel (NTK) region. We briefly summarize their results here: For a two-layer fully-connected neural network (with fixed second layer  $\{v_r\}$ ):
$$
f _ {F C} (\mathbf {W}, \mathbf {x}) \triangleq \frac {1}{\sqrt {m}} \sum_ {r = 1} ^ {m} v _ {r} \psi \left(\mathbf {w} _ {r} ^ {\top} \mathbf {x} / \sqrt {d}\right), \tag {21}
$$
considering the  $\ell_2$  training loss  $\ell_2(\hat{y}, y) \triangleq \frac{1}{2} (\hat{y} - y)^2$  and the ERM objective  $L_{\mathrm{ERM}}(\mathbf{W}) = \frac{1}{n}\sum_{i=1}^{n}\ell_2\big(f_{FC}(\mathbf{W}, \mathbf{x}_i), y_i\big)$ , when using GD:  $\mathbf{W}^{t+1} = \mathbf{W}^t - \boldsymbol{\eta} \cdot \nabla L_{\mathrm{ERM}}(\mathbf{W}^t)$  to minimize the ERM objective, the following holds.
Theorem D.4 (Theorem 3.2 of [30]). Let  $\alpha_{nl} \in (0, \frac{1}{4})$  be a fixed constant, and  $\psi(\cdot)$  be a smooth (with bounded first and second derivatives) or piece-wise linear activation function. Suppose that  $n$  and  $m$  satisfy  $n = \Omega(d^{1 + \alpha_{nl}})$  and  $m = \Omega(d^{1 + \alpha_{nl}})$ . Suppose that  $\eta \ll d$ . Then there exists a universal constant  $c > 0$  such that with high probability, for all  $t = O\left(\frac{d}{\eta}\log d\right)$  simultaneously, the learned neural network  $f_{FC}^t$  and the linear model  $f_{lin}^t$  (defined below) at iteration  $t$  are close on average on the training data:
$$
\frac {1}{n} \sum_ {i = 1} ^ {n} \left(f _ {F C} ^ {t} (\mathbf {x} _ {i}) - f _ {l i n} ^ {t} (\mathbf {x} _ {i})\right) ^ {2} = O (d ^ {- \Omega (\alpha_ {n l})}).
$$
The linear model  $f_{lin}(\beta, \mathbf{x}) = \beta^\top R(\mathbf{x})$  is a linear function of the transformed data  $R(\mathbf{x}) = \frac{1}{\sqrt{d}} \begin{bmatrix} \zeta \mathbf{x} \\ \nu \end{bmatrix}$ , where  $\zeta$  and  $\nu$  are constants related to  $\psi'$  and the dataset distribution (see (5) in [30] for formal definitions).
We show that we can relate our data model to the dataset setup in [30], and thus by analyzing the feature learning terms for the linear model  $f_{lin}(\beta, \mathbf{x})$  similar to the analysis<sup>4</sup> in Appendix D.2, we obtain similar results as in Theorem D.1 in the early stage of GD training, but with an error of  $O(d^{-\Omega (\alpha_{nl})})$ .
Recall that our CNN model is  $f(\mathbf{W}, \mathbf{x}) = F_{+1}(\mathbf{W}_{+1}, \mathbf{x}) - F_{-1}(\mathbf{W}_{-1}, \mathbf{x})$ , where  $F_{+1}(\mathbf{W}_{+1}, \mathbf{x})$  and  $F_{-1}(\mathbf{W}_{-1}, \mathbf{x})$  are defined as follows:
$$
F _ {j} \left(\mathbf {W} _ {j}, \mathbf {x}\right) = \frac {1}{m} \sum_ {r = 1} ^ {m} \left[ \psi \left(\mathbf {w} _ {j, r} ^ {\top} \mathbf {x} _ {1}\right) + \psi \left(\mathbf {w} _ {j, r} ^ {\top} \mathbf {x} _ {2}\right) \right], j \in \{\pm 1 \}.
$$
We can cast this CNN model into an instance of the two-layer fully connected neural network defined at (21) by specifying the values of  $\{v_{r} = \pm \frac{1}{\sqrt{m}}\}$  and transforming the dataset as  $\left\{\sqrt{d}\left[y\cdot \mathrm{Rad}(\alpha)\cdot \mathbf{v}_1 + y\cdot \mathrm{Rad}(\beta)\cdot \mathbf{v}_2\right],\sqrt{d}\xi \right\}$ . Then by tuning the norms of  $\mathbf{v}_1,\mathbf{v}_2$  and  $\pmb{\xi}$ , we obtain a dataset that satisfies the input assumptions in [30]. Note that this cast drops the shared variable of our CNN model and thus might lead to a slightly different training dynamic. To fix such gap, we can leverage Proposition 6.4.1 in [29] for the early stage behavior of training a CNN model.
Based on the above ideas, to formalize the convergence results of the feature learning terms in the non-linear case, it remains to re-derive the analysis in Appendix D.2 based on  $\ell_2$  loss function, which follows a similar line of proofs and has a simpler dynamic.
Note that when  $\psi (x) = x$ , our CNN model can be viewed as a linear model with re-parameterized weight matrices. Thus, the discussion in Appendix D.2 can be viewed as studying the feature learning terms for a linear model with logistic loss function.
D.3 Proof for Theorem 4.2
Theorem D.5 (Restatement of Theorem 4.2). Consider training a CNN model with the same data as in Theorem 4.1, define
$$
\mathbf {c} (t) \triangleq \left[ C _ {I R M v I} ^ {1} (\mathbf {W}, t), C _ {I R M v I} ^ {2} (\mathbf {W}, t), \dots , C _ {I R M v I} ^ {| \mathcal {E} _ {t r} |} (\mathbf {W}, t) \right],
$$
and  $\lambda_0 = \lambda_{\min}(\mathbf{H}^\infty)$ , where we define
$$
\mathbf {H} _ {e, e ^ {\prime}} ^ {\infty} \triangleq \frac {1}{2 m n _ {e} n _ {e ^ {\prime}}} \sum_ {i = 1} ^ {n _ {e}} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (0), \mathbf {x} _ {1, i} ^ {e} \rangle) \mathbf {x} _ {1, i} ^ {e \top} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (0), \mathbf {x} _ {1, i ^ {\prime}} ^ {e ^ {\prime}} \rangle) \mathbf {x} _ {1, i ^ {\prime}} ^ {e ^ {\prime}}.
$$
Suppose that activation function is smooth,  $\psi'(0) \leq \beta$ ,  $|\psi'(x) - \psi'(x')| < \beta |x - x'|$  and Lipschitz  $|\psi(0)| < L$ ,  $|\psi(x) - \psi(x')| < L|x - x'|$ . Assume that dimension  $d = \Omega (\log(m / \delta))$ , network width  $m = \Omega(1 / \delta)$ , regularization factor  $\lambda \geq 1 / (\sigma_0 \sqrt{|\mathcal{E}_{tr}|^3})$ , noise variance  $\sigma_p = O(d^{-2})$ , weight initial scale  $\sigma_0 = O(\frac{|\mathcal{E}_{tr}|^{7/2} \beta^3 L}{d^{1/2} m^2 \lambda_0^2 \log(1 / \epsilon)})$ , then with probability at least  $1 - \delta$ , after training time  $T = \Omega\left(\frac{\log(1 / \epsilon)}{\eta \lambda \lambda_0}\right)$ , we have:
$$
\| \mathbf {c} (T) \| _ {2} \leq \epsilon , \quad \gamma_ {j, r} ^ {i n v} (T) = o (1), \quad \gamma_ {j, r} ^ {s p u} (T) = o (1).
$$
Before proving Theorem D.5, we first provide some useful lemmas as follows:
Lemma D.6 ([12]). Suppose that  $\delta > 0$  and  $d = \Omega (\log (4n / \delta))$ . Then with probability at least  $1 - \delta$ ,
$$
\sigma_ {p} ^ {2} d / 2 \leq \| \boldsymbol {\xi} _ {i} \| _ {2} ^ {2} \leq 3 \sigma_ {p} ^ {2} d / 2
$$
for all  $i, i' \in [n]$ .
Lemma D.7 ([12]). Suppose that  $d \geq \Omega(\log(mn/\delta))$ ,  $m = \Omega(\log(1/\delta))$ . Then with probability at least  $1 - \delta$ ,
$$
\left| \left\langle \mathbf {w} _ {j, r} ^ {(0)}, \mathbf {v} _ {1} \right\rangle \right| \leq \sqrt {2 \log (8 m / \delta)} \cdot \sigma_ {0} \| \mathbf {v} _ {1} \| _ {2},
$$
$$
| \langle \mathbf {w} _ {j, r} ^ {(0)}, \mathbf {v} _ {2} \rangle | \leq \sqrt {2 \log (8 m / \delta)} \cdot \sigma_ {0} \| \mathbf {v} _ {2} \| _ {2},
$$
$$
| \langle \mathbf {w} _ {j, r} ^ {(0)}, \boldsymbol {\xi} _ {i} \rangle | \leq 2 \sqrt {\log (8 m n / \delta)} \cdot \sigma_ {0} \sigma_ {p} \sqrt {d}
$$
Lemma D.8. Suppose that  $\delta > 0$  and  $d = \Omega (\log (4m / \delta))$ . Then with probability at least  $1 - \delta$ , for all  $r \in [m]$  and  $j \in \{-1,1\}$ , we have
$$
\sigma_ {0} ^ {2} d / 2 \leq \| \mathbf {w} _ {j, r} (0) \| _ {2} ^ {2} \leq 3 \sigma_ {0} ^ {2} d / 2.
$$
Proof of Lemma D.8. By Bernstein's inequality, with probability at least  $1 - \delta /(2m)$  we have
$$
\left| \left\| \mathbf {w} _ {j, r} (0) \right\| _ {2} ^ {2} - \sigma_ {0} ^ {2} d \right| = O \left(\sigma_ {0} ^ {2} \cdot \sqrt {d \log (4 m / \delta)}\right).
$$
Therefore, as long as  $d = \Omega (\log (4m / \delta))$ , we have
Proof of Theorem D.5. The proof is by induction method. First we show the gradient flow of weights by IRMv1 objective function (5):
$$
\begin{array}{l} \frac {d \mathbf {w} _ {j , r} (t)}{d t} = - \eta \cdot \nabla_ {\mathbf {w} _ {j, r}} L _ {\mathrm {I R M v 1}} (\mathbf {W} (t)) \\ = - \frac {\eta}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} \rangle) \cdot j \mathbf {v} _ {i} ^ {e} - \frac {\eta}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), \pmb {\xi} _ {i} \rangle) \cdot j y _ {i} ^ {e} \pmb {\xi} _ {i} \\ - \frac {2 \eta \lambda}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} C _ {\mathrm {I R M v 1}} ^ {e} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} \rangle) j y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} - \frac {2 \eta \lambda}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} C _ {\mathrm {I R M v 1}} ^ {e} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} \psi^ {\prime} (\langle \mathbb {w} _ {j, r} (t), \boldsymbol {\xi} _ {i} \rangle) j \boldsymbol {\xi} _ {i} \\ - \frac {2 \eta \lambda}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} C _ {\mathrm {I R M v 1}} ^ {e} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} \rangle) j \mathbf {v} _ {i} ^ {e} - \frac {2 \eta \lambda}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} C _ {\mathrm {I R M v 1}} ^ {e} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t) , \boldsymbol {\xi} _ {i} \rangle) j y _ {i} ^ {e} \boldsymbol {\xi} _ {i} \\ = - \frac {\eta}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} (1 + 2 \lambda C _ {\mathrm {I R M v 1}} ^ {e} (t)) \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} \rangle) \cdot j \mathbf {v} _ {i} ^ {e} \\ - \frac {\eta}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} (1 + 2 \lambda C _ {\mathrm {I R M v l}} ^ {e} (t)) \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), \boldsymbol {\xi} _ {i} \rangle) \cdot j y _ {i} ^ {e} \boldsymbol {\xi} _ {i} \\ - \frac {2 \eta \lambda}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} C _ {\mathrm {I R M v 1}} ^ {e} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \dot {y} _ {i} ^ {e} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} \rangle) j y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} - \frac {2 \eta \lambda}{n m} \sum_ {e \in \mathcal {E} _ {\mathrm {t r}}} C _ {\mathrm {I R M v 1}} ^ {e} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \dot {y} _ {i} ^ {e} \psi^ {\prime} (\langle \mathbb {w} _ {j, r} (t), \pmb {\xi} _ {i} \rangle) j \pmb {\xi} _ {i}, \\ \end{array}
$$
where  $C_{\mathrm{IRMv1}}^e = \frac{1}{n_e}\sum_{i = 1}^{n_e}\ell_i'^e\hat{y}_i^e y_i^e$  and  $\mathbf{v}_i^e = \operatorname {Rad}(\alpha)_i\cdot \mathbf{v}_1 + \operatorname {Rad}(\beta_e)_i\cdot \mathbf{v}_2$ . Note that  $\ell ''$  has the opposite sign to  $\ell^{\prime}$ .
Then we look at the dynamics of  $C_{\mathrm{IRMv1}}^e (t)$  according to the gradient flow update rule:
$$
\begin{array}{l} \frac {d C _ {\mathrm {I R M v 1}} ^ {e} (\mathbf {W} , t)}{d t} = \sum_ {j = \pm 1} \sum_ {r = 1} ^ {m} \left\langle \frac {\partial C _ {\mathrm {I R M v 1}} ^ {e} (\mathbf {W} , t)}{\partial \mathbf {w} _ {j , r} (t)}, \frac {d \mathbf {w} _ {j , r} (t)}{d t} \right\rangle \\ = \sum_ {e ^ {\prime}} 2 \lambda C _ {\mathrm {I R M v 1}} ^ {e ^ {\prime}} (\mathbf {W}, t) \sum_ {j} \sum_ {r = 1} ^ {m} \left\langle \frac {\partial C _ {\mathrm {I R M v 1}} ^ {e} (\mathbf {W} , t)}{\partial \mathbf {w} _ {j , r} (t)}, \frac {\partial C _ {\mathrm {I R M v 1}} ^ {e ^ {\prime}} (\mathbf {W} , t)}{\partial \mathbf {w} _ {j , r} (t)} \right\rangle \\ + \sum_ {j = \pm 1} \sum_ {r = 1} ^ {m} \left\langle \frac {\partial C _ {\mathrm {I R M v 1}} ^ {e} (\mathbf {W} , t)}{\partial \mathbf {w} _ {j , r} (t)}, \frac {\partial L (\mathbf {W} , t)}{\partial \mathbf {w} _ {j , r} (t)} \right\rangle \\ = 2 \lambda \sum_ {e ^ {\prime}} C _ {\mathrm {I R M v 1}} ^ {e ^ {\prime}} (\mathbf {W}, t) \cdot \mathbf {H} _ {e, e ^ {\prime}} (t) + \mathbf {g} _ {e} (t), \\ \end{array}
$$
where we define  $\mathbf{H}_{e,e^{\prime}}(t) = \sum_{j}\sum_{r = 1}^{m}\left\langle \frac{\partial C_{\mathrm{IRMv1}}^{\epsilon}(\mathbf{W},t)}{\partial\mathbf{w}_{j,r}(t)},\frac{\partial C_{\mathrm{IRMv1}}^{\epsilon^{\prime}}(\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_{\mathrm{IRMv1}}^{\epsilon}(\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}_{\mathrm{tr}}| \times |\mathcal{E}_{\mathrm{tr}}|$  matrix. We can write the dynamics of  $\mathbf{c}(t) = \left[C_{\mathrm{IRMv1}}^{1}(\mathbf{W},t),C_{\mathrm{IRMv1}}^{2}(\mathbf{W},t),\dots ,C_{\mathrm{IRMv1}}^{|\mathcal{E}_{\mathrm{tr}}|}(\mathbf{W},t)\right]$  in a compact way:
$$
\frac {d \mathbf {c} (t)}{d t} = 2 \lambda \cdot \mathbf {H} (t) \mathbf {c} (t) + \mathbf {g} (t). \tag {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_{\mathrm{IRMv1}}^e(\mathbf{W},t)}{\partial\mathbf{w}_{j,r}(t)}\in \mathbb{R}^d$  ..
$$
\begin{array}{l} \frac {\partial C _ {\mathrm {I R M v 1}} ^ {e} (\mathbf {W} (t))}{\partial \mathbf {w} _ {j , r} (t)} = \frac {1}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime} (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_ {i} ^ {\prime} (t) \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), \pmb {\xi} _ {i} \rangle) \cdot j y _ {i} ^ {e} \pmb {\xi} _ {i} \\ + \frac {1}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} \rangle) \cdot j y _ {i} ^ {e} \mathbf {v} _ {i} ^ {e} + \frac {1}{n _ {e} m} \sum_ {i = 1} ^ {n _ {e}} \ell_ {i} ^ {\prime \prime} \hat {y} _ {i} ^ {e} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (t), \pmb {\xi} _ {i} \rangle) \cdot j \pmb {\xi} _ {i}. \\ \end{array}
$$
When we consider non-linear activation function  $\psi (x)$ , the entry of matrix  $\mathbf{H}(t)$  can be computed as follows:
The matrix  $\mathbf{H}$  is composed of eight elements. In addition, we define
$$
\begin{array}{l} \mathbf {H} _ {e, e ^ {\prime}} ^ {1, \infty} = \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} _ {i} ^ {e} \rangle) j \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} - \frac {1}{2} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (0), \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle) j \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \right] \\ = \frac {1}{2 m n _ {e} n _ {e ^ {\prime}}} \sum_ {i = 1} ^ {n _ {e}} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (0), \mathbf {v} _ {i} ^ {e} \rangle) \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \psi^ {\prime} (\langle \mathbf {w} _ {j, r} (0), \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \rangle) \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}}. \\ \end{array}
$$
Then we can show that:
where we calculate each item as follows:
$$
\begin{array}{l} I _ {1} = \frac {2}{m n _ {e} n _ {e ^ {\prime}}} \left| \sum_ {i = 1} ^ {n _ {e}} (\psi^ {\prime} (t) - \psi^ {\prime} (0)) \ell_ {i} ^ {\prime} (t) \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \ell_ {i ^ {\prime}} ^ {\prime} (t) \psi^ {\prime} (t) \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \right| \\ \stackrel {(a)} {\leq} \frac {2}{m n _ {e} n _ {e ^ {\prime}}} \left| \sum_ {i = 1} ^ {n _ {e}} \beta \| \mathbf {w} _ {j, r} (t) - \mathbf {w} _ {j, r} (0) \| _ {2} \| \mathbf {v} _ {i} ^ {e} \| _ {2} \ell_ {i ^ {\prime}} ^ {\prime} (t) \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \ell_ {i ^ {\prime}} ^ {\prime} (t) \psi^ {\prime} (t) \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \right| \\ \stackrel {(b)} {\leq} \frac {2 \beta^ {2}}{m n _ {e} n _ {e ^ {\prime}}} \sum_ {i = 1} ^ {n _ {e}} \| \mathbf {w} _ {j, r} (t) - \mathbf {w} _ {j, r} (0) \| _ {2} \| \mathbf {v} _ {i} ^ {e} \| _ {2} ^ {2} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \| \mathbf {w} _ {j, r} (t) \| _ {2} \| \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \| _ {2} ^ {2} \\ \stackrel {(c)} {\leq} \frac {3 2 \beta^ {2} R (R + \frac {3}{2} \sigma_ {0} d)}{m}, \\ \end{array}
$$
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_i^e |\leq 1$  for all  $i\in n_e$  and  $e\in \mathcal{E}_{\mathrm{all}}$ ; inequality (c) is by the fact that  $\| \mathbf{v}_i^e\| _2\leq 2$  for all  $i\in n_e$  and  $e\in \mathcal{E}_{\mathrm{all}}$  and Lemma D.8.
Similarly, we calculate the upper bound for  $I_{2}$  as follows:
$$
\begin{array}{l} I _ {2} = \frac {2}{m n _ {e} n _ {e ^ {\prime}}} \left| \sum_ {i = 1} ^ {n _ {e}} \psi^ {\prime} (t) \ell_ {i} ^ {\prime} (t) \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} (\psi^ {\prime} (t) - \psi^ {\prime} (0)) \frac {1}{2} \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \right| \\ \leq \frac {3 2 \beta^ {2} R (R + \frac {3}{2} \sigma_ {0} d)}{m}. \\ \end{array}
$$
Next, we give the upper bound of  $I_3$ :
$$
\begin{array}{l} I _ {3} = \frac {2}{m n _ {e} n _ {e ^ {\prime}}} \left| \sum_ {i = 1} ^ {n _ {e}} \psi^ {\prime} (0) \ell_ {i} ^ {\prime} (t) \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \psi^ {\prime} (0) \left(\ell_ {i ^ {\prime}} ^ {\prime} (t) + \frac {1}{2}\right) \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \right| \\ \leq \frac {2}{m n _ {e} n _ {e ^ {\prime}}} \left| \sum_ {i = 1} ^ {n _ {e}} \beta \| \mathbf {w} _ {j, r} (0) \| _ {2} \| \mathbf {v} _ {i} ^ {e} \| _ {2} \ell_ {i} ^ {\prime} (t) \mathbf {v} _ {i} ^ {e ^ {\top}} \sum_ {i ^ {\prime} = 1} ^ {n _ {e ^ {\prime}}} \beta \| \mathbf {w} _ {j, r} (0) \| _ {2} \| \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \| _ {2} \left(\ell_ {i ^ {\prime}} ^ {\prime} (t) + \frac {1}{2}\right) \mathbf {v} _ {i ^ {\prime}} ^ {e ^ {\prime}} \right| \\ \leq \frac {6 4 \beta^ {2} L R \left(\frac {3}{2} \sigma_ {0} d\right) ^ {2}}{m}, \\ \end{array}
$$
where we have used  $\gamma$  which is defined as follows:
$$
\begin{array}{l} | \hat {y} _ {i} ^ {e} (t) | = \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| \\ \stackrel {(a)} {\leq} 2 L R, \\ \end{array}
$$
where inequality (a) is by the Lipschitz property of non-linear activation function and we have used the bound for  $\ell_i^\prime (t) + \frac{1}{2}$ :
$$
\begin{array}{l} \left| \ell_ {i} ^ {\prime} (t) + \frac {1}{2} \right| = \left| - \frac {\exp \left(- y _ {i} ^ {e} \cdot f (\mathbf {W} , \mathbf {x} _ {i} , t)\right)}{1 + \exp \left(- y _ {i} ^ {e} \cdot f (\mathbf {W} , \mathbf {x} _ {i} , t)\right)} + \frac {1}{2} \right| \\ = \left| \frac {1}{2} - \frac {1}{1 + \exp \left(y _ {i} ^ {e} \cdot f (\mathbf {W} , \mathbf {x} _ {i} , t)\right)} \right| \\ \leq \max  \left\{\left| \frac {1}{2} - \frac {1}{1 + \exp (2 L R)} \right|, \left| \frac {1}{2} - \frac {1}{1 + \exp (- 2 L R)} \right| \right\} \\ \leq \max  \left\{\left| \frac {1}{2} - \frac {1}{2 + \frac {7}{4} 2 L R} \right|, \left| \frac {1}{2} - \frac {1}{2 - 2 L R} \right| \right\} = \Theta (L R). \\ \end{array}
$$
and we provide the bound of  $\ell_i^{\prime \prime}(t) - \frac{1}{4}$ :