$\displaystyle\eta\biggr{\|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widehat{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}$
(83) $\displaystyle-$
$\displaystyle\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widehat{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}\biggr{\|}_{2}$
(84) $\displaystyle+$
$\displaystyle\eta\biggr{\|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widehat{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}$
(85) $\displaystyle-$
$\displaystyle\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\accentset{\circ}{\bm{v}}}(\bm{X}))\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{b}\biggr{]}\biggr{\|}_{2}.$
(86)

Remark that (86) is equal to zero because $\mathbb{E}[\mathbf{P}\bm{X}_{u}]=0$
for all $u\in[D].$

##### Summand 1: $\|\eqref{eq:fewpwpe}-\eqref{eq:fewefewpe}\|_{2}$.

Using the matrix Hoeffding inequality, we have with high probability,
$\|\eqref{eq:fewpwpe}-\eqref{eq:fewefewpe}\|_{2}\leq
16\sqrt{\log(d)\sum_{i=1}^{N}M_{i}^{2}},$ where
$\Big{\|}\frac{\eta
Dy[i]}{N}\mathfrak{S}(-y[i]F_{\widehat{\bm{v}}}(\bm{X}[i]))\sum_{m=1}^{D}\left(D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{r}[i]\rangle\right)^{p-1}\sum_{u=1}^{D}\widehat{S}_{m,u}^{(t)}\mathbf{P}\bm{X}_{u}[i]\Big{\|}_{2}^{2}\leq
M_{i}^{2}.$ Induction Hypothesis B.1, Lemma D.11 and
$\|\mathbf{P}\bm{X}_{u}[i]\|_{2}\leq\sigma^{2}d\log(d)$ imply
$M_{i}^{2}\leq\frac{\eta^{2}\mathrm{poly}(D)}{N^{2}}.$ We deduce that
$\|\eqref{eq:fewpwpe}-\eqref{eq:fewefewpe}\|_{2}\leq\eta\frac{\mathrm{poly}(D)}{\sqrt{N}}.$

##### Summand 2: $\|\eqref{eq:woesddsfnfeo}-\eqref{eq:woenfeofrwed}\|_{2}$.

The function $x\mapsto x^{p-1}$ is $(p-1)M^{p-2}$ Lipschitz on a bounded
domain $[0,M]$. We apply this property and
$(p-1)\max_{m\in[D]}\big{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{r}\rangle\big{)}^{p-2}\leq(\widehat{\alpha}^{(t)})^{p-2}\mathrm{poly}(D)$
to get:

$\displaystyle\|\eqref{eq:woesddsfnfeo}-\eqref{eq:woenfeofrwed}\|_{2}$
$\displaystyle\leq$
$\displaystyle\eta(\widehat{\alpha}^{(t)})^{p-2}\mathrm{poly}(D)\mathbb{E}\Biggr{[}\sum_{m=1}^{D}\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}|\langle\widehat{\bm{v}}^{(t)}-\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle|\cdot\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\|\mathbf{P}\bm{X}_{b}\|_{2}\Biggr{]}$
$\displaystyle\leq$
$\displaystyle\eta(\widehat{\alpha}^{(t)})^{p-2}\mathrm{poly}(D)\mathbb{E}\Biggr{[}\sum_{m=1}^{D}\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\varepsilon_{\bm{v}}^{(t)}|\langle\bm{u}^{(t)},\bm{X}_{r}\rangle|\cdot\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\|\mathbf{P}\bm{X}_{b}\|_{2}\Biggr{]}.$
(87)

With high probability,
$\|\bm{\xi}_{b}\|_{2}\leq\sigma\sqrt{\log(d)}\leq\sqrt{\log(d)/d}$,
$|\langle\bm{u}^{(t)},\bm{\xi}_{r}\rangle|\leq\sqrt{\log(d)}\sigma$. We thus
get:

$\displaystyle\|\eqref{eq:woesddsfnfeo}-\eqref{eq:woenfeofrwed}\|_{2}$
$\displaystyle\leq\eta(\widehat{\alpha}^{(t)})^{p-2}\frac{\mathrm{poly}(D)}{\sqrt{d}}\varepsilon_{\bm{v}}^{(t)}.$
(88)

##### Summand 3: $\|\eqref{eq:woenfeo}-\eqref{eq:ewjfnweonj}\|_{2}$.

We have:

$\displaystyle\|\eqref{eq:woenfeo}-\eqref{eq:ewjfnweonj}\|_{2}$ (89)
$\displaystyle\leq$
$\displaystyle\eta\Biggr{\|}\mathbb{E}\Biggr{[}Dy\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{b=1}^{D}\widehat{S}_{m,b}^{(t)}\mathbf{P}\bm{X}_{u}\cdot$
$\displaystyle\biggr{[}\mathfrak{S}\biggr{(}-y\sum_{a=1}^{D}\big{(}D\sum_{a^{\prime}=1}^{D}\widehat{S}_{a,a^{\prime}}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{a^{\prime}}\rangle\big{)}^{p}\biggr{)}-\mathfrak{S}\biggr{(}-y\sum_{a=1}^{D}\big{(}D\sum_{a^{\prime}=1}^{D}\widehat{S}_{a,a^{\prime}}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{a^{\prime}}\rangle\big{)}^{p}\biggr{)}\biggr{]}\Biggr{]}\Biggr{\|}_{2}.$

We apply Lemma E.12 to bound the local change of the sigmoid in (89) which
yields:

$\displaystyle\|\eqref{eq:woenfeo}-\eqref{eq:ewjfnweonj}\|_{2}$
$\displaystyle\leq$ $\displaystyle\Theta(\eta
D)\mathbb{E}\biggr{[}\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{a=1}^{D}\biggr{|}D\sum_{a^{\prime}=1}^{D}\widehat{S}_{a,a^{\prime}}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)}-\widehat{\bm{v}}^{(t)},\bm{X}_{a^{\prime}}\rangle\biggr{|}\cdot\sum_{b=1}^{D}\widehat{S}_{m,u}^{(t)}\|\mathbf{P}\bm{X}_{b}\|_{2}\biggr{]}$
$\displaystyle\leq$ $\displaystyle\Theta(\eta
D^{2})\mathbb{E}\biggr{[}\sum_{m=1}^{D}\biggr{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(t)}\langle\accentset{\circ}{\bm{v}}^{(t)},\bm{X}_{r}\rangle\biggr{)}^{p-1}\sum_{a=1}^{D}\sum_{a^{\prime}=1}^{D}\widehat{S}_{a,a^{\prime}}^{(t)}\varepsilon_{\bm{v}}^{(t)}|\langle\bm{u}^{(t)},\bm{X}_{a^{\prime}}\rangle|\sum_{b=1}^{D}\widehat{S}_{m,u}^{(t)}\|\mathbf{P}\bm{X}_{b}\|_{2}\biggr{]}.$
(90)

We apply Induction Hypothesis B.1 to bound the softmax terms in (90). Besides,
with high probability, we have
$\|\bm{\xi}_{b}\|_{2}\leq\sigma\sqrt{d\log(d)}$,
$|\langle\bm{u}^{(t)},\bm{\xi}_{a^{\prime}}\rangle|\leq\sqrt{\log(d)}\sigma$.
Thus, we have:

$\displaystyle\|\eqref{eq:woenfeo}-\eqref{eq:ewjfnweonj}\|_{2}$
$\displaystyle\leq\eta\cdot\mathrm{poly}(D)(\widehat{\alpha}^{(t)})^{p-1}\varepsilon_{\bm{v}}^{(t)}\sigma\leq\eta\frac{\mathrm{poly}(D)}{\sqrt{d}}(\widehat{\alpha}^{(t)})^{p-1}\varepsilon_{\bm{v}}^{(t)}.$
(91)

We combine the bounds on the three summands to obtain the recursion of
$\varepsilon_{\bm{v}}^{(t)}$. ∎

We now prove Lemma E.11 that gives the final bound on
$\varepsilon_{\bm{v}}^{(t)}$ for $t\leq T.$

######  Lemma E.11.

For all $t\leq T$, $\varepsilon_{\bm{v}}^{(t)}\leq
O(\frac{\mathrm{poly}(d)}{\sqrt{N}}).$ By setting $N=\mathrm{poly}(d)$,
$\varepsilon_{\bm{v}}^{(t)}\leq\frac{1}{\mathrm{poly}(d)}.$

###### Proof of Lemma E.11.

We bound $\varepsilon_{\bm{v}}^{(t)}$ in the following two regimes:
$t\in[\mathscr{T},\mathscr{T}+\widecheck{\mathcal{T}}_{0}]$ and
$t\in[\mathscr{T}+\widecheck{\mathcal{T}}_{0},T].$

##### First phase:
$t\in[\mathscr{T},\mathscr{T}+\widecheck{\mathcal{T}}_{0}]$.

Unraveling Lemma E.10 for
$t=\mathscr{T},\dots,\mathscr{T}+\widecheck{\mathcal{T}}_{0}$ leads to:

$\displaystyle\varepsilon_{\bm{v}}^{(t)}$
$\displaystyle\leq\Big{[}\varepsilon_{\bm{v}}^{(\mathscr{T})}+\eta\zeta\widecheck{\mathcal{T}}_{0}\Big{]}\prod_{\tau=\mathscr{T}}^{\mathscr{T}+\widecheck{\mathcal{T}}_{0}}\biggr{(}1+\eta(\widehat{\alpha}^{(\tau)})^{p-1}\frac{\mathrm{poly}(D)}{\sqrt{d}}\biggr{)}$
(92)
$\displaystyle\leq[\varepsilon_{\bm{v}}^{(\mathscr{T})}+\eta\zeta\widecheck{\mathcal{T}}_{0}]\prod_{\tau=0}^{\widecheck{\mathcal{T}}_{0}-1}\biggr{(}1+\eta(\widehat{\alpha}^{(\tau)})^{p-1}\frac{\mathrm{poly}(D)}{\sqrt{d}}\biggr{)}.$

Lemma E.20 provides the update of $\widehat{\alpha}^{(t)}$ during this time
phase. We thus apply Lemma K.2 to bound the product term in (92).

$\displaystyle\prod_{\tau=0}^{\widecheck{\mathcal{T}}_{0}-1}\biggr{(}1+\eta(\widehat{\alpha}^{(\tau)})^{p-1}\frac{\mathrm{poly}(D)}{\sqrt{d}}\biggr{)}$
$\displaystyle\leq\biggr{(}1+\frac{\Theta(1)}{C^{2(p-2)}\lambda_{0}^{p-2}}\biggr{)}^{\frac{\eta\mathrm{poly}(D)}{\sqrt{d}}}\leq
O(1).$ (93)

Plugging (93) in (92) yields a bound on
$\varepsilon_{\bm{v}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{0})}.$

$\displaystyle\varepsilon_{\bm{v}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{0})}$
$\displaystyle\leq
O(1)\big{(}\varepsilon_{\bm{v}}^{(\mathscr{T})}+\eta\zeta\widecheck{\mathcal{T}}_{0}\big{)}.$
(94)

##### Second phase: $t\in(\mathscr{T}+\widecheck{\mathcal{T}}_{0},T]$.

Lemma E.23 shows that $\widehat{\alpha}^{(t)}$ gets updated until
$t=\widecheck{\mathcal{T}}_{2}<T.$ Therefore, we have
$\varepsilon_{\bm{v}}^{(T)}=\varepsilon_{\bm{v}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{2})}.$
Unraveling Lemma E.10 for
$t=\mathscr{T}+\widecheck{\mathcal{T}}_{0},\dots,\mathscr{T}+\widecheck{\mathcal{T}}_{2}$
and using
$\widehat{\alpha}^{(t)}\leq\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}$ for
$t\leq T$ leads to:

$\displaystyle\varepsilon_{\bm{v}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{2})}$
$\displaystyle\leq\biggr{(}1+\frac{\eta\mathrm{polylog}(d)}{\lambda_{0}^{p-1}}\frac{\mathrm{poly}(D)}{\sqrt{d}}\biggr{)}^{\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}}\varepsilon_{\bm{v}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{0})}+\eta\zeta\frac{\biggr{(}1+\frac{\eta\mathrm{polylog}(d)}{\lambda_{0}^{p-1}}\frac{\mathrm{poly}(D)}{\sqrt{d}}\biggr{)}^{\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}}-1}{\frac{\eta\mathrm{polylog}(d)}{\lambda_{0}^{p-1}}\frac{\mathrm{poly}(D)}{\sqrt{d}}}.$
(95)

Since
$\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}\leq\frac{\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}$, we have
$\Big{(}1+\frac{\eta\mathrm{polylog}(d)}{\lambda_{0}^{p-1}}\frac{\mathrm{poly}(D)}{\sqrt{d}}\Big{)}^{\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}}\leq
O(1).$ Simplifying (95) yields:

$\displaystyle\varepsilon_{\bm{v}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{2})}$
$\displaystyle\leq O\biggr{(}\frac{\mathrm{poly}(d)}{\sqrt{N}}\biggr{)}.$ (96)

(96) implies that we need $N=\mathrm{poly}(d)$ samples to have
$\varepsilon_{\bm{v}}^{(T)}\leq 1/\mathrm{poly}(d).$ ∎

####  E.2.3 Auxiliary lemmas

In this section, we prove the Lipschitzness of the function appearing in the
proof of Lemma E.10.

######  Lemma E.12.

Let $\psi\colon\mathbb{R}^{d}\rightarrow\mathbb{R}$ defined as
$\psi(\bm{x}):=\mathfrak{S}(-\sum_{m=1}^{D}x_{m}^{p})$ and $p\geq 3$ be an odd
integer. Assume that $x_{m}^{p}\geq-1/D$. Then, $\psi$ is $p$-Lipschitz i.e.
for all $\bm{x},\bm{y}\in\mathbb{R}^{d},$ $|\psi(\bm{x})-\psi(\bm{y})|\leq
p\|\bm{x}-\bm{y}\|_{1}.$

###### Proof of Lemma E.12.

Let $l\in[D].$ The derivative of $\psi$ with respect to a variable $x_{l}$ is:

$\displaystyle\left|\frac{\partial\psi}{\partial x_{l}}\right|$
$\displaystyle=p\frac{x_{l}^{p-1}\exp(\sum_{m=1}^{D}x_{m}^{p})}{(1+\exp(\sum_{m=1}^{D}x_{m}^{p}))^{2}}\leq\frac{px_{l}^{p-1}}{1+\exp(\sum_{m=1}^{D}x_{m}^{p})}.$
(97)

(97) implies a bound on $\|\nabla\psi(\bm{x})\|_{1}$. Indeed, since
$\sum_{m=1}^{D}x_{m}^{p}\geq-1$, we have:

$\displaystyle\left\|\nabla\psi(\bm{x})\right\|_{1}$ $\displaystyle\leq
p\frac{\sum_{m=1}^{D}(yx_{m})^{p-1}}{1+\exp(\sum_{m=1}^{D}(yx_{m})^{p})}\leq
p.$ (98)

(98) shows that $\psi$ is $p$-Lipschitz. ∎

####  E.2.4 $\Delta_{\bm{A}}^{(t)}$ stays small during the learning process

Here, we bound the gap in attention coefficients between the realistic and
semi-idealized cases.

######  Lemma E.13.

Assume that we run GD on the empirical risk (E) for $T$ iterations with
parameters set as in Parametrization 3.1 and the number of samples is
$N=\mathrm{poly}(d).$ Then, the attention matrix in the realistic case is very
close to the semi-idealized one i.e. for $t\in[\mathscr{T},T]$,

$\displaystyle\Delta_{\bm{A}}^{(t)}:=\max_{i\neq
j}|\widehat{A}_{i,j}^{(t)}-A_{i,j}^{(t)}|\leq\frac{1}{\mathrm{poly}(d)}.$

We now detail the steps to prove Lemma E.13. We first provide the recursion
that $\Delta_{\bm{A}}^{(t)}$ satisfies.

######  Lemma E.14.

Assume that we run GD on the empirical risk (E) for $T$ iterations with
parameters set as in Parametrization 3.1. Then, the discrepancy
$\Delta_{\bm{A}}$ satisfies for $t\in(\mathscr{T},T]$,

$\displaystyle\Delta_{\bm{A}}^{(t+1)}$ $\displaystyle\leq\big{(}1+\eta
R^{(t)}\big{)}\Delta_{\bm{A}}^{(t)}+\eta\mathrm{poly}(D)\sigma\varepsilon_{\bm{v}}^{(t)}+\eta\zeta,$

