SlowGuess/NeurIPS_2023_OCR huggingface dataset:
is positive semi-definite, i.e., the loss is convex with respect to trainable parameters.
(2) For  $\tilde{\mathcal{L}} (\bar{\theta};\lambda (\cdot))$  , notice that
$$
\begin{array}{l} \| \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}}} (\boldsymbol {x}) \| _ {2} ^ {2} = \mathbb {E} _ {(\boldsymbol {w}, \boldsymbol {u}) \sim \rho_ {0}} \left[ \sigma (\boldsymbol {w} ^ {\top} \boldsymbol {x} + \boldsymbol {u} ^ {\top} \boldsymbol {e} (t)) \boldsymbol {a} ^ {\top} (\boldsymbol {w}, \boldsymbol {u}) \right] \mathbb {E} _ {(\boldsymbol {w} ^ {\prime}, \boldsymbol {u} ^ {\prime}) \sim \rho_ {0}} \left[ \boldsymbol {a} (\boldsymbol {w} ^ {\prime}, \boldsymbol {u} ^ {\prime}) \sigma (\boldsymbol {w} ^ {\prime \top} \boldsymbol {x} + \boldsymbol {u} ^ {\prime \top} \boldsymbol {e} (t)) \right] \\ = \mathbb {E} _ {\left(\boldsymbol {w}, \boldsymbol {u}\right), \left(\boldsymbol {w} ^ {\prime}, \boldsymbol {u} ^ {\prime}\right) \sim \rho_ {0}} \left[ \boldsymbol {a} ^ {\top} (\boldsymbol {w}, \boldsymbol {u}) \sigma \left(\boldsymbol {w} ^ {\top} \boldsymbol {x} + \boldsymbol {u} ^ {\top} \boldsymbol {e} (t)\right) \sigma \left(\boldsymbol {w} ^ {\prime \top} \boldsymbol {x} + \boldsymbol {u} ^ {\prime \top} \boldsymbol {e} (t)\right) \boldsymbol {a} \left(\boldsymbol {w} ^ {\prime}, \boldsymbol {u} ^ {\prime}\right) \right]. \\ \end{array}
$$
Let  $\pmb{v} \coloneqq (\pmb{w}, \pmb{u}), \pmb{v}' \coloneqq (\pmb{w}', \pmb{u}')$ ,  $z(t) \coloneqq (\pmb{x}^{\top}(t), e^{\top}(t))^{\top}$ , we get
$$
\begin{array}{l} \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}}; \lambda (\cdot)) = \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \left(\| \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}}} (\boldsymbol {x} (t)) \| _ {2} ^ {2} - 2 \boldsymbol {s} _ {t} ^ {\top} (\boldsymbol {x} (t)) \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}}} (\boldsymbol {x} (t)) + \| \boldsymbol {s} _ {t} (\boldsymbol {x} (t)) \| _ {2} ^ {2}\right) \right] \\ = \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \left(\mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ \boldsymbol {a} ^ {\top} (\boldsymbol {v}) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \sigma (\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)) \boldsymbol {a} (\boldsymbol {v} ^ {\prime}) \right] \right. \right. \\ \left. - 2 \boldsymbol {s} _ {t} ^ {\top} (\boldsymbol {x} (t)) \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} \left[ \boldsymbol {a} (\boldsymbol {v}) \sigma \left(\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)\right) \right]\right) \bigg ] + \text {c o n s t a n t} \\ = \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim p _ {0}} \left[ \boldsymbol {a} ^ {\top} (\boldsymbol {v}) \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \sigma \left(\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)\right) \sigma \left(\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)\right) \right]\right) \boldsymbol {a} \left(\boldsymbol {v} ^ {\prime}\right) \right] \\ - 2 \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} \left[ \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \boldsymbol {s} _ {t} ^ {\top} (\boldsymbol {x} (t)) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \right]\right) \boldsymbol {a} (\boldsymbol {v}) \right] + \text {c o n s t a n t}. \\ \end{array}
$$
<sup>7</sup>In fact, if  $A \in \mathbb{R}^{n \times n}$ ,  $B \in \mathbb{R}^{m \times m}$  have the eigenvalues  $\{\nu_i\}_{i=1}^n$ ,  $\{\mu_j\}_{j=1}^m$ , respectively, the eigenvalues of  $A \otimes B$  are  $\nu_i \mu_j$ ,  $i = 1, \dots, n$ ,  $j = 1, \dots, m$ . This is due to the following argument. By Jordan-Chevalley decomposition, there exist invertible matrices  $P, Q$  such that  $A = P \Lambda P^{-1}$ ,  $B = Q \Delta Q^{-1}$ , where  $\Lambda, \Delta$  are upper triangular matrices. Therefore, we have
$\pmb{A} \otimes \pmb{B} = (\pmb{P} \pmb{\Lambda} \pmb{P}^{-1}) \otimes (\pmb{Q} \pmb{\Delta} \pmb{Q}^{-1}) = (\pmb{P} \otimes \pmb{Q})(\pmb{\Lambda} \otimes \pmb{\Delta})(\pmb{P}^{-1} \otimes \pmb{Q}^{-1}) = (\pmb{P} \otimes \pmb{Q})(\pmb{\Lambda} \otimes \pmb{\Delta})(\pmb{P} \otimes \pmb{Q})^{-1}$ .
That is, the matrices  $\mathbf{A} \otimes \mathbf{B}$  and  $\boldsymbol{\Lambda} \otimes \boldsymbol{\Delta}$  are similar. Notice that  $\boldsymbol{\Lambda} \otimes \boldsymbol{\Delta}$  is still an upper triangular matrix with diagonal elements  $\nu_{i}\mu_{j}, i = 1,\dots ,n, j = 1,\dots ,m$ , we obtain the desired result.
Then for any  $\phi \in L^{2}(\rho_{0})$ , by symmetry we have
$$
\begin{array}{l} \delta \bar {\tilde {\mathcal {L}}} (\cdot ; \lambda (\cdot)) [ \bar {\boldsymbol {\theta}}, \phi ] = \lim  _ {\epsilon \rightarrow 0} \frac {1}{\epsilon} \left(\bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} + \epsilon \phi ; \lambda (\cdot)) - \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}}; \lambda (\cdot))\right) \\ \left. \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ \phi^ {\top} (\boldsymbol {v}) \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \sigma \left(\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)\right) \sigma \left(\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)\right) \right]\right) \boldsymbol {a} \left(\boldsymbol {v} ^ {\prime}\right) \right. \right. \\ \left. + \boldsymbol {a} ^ {\top} (\boldsymbol {v}) \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \sigma \left(\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)\right) \sigma \left(\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)\right) \right]\right) \phi \left(\boldsymbol {v} ^ {\prime}\right) \right] \\ - 2 \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} \left[ \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \boldsymbol {s} _ {t} ^ {\top} (\boldsymbol {x} (t)) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \right]\right) \phi (\boldsymbol {v}) \right] \\ = 2 \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim p _ {0}} \left[ \boldsymbol {a} ^ {\top} (\boldsymbol {v} ^ {\prime}) \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \sigma (\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)) \right]\right) \phi (\boldsymbol {v}) \right] \\ - 2 \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} \left[ \left(\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \boldsymbol {s} _ {t} ^ {\top} (\boldsymbol {x} (t)) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \right]\right) \boldsymbol {\phi} (\boldsymbol {v}) \right] \\ = 2 \left\langle \mathbb {E} _ {\boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) \boldsymbol {a} (\boldsymbol {v} ^ {\prime}) \right] - \boldsymbol {k} (\boldsymbol {v}; \lambda (\cdot)), \phi (\boldsymbol {v}) \right\rangle_ {L ^ {2} \left(\rho_ {0}\right)}, \\ \end{array}
$$
$$
K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) := \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \sigma (\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)) \right],
$$
$$
\boldsymbol {k} (\boldsymbol {v}; \lambda (\cdot)) := \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \boldsymbol {s} _ {t} (\boldsymbol {x} (t)) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \right].
$$
$$
\frac {\delta \tilde {\bar {\mathcal {L}}} (\cdot ; \lambda (\cdot))}{\delta \bar {\boldsymbol {\theta}}} = 2 \mathbb {E} _ {\boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) \boldsymbol {a} \left(\boldsymbol {v} ^ {\prime}\right) \right] - 2 \boldsymbol {k} (\boldsymbol {v}; \lambda (\cdot)), \tag {24}
$$
$$
\begin{array}{l} \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} _ {1}; \lambda (\cdot)) + \left\langle \frac {\delta \bar {\tilde {\mathcal {L}}} (\cdot ; \lambda (\cdot))}{\delta \bar {\boldsymbol {\theta}}} \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} _ {1}}, \bar {\boldsymbol {\theta}} _ {2} - \bar {\boldsymbol {\theta}} _ {1} \Bigg \rangle_ {L ^ {2} (\rho_ {0})} - \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} _ {2}; \lambda (\cdot)) \\ = \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ \boldsymbol {a} _ {1} ^ {\top} (\boldsymbol {v}) K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) \boldsymbol {a} _ {1} (\boldsymbol {v} ^ {\prime}) \right] - 2 \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} \left[ \boldsymbol {k} ^ {\top} (\boldsymbol {v}; \lambda (\cdot)) \boldsymbol {a} _ {1} (\boldsymbol {v}) \right] \\ + 2 \left\langle \mathbb {E} _ {\boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) \boldsymbol {a} _ {1} \left(\boldsymbol {v} ^ {\prime}\right) \right] - \boldsymbol {k} (\boldsymbol {v}; \lambda (\cdot)), \boldsymbol {a} _ {2} (\boldsymbol {v}) - \boldsymbol {a} _ {1} (\boldsymbol {v}) \right\rangle_ {L ^ {2} \left(\rho_ {0}\right)} \\ - \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ \boldsymbol {a} _ {2} ^ {\top} (\boldsymbol {v}) K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) \boldsymbol {a} _ {2} (\boldsymbol {v} ^ {\prime}) \right] + 2 \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} \left[ \boldsymbol {k} ^ {\top} (\boldsymbol {v}; \lambda (\cdot)) \boldsymbol {a} _ {2} (\boldsymbol {v}) \right] \\ = - \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ \left(\boldsymbol {a} _ {2} (\boldsymbol {v}) - \boldsymbol {a} _ {1} (\boldsymbol {v})\right) ^ {\top} K (\boldsymbol {v}, \boldsymbol {v} ^ {\prime}; \lambda (\cdot)) \left(\boldsymbol {a} _ {2} (\boldsymbol {v} ^ {\prime}) - \boldsymbol {a} _ {1} (\boldsymbol {v} ^ {\prime})\right) \right] \\ = - \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \mathbb {E} _ {\boldsymbol {v}, \boldsymbol {v} ^ {\prime} \sim \rho_ {0}} \left[ (\boldsymbol {a} _ {2} (\boldsymbol {v}) - \boldsymbol {a} _ {1} (\boldsymbol {v})) ^ {\top} \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) \sigma (\boldsymbol {v} ^ {\prime \top} \boldsymbol {z} (t)) (\boldsymbol {a} _ {2} (\boldsymbol {v} ^ {\prime}) - \boldsymbol {a} _ {1} (\boldsymbol {v} ^ {\prime})) \right] \right] \\ = - \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} [ \lambda (t) \| \mathbb {E} _ {\boldsymbol {v} \sim \rho_ {0}} [ (\boldsymbol {a} _ {2} (\boldsymbol {v}) - \boldsymbol {a} _ {1} (\boldsymbol {v})) \sigma (\boldsymbol {v} ^ {\top} \boldsymbol {z} (t)) ] \| _ {2} ^ {2} ] \leq 0, \\ \end{array}
$$
hence the functional  $\tilde{\mathcal{L}} (\bar{\theta};\lambda (\cdot))$  is convex with respect to  $\bar{\theta}$  (given any positive weighting function  $\lambda (\cdot)$ ). The proof is completed.
Since the weighting function is fixed in our analysis, we omit the notation  $\lambda(\cdot)$  without ambiguity in the following contents.
Lemma 4. For any  $\tau >0$  and  $\theta ,\bar{\theta}$  we have
$$
\bar {\tilde {\mathcal {L}}} \left(\bar {\boldsymbol {\theta}} (\tau)\right) - \tilde {\bar {\mathcal {L}}} \left(\bar {\boldsymbol {\theta}}\right) \lesssim \frac {\left\| \bar {\boldsymbol {s}} _ {0 , \bar {\boldsymbol {\theta}} _ {0}} \right\| _ {\mathcal {H}} ^ {2} + \left\| \bar {\boldsymbol {s}} _ {0 , \bar {\boldsymbol {\theta}}} \right\| _ {\mathcal {H}} ^ {2}}{\tau}, \quad \tilde {\mathcal {L}} \left(\boldsymbol {\theta} (\tau)\right) - \tilde {\mathcal {L}} \left(\boldsymbol {\theta}\right) \lesssim \frac {\left\| \boldsymbol {s} _ {0 , \boldsymbol {\theta} _ {0}} \right\| _ {\mathcal {H}} ^ {2} + \left\| \boldsymbol {s} _ {0 , \boldsymbol {\theta}} \right\| _ {\mathcal {H}} ^ {2}}{\tau}.
$$
Proof. (1) For the loss objective  $\tilde{\mathcal{L}}$ , we define the Lyapunov function
$$
\bar {E} (\tau) := \tau \left(\bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) - \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}})\right) + \frac {1}{2} \| \boldsymbol {a} _ {\tau} - \boldsymbol {a} \| _ {L ^ {2} (\rho_ {0})} ^ {2}.
$$
$$
\begin{array}{l} \frac {d}{d \tau} \bar {E} (\tau) = \left(\bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) - \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}})\right) + \tau \cdot \frac {d}{d \tau} \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) + \left\langle \boldsymbol {a} _ {\tau} - \boldsymbol {a}, \frac {d}{d \tau} \boldsymbol {a} _ {\tau} \right\rangle_ {L ^ {2} (\rho_ {0})} \\ \leq \left(\bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) - \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}})\right) - \left\langle \boldsymbol {a} _ {\tau} - \boldsymbol {a}, \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \Big | _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)} \right\rangle_ {L ^ {2} (\rho_ {0})}, \\ \end{array}
$$
where the last inequality holds since
$$
\frac {d}{d \tau} \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) = \left\langle \left. \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)}, \frac {d}{d \tau} \bar {\boldsymbol {\theta}} (\tau) \right\rangle_ {L ^ {2} (\rho_ {0})} = - \left\langle \left. \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)}, \left. \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)} \right\rangle_ {L ^ {2} (\rho_ {0})} \leq 0.
$$
By convexity, for any  $\tau_{1},\tau_{2}$  , it holds that
$$
\left. \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau_ {1})) + \left\langle \boldsymbol {a} _ {\tau_ {2}} - \boldsymbol {a} _ {\tau_ {1}}, \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau_ {1})} \right\rangle_ {L ^ {2} (\rho_ {0})} \leq \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau_ {2})),
$$
hence  $\frac{d}{d\tau}\bar{E} (\tau)\leq 0$  . We conclude that  $\bar{E} (\tau)\leq \bar{E} (0)$  , or equivalently
$$
\tau \left(\bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) - \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}})\right) + \frac {1}{2} \| \boldsymbol {a} _ {\tau} - \boldsymbol {a} \| _ {L ^ {2} (\rho_ {0})} ^ {2} \leq \frac {1}{2} \| \boldsymbol {a} _ {0} - \boldsymbol {a} \| _ {L ^ {2} (\rho_ {0})} ^ {2}.
$$
Therefore, note that  $\left\| \bar{s}_{0,\bar{\theta}}\right\|_{\mathcal{H}}^2 \coloneqq \mathbb{E}_{\rho_0}\left[\| \boldsymbol{a}\| _2^2\right]$  (let  $\mathcal{H} \coloneqq \mathcal{H}_{k_{\rho_0}}$ ), we obtain
$$
\bar {\tilde {\mathcal {L}}} \left(\bar {\boldsymbol {\theta}} (\tau)\right) - \bar {\tilde {\mathcal {L}}} \left(\bar {\boldsymbol {\theta}}\right) \lesssim \frac {\| \boldsymbol {a} _ {0} - \boldsymbol {a} \| _ {L ^ {2} \left(\rho_ {0}\right)} ^ {2}}{\tau} \lesssim \frac {\left\| \bar {\boldsymbol {s}} _ {0 , \bar {\boldsymbol {\theta}} _ {0}} \right\| _ {\mathcal {H}} ^ {2} + \left\| \bar {\boldsymbol {s}} _ {0 , \bar {\boldsymbol {\theta}}} \right\| _ {\mathcal {H}} ^ {2}}{\tau},
$$
which gives the desired estimate.
(2) For the loss objective  $\tilde{\mathcal{L}}$ , the argument is almost the same, except replacing the Lyapunov function by
$$
E (\tau) := \tau \left(\tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau)) - \tilde {\mathcal {L}} (\boldsymbol {\theta})\right) + \frac {1}{2 m} \| \boldsymbol {A} _ {\tau} - \boldsymbol {A} \| _ {F} ^ {2},
$$
and  $\langle \cdot ,\cdot \rangle_{L^2 (\rho_0)}$  by  $\langle \cdot ,\cdot \rangle$  . Note that  $\pmb {\theta} = \operatorname {vec}(\pmb {A}) / \sqrt{m}$  and  $\| s_{0,\pmb{\theta}}\|_{\mathcal{H}}^{2} = \| \pmb {A}\|_{F}^{2} / m$  , we obtain
$$
\tilde {\mathcal {L}} (\pmb {\theta} (\tau)) - \tilde {\mathcal {L}} (\pmb {\theta}) \lesssim \frac {\| \pmb {A} _ {0} - \pmb {A} \| _ {F} ^ {2}}{m \tau} \lesssim \frac {\| \pmb {s} _ {0 , \pmb {\theta} _ {0}} \| _ {\mathcal {H}} ^ {2} + \| \pmb {s} _ {0 , \pmb {\theta}} \| _ {\mathcal {H}} ^ {2}}{\tau},
$$
Lemma 5. Suppose that the loss objectives  $\tilde{\mathcal{L}},\tilde{\mathcal{L}}^{(n)},\tilde{\mathcal{L}},\tilde{\mathcal{L}}^{(n)}$  are bounded at the initialization, then for any  $\tau >0$ , we have
$$
\left\| \boldsymbol {s} _ {0, \boldsymbol {\theta} (\tau)} \right\| _ {\mathcal {H}}, \left\| \boldsymbol {s} _ {0, \hat {\boldsymbol {\theta}} _ {n} (\tau)} \right\| _ {\mathcal {H}} \lesssim \left\| \boldsymbol {s} _ {0, \boldsymbol {\theta} _ {0}} \right\| _ {\mathcal {H}} + \sqrt {\frac {\tau}{m}}, \left\| \overline {{\boldsymbol {s}}} _ {0, \bar {\boldsymbol {\theta}} (\tau)} \right\| _ {\mathcal {H}}, \left\| \overline {{\boldsymbol {s}}} _ {0, \hat {\bar {\boldsymbol {\theta}}} _ {n} (\tau)} \right\| _ {\mathcal {H}} \lesssim \left\| \overline {{\boldsymbol {s}}} _ {0, \bar {\boldsymbol {\theta}} _ {0}} \right\| _ {\mathcal {H}} + \sqrt {\tau}.
$$
Proof. It's sufficient to prove  $\left\| \overline{s}_{0,\bar{\theta} (\tau)}\right\|_{\mathcal{H}}\lesssim \sqrt{\tau}$ , and the rest part follows similarly.
$$
\left\| \boldsymbol {a} _ {\tau} \right\| _ {L ^ {2} (\rho_ {0})} \frac {d}{d \tau} \left\| \boldsymbol {a} _ {\tau} \right\| _ {L ^ {2} (\rho_ {0})} = \frac {1}{2} \frac {d}{d \tau} \left\| \boldsymbol {a} _ {\tau} \right\| _ {L ^ {2} (\rho_ {0})} ^ {2} = \left\langle \boldsymbol {a} _ {\tau}, \frac {d}{d \tau} \boldsymbol {a} _ {\tau} \right\rangle_ {L ^ {2} (\rho_ {0})} = \left\langle \boldsymbol {a} _ {\tau}, - \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)} \Bigg \rangle_ {L ^ {2} (\rho_ {0})},
$$
applying Cauchy-Schwartz inequality yields
$$
\begin{array}{l} \frac {d}{d \tau} \| \boldsymbol {a} _ {\tau} \| _ {L ^ {2} \left(\rho_ {0}\right)} = \left\langle \right. \frac {\boldsymbol {a} _ {\tau}}{\| \boldsymbol {a} _ {\tau} \| _ {L ^ {2} \left(\rho_ {0}\right)}}, - \frac {\delta \tilde {\bar {\boldsymbol {\mathcal {L}}}}}{\delta \bar {\boldsymbol {\mathcal {B}}}} \left. \right| _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)} \left. \right\rangle_ {L ^ {2} \left(\rho_ {0}\right)} \\ \leq \left\| \frac {\delta \bar {\tilde {\mathcal {L}}}}{\delta \bar {\boldsymbol {\theta}}} \Big | _ {\bar {\boldsymbol {\theta}} = \bar {\boldsymbol {\theta}} (\tau)} \right\| _ {L ^ {2} (\rho_ {0})} = \sqrt {- \frac {d}{d \tau} \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau))}. \\ \end{array}
$$
Thus, again by Cauchy-Schwartz inequality, for any  $\tau >\tau_0\geq 0$  , we have
$$
\begin{array}{l} \left\| \boldsymbol {a} _ {\tau} \right\| _ {L ^ {2} \left(\rho_ {0}\right)} - \left\| \boldsymbol {a} _ {\tau_ {0}} \right\| _ {L ^ {2} \left(\rho_ {0}\right)} \leq \int_ {\tau_ {0}} ^ {\tau} \sqrt {- \frac {d}{d s} \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (s))} d s \\ \leq \sqrt {\tau - \tau_ {0}} \sqrt {- \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) + \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau_ {0}))} \\ \leq \sqrt {\tau} \sqrt {\bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (0))}. \\ \end{array}
$$
By choosing  $\tau_0 = 0$ , we have  $\| \pmb{a}_{\tau}\|_{L^2 (\rho_0)}\lesssim \left\| \overline{\pmb{s}}_{0,\overline{\pmb{\theta}} (0)}\right\|_{\mathcal{H}} + \sqrt{\tau}$ , hence  $\left\| \overline{\pmb{s}}_{0,\overline{\pmb{\theta}} (\tau)}\right\|_{\mathcal{H}}\lesssim \left\| \overline{\pmb{s}}_{0,\overline{\pmb{\theta}} (0)}\right\|_{\mathcal{H}}+$ $\sqrt{\tau}$ , which completes the proof.
Lemma 6 (Monte Carlo estimates). Define the Monte Carlo error
$$
\operatorname {E r r} _ {\mathrm {M C}} = \operatorname {E r r} _ {\mathrm {M C}} (\boldsymbol {\theta}, \bar {\boldsymbol {\theta}}; T, \lambda (\cdot)) := \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \left[ \lambda (t) \cdot \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \| \boldsymbol {s} _ {t, \boldsymbol {\theta}} (\boldsymbol {x} (t)) - \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}}} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right]. \tag {25}
$$
Suppose that  $\| \pmb{x}(0) \|_{\infty} \leq 1$ , and the trainable parameter  $\pmb{a}$  and embedding function  $e(\cdot)$  are both bounded. Then, given any  $\pmb{\theta}$ , for any  $\delta > 0$ ,  $\delta \ll 1$ , with the probability of at least  $1 - \delta$ , there exists  $\pmb{\theta}$  such that
$$
\operatorname {E r r} _ {\mathrm {M C}} \lesssim \frac {\log^ {2} \left(1 / \delta^ {2}\right)}{m} d, \tag {26}
$$
where  $\lesssim$  hides universal positive constants only depending on  $T$ .
Proof. Fix any  $\bar{\theta}$ . According to Lemma 1, for any  $\delta > 0$ ,  $\delta \ll 1$ , with the probability of at least  $1 - \delta$ , we have
$$
\left\| \boldsymbol {x} (t) \right\| _ {\infty} \lesssim C _ {T} \left(1 + \sqrt {\log \left(1 / \delta^ {2}\right)}\right) \triangleq C _ {T, \delta}. \tag {27}
$$
Based on the representation (13), for any  $\pmb{W} = (\pmb{w}_1, \dots, \pmb{w}_m)^\top \in \mathbb{R}^{m \times d}$ ,  $\pmb{U} = (\pmb{u}_1, \dots, \pmb{u}_m)^\top \in \mathbb{R}^{m \times d_e}$  with  $(\pmb{w}_i, \pmb{u}_i) \sim \rho_0$ ,  $i = 1, \dots, m$ , let  $\pmb{A} := (\pmb{a}_1, \dots, \pmb{a}_m) \in \mathbb{R}^{d \times m}$  with  $\pmb{a}_i := \pmb{a}(\pmb{w}_i, \pmb{u}_i)$  for  $i = 1, \dots, m$ , and
$$
\boldsymbol {s} _ {t, \boldsymbol {\theta}} (\boldsymbol {x}) := \frac {1}{m} \boldsymbol {A} \sigma (\boldsymbol {W} \boldsymbol {x} + \boldsymbol {U} \boldsymbol {e} (t)) = \frac {1}{m} \sum_ {i = 1} ^ {m} \boldsymbol {a} _ {i} \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right), \tag {28}
$$
then  $\mathbb{E}_{\pmb{W},\pmb{U}}[\pmb{s}_{t,\pmb{\theta}}(\pmb{x})] = \mathbb{E}_{(\pmb{w},\pmb{u})\sim \rho_0}\left[\pmb{a}(\pmb{w},\pmb{u})\sigma (\pmb{w}^\top \pmb{x} + \pmb{u}^\top \pmb{e}(t))\right] = \bar{\pmb{s}}_{t,\pmb{\theta}}(\pmb{x})$  . For  $k = 1,\dots ,d$  , let
$$
\begin{array}{l} Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) := \left\| s _ {t, \boldsymbol {\theta}, k} (\boldsymbol {x}) - \bar {s} _ {t, \bar {\boldsymbol {\theta}}, k} (\boldsymbol {x}) \right\| _ {L ^ {2} (p _ {t})} = \mathbb {E} _ {\boldsymbol {x} \sim p _ {t}} ^ {1 / 2} \left[ \left| s _ {t, \boldsymbol {\theta}, k} (\boldsymbol {x}) - \bar {s} _ {t, \bar {\boldsymbol {\theta}}, k} (\boldsymbol {x}) \right| ^ {2} \right] \\ = \mathbb {E} _ {\boldsymbol {x} \sim p _ {t}} ^ {1 / 2} \left[ \left| \frac {1}{m} \sum_ {i = 1} ^ {m} a _ {i, k} \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) - \mathbb {E} _ {(\boldsymbol {w}, \boldsymbol {u}) \sim \rho_ {0}} [ a _ {k} (\boldsymbol {w}, \boldsymbol {u}) \sigma (\boldsymbol {w} ^ {\top} \boldsymbol {x} + \boldsymbol {u} ^ {\top} \boldsymbol {e} (t)) ] \right| ^ {2} \right]. \\ \end{array}
$$
If  $(\tilde{W},\tilde{U})$  is different from  $(W,U)$  at only one component indexed by  $i$ , we have