where
$\Delta_{\bm{A}}^{(\mathscr{T})}\leq|\widehat{A}_{i,j}^{(\mathscr{T})}|$,
$R^{(t)}=O(C)(G^{(t)})^{p-1}\widecheck{\Gamma}^{(t)}(\widehat{\alpha}^{(t)})^{p}$
and $\zeta=\frac{\mathrm{poly}(D)}{\sqrt{N}}.$

###### Proof of Lemma E.14.

In this proof, we maintain the hypothesis that $\Delta_{\bm{A}}^{(t)}$ is
small. We will eventually prove this statement in Lemma E.15. Let $a,b\in[D]$
such that $a\neq b$. Using GD,
$\widehat{A}_{a,b}^{(t+1)}-\widecheck{A}_{a,b}^{(t+1-\mathscr{T})}$ satisfies:

$\displaystyle\big{|}\widehat{A}_{a,b}^{(t+1)}-\widecheck{A}_{a,b}^{(t+1-\mathscr{T})}\big{|}\leq\big{|}\widehat{A}_{a,b}^{(t)}-\widecheck{A}_{a,b}^{(t-\mathscr{T})}\big{|}$
$\displaystyle+\eta\biggr{|}\frac{D}{N}\sum_{i=1}^{N}y[i]\mathfrak{S}\big{(}-y[i]F_{\widehat{\bm{v}},\widehat{\bm{A}}}(\bm{X}[i])\big{)}\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{c}[i]\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{b}[i]-\bm{X}_{r}[i]\rangle$
(99) $\displaystyle-$
$\displaystyle\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widehat{\bm{v}},\widehat{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}$
(100)
$\displaystyle+\eta\biggr{|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widehat{\bm{v}},\widehat{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}$
(101) $\displaystyle-$
$\displaystyle\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widehat{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t-\mathscr{T})},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t-\mathscr{T})},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}$
(102)
$\displaystyle+D\eta\biggr{|}\mathbb{E}\biggr{[}y\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widehat{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t-\mathscr{T})},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}$
(103)
$\displaystyle-\mathbb{E}\biggr{[}y\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t-\mathscr{T})},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widecheck{S}_{a,b}^{(t-\mathscr{T})}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t-\mathscr{T})},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}.$
(104)

We now bound the three summands above.

##### Summand 1: $|\eqref{eq:wekenew}-\eqref{eq:oewjwdw}|$.

We apply the Hoeffding inequality. With high probability, we have:
$|\eqref{eq:wekenew}-\eqref{eq:oewjwdw}|\leq\eta\frac{\mathrm{poly}(D)}{\sqrt{N}}$.

##### Summand 2: $|\eqref{eq:odjw}-\eqref{eq:ofjnewend}|$.

Since $\varepsilon_{\bm{v}}^{(t)}$ is small (Lemma E.11), we can show that:

$\displaystyle|\eqref{eq:odjw}-\eqref{eq:ofjnewend}|\leq\eta
D\mathbb{E}\biggr{[}\Big{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\Big{)}^{p-1}\widehat{S}_{a,b}^{(t)}\cdot\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}|\langle\bm{u}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\biggr{]}\varepsilon_{\bm{v}}^{(t)}.$
(105)

With high probability, we have
$\langle\bm{u}^{(t)},\bm{\xi}_{b}\rangle\leq\sigma\sqrt{\log(d)}$. Using this
fact along with
$D\Big{|}\Big{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\Big{)}^{p-1}\widehat{S}_{a,b}^{(t)}\cdot\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\Big{|}\leq\mathrm{poly}(D)$, we further bound (105)
as:

$\displaystyle|\eqref{eq:odjw}-\eqref{eq:ofjnewend}|$
$\displaystyle\leq\eta\mathrm{poly}(D)\varepsilon_{\bm{v}}^{(t)}\sigma\sqrt{\log(d)}.$
(106)

##### Summand 3: $|\eqref{eq:ofnewow}-\eqref{eq:fijoffw}|$.

We have the following decomposition.

$\displaystyle|\eqref{eq:ofnewow}-\eqref{eq:fijoffw}|$
$\displaystyle\leq\eta\biggr{|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widehat{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}$
(107)
$\displaystyle-\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}$
(108)
$\displaystyle+\eta\biggr{|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}$
(109)
$\displaystyle-\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}$
(110)
$\displaystyle+\eta\biggr{|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}$
(111)
$\displaystyle-\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}S_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}$
(112)
$\displaystyle+\eta\biggr{|}\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}$
(113)
$\displaystyle-\mathbb{E}\biggr{[}Dy\mathfrak{S}(-yF_{\widecheck{\bm{v}},\widecheck{\bm{A}}}(\bm{X}))\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widecheck{S}_{a,b}^{(t-\mathscr{T})}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\biggr{]}\biggr{|}.$
(114)

We need to distinguish two sub-cases: $a,b\in\mathcal{S}_{\ell}$ and
$a\in\mathcal{S}_{\ell}$, $b\in\mathcal{S}_{m}$ with $\ell\neq m.$

#####  Subcase 1: $a,b\in\mathcal{S}_{\ell}$.

The proof of Lemma G.1 highlights that when $a,b\in\mathcal{S}_{\ell}$ the
event with largest gradient is event a: "$\ell=\ell(\pi(\bm{X}))$" which
happens with probability $1/L.$ Therefore, to simplify the calculations, we
will only take into account this event. We first bound
$|\eqref{eq:efrpefr}-\eqref{eq:jwejwfej}|$. We successively apply Lemma E.12
(Lipschitzness of sigmoid) and Lemma E.17 (Lipschitzness of softmax) and get:

$\displaystyle|\eqref{eq:efrpefr}-\eqref{eq:jwejwfej}|$ $\displaystyle\leq$
$\displaystyle\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{c=1}^{D}\sum_{c^{\prime}=1}^{D}|\widecheck{S}_{c,c^{\prime}}^{(t-\mathscr{T})}-\widehat{S}_{c,c^{\prime}}^{(t)}|\cdot|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c^{\prime}}\rangle|\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\;\biggr{|}\;\mathbf{a}\biggr{]}$
$\displaystyle\leq$ $\displaystyle\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{c=1}^{D}\sum_{c^{\prime}=1}^{D}\widecheck{S}_{c,c^{\prime}}^{(t-\mathscr{T})}|e^{2\Delta_{\bm{A}}^{(t)}}-1||\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c^{\prime}}\rangle|\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\;\biggr{|}\;\mathbf{a}\biggr{]}$
$\displaystyle\leq$ $\displaystyle\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{c=1}^{D}\sum_{c^{\prime}=1}^{D}\widecheck{S}_{c,c^{\prime}}^{(t-\mathscr{T})}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c^{\prime}}\rangle|\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\;\biggr{|}\;\mathbf{a}\biggr{]}\Delta_{\bm{A}}^{(t)}$
(115) $\displaystyle\ \leq$
$\displaystyle\frac{\Theta(\eta\widehat{\alpha}^{(t)}\lambda_{0}D^{2})}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widehat{S}_{a,c}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\;\biggr{|}\;\mathbf{a}\biggr{]}\Delta_{\bm{A}}^{(t)}.$
(116)

where we used $e^{\Delta_{\bm{A}}^{(t)}}-1\leq 2\Delta_{\bm{A}}^{(t)}$ in
(115) and Induction Hypothesis B.1 in (116). Using Lipschitz inequalities, we
can further expand (116) as a function of the coefficients from
$\widecheck{S}^{(t)}$ and $\Delta_{\bm{A}}^{(t)}$. However,
$\Delta_{\bm{A}}^{(t)}$ is small and we only want terms of order 1 in
$\Delta_{\bm{A}}^{(t)}$ in (116). Therefore, the only term of order 1 that
remains is:

$\displaystyle|\eqref{eq:efrpefr}-\eqref{eq:jwejwfej}|$ (117)
$\displaystyle\leq\frac{\Theta(\eta\widehat{\alpha}^{(t)}\lambda_{0}D^{2})}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widecheck{S}_{a,b}^{(t-\mathscr{T})}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\;\biggr{|}\;\mathbf{a}\biggr{]}\Delta_{\bm{A}}^{(t)}.$

Bounding the expectation in (LABEL:eq:wefcerffer) as in the proof of Lemma G.1
yields $|\eqref{eq:efrpefr}-\eqref{eq:jwejwfej}|\leq\eta
R^{(t)}\Delta_{\bm{A}}^{(t)}$. We now bound
$|\eqref{eq:ojdeojedoj}-\eqref{eq:weejiwjwei}|$. We therefore apply (Lemma
E.16) and get:

$\displaystyle|\eqref{eq:ojdeojedoj}-\eqref{eq:weejiwjwei}|$ (118)
$\displaystyle\leq$ $\displaystyle\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}|e^{2(p-1)\Delta_{\bm{A}}^{(t)}}-1|\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\;\biggr{|}\;\mathbf{a}\biggr{]}$
(119) $\displaystyle\leq$ $\displaystyle\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widehat{S}_{a,r}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle\;\biggr{|}\;\mathbf{a}\biggr{]}\Delta_{\bm{A}}^{(t)}.$
(120)

where we used $e^{2(p-1)\Delta_{\bm{A}}^{(t)}}-1\leq
4(p-1)\Delta_{\bm{A}}^{(t)}$ in (120). We can further expand (120), keep the
terms of first order in $\Delta_{\bm{A}}$ and get
$|\eqref{eq:kmnkknk}-\eqref{eq:evfejr}|\leq\eta R^{(t)}\Delta_{\bm{A}}^{(t)}.$

We now bound $|\eqref{eq:kmnkknk}-\eqref{eq:evfejr}|$. Using the Lipschitz
property of the softmax (Lemma E.17), we have:

$\displaystyle|\eqref{eq:kmnkknk}-\eqref{eq:evfejr}|$ $\displaystyle\leq$
$\displaystyle\frac{\eta
D}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}|e^{2\Delta_{\bm{A}}^{(t)}}-1|\cdot\widecheck{S}_{a,r}^{(t-\mathscr{T})}\cdot|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\Big{|}\textbf{a}\biggr{]}$
$\displaystyle\leq$ $\displaystyle\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widehat{S}_{a,b}^{(t)}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\Big{|}\textbf{a}\biggr{]}\Delta_{\bm{A}}^{(t)},$
(121)

where we used $|e^{2\Delta_{\bm{A}}^{(t)}}-1|\leq 4\Delta_{\bm{A}}^{(t)}$ in
(121). Using the same arguments as above, we obtain
$|\eqref{eq:kmnkknk}-\eqref{eq:evfejr}|\leq\eta R^{(t)}\Delta_{\bm{A}}^{(t)}.$

The bound on $|\eqref{eq:rfjrf}-\eqref{eq:ofeojeo}|$ can be derived as above.
We again use the Lipschitz property of softmax (Lemma E.17) which leads to

$\displaystyle|\eqref{eq:rfjrf}-\eqref{eq:ofeojeo}|$
$\displaystyle\leq\frac{\Theta(\eta
D)}{L}\mathbb{E}\biggr{[}\biggr{(}D\sum_{c=1}^{D}\widecheck{S}_{a,c}^{(t-\mathscr{T})}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{c}\rangle\biggr{)}^{p-1}\widecheck{S}_{a,b}^{(t-\mathscr{T})}\sum_{r\neq
b}\widecheck{S}_{a,r}^{(t-\mathscr{T})}|\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{b}-\bm{X}_{r}\rangle|\Big{|}\textbf{a}\biggr{]}\Delta_{\bm{A}}^{(t)}$
$\displaystyle\leq\eta R^{(t)}\Delta_{\bm{A}}^{(t)}.$ (122)

#####  Subcase 2: $a\in\mathcal{S}_{\ell}$ and $b\in\mathcal{S}_{m}$ with
$\ell\neq m$.

The proof is analogous to the Subcase 1. We only take into account event a:
"$\ell=\ell(\pi(\bm{X}))$ and $\delta_{j}=0$" and Event e:
"$\ell,m\neq\ell(\pi(\bm{X}))$ and $\delta_{j}=0$ and $\delta_{s}=0$ and show
that
$|\eqref{eq:ofnewow}-\eqref{eq:fijoffw}|\leq\eta\tilde{R}^{(t)}\Delta_{\bm{A}}^{(t)}$
where
$\tilde{R}^{(t)}:=O(\widehat{\alpha}^{(t)})\left(D\widehat{\alpha}^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\right)^{p-1}\Xi^{(t)}\frac{\lambda_{0}}{D}+O(\widehat{\alpha}^{(t)})\left(D\widehat{\alpha}^{(t)}qD\log(d)\Xi^{(t)}\right)^{p-1}\Xi^{(t)}qD\log(d)$.

#####  Putting all the pieces together.

Given the value of the parameters, we know that $\tilde{R}^{(t)}\leq R^{(t)}$
for all $t\in[T].$ Therefore, Summand 3 is bounded as:

$\displaystyle|\eqref{eq:ofnewow}-\eqref{eq:fijoffw}|\leq\eta
R^{(t)}\Delta_{\bm{A}}^{(t)}.$ (123)

##### Conclusion.

Plugging the bounds on Summands 1, 2 and 3 in the original decomposition of
$\big{|}\widehat{A}_{a,b}^{(t+1)}-\widecheck{A}_{a,b}^{(t+1-\mathscr{T})}\big{|}$
yields the bound on $\Delta_{\bm{A}}^{(t)}$. The second part of the lemma is
obtained using Lemma E.15.

∎

######  Lemma E.15.

Let $N=\mathrm{poly}(d)$. Then, for all $t\leq T$,
$\Delta_{\bm{A}}^{(t)}\leq\frac{1}{\mathrm{poly}(d)}.$

###### Proof of Lemma E.15.

Let $\mathcal{E}_{\bm{v}}>0$ such that
$\varepsilon_{\bm{v}}^{(t)}\leq\mathcal{E}_{\bm{v}}$ for $t\in[T]$ – we proved
the existence of $\mathcal{E}_{\bm{v}}$ in Lemma E.11. We bound
$\Delta_{\bm{A}}^{(t)}$ when
$t\in[\mathscr{T},\mathscr{T}+\widecheck{\mathcal{T}}_{0}]$ and
$t\in[\mathscr{T}+\widecheck{\mathcal{T}}_{0},T].$

##### First phase:
$t\in[\mathscr{T},\mathscr{T}+\widecheck{\mathcal{T}}_{0}]$.

Unraveling Lemma E.14 for
$t=\mathscr{T},\dots,\mathscr{T}+\widecheck{\mathcal{T}}_{0}$ leads to:

$\displaystyle\Delta_{\bm{A}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{0})}$
(124)
$\displaystyle\leq\Big{(}\Delta_{\bm{A}}^{(\mathscr{T})}+\eta\widecheck{\mathcal{T}}_{0}\big{(}\mathcal{E}_{\bm{v}}\mathrm{poly}(D)\sigma+\zeta)\Big{)}\prod_{\tau=\mathscr{T}}^{\mathscr{T}+\widecheck{\mathcal{T}}_{0}-1}\biggr{(}1+\frac{\eta
D(\widehat{\alpha}^{(\tau)})^{p}}{L}O\big{(}D(\widecheck{\Lambda}^{(t)}+(C-1)\widecheck{\Gamma}^{(t)})\big{)}^{p-1}\widecheck{\Gamma}^{(t)}\biggr{)}.$

We now apply Induction Hypothesis B.1 to simplify (LABEL:eq:jfejfe) and get:

$\displaystyle\Delta_{\bm{A}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{0})}$
$\displaystyle\leq\Big{(}\Delta_{\bm{A}}^{(\mathscr{T})}+\eta\widecheck{\mathcal{T}}_{0}\big{(}\mathcal{E}_{\bm{v}}\mathrm{poly}(D)\sigma+\zeta)\Big{)}\prod_{\tau=\mathscr{T}}^{\mathscr{T}+\widecheck{\mathcal{T}}_{0}-1}\biggr{(}1+\frac{O(\eta\lambda_{0}^{p})}{L}(\widehat{\alpha}^{(t)})^{p}\biggr{)}.$
(125)

We then apply Lemma K.2 to bound the product term in (125). We obtain:

$\displaystyle\Delta_{\bm{A}}^{(\mathscr{T}+\mathcal{T}_{0})}$
$\displaystyle\leq\Big{(}\Delta_{\bm{A}}^{(\mathscr{T})}+\eta\widecheck{\mathcal{T}}_{0}\big{(}\mathcal{E}_{\bm{v}}\mathrm{poly}(D)\sigma+\zeta)\Big{)}\biggr{(}1+\frac{\Theta(1)}{C^{2(p-2)}\lambda_{0}^{p-2}}\biggr{)}^{\frac{\eta\lambda_{0}^{2p}}{L}}$
$\displaystyle\leq
O\Big{(}\Delta_{\bm{A}}^{(\mathscr{T})}+\eta\widecheck{\mathcal{T}}_{0}\big{(}\mathcal{E}_{\bm{v}}\mathrm{poly}(D)\sigma+\zeta)\Big{)}.$
(126)

##### Second phase: $t\in[\mathscr{T}+\widecheck{\mathcal{T}}_{0},T]$.

Unraveling Lemma E.14 for
$t=\mathscr{T}+\widecheck{\mathcal{T}}_{0},\dots,\mathscr{T}+\widecheck{\mathcal{T}}_{2}$
and and using
$\widehat{\alpha}^{(t)}\leq\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}$ for
$t\leq T$ leads to:

$\displaystyle\Delta_{\bm{A}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{2})}$
$\displaystyle\leq\biggr{(}1+\frac{\eta\mathrm{polylog}(d)}{L}\biggr{)}^{\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}}\Big{(}\Delta_{\bm{A}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{0})}+L\cdot\mathrm{polylog}(d)\big{(}\mathcal{E}_{v}\mathrm{poly}(D)\sigma+\zeta\big{)}\Big{)}.$
(127)

Using
$\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}\leq\frac{\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}$, we have
$\Big{(}1+\frac{\eta\mathrm{polylog}(d)}{L}\Big{)}^{\widecheck{\mathcal{T}}_{2}-\widecheck{\mathcal{T}}_{0}}\leq
O(1)$. We thus bound (126) as:

$\displaystyle\Delta_{\bm{A}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{2})}$
$\displaystyle\leq
O(\Delta_{\bm{A}}^{(\mathscr{T})})+\mathrm{polylog}(d)\biggr{[}\mathcal{E}_{\bm{v}}\sigma+\frac{1}{\sqrt{N}}\Biggr{]}\Biggr{[}\frac{1}{C(\widehat{\alpha}^{(\mathscr{T})})^{p-1}e^{\mathrm{polyloglog}(d)}}+L\biggr{]}\leq
O\Big{(}\frac{\nu}{\sqrt{N}}\Big{)}.$ (128)

We deduce that setting $N=\mathrm{poly}(d)$ yields
$\Delta_{\bm{A}}^{(T)}=\Delta_{\bm{A}}^{(\mathscr{T}+\widecheck{\mathcal{T}}_{2})}\leq
1/\mathrm{poly}(d).$ ∎

####  E.2.5 Auxiliary lemmas

######  Lemma E.16.

Let $\psi\colon\mathbb{R}^{D}\rightarrow\mathbb{R}$ defined as
$\psi(\bm{x}):=(D\sum_{m=1}^{D}S_{m}\langle\bm{v},\bm{X}_{m}\rangle)^{p-1}$
where $S_{m}=(\mathrm{softmax}(x_{1},\dots,x_{D}))_{m}$,
$\\{\bm{X}_{m}\\}_{m=1}^{D}$ and $\bm{v}$ are fixed vectors and $p\geq 3$ is
an odd integer. Then, we have:

$\displaystyle|\psi(\bm{x})-\psi(\bm{y})|\leq\psi(\bm{x})\big{|}e^{2(p-1)\max_{c\in[D]}|x_{c}-y_{c}|}-1\big{|}.$

###### Proof of Lemma E.16.

Let $S_{c}=(\mathrm{softmax}(x_{1},\dots,x_{D}))_{c}$,
$\widehat{S}_{c}=(\mathrm{softmax}(y_{1},\dots,y_{D}))_{c}$. We have:

$\displaystyle|\psi(\bm{x})-\psi(\bm{y})|$
$\displaystyle\leq\psi(\bm{x})\Bigg{|}\Bigg{[}\frac{\sum_{c=1}^{D}S_{c}\langle\bm{v},\bm{X}_{c}\rangle}{\sum_{c=1}^{D}\widehat{S}_{c}\langle\bm{v},\bm{X}_{c}\rangle}\Bigg{]}^{p-1}-1\Bigg{|}$
(129)
$\displaystyle\leq\psi(\bm{x})\Bigg{|}\Bigg{[}\frac{\sum_{c=1}^{D}e^{x_{c}}\langle\bm{v},\bm{X}_{c}\rangle}{\sum_{c=1}^{D}e^{y_{c}}\langle\bm{v},\bm{X}_{c}\rangle}\cdot\frac{\sum_{r=1}^{D}e^{y_{r}}}{\sum_{r=1}^{D}e^{x_{r}}}\Bigg{]}^{p-1}-1\Bigg{|},$
(130)

We finally apply the generalized mediant inequality in (130) and get:

$\displaystyle|\psi(\bm{x})-\psi(\bm{y})|$
$\displaystyle\leq\psi(\bm{x})\Big{|}\Big{(}\max_{c\in[D]}e^{|x_{c}-y_{c}|}\cdot\max_{r\in[D]}e^{|x_{r}-y_{r}|}\Big{)}^{p-1}-1\Big{|}$
$\displaystyle\leq\psi(\bm{x})\big{|}e^{2(p-1)\max_{c\in[D]}|x_{c}-y_{c}|}-1\big{|}.$

∎

######  Lemma E.17 (Lipschitzness of Softmax).

Let $\bm{a},\bm{b}\in\mathbb{R}^{D}$. For all $i\in[D]$, we have:

$\displaystyle|\mathrm{softmax}(\bm{a})_{i}-\mathrm{softmax}(\bm{b})_{i}|\leq\big{|}e^{2\max_{k\in[D]}|b_{k}-a_{k}|}-1\big{|}\cdot\mathrm{softmax}(\bm{a})_{i}.$

###### Proof of Lemma E.17.

Let $i\in[D]$. The difference of softmax is bounded as:

$\displaystyle|\mathrm{softmax}(\bm{a})_{i}-\mathrm{softmax}(\bm{b})_{i}|$
$\displaystyle\leq\frac{e^{a_{i}}}{\sum_{j=1}^{D}e^{a_{j}}}\cdot\bigg{|}\frac{e^{b_{i}}}{e^{a_{i}}}\frac{\sum_{j=1}^{D}e^{a_{j}}}{\sum_{j=1}^{D}e^{b_{j}}}-1\bigg{|}$
$\displaystyle\leq\frac{e^{a_{i}}}{\sum_{j=1}^{D}e^{a_{j}}}\cdot\Big{|}e^{b_{i}-a_{i}}\max_{k\in[D]}e^{a_{k}-b_{k}}-1\Big{|},$
(131)

where we used the mediant inequality in the last inequality of (131). Since
the exponential function is non-decreasing, we deduce:

$\displaystyle|\mathrm{softmax}(\bm{a})_{i}-\mathrm{softmax}(\bm{b})_{i}|$
$\displaystyle\leq\big{|}e^{2\max_{k\in[D]}|b_{k}-a_{k}|}-1\big{|}\cdot\mathrm{softmax}(\bm{a})_{i}.$
(132)

∎

####  E.2.6 Dynamics of $\widehat{\bm{v}}^{(t)}$

Lastly, since $\Delta_{\bm{A}}^{(t)}$ remains small and
$\widehat{\bm{v}}^{(t)}$ mainly lies in $\mathrm{span}(\bm{w}^{*})$, we show
that $\widehat{\alpha}^{(t)}$ satisfies the same updates as the ideal
$\alpha^{(t)}$ (up to some constant factors).

######  Lemma E.18.

Assume that we run GD on the empirical risk (E) for $T$ iterations with
parameters set as in Parametrization 3.1 and the number of samples is
$N=\mathrm{poly}(d).$ Then, there exist times
$\mathscr{T},\widehat{\mathcal{T}}_{0},\widehat{\mathcal{T}}_{1}>0$ such that

  1. 1.

Analog of Event I (Lemma D.2):
$\widehat{\alpha}^{(t+1)}=\widehat{\alpha}^{(t)}+\Theta(\eta)(\widehat{G}^{(t)})^{p}(\widehat{\alpha}^{(t)})^{p-1}$
for $t\in[\mathscr{T},\mathscr{T}+\widehat{\mathcal{T}}_{0}]$.

  2. 2.

Analog of Event III (Lemma D.11):
$\widehat{\alpha}^{(t+1)}=\widehat{\alpha}^{(t)}+\Theta(\eta)(\widehat{G}^{(t)})^{p}(\widehat{\alpha}^{(t)})^{p-1}$
for $t\in[\mathscr{T}+\widehat{\mathcal{T}}_{1},T]$.

Consequently, $\widehat{\alpha}^{(t)}$ is non-decreasing and eventually,
$\widehat{\alpha}^{(T)}=\alpha^{(T)}=\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}.$

These three lemmas imply that at time $T$, the realistic iterates are very
close to the ideal ones. Therefore, they incur nearby test loss and thus the
realistic model generalizes. We now proceed to the proof of

In order to analyze the dynamics of $\widehat{\bm{v}}^{(t)}$, we first show
that the gradient (with respect to $\widehat{\bm{v}}$) in the realistic
learning process is very close to the one in the semi-idealized one.

######  Lemma E.19.

Let $t\in[\mathscr{T},T]$. With high probability, we have

$\displaystyle\|\nabla_{\widehat{\bm{v}}}\widehat{\mathcal{L}}(\widehat{\bm{v}}^{(t)},\widehat{\bm{A}}^{(t)})-\nabla_{\widecheck{\bm{v}}}\widecheck{\mathcal{L}}(\widecheck{\bm{v}}^{(t-\mathscr{T})},\widecheck{\bm{A}}^{(t-\mathscr{T})})\|_{2}\leq\mathrm{poly}(D)\Big{(}\frac{1}{\sqrt{N}}+\varepsilon_{\bm{v}}^{(t)}+\Delta_{\bm{A}}^{(t)}\Big{)}.$
(133)

By choosing $N=\mathrm{poly}(d)$, we have
$\|\nabla_{\widehat{\bm{v}}}\widehat{\mathcal{L}}(\widehat{\bm{v}}^{(t)},\widehat{\bm{A}}^{(t)})-\nabla_{\widecheck{\bm{v}}}\widecheck{\mathcal{L}}(\widecheck{\bm{v}}^{(t-\mathscr{T})},\widecheck{\bm{A}}^{(t-\mathscr{T})})\|_{2}\leq
1/\mathrm{poly}(d).$

###### Proof of Lemma E.19.

We have:

$\displaystyle\|\nabla_{\widehat{\bm{v}}}\widehat{\mathcal{L}}(\widehat{\bm{v}},\widehat{\bm{A}})-\nabla_{\widecheck{\bm{v}}}\widecheck{\mathcal{L}}(\widecheck{\bm{v}},\widecheck{\bm{A}})\|_{2}$
(134)
$\displaystyle=\Big{|}\frac{D}{N}\sum_{i=1}^{N}y[i]\mathfrak{S}\big{(}-y[i]F_{\widehat{\bm{A}}^{(t)},\widehat{\bm{v}}^{(t)}}(\bm{X}[i])\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{k}[i]\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{X}_{r}[i],\bm{w}^{*}\rangle$
(135)
$\displaystyle-D\mathbb{E}\Big{[}y\mathfrak{S}\big{(}-yF_{\widehat{\bm{A}}^{(t)},\widehat{\bm{v}}^{(t)}}(\bm{X})\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{k}\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{X}_{r},\bm{w}^{*}\rangle\Big{]}\Big{|}$
(136) $\displaystyle+$ $\displaystyle
D\Bigg{|}\mathbb{E}\Big{[}y\mathfrak{S}\big{(}-yF_{\widehat{\bm{A}}^{(t)},\widehat{\bm{v}}^{(t)}}(\bm{X})\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{k}\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{X}_{r},\bm{w}^{*}\rangle\Big{]}$
(137)
$\displaystyle-\mathbb{E}\Big{[}y\mathfrak{S}\big{(}-yF_{\widehat{\bm{A}}^{(t)},\widecheck{\bm{v}}^{(t)}}(\bm{X})\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{k}\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{X}_{r},\bm{w}^{*}\rangle\Big{]}\Big{|}$
(138) $\displaystyle+$ $\displaystyle
D\Big{|}\mathbb{E}\Big{[}y\mathfrak{S}\big{(}-yF_{\widehat{\bm{A}}^{(t)},\widecheck{\bm{v}}^{(t)}}(\bm{X})\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{k}\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{X}_{r},\bm{w}^{*}\rangle\Big{]}$
(139)
$\displaystyle-\mathbb{E}\Big{[}y\mathfrak{S}\big{(}-yF_{\widecheck{\bm{A}}^{(t)},\widecheck{\bm{v}}^{(t)}}(\bm{X})\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widecheck{S}_{j,k}^{(t)}\langle\widecheck{\bm{v}}^{(t)},\bm{X}_{k}\rangle\Big{)}\sum_{r=1}^{D}\widecheck{S}_{j,r}^{(t)}\langle\bm{X}_{r},\bm{w}^{*}\rangle\Big{]}\Big{|}$
(140)
$\displaystyle+\varepsilon_{\bm{v}}^{(t)}\Big{|}\frac{D}{N}\sum_{i=1}^{N}y[i]\mathfrak{S}\big{(}-y[i]F_{\widehat{\bm{A}}^{(t)},\widehat{\bm{v}}^{(t)}}(\bm{X}[i])\big{)}\sum_{j=1}^{D}\sigma^{\prime}\Big{(}\sum_{k=1}^{D}\widehat{S}_{j,k}^{(t)}\langle\widehat{\bm{v}}^{(t)},\bm{X}_{k}[i]\rangle\Big{)}\sum_{r=1}^{D}\widehat{S}_{j,r}^{(t)}\langle\bm{X}_{r}[i],\bm{u}^{(t)}\rangle\Big{|}.$
(141)

We bound each of the terms above using concentration or lipschitz
inequalities. Using the same arguments as in the proof of Lemma E.10, we have
$|\eqref{eq:oewdojfwe}-\eqref{eq:fnewfwe}|\leq\mathrm{poly}(D)/\sqrt{N}$ and
$|\eqref{eq:wfojw}-\eqref{eq:ojerofjr}|\leq\mathrm{poly}(D)\varepsilon_{\bm{v}}^{(t)}$.
Using the same steps as in the proof of Lemma E.14, we have
$|\eqref{eq:fojoe}-\eqref{eq:qwdwq}|\leq\mathrm{poly}(D)\Delta_{\bm{A}}^{(t)}$.
Lastly, $|\eqref{eq:kewffkpd}|\leq\mathrm{poly}(D)\varepsilon_{\bm{v}}^{(t)}.$
Summing up all these terms yields the aimed result.

∎

Lemma E.19 shows that we can use the gradient from the semi-idealized process
to analyze the dynamics of $\widehat{\alpha}^{(t)}$ in the real process.
Therefore, we can derive similar updates for $\widehat{\alpha}^{(t)}$ as in
Lemma D.2, Lemma D.8 and Lemma D.11.

######  Lemma E.20.

Let $\widecheck{\mathcal{T}}_{0}=\Theta\left(\frac{1}{\eta
C(\widehat{\alpha}^{(\mathscr{T})})^{p-1}}\right)$. Therefore,
$\widehat{\alpha}^{(t)}$ is updated as

$\displaystyle\widehat{\alpha}^{(t+1)}$
$\displaystyle=\widehat{\alpha}^{(t)}+\Theta(\eta
C)(\widecheck{G}^{(t)})^{p}(\widehat{\alpha}^{(t)})^{p-1}.$

Consequently, $\widehat{\alpha}^{(t)}$ is non-decreasing and after
$\widecheck{\mathcal{T}}_{0}$ iterations, we have
$\widehat{\alpha}^{(t)}\geq\frac{\Omega(1)}{C^{2}\lambda_{0}}$ for
$t\geq\widecheck{\mathcal{T}}_{0}.$

######  Lemma E.21.

Let
$\widecheck{\mathcal{T}}_{1}=\widecheck{\mathcal{T}}_{0}+\Theta\Big{(}\frac{\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}\Big{)}$. For all
$t\in[\widecheck{\mathcal{T}}_{1},T]$, we have
$\widecheck{\Gamma}^{(t)}\geq\frac{\Omega(\lambda_{0})}{D}$.

######  Lemma E.22.

Let
$\widecheck{\mathcal{T}}_{1}=\widecheck{\mathcal{T}}_{0}+\Theta\Big{(}\frac{\lambda_{0}^{p}}{\eta
e^{\mathrm{polyloglog}(d)}}\Big{)}$ and $t\in[\widecheck{\mathcal{T}}_{1},T]$.
$\widehat{\alpha}^{(t)}$ updates as

$\displaystyle\widehat{\alpha}^{(t+1)}$
$\displaystyle=\widehat{\alpha}^{(t)}+\Theta(\eta
C)(\widecheck{G}^{(t)})^{p}(\widehat{\alpha}^{(t)})^{p-1}.$ (142)

Consequently, $\widehat{\alpha}^{(t)}$ is non-decreasing and eventually
$\widehat{\alpha}^{(T)}=\frac{\mathrm{polylog}(d)}{C^{2}\lambda_{0}}.$

##### Auxiliary lemma.

The following lemma is useful to prove Lemma E.10 and Lemma E.14.

######  Lemma E.23.

The time at which $\widehat{\alpha}^{(t)}$ stops increasing is
$\widecheck{\mathcal{T}}_{2}=\widecheck{\mathcal{T}}_{1}+\Theta\left(\frac{\lambda_{0}^{p-2}}{\eta\mathrm{polylog}(d)}\right).$

###### Proof of Lemma E.23.

The result is obtained by applying Lemma K.1 to (142). We have:

$\displaystyle\widecheck{\mathcal{T}}_{2}=\widecheck{\mathcal{T}}_{1}+\frac{3(\lambda_{0})^{p-2}}{\eta\mathrm{polylog}(d)e^{\mathrm{polyloglog}(d)}}+\frac{2^{p}\lambda_{0}^{p}}{e^{\mathrm{polyloglog}(d)}}\log\log(d).$

∎

####  E.2.7 The realistic model fits the labeling function

######  Lemma E.24.

In the realistic case, the model fits the labeling function i.e.

$\displaystyle\mathbb{P}_{\mathcal{D}}[f^{*}(\bm{X})F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})>0]\geq
1-o(1).$ (143)

###### Proof of Lemma E.24.

We bound the population risk
$\mathcal{L}(\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)})$. We have:

$\displaystyle\mathcal{L}(\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)})$
$\displaystyle\leq\Big{|}\mathcal{L}(\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)})-\mathcal{L}(\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)})\Big{|}+\mathcal{L}(\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}).$
(144)

Using Lemma E.25, we have
$\mathcal{L}(\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)})\leq
1/\mathrm{poly}(d).$ We now bound the first summand in (144) using the
1-Lipschitzness of the logistic function and get:

$\displaystyle\Big{|}\mathcal{L}(\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)})-\mathcal{L}(\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)})\Big{|}$
$\displaystyle\leq\mathbb{E}\Big{[}\big{|}F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})-F_{\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})\big{|}\Big{]}$
(145)

Using Lemma E.26, we have
$\mathbb{E}\Big{[}\big{|}F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})-F_{\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})\big{|}\Big{]}\leq
1/\mathrm{poly}(d).$ Therefore, we deduce that
$\mathcal{L}(\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)})\leq
1/\mathrm{poly}(d).$ Since the 0-1 loss is a convex surrogate, we have

$\displaystyle\mathbb{P}_{(\bm{X},y)\sim\mathcal{D}}[yF_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})<0]\leq\mathcal{L}(\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)})\leq
1/\mathrm{poly}(d).$ (146)

We can further expand (146) as in the proof of Lemma D.15 and deduce the aimed
result. ∎

To prove Lemma E.24, we use the following auxiliary lemma.

######  Lemma E.25.

After $T$ iterations, the population risk in the semi-idealized case converges
i.e.
$\widecheck{\mathcal{L}}(\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)})\leq
o(1).$

###### Proof of Lemma E.25.

The proof is similar to the one of Lemma D.15. ∎

######  Lemma E.26.

For all $\bm{X}$ sampled from $\mathcal{D}$, we have

$\displaystyle\big{|}F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})-F_{\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})\big{|}\leq
1/\mathrm{poly}(d).$

###### Proof of Lemma E.26.

We have:

$\displaystyle\big{|}F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})-F_{\bm{Q}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})\big{|}$
$\displaystyle\leq\big{|}F_{\widehat{\bm{A}}^{(T)},\widehat{\bm{v}}^{(T)}}(\bm{X})-F_{\widehat{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})\big{|}$
(147)
$\displaystyle+\big{|}F_{\widehat{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})-F_{\widecheck{\bm{A}}^{(T)},\widecheck{\bm{v}}^{(T)}}(\bm{X})\big{|}.$
(148)

We now separately bound (147) and (148).

##### Bound on (147).

Since $x\mapsto x^{p}+\nu x$ is Lipschitz on a bounded domain, we have:

$\displaystyle\eqref{eq:fweojowjeef}$
$\displaystyle\leq\sum_{m=1}^{D}\sigma^{\prime}\Big{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(T)}|\langle\widehat{\bm{v}}^{(T)},\bm{X}_{r}\rangle|\Big{)}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(T)}|\langle\widehat{\bm{v}}^{(T)}-\widecheck{\bm{v}}^{(T)},\bm{X}_{r}\rangle|$
(149)
$\displaystyle\leq\varepsilon_{\bm{v}}^{(T)}\sum_{m=1}^{D}\sigma^{\prime}\Big{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(T)}|\langle\widehat{\bm{v}}^{(T)},\bm{X}_{r}\rangle|\Big{)}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(T)}|\langle\bm{u}^{(t)},\bm{X}_{r}\rangle|,$
(150)

since
$\langle\widehat{\bm{v}}^{(T)},\bm{w}^{*}\rangle=\langle\widecheck{\bm{v}}^{(T)},\bm{w}^{*}\rangle$.
Using Cauchy-Schwarz inequality, (149) simplifies as:

$\displaystyle\eqref{eq:fweojowjeef}$
$\displaystyle\leq\sum_{m=1}^{D}\sigma^{\prime}\Big{(}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(T)}|\langle\widehat{\bm{v}}^{(T)},\bm{X}_{r}\rangle|\Big{)}^{p-1}D\sum_{r=1}^{D}\widehat{S}_{m,r}^{(T)}\varepsilon_{\bm{v}}^{(T)}\leq\mathrm{poly}(D)\varepsilon_{\bm{v}}^{(T)}.$
(151)

Using Lemma E.10, we conclude that $\eqref{eq:fweojowjeef}\leq
1/\mathrm{poly}(d).$

##### Bound on (148).

We again use the Lipschitzness of the power function and get:

$\displaystyle\eqref{eq:fwejoojrw}$
$\displaystyle\leq\sum_{m=1}^{D}\sigma^{\prime}\Big{(}D\sum_{r=1}^{D}\max\\{\widecheck{S}_{m,r}^{(T)},\widehat{S}_{m,r}^{(T)}\\}|\langle\bm{v}^{(T)},\bm{X}_{r}\rangle|\Big{)}D\sum_{r=1}^{D}|\widehat{S}_{m,r}^{(T)}-\widecheck{S}_{m,r}^{(T)}|\cdot|\langle\widehat{\bm{v}}^{(T)},\bm{X}_{r}\rangle|.$
(152)

We apply Lemma E.17 in (152) to get
$\eqref{eq:fwejoojrw}\leq\mathrm{poly}(D)\Delta_{\bm{A}}^{(T)}.$ Finally, we
apply Lemma E.14 to get $\eqref{eq:fwejoojrw}\leq 1/\mathrm{poly}(d).$

∎

##  Appendix F Transfer Learning

In this section, we show that a transformer that has been pre-trained on a
structured dataset require a few samples to generalize in a new dataset
sharing the same structure.

See 5.1

###### Proof of Theorem 5.1.

Actually, even one step of the update using normalized gradient descent on
$\bm{v}$ can already achieve test accuracy $\geq 1-o(1).$ We know that for a
datum $(\bm{X},y)$, the gradient of $L(\bm{X})$ with respect to $\bm{v}$ is

$\displaystyle\nabla_{\bm{v}}L(\bm{X})=-y\mathfrak{S}(-yF(\bm{X}))\sum_{m=1}^{D}\sigma^{\prime}(\langle\bm{O}_{m}^{(t)},\widetilde{\bm{v}}^{(t)}\rangle)\bm{O}_{m}^{(t)}$
(153)

Since $\widetilde{\bm{v}}^{(0)}=\bm{0},$ we have
$F_{\widetilde{\bm{v}}^{(0)}}(\bm{X})=0$ and
$\sigma^{\prime}(\langle\bm{O}_{m}^{(0)},\widetilde{\bm{v}}^{(0)}\rangle)=\nu.$
Thus, the gradient (153) simplifies to

$\displaystyle\nabla_{\bm{v}}L(\bm{X})=-\frac{\nu
y}{2}\sum_{j=1}^{D}\bm{X}_{j}\sum_{m=1}^{D}\widehat{S}_{m,j}^{(t)}$ (154)

By symmetry of the $\widehat{S}_{j,m}^{(t)}$ , we know that

$\displaystyle\Big{|}\sum_{m=1}^{D}\widehat{S}_{m,j}^{(t)}-D\Big{|}\leq\frac{D}{\mathrm{poly}(d)},$
(155)

where $\frac{1}{\mathrm{poly}(d)}$ comes from the $\varepsilon_{\bm{Q}}$ part
in the previous section. Moreover, since the noise and feature noise has mean
zero independent of $y$, we know that there exists some value $c_{0}>0$
(roughly equal to $\alpha DC$) ) such that:

$\displaystyle\mathbb{E}[\nabla_{\bm{v}}L(\bm{X})]$
$\displaystyle=c_{0}\widetilde{\bm{w}}^{*}.$ (156)

Now, by standard concentration inequality, we know that for $N$ i.i.d. samples
$\bm{X}[i],y[i]$, with high probability

$\displaystyle\bm{\nabla}_{0}:=\sum_{i=1}^{N}\nabla_{\bm{v}}L(\bm{X}[i])=c_{0}\widetilde{\bm{w}}^{*}+\varepsilon_{0}\widetilde{\bm{w}}^{*}+\bm{\chi}_{0},$
(157)

where $\varepsilon_{0}$ comes from the feature noise

$\displaystyle|\varepsilon_{0}|\leq\frac{c_{0}qD\log(d)}{\sqrt{N}},$ (158)

and $\bm{\chi}_{0}$ comes from the noise:

$\displaystyle\|\bm{\chi}_{0}\|_{2}\leq\frac{c_{0}\sigma\sqrt{d}\sqrt{D}}{\sqrt{N}}.$
(159)

Therefore, if we update using normalized GD:

$\displaystyle\widetilde{\bm{v}}^{(1)}$
$\displaystyle=\widetilde{\bm{v}}^{(0)}+\frac{\bm{\nabla}_{0}}{\|\bm{\nabla}_{0}\|_{2}}$
(160)

we have that:

$\displaystyle\widetilde{\bm{v}}^{(1)}$
$\displaystyle=c_{1}\widetilde{\bm{w}}^{*}+\bm{\chi}_{1},$ (161)

where $\bm{\chi}_{1}\perp\widetilde{\bm{w}}^{*}$, $\|\bm{\chi}_{1}\|_{2}=O(1)$
and $c_{1}\geq\tilde{\Omega}(1)/\sqrt{D}.$ Now, for a new datum $\bm{X}^{new}$
with noises $\\{\bm{\xi}_{i}^{new}\\}_{i=1}^{D}$, we know that w.h.p

$\displaystyle|\langle\bm{\xi}_{i}^{new},\widetilde{\bm{v}}^{(1)}\rangle|\leq\frac{\tilde{O}(1)}{\sqrt{d}},\qquad|\langle\widetilde{\bm{w}}^{*},\widetilde{\bm{v}}^{(1)}\rangle|\geq\frac{\tilde{\Omega}(1)}{\sqrt{D}}.$
(162)

We can prove the test accuracy is small using the same proof as in Lemma D.15,
where we show that:

$\displaystyle
y\sigma(\langle\widetilde{\bm{v}}^{(1)},\bm{O}_{m}^{(1)}\rangle)\geq 0,$ (163)

for $m\in\mathcal{S}_{\ell(\bm{X})}$ and it dominates the other
$y\sigma(\langle\widetilde{\bm{v}}^{(1)},\bm{O}_{j}^{(1)}\rangle)$ for
$j\in\mathcal{S}_{\ell}$ with $\ell\neq\ell(\bm{X}).$ Therefore, we prove that
$\mathbb{P}_{\widetilde{\mathcal{D}}}[yF(\bm{X})>0]\geq 1-o(1)$ which implies
the aimed result. ∎

See 5.2

###### Proof of Theorem 5.2.

Let $\mathcal{A}$ be an algorithm. Assume that at training time the algorithm
$\mathcal{A}$ has access to $D^{o(1)}$ training data, Since each input $X$ is
made of $D$ patches, this means that there exist at least $\Omega(1)$ fraction
of $k\in[D]$ such that $\mathcal{A}$ has not seen training samples with
$\ell(X)=k$. Consider the following two distributions over $\\{-1,0,1\\}^{m}$:

  1. 1.

$\mathcal{D}_{1}$: Sample $z\in\\{0,1\\}^{m}$ where each $z_{i}$ i.i.d. $1$
w.p. $q/2$, $-1$ w.p. $q/2$ and $0$ otherwise.

  2. 2.

$\mathcal{D}_{2}$: Sample a set $\mathcal{S}$ uniformly at random from $[m]$
of size $C$, set all $z_{i}=1$ for $i\in\mathcal{S}$, and sample other $z_{j}$
i.i.d. $1$ w.p. $q/2$, $-1$ w.p. $q/2$ and $0$ otherwise.

We can easily see that as long as $qm=\text{poly}(C)$, then

$\textbf{TV}(\mathcal{D}_{1},\mathcal{D}_{2})=o(1)$

This implies that $\mathcal{A}$ must have bad generalization error
($\Omega(1)$) on $\tilde{\mathcal{D}}$.

∎

##  Appendix G Gradient descent updates in the idealized process

In this section, we derive the gradient descent updates of $A_{i,j}$ in the
idealized learning process.

###  G.1 Indices in the same set: $i,j\in\mathcal{S}_{\ell}$

######  Lemma G.1.

Let $T>0$ be the time where the population loss is at most $o(1)$ and
$t\in[0,T].$ Then, $\gamma^{(t)}$ satisfies the update

$\displaystyle\gamma^{(t+1)}$
$\displaystyle=\gamma^{(t)}+\eta\Theta(C)\alpha^{(t)}\Big{(}D\alpha^{(t)}\big{[}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\big{]}\Big{)}^{p-1}\Gamma^{(t)}.$

###### Proof of Lemma G.1.

Let $\ell\in[L]$ and $i,j\in\mathcal{S}_{\ell}$ with $i\neq j.$ The main idea
of the proof is to bound the gradient of $L(\bm{X})$ with respect to
$A_{i,j}$. This gradient is given by Lemma H.2 and is made of two terms: the
$\sigma^{\prime}$ term and the sum outside $\sigma^{\prime}$. We distinguish
the following cases and bound these two terms.

1. $\bm{\ell=\ell(X)}$. We first bound the outside sum. Using Lemma K.3, we have

$\displaystyle\Big{|}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}(1-y\delta_{r})-(D-C)\Big{|}$
$\displaystyle\leq\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|\leq
qD\log(d).$ (164)

Since $(D-C)/2-qD\log(d)\geq 0$, we rewrite (164) as:

$\displaystyle\frac{1}{2}(D-C)\leq\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}(1-y\delta_{r})\leq
2(D-C).$ (165)

Regarding the sum inside $\sigma^{\prime}$, we use Lemma G.3 which shows:

$\displaystyle\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}\delta_{h,r}=\Theta(1)\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}\right).$
(166)

By using (165) and (166), we finally obtain:

$\displaystyle-\frac{\partial L(\bm{X})}{\partial A_{i,j}}$
$\displaystyle=\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\left(D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\right)^{p-1}\Gamma^{(t)}(D-C)\Xi^{(t)}.$
(167)

2. $\bm{\ell\neq\ell(X)}$ and $\bm{\delta_{s}=0}$ for all $\bm{s\in\mathcal{S}_{\ell}}$. We first bound the outside sum. Since $\delta_{s}=0$, the only non-zero term is the one with factor $\Xi^{(t)}.$ Using triangle inequality, we have:

$\displaystyle\Xi^{(t)}\Big{|}Cy+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{|}$
$\displaystyle\leq\Xi^{(t)}\bigg{(}C+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|\bigg{)}=\Xi^{(t)}\kappa(\bm{X}).$
(168)

We now bound the sum inside $\sigma^{\prime}$. This sum is actually equal to
the outside sum and we can therefore use the bound (168). Therefore, the
overall gradient is bounded as:

$\displaystyle\Big{|}\frac{\partial L(\bm{X})}{\partial A_{i,j}}\Big{|}$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\left(D\Xi^{(t)}\alpha^{(t)}\kappa(\bm{X})\right)^{p-1}\Gamma^{(t)}\Xi^{(t)}\kappa(\bm{X}).$
(169)

3. $\bm{\ell\neq\ell(\bm{X})}$ and at least one  $\bm{\delta_{s}\neq 0}$ and  $\bm{\delta_{i}=0}$. We first bound the outside sum. Using Lemma K.3 and $\Lambda^{(t)}+\Gamma^{(t)}+\Xi^{(t)}=1$, we have:

$\displaystyle\alpha^{(t)}\bigg{|}\Lambda^{(t)}\delta_{j}+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}(\delta_{j}-\delta_{r})+\Xi^{(t)}C(\delta_{j}-y)+\Xi^{(t)}\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{m\in\mathcal{S}_{h}}(\delta_{j}-\delta_{m})\bigg{|}$
(170) $\displaystyle\leq$
$\displaystyle\alpha^{(t)}\left[(C-1)\Gamma^{(t)}+\Xi^{(t)}\left(C+\Theta(qD)\log(d)\right)+\mathbf{1}_{\delta_{j}\neq
0}\right].$

We lastly apply Induction Hypothesis B.1 to show that (170) is less or equal
to $\Theta(\alpha^{(t)}).$ We now bound the sum inside $\sigma^{\prime}$.

$\displaystyle\alpha^{(t)}\bigg{|}\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i,j\\}}\delta_{r}+\Xi^{(t)}\Big{(}Cy+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{m\in\mathcal{S}_{h}}\delta_{m}\Big{)}\bigg{|}$
$\displaystyle\leq\alpha^{(t)}\bigg{(}\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i,j\\}}|\delta_{r}|+\Xi^{(t)}\kappa(\bm{X})\bigg{)}$
(171)
$\displaystyle\leq\alpha^{(t)}\left((C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})\right).$

The overall bound on the derivative is:

$\displaystyle\Big{|}\frac{\partial L(\bm{X})}{\partial A_{i,j}}\Big{|}$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\Big{(}D\alpha^{(t)}\big{(}(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(X)\big{)}\Big{)}^{p-1}\Gamma^{(t)}.$
(172)

4. $\bm{\ell\neq\ell(X)}$ where  $\bm{\delta_{s}\neq 0}$ for all $s$. We first bound the outside sum as follows.

$\displaystyle\alpha^{(t)}\biggr{|}\Lambda^{(t)}(\delta_{j}-\delta_{i})+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i,j\\}}(\delta_{j}-\delta_{r})+\Xi^{(t)}\Big{[}C(\delta_{j}-y)+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{m\in\mathcal{S}_{h}}(\delta_{j}-\delta_{m})\Big{]}\bigg{|}$
(173)
$\displaystyle\leq\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{\delta_{j}\neq
0}).$

We now bound the sum inside $\sigma^{\prime}$.

$\displaystyle\alpha^{(t)}\bigg{|}\Lambda^{(t)}\delta_{i}+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i,j\\}}\delta_{r}+\Xi^{(t)}\Big{[}Cy+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{m\in\mathcal{S}_{h}}\delta_{m}\Big{]}\bigg{|}$
(174) $\displaystyle\leq$
$\displaystyle\alpha^{(t)}\Big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})\Big{)}.$

We lastly apply Lemma G.3 to show that (LABEL:eq:jewnwcdw) is bounded by
$\Theta(\alpha^{(t)})(\Lambda^{(t)}+(C-1)\Gamma^{(t)}).$ Therefore, the
overall gradient is bounded as

$\displaystyle\Big{|}\frac{\partial L(\bm{X})}{\partial A_{i,j}}\Big{|}$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\left(D\alpha^{(t)}((C-1)\Gamma^{(t)}+\Lambda^{(t)})\right)^{p-1}\Gamma^{(t)}\cdot$
(175) $\displaystyle\hskip
14.22636pt\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{\delta_{j}\neq
0}\right).$

##### Putting all the pieces together.

We now bound the derivative of the population loss. Using Tower property and
Lemma I.1, we have:

$\displaystyle\mathbb{E}_{\bm{X}}\bigg{[}\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\bigg{]}=\mathbb{E}_{\bm{X}}\bigg{[}\frac{\partial
L(\pi(\bm{X}))}{\partial
A_{i,j}}\bigg{]}=\mathbb{E}_{\bm{X}}\bigg{[}\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\Big{|}\bm{X}\Big{]}\bigg{]}.$ (176)

For a fixed $\bm{X}$, we now bound the derivative of the loss evaluated in
$\bm{X}$. For $i,j\in\mathcal{S}_{\ell}$, we distinguish the four possible
events depending on the randomness of $\pi.$

  * –

Event a: "$\ell=\ell(\pi(\bm{X}))$" occurs with probability $1/L.$

  * –

Event b: "$\ell\neq\ell(\pi(\bm{X}))$ and $\delta_{s}=0$ for all
$s\in\mathcal{S}_{\ell}$" occurs with probability $(1-1/L)(1-q)^{C}$.

  * –

Event ck: "$\ell\neq\ell(\pi(\bm{X}))$ and $\\#\\{s:\delta_{s}\neq 0\\}=k$ for
$1\leq k\leq C-1$ and $\delta_{i}=0$" occurs with probability
$(1-1/L)\binom{C-1}{k}q^{k}(1-q)^{C-k}$.

  * –

Event d: "$\ell\neq\ell(\pi(\bm{X}))$ and $\delta_{s}\neq 0$ for all $s$"
occurs with probability $(1-1/L)\cdot q^{C}$.

Therefore, the derivative of the loss in $\bm{X}$ is:

$\displaystyle\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\Big{|}\bm{X}\Big{]}$
$\displaystyle=\frac{1}{L}\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\;\Big{|}\;\bm{X},\textbf{a}\Big{]}$ (177)
$\displaystyle+$
$\displaystyle\left(1-\frac{1}{L}\right)(1-q)^{C}\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\;\Big{|}\;\bm{X},\textbf{b}\Big{]}$
$\displaystyle+$
$\displaystyle\left(1-\frac{1}{L}\right)\sum_{k=1}^{C-1}\binom{C-1}{k}q^{k}(1-q)^{C-k}\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\;\Big{|}\;\bm{X},\textbf{c}_{k}\Big{]}$
$\displaystyle+$
$\displaystyle\left(1-\frac{1}{L}\right)q^{C}\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\;\Big{|}\;\bm{X},\textbf{d}\Big{]}.$

Event a is the event that is the most likely to happen. Therefore, we only
take into account $\mathbb{E}_{\pi_{1},\pi_{2}}\left[\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\;\middle|\;\textbf{a}\right]$ in (177) and
obtain:

$\displaystyle\mathbb{E}_{\bm{X}}\bigg{[}\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\bigg{]}$ $\displaystyle=\
\frac{\Theta(D\alpha^{(t)})}{L}\mathbb{E}_{X}[\mathfrak{S}(-yF(\bm{X}))]\left(D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)}))\right)^{p-1}\Gamma^{(t)}.$
(178)

Since the population loss is a $\Omega(1)$ for $t\leq T$, this implies that
$\mathbb{E}_{\bm{X}}[\mathfrak{S}(-yF(\bm{X}))]=\Theta(1)$ (Lemma D.10). We
thus plug (178) in the update of $\gamma^{(t)}$ to obtain the desired result.
∎

######  Corollary G.1.

Let $T>0$ be the time where the population loss is $o(1)$ and $t\leq T.$ Let
$G^{(t)}:=D(\Lambda^{(t)}+(C-1)\Gamma^{(t)})$. The update of $\gamma^{(t)}$
satisfies:

$\displaystyle\gamma^{(t+1)}$
$\displaystyle=\gamma^{(t)}+\Theta(C\eta)(\alpha^{(t)})^{p}\Gamma^{(t)}(G^{(t)})^{p-1}.$

###### Proof of Corollary G.1.

Lemma G.1 provides the update rule of $\gamma^{(t)}$.

$\displaystyle\gamma^{(t+1)}$
$\displaystyle=\gamma^{(t)}+\eta\Theta(C)(\alpha^{(t)})^{p}\Gamma^{(t)}(G^{(t)})^{p-1}.$

∎

###  G.2 Update for $i\in\mathcal{S}_{\ell}$ and $j\in\mathcal{S}_{m}$

######  Lemma G.2.

Let $T>0$ be the time where the population loss is at most $o(1)$ and
$t\in[0,T].$ Then, $\rho^{(t)}$ satisfies the update

$\displaystyle\left|\frac{\rho^{(t+1)}-\rho^{(t)}}{\eta}\right|$
$\displaystyle\leq\Theta(C^{2}\alpha^{(t)})\left(D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\right)^{p-1}\Xi^{(t)}\frac{\lambda_{0}}{D}$
$\displaystyle+\Theta(\alpha^{(t)})\left(D\alpha^{(t)}qD\log(d)\Xi^{(t)}\right)^{p-1}\Xi^{(t)}qD\log(d).$

###### Proof of Lemma G.1.

Let $\ell,m\in[L]$ such that $\ell\neq m$ and $i\in\mathcal{S}_{\ell}$,
$j\in\mathcal{S}_{m}$. The main idea of the proof is to bound the gradient of
$L(\bm{X})$ with respect to $A_{i,j}$. This gradient is given by Lemma H.2 and
is made of two terms: the $\sigma^{\prime}$ term and the sum outside
$\sigma^{\prime}$. We distinguish the following cases and bound these two
terms.

1. $\bm{\ell=\ell(X)}$ and $\bm{\delta_{j}=0}$. We first bound the outside sum. We apply Lemma G.3 to obtain:

$\displaystyle\alpha^{(t)}\Big{|}\big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\big{)}y+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{|}\leq
2\alpha^{(t)}\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}\right).$ (179)

We now bound the sum inside $\sigma^{\prime}$. This sum is actually equal to
the outside sum and we can therefore use the bound (179). Therefore, the
overall gradient is bounded as:

$\left|\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right|\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\left(D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\right)^{p-1}\Xi^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)}).$
(180)

2. $\bm{\ell=\ell(X)}$ and $\bm{\delta_{j}\neq 0}$. We first bound the outside sum. We successively apply $\Lambda^{(t)}+\Gamma^{(t)}+\Xi^{(t)}=1$, Lemma K.3 and Induction Hypothesis B.1 to obtain:

$\displaystyle\alpha^{(t)}\bigg{|}\Lambda^{(t)}(\delta_{j}-y)+(C-1)\Gamma^{(t)}(\delta_{j}-y)+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}(\delta_{j}-\delta_{r})\bigg{|}$
$\displaystyle\leq\;$
$\displaystyle\alpha^{(t)}\bigg{[}\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|+\mathbf{1}_{\delta_{j}\neq
0}\bigg{]}$ $\displaystyle\leq\;$
$\displaystyle\alpha^{(t)}\left[\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}qD\log(d)+\mathbf{1}_{\delta_{j}\neq
0}\right]$ $\displaystyle\leq\;$
$\displaystyle\alpha^{(t)}\left[\frac{e^{\mathrm{polyloglog}(d)}+\lambda_{0}}{D}+\Theta(q\log(d))+\mathbf{1}_{\delta_{j}\neq
0}\right]=\Theta(\alpha^{(t)}).$ (181)

We now bound the sum inside $\sigma^{\prime}$. We successively apply Lemma
K.3, triangle inequality and Lemma G.3 to obtain:

$\displaystyle\alpha^{(t)}\bigg{|}y(\Lambda^{(t)}+(C-1)\Gamma^{(t)})+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\bigg{|}\leq
2\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)}).$ (182)

Thus, we use (181) and (182) to obtain a bound on the derivative.

$\displaystyle\left|\frac{\partial L(\bm{X})}{\partial A_{i,j}}\right|$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\Big{(}D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\Big{)}^{p-1}\Xi^{(t)}.$
(183)

3. $\bm{\ell\neq\ell(X)}$ and $\bm{\delta_{s}=0}$ for all $\bm{s}$ and $\bm{m=\ell(X)}$. We first bound the outside sum. Since $\Lambda^{(t)}+(C-1)\Gamma^{(t)}+(D-C)\Xi^{(t)}=1$, we have:

$\displaystyle\Xi^{(t)}\bigg{|}y\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}\right)+\Xi^{(t)}\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}(y-\delta_{r})\bigg{|}\leq\Xi^{(t)}\left(\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{m=\ell(\bm{X})}\right).$
(184)

We now bound the sum inside $\sigma^{\prime}$.

$\displaystyle\Xi^{(t)}\bigg{|}y+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}(X)\bigg{|}\leq\Xi^{(t)}\kappa(\bm{X}).$
(185)

Using (184) and (185), we obtain a bound on the derivative.

$\displaystyle\left|\frac{\partial L(\bm{X})}{\partial A_{i,j}}\right|$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\Big{(}D\alpha^{(t)}\kappa(\bm{X})\Xi^{(t)}\Big{)}^{p-1}\Xi^{(t)}\left(\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{m=\ell(\bm{X})}\right).$
(186)

4. $\bm{\ell\neq\ell(\bm{X})}$ and $\bm{\delta_{s}\neq 0}$ for some $\bm{s}$ and $\bm{m=\ell(X)}$. We first bound the outside sum. We apply $\Lambda^{(t)}+(C-1)\Gamma^{(t)}+(D-C)\Xi^{(t)}=1$ and Lemma G.4 to get:

$\displaystyle\Xi^{(t)}\left|\Lambda^{(t)}(y-\delta_{i})+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}(y-\delta_{r})+\Xi^{(t)}\sum_{h\neq\\{\ell(X),\ell\\}}\sum_{r\in\mathcal{S}_{h}}(y-\delta_{r})\right|$
$\displaystyle\leq\;$
$\displaystyle\Xi^{(t)}\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|+\mathbf{1}_{m=\ell(\bm{X})}\right)$
$\displaystyle\leq\;$ $\displaystyle\Theta(\Xi^{(t)}).$ (187)

We now bound the sum inside $\sigma^{\prime}$. We apply Lemma G.3 to obtain:

$\displaystyle\left|\Lambda^{(t)}\delta_{i}(X)+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}\delta_{r}(X)+\Xi^{(t)}\left(Cy+\sum_{h\in[L]\backslash\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{h,r}(X)\right)\right|$
(188) $\displaystyle\leq$
$\displaystyle\;\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})$
$\displaystyle\leq$
$\displaystyle\;2\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}\right).$

We combine (187) and (LABEL:eq:fjqjeq) and obtain:

$\displaystyle\left|\frac{\partial L(\bm{X})}{\partial A_{i,j}}\right|$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\Big{(}D\alpha^{(t)}\big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\big{)}\Big{)}^{p-1}\Xi^{(t)}.$
(189)

5. $\bm{\ell,m\neq\ell(X)}$ and $\bm{\delta_{s}=0}$ for all $s$ and $\bm{\delta_{j}=0}$. We first bound the outside sum.

$\displaystyle\Xi^{(t)}\Big{|}Cy+\sum_{h\neq\\{\ell,\ell(\bm{X})\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{|}\leq\Xi^{(t)}\kappa(\bm{X}).$
(190)

We now bound the sum inside $\sigma^{\prime}$. This sum is actually equal to
the outside sum outside and we can therefore use the bound
(LABEL:eq:frekfwew). Thus, the derivative is bounded as:

$\displaystyle\left|\frac{\partial L(\bm{X})}{\partial A_{i,j}}\right|$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\left(D\alpha^{(t)}\kappa(\bm{X})\Xi^{(t)}\right)^{p-1}(\Xi^{(t)})^{2}\kappa(\bm{X}).$
(191)

6. $\bm{\ell,m\neq\ell(X)}$ and $\bm{\delta_{s}\neq 0}$ for some $s$ and $\bm{\delta_{j}=0}$. We first bound the outside sum. We apply Lemma G.3 and obtain:

$\displaystyle\alpha^{(t)}\Big{|}\Lambda^{(t)}\delta_{i}+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}\delta_{r}+\Xi^{(t)}\big{[}Cy+\sum_{h\neq\\{\ell,\ell(\bm{X})\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\big{]}\Big{|}$
(192) $\displaystyle\leq$
$\displaystyle\;\alpha^{(t)}\Big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\big{[}C+\sum_{h\neq\\{\ell,\ell(\bm{X})\\}}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|\big{]}\Big{)}$
$\displaystyle\leq$
$\displaystyle\;2\alpha^{(t)}\left(\Lambda^{(t)}+(C-1)\Gamma^{(t)}\right).$

We now bound the sum inside the power term. This sum is actually equal to the
sum outside the power term and we can therefore use the bound
(LABEL:eq:redeww). Thus, the derivative is bounded as:

$\displaystyle\left|\frac{\partial L(\bm{X})}{\partial A_{i,j}}\right|$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\left(D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\right)^{p-1}\Xi^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)}).$
(193)

7. $\bm{\ell,m\neq\ell(X)}$ and $\bm{\delta_{s}=0}$ for all $s$ and $\bm{\delta_{j}\neq 0}$. We first bound the outside sum.

$\displaystyle\Xi^{(t)}\bigg{|}\Lambda^{(t)}\delta_{j}+(C-1)\Gamma^{(t)}\delta_{j}+\Xi^{(t)}\Big{[}C(\delta_{j}-y)+\sum_{h\neq\\{\ell(X),\ell\\}}\sum_{r\in\mathcal{S}_{h}}(\delta_{j}-\delta_{r})\Big{]}\bigg{|}$
(194) $\displaystyle\leq$
$\displaystyle\Xi^{(t)}\Big{(}\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{\delta_{j}\neq
0}\Big{)}.$

We now bound the sum inside $\sigma^{\prime}$.

$\displaystyle\Xi^{(t)}\Big{|}Cy+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{|}\leq\Xi^{(t)}\kappa(\bm{X}).$
(195)

Using (LABEL:eq:jfnrorjw3) and (195), the bound on the derivative is:

$\displaystyle\Big{|}\frac{\partial L(\bm{X})}{\partial A_{i,j}}\Big{|}$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\Big{(}D\alpha^{(t)}\Xi^{(t)}\kappa(\bm{X})\Big{)}^{p-1}\Xi^{(t)}\Big{(}\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{\delta_{j}\neq
0}\Big{)}.$ (196)

8. $\bm{\ell,m\neq\ell(X)}$ and $\bm{\delta_{s}(X)\neq 0}$ for some $s$ and $\bm{\delta_{j}\neq 0}$: We first bound the outside sum. We apply Lemma G.4 to get:

$\displaystyle\Xi^{(t)}\bigg{|}\Lambda^{(t)}(\delta_{j}-\delta_{i})+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}(\delta_{j}-\delta_{r})+\Xi^{(t)}\Big{[}C(\delta_{j}-y)+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}(\delta_{j}-\delta_{r})\Big{]}\bigg{|}$
(197) $\displaystyle\leq$
$\displaystyle\Xi^{(t)}\Big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})+\mathbf{1}_{\delta_{j}\neq
0}\Big{)}$ $\displaystyle\leq$ $\displaystyle\Theta(\Xi^{(t)}).$

We now bound the sum inside the power term. We apply Lemma G.3 and get:

$\displaystyle\bigg{|}\Lambda^{(t)}\delta_{i}+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}\delta_{r}+\Xi^{(t)}\Big{(}Cy+\sum_{h\neq\\{\ell(\bm{X}),\ell\\}}\sum_{r\in\mathcal{S}_{h}}\delta_{r}\Big{)}\bigg{|}$
$\displaystyle\leq\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\kappa(\bm{X})$
(198) $\displaystyle\leq\Theta\Big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\Big{)}.$

We plug (LABEL:eq:fjeiwbf) and (198) to obtain the derivative.

$\displaystyle\left|\frac{\partial L(\bm{X})}{\partial A_{i,j}}\right|$
$\displaystyle\leq\Theta(D\alpha^{(t)})\mathfrak{S}(-yF(\bm{X}))\Big{(}D\alpha^{(t)}\big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\big{)}\Big{)}^{p-1}\Xi^{(t)}.$
(199)

##### Putting all the pieces together.

We now bound the derivative of the population loss. Using Tower property and
and Lemma I.1, we have:

$\displaystyle\mathbb{E}_{X}\left[\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right]=\mathbb{E}_{\bm{X}}\left[\mathbb{E}_{\pi_{1},\pi_{2}}\left[\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\middle|\bm{X}\right]\right].$ (200)

For a fixed $\bm{X}$, we now bound the derivative of the loss evaluated in
$\bm{X}$. For $i\in\mathcal{S}_{\ell}$ and $j\in\mathcal{S}_{m}$, we
distinguish the eight possible events depending on the randomness of $\pi.$

  * –

Event a: "$\ell=\ell(\pi(\bm{X}))$ and $\delta_{j}=0$" occurs with probability
$(1-q)/L.$

  * –

Event b: "$\ell=\ell(\pi(\bm{X}))$ and $\delta_{j}\neq 0$" occurs with
probability $q/L.$

  * –

Event c: "$\ell\neq\ell(\pi(\bm{X}))$ and $\delta_{s}=0$ for all $s$ and
$m=\ell(\pi(\bm{X}))$" occurs with probability $(1-q)^{C}/L$.

  * –

Event dk: "$\ell\neq\ell(\pi(\bm{X}))$ and $\\#\\{s:\delta_{s}\neq 0\\}=k$ for
$1\leq k\leq C$ and $m=\ell(\pi(\bm{X}))$" occurs with probability
$\binom{C}{k}q^{k}(1-q)^{C-k}/L$.

  * –

Event e: "$\ell,m\neq\ell(\pi(\bm{X}))$ and $\delta_{j}=0$ and $\delta_{s}=0$
for all $s$" occurs with probability $(1-2/L)(1-q)^{C+1}.$

  * –

Event fk: "$\ell,m\neq\ell(\pi(\bm{X}))$ and $\\#\\{s:\delta_{s}\neq 0\\}=k$
for $1\leq k\leq C$ and $\delta_{j}=0$" occurs with probability
$(1-2/L)(1-q)\binom{C}{k}q^{k}(1-q)^{C-k}$.

  * –

Event g: "$\ell,m\neq\ell(\pi(\bm{X}))$ and $\delta_{s}=0$ for all $s$ and
$\delta_{j}\neq 0$" occurs with probability $(1-2/L)(1-q)^{C}q.$

  * –

Event hk: "$\ell,m\neq\ell(\pi(\bm{X}))$ and $\\#\\{s:\delta_{s}\neq 0\\}=k$
for $0\leq k\leq C$ and $\delta_{j}\neq 0$" occurs with probability
$(1-2/L)q\binom{C}{k}q^{k}(1-q)^{C-k}$

Since events a and e are the ones with highest probabilities, the derivative
of the loss is bounded by the expectations conditioned on these events. We
have:

$\displaystyle\mathbb{E}_{\pi_{1},\pi_{2}}\Big{[}\Big{|}\frac{\partial
L(\pi(\bm{X}))}{\partial A_{i,j}}\Big{|}\Big{]}$ (201) $\displaystyle\leq$
$\displaystyle\frac{\Theta(D\alpha^{(t)})}{L}\Xi^{(t)}\mathfrak{S}(-yF(\bm{X}))\big{(}D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\big{)}^{p-1}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})$
$\displaystyle+$
$\displaystyle\Theta(D\alpha^{(t)})\Xi^{(t)}\mathfrak{S}(-yF(\bm{X}))\mathbb{E}_{\pi_{1},\pi_{2}}\big{[}\big{(}D\alpha^{(t)}\kappa(\pi(\bm{X}))\Xi^{(t)}\big{)}^{p-1}\Xi^{(t)}\kappa(\pi(\bm{X}))\Big{|}\mathbf{e}\Big{]}.$

We now apply Induction Hypothesis B.1 and Lemma K.3 and finally obtain:

$\displaystyle\mathbb{E}_{\bm{X}}\Big{[}\Big{|}\frac{\partial
L(\bm{X})}{\partial A_{i,j}}\Big{|}\Big{]}$
$\displaystyle\leq\Theta(C^{2}\alpha^{(t)})\mathbb{E}_{\bm{X}}[\mathfrak{S}(-yF(\bm{X}))]\big{(}D\alpha^{(t)}(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\big{)}^{p-1}\Xi^{(t)}\frac{\lambda_{0}}{D}$
(202) $\displaystyle+$
$\displaystyle\Theta(\alpha^{(t)})\mathbb{E}_{\bm{X}}[\mathfrak{S}(-yF(\bm{X}))]\big{(}D\alpha^{(t)}qD\log(d)\Xi^{(t)}\big{)}^{p-1}\Xi^{(t)}qD\log(d).$

Since $\mathbb{E}_{\bm{X}}[\mathfrak{S}(-yF(\bm{X}))]\leq 1$, we thus plug
(202) in the update of $\rho^{(t)}$ to obtain the aimed result.

∎

######  Corollary G.2.

Let $T>0$ be the time where the population loss is $o(1)$ and $t\leq T.$ The
update of $\rho^{(t)}$ satisfies:

$\displaystyle|\rho^{(t+1)}|$
$\displaystyle\leq|\rho^{(t)}|+\eta\cdot\mathrm{polylog}(d)(\alpha^{(t)})^{p}\left(\frac{1}{D}+\frac{\lambda_{0}}{D}\Xi^{(t)}(G^{(t)})^{p-1}\right).$

###### Proof of Corollary G.2.

Using Lemma G.2, $C=\mathrm{polylog}(d)$ and $qD\leq 1$, $\rho^{(t)}$’s update
is:

$\displaystyle|\rho^{(t+1)}|$
$\displaystyle\leq|\rho^{(t)}|+\eta(\alpha^{(t)})^{p}\mathrm{polylog}(d)\left((D\Xi^{(t)})^{p-1}\Xi^{(t)}+\frac{\lambda_{0}}{D}(G^{(t)})^{p-1}\Xi^{(t)}\right).$
(203)

Lastly, we apply Induction Hypothesis B.1 to replace $\Xi^{(t)}$ by its value
in (203) and thus obtain the aimed result. ∎

###  G.3 Auxiliary lemmas

######  Lemma G.3.

Let $t>0.$ In the idealized learning process, with high probability, we have:

$\displaystyle\Lambda^{(t)}+(C-1)\Gamma^{(t)}+y\Xi^{(t)}\sum_{\ell\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{\ell}}\delta_{r}=\Theta\big{(}\Lambda^{(t)}+(C-1)\Gamma^{(t)}\big{)}.$

###### Proof of Lemma G.3.

We first bound the sum with factor $\Xi^{(t)}$. Using Induction Hypothesis B.1
and Lemma K.3, we have:

$\displaystyle\Big{|}y\Xi^{(t)}\sum_{\ell\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{\ell}}\delta_{r}\Big{|}\leq\Theta(q)\log(d)=\frac{1}{D}\mathrm{poly}(C)\log(d).$
(204)

We now bound $\Lambda^{(t)}+(C-1)\Gamma^{(t)}$. Using Induction Hypothesis
B.1, we have:

$\displaystyle\frac{1}{D}\left(e^{\mathrm{polyloglog}(d)}+\Omega(C)\right)\leq(\Lambda^{(t)}+(C-1)\Gamma^{(t)})\leq\frac{1}{D}\left(e^{\mathrm{polyloglog}(d)}+\lambda_{0}\right).$
(205)

Since $\mathrm{poly}(C)\log(d)\leq e^{\mathrm{polyloglog}(d)}/2$, we combine
(204) and (205) to get the aimed result. ∎

######  Lemma G.4.

Let $t>0$. In the idealized learning process, we have with high probability:

$\displaystyle\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|\leq\Theta(1).$

###### Proof of Lemma G.4.

We successively apply Lemma K.3 and Induction Hypothesis B.1 to get the
desired bound. Indeed, we have:

$\displaystyle\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}\sum_{h\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{h}}|\delta_{r}|$
$\displaystyle\leq\Lambda^{(t)}+(C-1)\Gamma^{(t)}+\Xi^{(t)}qD\log(d)$
$\displaystyle\leq\frac{1}{D}\left(e^{\mathrm{polyloglog}(d)}+\lambda_{0}+qD\log(d)\right)$
$\displaystyle=\Theta(1).$

∎

##  Appendix H Gradients

In this section, we present the gradients of the loss $\mathcal{L}$ with
respect to $\bm{v}$ and $\bm{A}_{i,j}$.

######  Lemma H.1.

Let $(\bm{X},y)$ be a data-point. Then, the gradient of $L(\bm{X})$ with
respect to $\bm{v}$ is:

$\displaystyle-\nabla_{v}L(\bm{X})$
$\displaystyle=Dy\mathfrak{S}(-yF(\bm{X}))\sum_{i=1}^{D}\sigma^{\prime}\bigg{(}D\sum_{r\in[D]}\bm{S}_{i,r}\langle\bm{v},\bm{X}_{r}\rangle\bigg{)}\sum_{j=1}^{D}\bm{S}_{i,j}\bm{X}_{j}.$

######  Lemma H.2.

Let $(\bm{X},y)$ be a data-point and $i,j\in[D]$. The derivative of
$L(\bm{X})$ with respect to $A_{i,j}$ is:

$\displaystyle-\frac{\partial L(\bm{X})}{\partial A_{i,j}}$
$\displaystyle=pDy\mathfrak{S}(-yF(\bm{X}))\bigg{(}D\sum_{r\in[D]}\bm{S}_{i,r}\langle\bm{v},\bm{X}_{r}\rangle\bigg{)}^{p-1}\sum_{r\neq
j}\bm{S}_{i,r}\langle\bm{v},\bm{X}_{j}-\bm{X}_{r}\rangle.$

##  Appendix I Invariance of the problem

###  I.1 Invariance of the parameters

See 4.2

###### Proof of Lemma 4.2.

The proof is by induction. Our induction hypothesis is
$\bm{v}^{(t)}=\alpha^{(t)}\bm{w}^{*}$ for all $t\geq 0.$ For $t=0$, we know
that $\bm{v}^{(0)}=\alpha^{(0)}\bm{w}^{*}\in\mathrm{span}(\bm{w}^{*}).$ Assume
that $\bm{v}^{(t)}=\alpha^{(t)}\bm{w}^{*}$. Let’s show that there exists
$\alpha^{(t+1)}\in\mathbb{R}$ such that
$\bm{v}^{(t+1)}=\alpha^{(t+1)}\bm{w}^{*}$. From the update rule, we have:

$\displaystyle\bm{v}^{(t+1)}$
$\displaystyle=\bm{v}^{(t)}+\eta\mathbb{E}\Big{[}y\mathfrak{S}(-yF(\bm{X}))\sum_{i\in[D]}\sigma^{\prime}(\langle\bm{v}^{(t)},\bm{O}_{i}^{(t)}\rangle)\sum_{r\in[D]}S_{i,r}\bm{X}_{r}\Big{]}$
$\displaystyle=\bm{v}^{(t)}+\eta\mathbb{E}\Big{[}y\mathfrak{S}(-yF(\bm{X}))\sum_{i\in[D]}\sigma^{\prime}(\langle\bm{v}^{(t)},\bm{O}_{i}^{(t)}\rangle)\sum_{r\in[D]}S_{i,r}\mathbb{E}_{\bm{\xi}}[\bm{X}_{r}]\Big{]},$
(206)

where $\mathbb{E}_{\bm{\xi}}$ is the expectation with respect to noise vectors
$\bm{\xi}_{r}$. Using the definition of the data distribution and
$\mathbb{E}[\bm{\xi}_{r}]=0$, we simplify the update (206):

$\displaystyle\bm{v}^{(t+1)}$
$\displaystyle=\bm{v}^{(t)}+\eta\mathbb{E}\Big{[}\mathfrak{S}(-yF(\bm{X}))\sum_{i\in[D]}\sigma^{\prime}(\langle\bm{v}^{(t)},\bm{O}_{i}^{(t)}\rangle)\sum_{r\in\mathcal{S}_{\ell(X)}}S_{i,r}\Big{]}\bm{w}^{*}$
(207)
$\displaystyle+\eta\mathbb{E}\Big{[}y\mathfrak{S}(-yF(\bm{X}))\sum_{i\in[D]}\sigma^{\prime}(\langle\bm{v}^{(t)},\bm{O}_{i}^{(t)}\rangle)\sum_{h\neq\ell(X)}\sum_{r\in\mathcal{S}_{h}}S_{i,r}\delta_{r}\Big{]}\bm{w}^{*}.$

Since $\bm{v}^{(t)}=\alpha^{(t)}\bm{w}^{*}$, there exists
$\alpha^{(t+1)}\in\mathbb{R}$ such that
$\bm{v}^{(t+1)}=\alpha^{(t+1)}\bm{w}^{*}.$

∎

See 4.1

###### Proof of Lemma 4.1.

We initialize $A_{i,i}^{(0)}=\sigma_{\bm{A}}$ and do not update $A_{i,i}$,
Thus, for all $t\geq 0,$ we have $A_{i,i}^{(t)}=\sigma_{\bm{A}}.$
The remaining of the proof is by induction. For $i,j\in\mathcal{S}_{\ell}$, we
initialize $A_{i,j}^{(0)}=0=\gamma^{(0)}$. For $i\in\mathcal{S}_{\ell}$ and
$j\in\mathcal{S}_{m}$ with $\ell\neq m$, $A_{i,j}^{(0)}=0=\rho^{(0)}$. The
induction hypothesis is true for $t=0.$

We assume that for all $i,j\in\mathcal{S}_{\ell}$,
$A_{i,j}^{(t)}=\gamma^{(t)}$ and for $i\in\mathcal{S}_{\ell}$ and
$j\in\mathcal{S}_{m}$, $A_{i,j}^{(t)}=\rho^{(t)}$. Let’s first prove that
$A_{i,j}^{(t+1)}=A_{k,n}^{(t+1)}$ for $i,j\in\mathcal{S}_{\ell}$ and
$k,n\in\mathcal{S}_{\ell^{\prime}}$. Since the GD update is
$A_{i,j}^{(t+1)}=A_{i,j}^{(t)}-\eta\mathbb{E}_{\bm{X}}\left[\frac{\partial
L(\bm{X})}{\partial A_{i,j}}\right]$, it’s sufficient to prove that
$\mathbb{E}_{\bm{X}}\left[\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right]=\mathbb{E}_{\bm{X}}\left[\frac{\partial L(\bm{X})}{\partial
A_{k,n}}\right]$. Let $\pi_{1}\colon[L]\rightarrow[L]$,
$\pi_{2}\colon[C]\rightarrow[C]$ and $\pi=(\pi_{1},\pi_{2})$ be permutations.
From Lemma I.1, we know that $\bm{X}$ and $\pi(\bm{X})$ have the same
distribution which implies $\mathbb{E}_{\bm{X}}\left[\frac{\partial
L(\bm{X})}{\partial
A_{i,j}}\right]=\mathbb{E}_{\pi(\bm{X})}\left[\frac{\partial
L(\bm{X})}{\partial A_{i,j}}\right]$. Therefore, we have:

$\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right]$ (208) $\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{X}}\Bigg{[}y\mathfrak{S}(-yF(\pi(\bm{X})))\sigma^{\prime}\bigg{(}\sum_{h=1}^{L}\sum_{r\in\mathcal{S}_{h}}\frac{e^{\langle\bm{A}^{(t)}\bm{p}_{i},\;\bm{p}_{r}\rangle}}{\sum_{s=1}^{D}e^{\langle\bm{A}^{(t)}\bm{p}_{i},\;\bm{p}_{s}\rangle}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(h),\pi_{2}(r)}\rangle\bigg{)}\cdot$
$\displaystyle\hskip
11.38092pt\frac{e^{\langle\bm{A}^{(t)}\bm{p}_{i},\;\bm{p}_{j}\rangle}}{\sum_{s=1}^{D}e^{\langle\bm{A}^{(t)}\bm{p}_{i},\;\bm{p}_{s}\rangle}}\sum_{k=1}^{L}\sum_{m\in\mathcal{S}_{k}}\frac{e^{\langle\bm{A}^{(t)}\bm{p}_{i},\;\bm{p}_{m}\rangle}}{\sum_{s=1}^{D}e^{\langle\bm{A}^{(t)}\bm{p}_{i},\;\bm{p}_{s}\rangle}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(k),\pi_{2}(m)}-\bm{X}_{\pi_{1}(\ell),\pi_{2}(j)}\rangle\Bigg{]}$

Using Induction Hypothesis B.1, we simplify (LABEL:eq:fewjoejdwe) as

$\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial\mathcal{L}^{(t)}(\bm{X})}{\partial
A_{i,j}}\right]$ (209) $\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{X}}\Bigg{[}y\mathfrak{S}(-yF(\pi(\bm{X})))\sigma^{\prime}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(i)}\rangle+\Gamma^{(t)}\sum_{m\in\mathcal{S}_{\ell}\backslash\\{i\\}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(m)}\rangle$
$\displaystyle\hskip
8.5359pt+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{\pi(r)}\rangle\bigg{)}\Gamma^{(t)}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(i)}-\bm{X}_{\pi_{1}(\ell),\pi_{2}(j)}\rangle$
$\displaystyle\hskip
8.5359pt+\Gamma^{(t)}\sum_{m\in\mathcal{S}_{\ell}\backslash\\{i\\}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(m)}-\bm{X}_{\pi_{1}(\ell),\pi_{2}(j)}\rangle+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{\pi(r)}-\bm{X}_{\pi_{1}(\ell),\pi_{2}(j)}\rangle\bigg{)}\Bigg{]}.$

We now set $\pi_{1}$ and $\pi_{2}$ such that $\pi_{1}(\ell)=\ell^{\prime}$,
$\pi_{2}(i)=k$ and $\pi_{2}(j)=n$. Using this choice along with
$F(\bm{X})=F(\pi(\bm{X}))$ (Lemma I.1), we finally have in
(LABEL:eq:fewjoejvferdwe)

$\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right]$ (210) $\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{X}}\Bigg{[}y\mathfrak{S}(-yF(\pi(\bm{X})))\sigma^{\prime}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},k}\rangle+\Gamma^{(t)}\sum_{m\in\mathcal{S}_{\ell^{\prime}}\backslash\\{i\\}}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},m}\rangle$
$\displaystyle\hskip
8.5359pt+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell^{\prime}}}\langle\bm{v}^{(t)},\bm{X}_{r}\rangle\bigg{)}\Gamma^{(t)}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},k}-\bm{X}_{\ell^{\prime},n}\rangle$
$\displaystyle\hskip
8.5359pt+\Gamma^{(t)}\sum_{m\in\mathcal{S}_{\ell^{\prime}}\backslash\\{k\\}}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},m}-\bm{X}_{\ell^{\prime},n}\rangle+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{r}-\bm{X}_{\ell^{\prime},n}\rangle\bigg{)}\Bigg{]}$
$\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial\mathcal{L}^{(t)}(\bm{X})}{\partial
A_{k,n}}\right].$

Therefore, (LABEL:eq:fdvdvwe) implies that $A_{i,j}^{(t+1)}=A_{k,n}^{(t+1)}$
thus proving the induction hypothesis. Let’s now show that for
$i\in\mathcal{S}_{\ell},j\in\mathcal{S}_{m}$ and
$k\in\mathcal{S}_{\ell^{\prime}},j\in\mathcal{S}_{m^{\prime}}$, we have
$A_{i,j}^{(t+1)}=A_{k,n}^{(t+1)}$. We apply a similar argument as above. Using
Induction Hypothesis B.1, we have:

$\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right]$ (211) $\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{X}}\Bigg{[}y\mathfrak{S}(-yF(\pi(\bm{X})))\sigma^{\prime}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(i)}\rangle+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(r)}\rangle$
$\displaystyle\hskip
8.5359pt+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{\pi(r)}\rangle\bigg{)}\Xi^{(t)}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(i)}-\bm{X}_{\pi_{1}(m),\pi_{2}(j)}\rangle$
$\displaystyle\hskip
8.5359pt+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{i\\}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{1}(\ell),\pi_{2}(r)}-\bm{X}_{\pi_{1}(m),\pi_{2}(j)}\rangle+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{\pi(r)}-\bm{X}_{\pi_{1}(m),\pi_{2}(j)}\rangle\bigg{)}\Bigg{]}.$

We now set $\pi_{1}$ and $\pi_{2}$ such that $\pi_{1}(\ell)=\ell^{\prime}$,
$\pi_{1}(m)=m^{\prime}$, $\pi_{2}(i)=k$ and $\pi_{2}(j)=n$. Using this choice,
we finally have in (LABEL:eq:fvfejhhrdwe)

$\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial L(\bm{X})}{\partial
A_{i,j}}\right]$ (212) $\displaystyle=$
$\displaystyle\mathbb{E}_{\bm{X}}\Bigg{[}y\mathfrak{S}(-yF(\pi(\bm{X})))\sigma^{\prime}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},k}\rangle+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell^{\prime}}\backslash\\{k\\}}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},r}\rangle$
$\displaystyle\hskip
8.5359pt+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell^{\prime}}}\langle\bm{v}^{(t)},\bm{X}_{r}\rangle\bigg{)}\Xi^{(t)}\bigg{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},k}-\bm{X}_{m^{\prime},n}\rangle$
$\displaystyle\hskip
8.5359pt+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell^{\prime}}\backslash\\{k\\}}\langle\bm{v}^{(t)},\bm{X}_{\ell^{\prime},r}-\bm{X}_{m^{\prime},n}\rangle+\Xi^{(t)}\sum_{r\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{r}-\bm{X}_{m^{\prime},n}\rangle\bigg{)}\Bigg{]}$
$\displaystyle=$ $\displaystyle\mathbb{E}_{\bm{X}}\left[\frac{\partial
L(\bm{X})}{\partial A_{k,n}}\right].$

Therefore, (LABEL:eq:fe) implies that $A_{i,j}^{(t+1)}=A_{k,n}^{(t+1)}$ thus
proving the induction hypothesis.

∎

###  I.2 Invariance by permutation

######  Lemma I.1.

Let $\pi_{1}\colon[L]\rightarrow[L]$ and $\pi_{2}\colon[C]\rightarrow[C]$ be
two permutations and $\pi=(\pi_{1},\pi_{2})$. Let
$(\bm{X},\cdot)\sim\mathcal{D}$. Then, we have:

  1. 1.

permutation-invariant distribution: $\bm{X}$ has the same distribution as
$\pi(\bm{X}).$

  2. 2.

permutation-invariant model: $F(\bm{X})=F(\pi(\bm{X})).$

###### Proof of Lemma I.1.

To show that $\bm{X}$ and $\pi(\bm{X})$ have to same distribution, it is
sufficient to show that the items 1 to 6 hold in our definition of the data
distribution. We still have that the label $y$ is uniformly sampled on
$\\{-1,1\\}$. Let
$\pi(\bm{X})=(\widebar{\bm{X}}_{1},\dots,\widebar{\bm{X}}_{D})$ where
$\widebar{\bm{X}}_{i}$ for some $j.$ For 2, the number of tokens is still $D$
after permutation $\pi$. For 4, we define the same partition with
$\widebar{\mathcal{S}}_{l}=\widebar{\mathcal{S}}_{\pi_{1}(l)}$. Besides, we
have $\ell(\bm{\widebar{X}})=\pi_{1}(\ell(\bm{X}))$ that is also uniformly
sampled on $[L]$ since the permutation on a uniform distribution is also
uniform. For $i\in\mathcal{S}_{\ell(\bm{\widebar{X}})}$, we have that
$\bm{\widebar{X}}_{i}$ writes $\bm{X}_{\pi_{2}(k)}$ for some
$k\in\mathcal{S}_{\ell(\bm{X})}$ so that we do have
$\bm{\widebar{X}}_{i}=y\bm{w}^{*}+\bm{\xi}_{i}$. The same goes for 6 when
$\ell\neq\ell(\bm{\widebar{X}})$: with $\ell\neq\ell(\bm{\widebar{X}})$,
$\bm{\widebar{X}}_{j}=\delta_{\ell,j}\bm{w}^{*}+\bm{\xi}_{j}$, where
$\delta_{j}=1$ with probability $q/2$, $-1$ with the same probability and $0$
otherwise.

Let $\bm{X}$ be a data-point. Using Lemma 4.1, we rewrite $F(\bm{X})$ as

$\displaystyle F(\bm{X})$
$\displaystyle=\sum_{m=1}^{D}\sigma\Big{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{m}\rangle+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\ell}\backslash\\{m\\}}\langle\bm{v}^{(t)},\bm{X}_{r}\rangle+\Xi^{(t)}\sum_{h\in[L]\backslash\\{\ell\\}}\sum_{j\in\mathcal{S}_{h}}\langle\bm{v}^{(t)},\bm{X}_{j}\rangle\Big{)}.$
(213)

Let $\pi=(\pi_{1},\pi_{2})$ be a permutation. For a given $\ell\in[L]$ and
$m\in\mathcal{S}_{\ell}$, assume that $\pi_{1}(\ell^{\prime})=\ell$ and
$\pi_{2}(m^{\prime})=m.$ Using Lemma 4.1, we have:

$\displaystyle F(\pi(\bm{X}))$
$\displaystyle=\sum_{\ell=1}^{L}\sum_{m\in\mathcal{S}_{\ell}}\sigma\Big{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{\pi_{2}(m^{\prime})}\rangle$
$\displaystyle+\Gamma^{(t)}\sum_{r\in\mathcal{S}_{\pi_{1}(\ell^{\prime})}\backslash\\{\pi_{2}(m^{\prime})\\}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{2}(r)}\rangle+\Xi^{(t)}\sum_{h^{\prime}\in[L]\backslash\\{\ell^{\prime}\\}}\sum_{j\in\mathcal{S}_{h^{\prime}}}\langle\bm{v}^{(t)},\bm{X}_{\pi_{2}(j^{\prime})}\rangle\Big{)}$
$\displaystyle=\sum_{\ell=1}^{L}\sum_{m\in\mathcal{S}_{\ell}}\sigma\Big{(}\Lambda^{(t)}\langle\bm{v}^{(t)},\bm{X}_{m}\rangle$
$\displaystyle+\Gamma^{(t)}\sum_{r^{\prime\prime}\in\mathcal{S}_{\ell}\backslash\\{m\\}}\langle\bm{v}^{(t)},\bm{X}_{r^{\prime\prime}}\rangle+\Xi^{(t)}\sum_{j^{\prime\prime}\not\in\mathcal{S}_{\ell}}\langle\bm{v}^{(t)},\bm{X}_{j^{\prime\prime}}\rangle\Big{)}$
$\displaystyle=F(\bm{X}).$

∎

##  Appendix J Justification of our data distribution

In this section, we justify why the distribution $\mathcal{D}$ (Assumption 1)
is relevant. We first show that linear classifiers poorly generalize
(subsection J.1). We then show that there exists classifiers that generalize
without learning patch association (subsection J.2).

###  J.1 Generalized linear models poorly generalize

See 3.1

###### Proof of Theorem 3.1.

For every data point $\bm{X}$, consider
$\Delta(\bm{X}):=\sum_{j\in[D]}\delta_{j}$, it is very easy to see that for
every integer $p$, as long as $\Pr[\Delta=p]=\Omega(1/\text{polylog}(d))$, we
have that:

$\displaystyle\frac{\Pr[\Delta=p]}{\Pr[\Delta=p-2C]}=1-o(1)$ (214)

Consider two independently sampled data points, $\bm{X},\bm{X}^{\prime}$ with
label $1,-1$ respectively, consider the event when
$\Delta(\bm{X})=p-2C,\Delta(\bm{X}^{\prime})=p$ and all the noises
$\xi_{i},\xi_{i}^{\prime}$ of $\bm{X},\bm{X}^{\prime}$ satisfies
$\xi_{i}=\xi_{i}^{\prime}$, then we know that

$\displaystyle\sum_{j\in[D]}\langle
w_{j},\bm{X}_{j}\rangle=\sum_{j\in[D]}\langle
w_{j},\bm{X}_{j}^{\prime}\rangle$ (215)

By Eq (214) we also know that the density of $\bm{X}$ and $\bm{X}^{\prime}$
under the data-generation distribution satisfies

$p(\bm{X})=(1\pm o(1))p(\bm{X}^{\prime})$

Now, we know that

$\displaystyle P_{0}:=$ $\displaystyle\mathbb{P}[f^{*}(\bm{X})g(\bm{X})\leq
0\mid y(\bm{X})=1]=2\int_{\bm{X}}1_{f^{*}(\bm{X})g(\bm{X})\leq
0}p(\bm{X},y(\bm{X})=1)d\bm{X}$ (216)
$\displaystyle=2\int_{\bm{X}}1_{f^{*}(\bm{X})g(\bm{X})\leq
0}p(\bm{X}^{\prime},y(\bm{X}^{\prime})=-1)d\bm{X}\pm o(1)$ (217)
$\displaystyle=2\int_{\bm{X}}1_{f^{*}(\bm{X})g(\bm{X}^{\prime})\leq
0}p(\bm{X}^{\prime},y(\bm{X}^{\prime})=-1)d\bm{X}\pm o(1)$ (218)
$\displaystyle=2\int_{\bm{X}^{\prime}}1_{f^{*}(\bm{X}^{\prime})g(\bm{X}^{\prime})\geq
0}p(\bm{X}^{\prime},y(\bm{X}^{\prime})=-1)d\bm{X}^{\prime}\pm o(1)$ (219)
$\displaystyle=1-\Pr[f^{*}(\bm{X}^{\prime})g(\bm{X}^{\prime})<0\mid
y(\bm{X^{\prime}})=-1]\pm o(1)$ (220) $\displaystyle\geq
1-\Pr[f^{*}(\bm{X}^{\prime})g(\bm{X}^{\prime})\leq 0\mid
y(\bm{X^{\prime}})=-1]\pm o(1)$ (221)
$\displaystyle=1-\Pr[f^{*}(\bm{X})g(\bm{X})\leq 0\mid y(\bm{X})=-1]\pm o(1)$
(222)

Therefore, $\mathbb{P}[f^{*}(\bm{X})g(\bm{X})\leq 0\mid
y(\bm{X})=1]\geq\frac{1}{2}-o(1)$. ∎

###  J.2 Classifiers fitting the labelling function without patch association

See 3.2

###### Proof of Theorem 3.2.

We can consider a transformer in our setting, whose weights are defined as:
$\bm{v}=\bm{w}^{\star}$, $A_{i,j}=\beta>0$ for $i,j\in\mathcal{S}_{\ell}$.
$\bm{A}_{i,j}=2\beta$ for $i\in\mathcal{S}_{\ell},j\in\mathcal{S}_{\ell+1}$
(We denote $\mathcal{S}_{L+1}=\mathcal{S}_{1}$). For a sufficiently large
$\beta$, it is easy to check that
$\mathbb{P}[f^{*}(\bm{X})\mathcal{M}(\bm{X})\leq 0]=d^{-\omega(1)}$ but for
all $\ell\in[L]$, $i\in\mathcal{S}_{\ell}$,
$\mathrm{Top}_{C}\;\\{\langle\bm{p}_{i}^{(\mathcal{M})},\bm{p}_{j}^{(\mathcal{M})}\rangle\\}_{j=1}^{D}=\mathcal{S}_{\ell+1}\cap\mathcal{S}_{\ell}=\emptyset$.

∎

##  Appendix K Technical lemmas

In this section, we present the technical lemmas used in the paper.

###  K.1 Tensor Power Method

######  Lemma K.1.

Let $\\{z^{(t)}\\}_{t\geq 0}$ be a positive sequence defined by the following
recursions

$\displaystyle\begin{cases}z^{(t+1)}\geq z^{(t)}+m(z^{(t)})^{k}\\\
z^{(t+1)}\leq z^{(t)}+M(z^{(t)})^{k}\end{cases},$

where $z^{(0)}>0$ is the initialization, $k>1$ is an integer and $m,M>0$. Let
$\upsilon>0$ such that $z^{(0)}\leq\upsilon.$ Then, the time $\mathcal{T}$
such that $z^{(t)}\geq\upsilon$ for all $t\geq\mathcal{T}$ is:

$\displaystyle\mathcal{T}=\frac{3}{m(z^{(0)})^{k-1}}+\frac{2^{k+1}M}{m}\left\lceil\frac{\log(\upsilon/z^{(0)})}{\log(2)}\right\rceil.$

###### Proof of Lemma K.1.

Let $n\in\mathbb{N}^{*}$. Let $T_{n}$ be the time where $z^{(t)}\geq
2^{n}z^{(0)}$. This time exists because $z^{(t)}$ is a non-decreasing
sequence. We want to find an upper bound on this time. We start with the case
$n=1.$ By summing the recursion, we have:

$\displaystyle z^{(T_{1})}\geq z^{(0)}+m\sum_{s=0}^{T_{1}-1}(z^{(s)})^{k}.$
(223)

We use the fact that $z^{(s)}\geq z^{(0)}$ in (223) and obtain:

$\displaystyle T_{1}\leq\frac{z^{(T_{1})}-z^{(0)}}{m(z^{(0)})^{k}}.$ (224)

Now, we want to bound $z^{(T_{1})}-z^{(0)}$. Using again the recursion and
$z^{(T_{1}-1)}\leq 2z^{(0)}$, we have:

$\displaystyle z^{(T_{1})}\leq z^{(T_{1}-1)}+M(z^{(T_{1}-1)})^{k}\leq
2z^{(0)}+2^{k}M(z^{(0)})^{k}.$ (225)

Combining (224) and (225), we get a bound on $T_{1}.$

$\displaystyle T_{1}\leq\frac{1}{m(z^{(0)})^{k-1}}+\frac{2^{k}M}{m}.$ (226)

Now, let’s find a bound for $T_{n}$. Starting from the recursion and using the
fact that $z^{(s)}\geq 2^{n-1}z^{(0)}$ for $s\geq T_{n-1}$ we have:

$\displaystyle z^{(T_{n})}\geq
z^{(T_{n-1})}+m\sum_{s=T_{n-1}}^{T_{n}-1}(z^{(s)})^{k}\geq
z^{(T_{n-1})}+2^{k(n-1)}m(z^{(0)})^{k}(T_{n}-T_{n-1}).$ (227)

On the other hand, by using $z^{(T_{n}-1)}\leq 2^{n}z^{(0)}$ we upper bound
$z^{(T_{n})}$ as follows.

$\displaystyle z^{(T_{n})}$ $\displaystyle\leq
z^{(T_{n}-1)}+M(z^{(T_{n}-1)})^{k}\leq 2^{n}z^{(0)}+2^{kn}M(z^{(0)})^{k}.$
(228)

Besides, we know that $z^{(T_{n-1})}\geq 2^{n-1}z^{(0)}$. Therefore, we upper
bound $z^{(T_{n})}-z^{(T_{n-1})}$ as

$\displaystyle z^{(T_{n})}-z^{(T_{n-1})}\leq
2^{n-1}z^{(0)}+2^{kn}M(z^{(0)})^{k}.$ (229)

Combining (227) and (229) yields:

$\displaystyle T_{n}\leq
T_{n-1}+\frac{1}{2^{(k-1)(n-1)}m(z^{(0)})^{k-1}}+\frac{2^{k}M}{m}.$ (230)

We now sum (230) for $n=2,\dots,n$, use (226) and obtain:

$\displaystyle T_{n}\leq
T_{1}+\frac{2}{m(z^{(0)})^{k-1}}+\frac{2^{k}Mn}{m}\leq\frac{3}{m(z^{(0)})^{k-1}}+\frac{2^{k}M(n+1)}{m}\leq\frac{3}{m(z^{(0)})^{k-1}}+\frac{2^{k+1}Mn}{m}.$
(231)

Lastly, we know that $n$ satisfies $2^{n}z^{(0)}\geq\upsilon$ which implies
$n=\left\lceil\frac{\log(\upsilon/z_{0})}{\log(2)}\right\rceil$ in (231). ∎

######  Lemma K.2.

Let $\\{z^{(t)}\\}_{t\geq 0}$ be a positive sequence defined by the following
recursions

$\displaystyle\begin{cases}z^{(t+1)}\geq z^{(t)}+m(z^{(t)})^{k}\\\
z^{(t+1)}\leq z^{(t)}+M(z^{(t)})^{k}\end{cases},$

where $z^{(0)}>0$, $k>1$ is an integer and $m,M>0$. Let $\upsilon>0$ such that
$z^{(0)}\leq\upsilon$ and $\mathcal{T}$ be the time such that
$z^{(t)}\geq\upsilon$ for all $t\geq\mathcal{T}$. Assume that
$\frac{A\upsilon^{2}}{m}\ll 1$. Then, we have for $\kappa\in\\{1,2\\}$:

$\displaystyle\prod_{\tau=0}^{\mathcal{T}-1}\big{(}1+A(z^{(\tau)})^{k-1+\kappa}\big{)}$
$\displaystyle\leq\big{(}1+\upsilon^{k-1}\big{)}^{\frac{2^{\kappa+1}MA\upsilon^{\kappa}\log(\upsilon/z^{(0)})}{m}}.$

###### Proof of Lemma K.2.

Let $n\in\mathbb{N}^{*}$ and let $T_{n}$ be the time such that $z^{(t)}\geq
2^{n}z^{(0)}$ for $t\geq T_{n}$. Starting from the recursion, we have:

$\displaystyle\log(z^{(t+1)})$
$\displaystyle\geq\log(z^{(t)})+\log\big{(}1+m(z^{(t)})^{k-1}\big{)}.$ (232)

Since $z^{(t)}$ is a non-decreasing sequence, (232) satisfies:

$\displaystyle(z^{(t+1)})^{\kappa}\log(z^{(t+1)})$
$\displaystyle\geq(z^{(t)})^{\kappa}\log(z^{(t)})+(z^{(t)})^{\kappa}\log\big{(}1+m(z^{(t)})^{k-1}\big{)}.$
(233)

We now sum (233) for $t=T_{n-1},\dots,T_{n}$ and get:

$\displaystyle\frac{A}{m}(z^{(T_{n})})^{\kappa}\log(z^{(T_{n})})$
$\displaystyle\geq\frac{A}{m}(z^{(T_{n-1})})^{\kappa}\log(z^{(T_{n-1})})+\sum_{t=T_{n-1}}^{T_{n}-1}\frac{A(z^{(t)})^{\kappa}}{m}\log\big{(}1+m(z^{(t)})^{k-1}\big{)}.$
(234)

Since
$\frac{A(z^{(t)})^{\kappa}}{m}\leq\frac{A(2^{n}z^{(0)})^{\kappa}}{m}\leq\frac{A\upsilon^{\kappa}}{m}\ll
1$, we have
$\big{(}1+m(z^{(t)})^{k-1}\big{)}^{\frac{A(z^{(t)})^{\kappa}}{m}}\geq
1+A(z^{(t)})^{k-1+\kappa}$. We thus lower bound (234) as:

$\displaystyle\frac{A}{m}(z^{(T_{n})})^{\kappa}\log(z^{(T_{n})})$
$\displaystyle\geq\frac{A}{m}(z^{(T_{n-1})})^{\kappa}\log(z^{(T_{n-1})})+\sum_{t=T_{n-1}}^{T_{n}-1}\log\big{(}1+A(z^{(t)})^{k-1+\kappa}\big{)}.$
(235)

On the other hand, by using $z^{(T_{n}-1)}\leq 2^{n}z^{(0)}$ and
$z^{(T_{n})}\leq 2^{n+1}z^{(0)}$, we have the following upper bound.

$\displaystyle\frac{A}{m}(z^{(T_{n})})^{\kappa}\log(z^{(T_{n})})$ (236)
$\displaystyle\leq$
$\displaystyle\frac{A}{m}(z^{(T_{n})})^{\kappa}\log\big{(}z^{(T_{n}-1)})+\frac{MA}{m}(z^{(T_{n})})^{\kappa}\log(1+(z^{(T_{n}-1)})^{k-1}\big{)}$
$\displaystyle\leq$
$\displaystyle\frac{A\cdot(2^{n+1}z^{(0)})^{\kappa}}{m}\log\big{(}2^{n}z^{(0)})+\frac{MA\cdot(2^{n+1}z^{(0)})^{\kappa}}{m}\log\big{(}1+(2^{n}z^{(0)})^{k-1}\big{)}.$

Since $z^{(T_{n-1})}\geq 2^{n-1}z^{(0)}$, (236) is finally bounded as:

$\displaystyle\frac{A}{m}\Big{(}(z^{(T_{n})})^{\kappa}\log(z^{(T_{n})})-(z^{(T_{n-1})})^{\kappa}\log(z^{(T_{n-1})})\Big{)}$
(237) $\displaystyle\leq$
$\displaystyle\frac{\Theta(A)\cdot(2^{n}z^{(0)})^{\kappa}}{m}\log\big{(}2^{n}z^{(0)})+\frac{MA\cdot(2^{n+1}z^{(0)})^{\kappa}}{m}\log\big{(}1+(2^{n}z^{(0)})^{k-1}\big{)}.$

We combine (235) and (LABEL:eq:ewjeof) to obtain:

$\displaystyle\sum_{t=T_{n-1}}^{T_{n}-1}\log\big{(}1+m(z^{(t)})^{k}\big{)}$
$\displaystyle\leq\frac{\Theta(A)\cdot(2^{n}z^{(0)})^{\kappa}}{m}\log\big{(}2^{n}z^{(0)})$
(238)
$\displaystyle+\frac{MA\cdot(2^{n+1}z^{(0)})^{\kappa}}{m}\log\big{(}1+(2^{n}z^{(0)})^{k-1}\big{)}.$

We now sum (238) and get:

$\displaystyle\sum_{t=0}^{T_{n}-1}\log\big{(}1+m(z^{(t)})^{k}\big{)}$
$\displaystyle\leq
n(2^{n}z^{(0)})^{\kappa}\left(\frac{\Theta(A)}{m}\log\big{(}2^{n}z^{(0)})+\frac{2^{\kappa}MA}{m}\log\big{(}1+(2^{n}z^{(0)})^{k-1}\big{)}\right).$
(239)

We replace $2^{n}z^{(0)}$ by $\upsilon$ and $n$ by $\log(\upsilon/z^{(0)})$ in
(239) to get the aimed result. ∎

###  K.2 Probabilistic lemmas

######  Lemma K.3.

Let $\\{\delta_{r}\\}_{r=1}^{D-C}$ be i.i.d. random variables such that with
probability $q$ $\delta_{r}=\pm 1$ and zero otherwise. Then, with probability
at least $1-1/\mathrm{poly}(d)$, we have:

$\displaystyle\sum_{\ell\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{\ell}}|\delta_{r}|\leq\Theta(q(D-C))\log(d).$

###### Proof of Lemma K.3.

First, note that $|\delta_{r}|$ is a Bernoulli random variable with parameter
$q.$ Therefore,
$\sum_{\ell\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{\ell}}|\delta_{r}|$ is a
binomial random variable $\mathcal{B}(D-C,q).$ Therefore, we apply a Chernoff
bound to obtain:

$\displaystyle\mathbb{P}\left[\sum_{\ell\neq\ell(\bm{X})}\sum_{r\in\mathcal{S}_{\ell}}|\delta_{r}|\geq(1+\varepsilon)(D-C)q\right]\leq\exp\left(-\frac{(D-C)q\varepsilon^{2}}{3}\right).$
(240)

Setting $\varepsilon=\sqrt{\frac{3}{(D-C)q}\log(\mathrm{poly}(d))}$ in (240)
yields the desired result. ∎

######  Lemma K.4.

Let $\\{\delta_{r}\\}_{r=1}^{C}$ be i.i.d. random variables such that with
probability $q$ $\delta_{r}=\pm 1$ and zero otherwise. Then, with probability
at least $1-1/\mathrm{poly}(d)$, we have:

$\displaystyle\sum_{r=1}^{C}|\delta_{r}|\leq O(1).$

###### Proof of Lemma K.4.

Let $k\in\mathbb{N}$ and $\Delta:=\sum_{r=1}^{C}|\delta_{r}|.$ The tail bound
is bounded as:

$\displaystyle\mathbb{P}[\Delta\geq k]$
$\displaystyle=\sum_{j=k}^{C}\binom{C}{j}q^{j}(1-q)^{C-j}\leq
q^{k}\sum_{j=k}^{C}\binom{C}{j}\leq
q^{k}\sum_{j=0}^{C}\binom{C}{j}=2^{C}q^{k}.$ (241)

We want to find $k$ such that $2^{C}q^{k}\leq 1/\mathrm{poly}(d)$ which
implies $k\leq\log(d)/\log(D)\leq O(1)$. ∎

###  K.3 Logarithmic inequalities

######  Lemma K.5.

Let $a\in\mathbb{R}$ such that $C_{-}\leq a\leq C_{+}$, where $C_{+},C_{-}>0.$
Let $p\geq 3$ be an odd integer. Then, the following inequality holds:

$\displaystyle\frac{0.1}{C+}\log\left(1+\exp\left(-a^{p}\right)\right)\leq\frac{a^{p-1}}{1+\exp(a^{p})}\leq\frac{10}{C_{-}}\log\left(1+\exp\left(-a^{p}\right)\right).$

###### Proof of Lemma K.5.

We first remark that:

$\displaystyle\frac{a^{p-1}}{1+\exp(a^{p})}$
$\displaystyle=\frac{a^{p}}{a(1+\exp(a^{p}))}.$ (242)

##### Upper bound.

We upper bound (242) by applying $a\geq C_{-}$:

$\displaystyle\frac{a^{p-1}}{1+\exp(a^{p})}\leq\frac{a^{p}}{C_{-}(1+\exp(a^{p}))}$
(243)

We obtain the final bound by applying Lemma K.6 to (243).

##### Lower bound.

We lower bound (242) by using $a\leq C_{+}$:

$\displaystyle\frac{a^{p-1}}{1+\exp(a^{p})}\geq\frac{a^{p}}{C_{+}(1+\exp(a^{p}))}$
(244)

We obtain the final bound by applying Lemma K.6 to (244). ∎

######  Lemma K.6 (Connection between derivative and loss).

Let $x>0.$ Then, we have:

$\displaystyle 0.1\log(1+\exp(-x))\leq\mathfrak{S}(-x)\leq 10\log(1+\exp(-x))$
(245)

######  Lemma K.7.

Let $x,y>0.$ Assume that $y\leq x.$ Then, we have:

$\displaystyle\log(1+xy)\leq(1+y)\log(1+x).$