SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \left| Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) - Z _ {t, k} (\tilde {\boldsymbol {W}}, \tilde {\boldsymbol {U}}) \right| \\ = \left| \left\| s _ {t, \boldsymbol {\theta}, k} (\boldsymbol {x}) - \bar {s} _ {t, \bar {\boldsymbol {\theta}}, k} (\boldsymbol {x}) \right\| _ {L ^ {2} \left(p _ {t}\right)} - \left\| s _ {t, \tilde {\boldsymbol {\theta}}, k} (\boldsymbol {x}) - \bar {s} _ {t, \bar {\boldsymbol {\theta}}, k} (\boldsymbol {x}) \right\| _ {L ^ {2} \left(p _ {t}\right)} \right| \\ \end{array}
$$
$$
\begin{array}{l} \stackrel {(i)} {\leq} \left\| s _ {t, \boldsymbol {\theta}, k} (\boldsymbol {x}) - s _ {t, \tilde {\boldsymbol {\theta}}, k} (\boldsymbol {x}) \right\| _ {L ^ {2} (p _ {t})} \\ = \frac {1}{m} \left\| a _ {i, k} \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) - \tilde {a} _ {i, k} \sigma \left(\tilde {\boldsymbol {w}} _ {i} ^ {\top} \boldsymbol {x} + \tilde {\boldsymbol {u}} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right\| _ {L ^ {2} \left(p _ {t}\right)} \\ \leq \frac {1}{m} \left(| a _ {i, k} | \| \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) \| _ {L ^ {2} \left(p _ {t}\right)} + | \tilde {a} _ {i, k} | \| \sigma (\tilde {\boldsymbol {w}} _ {i} ^ {\top} \boldsymbol {x} + \tilde {\boldsymbol {u}} _ {i} ^ {\top} \boldsymbol {e} (t)) \| _ {L ^ {2} \left(p _ {t}\right)}\right) \\ \stackrel {(i i)} {\leq} \frac {1}{m} \left(\left| a _ {i, k} \right| \left(\left\| \boldsymbol {w} _ {i} \right\| _ {1} C _ {T, \delta} + \left\| \boldsymbol {u} _ {i} \right\| _ {1} \left\| \boldsymbol {e} (t) \right\| _ {\infty}\right) + \left| \tilde {a} _ {i, k} \right| \left(\left\| \tilde {\boldsymbol {w}} _ {i} \right\| _ {1} C _ {T, \delta} + \left\| \tilde {\boldsymbol {u}} _ {i} \right\| _ {1} \left\| \boldsymbol {e} (t) \right\| _ {\infty}\right)\right) \\ \stackrel {(i i i)} {\leq} \frac {1}{m} \left(\left| a _ {i, k} \right| + \left| \tilde {a} _ {i, k} \right|\right) \left(C _ {T, \delta} + \| e (t) \| _ {\infty}\right) \\ \stackrel {\mathrm {(i v)}} {\lesssim} \frac {1}{m} \left(C _ {T, \delta} + C _ {T, e}\right), \\ \end{array}
$$
where (i) is from the triangle inequality, (ii) is due to the fact that  $|\sigma(y)| = |\mathrm{ReLU}(y)| \leq |y|$  for any  $y \in \mathbb{R}$ , the triangle and Hölder's inequality and (27), (iii) follows from the positive homogeneity property of the ReLU activation, and (iv) is due to the boundness of the trainable parameter  $a$  and
embedding function  $e(\cdot)$ . By McDiarmid's inequality (see e.g., Lemma 26.4 in [43]), for any  $\delta > 0$ , with the probability of at least  $1 - \delta$ , we have
$$
\begin{array}{l} \left| Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) - \mathbb {E} _ {\boldsymbol {W}, \boldsymbol {U}} \left[ Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) \right] \right| \lesssim \frac {1}{m} \left(C _ {T, \delta} + C _ {T, \boldsymbol {e}}\right) \sqrt {m \log (2 / \delta) / 2} (29) \\ \lesssim \left(C _ {T, \delta} + C _ {T, e}\right) \sqrt {\frac {\log (1 / \delta)}{m}}. (30) \\ \end{array}
$$
where (v) is due to Fubini's theorem, and (ii), (iii), (iv) is the same as before. By the triangle inequality, Jensen's inequality, (29) and (31), we obtain
$$
\mathbb {E} _ {\boldsymbol {x} \sim p _ {t}} \left[ \| \boldsymbol {s} _ {t, \boldsymbol {\theta}} (\boldsymbol {x}) - \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}}} (\boldsymbol {x}) \| _ {2} ^ {2} \right] = \sum_ {k = 1} ^ {d} \mathbb {E} _ {\boldsymbol {x} \sim p _ {t}} \left[ \left| s _ {t, \boldsymbol {\theta}, k} (\boldsymbol {x}) - \bar {s} _ {t, \bar {\boldsymbol {\theta}}, k} (\boldsymbol {x}) \right| ^ {2} \right]
$$
$$
\begin{array}{l} = \sum_ {k = 1} ^ {d} Z _ {t, k} ^ {2} (\boldsymbol {W}, \boldsymbol {U}) \\ \leq \sum_ {k = 1} ^ {d} \left(\left| Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) - \mathbb {E} _ {\boldsymbol {W}, \boldsymbol {U}} \left[ Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) \right] \right| + \left| \mathbb {E} _ {\boldsymbol {W}, \boldsymbol {U}} \left[ Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) \right] \right|\right) ^ {2} \\ \lesssim \sum_ {k = 1} ^ {d} \left(\left| Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) - \mathbb {E} _ {\boldsymbol {W}, \boldsymbol {U}} \left[ Z _ {t, k} (\boldsymbol {W}, \boldsymbol {U}) \right] \right| ^ {2} + \mathbb {E} _ {\boldsymbol {W}, \boldsymbol {U}} \left[ Z _ {t, k} ^ {2} (\boldsymbol {W}, \boldsymbol {U}) \right]\right) \\ \lesssim \sum_ {k = 1} ^ {d} \left(C _ {T, \delta} + C _ {T, e}\right) ^ {2} \frac {\log (1 / \delta)}{m} = \left(C _ {T, \delta} + C _ {T, e}\right) ^ {2} \frac {\log (1 / \delta)}{m} d, \tag {33} \\ \end{array}
$$
$$
\begin{array}{l} \operatorname {E r r} _ {\mathrm {M C}} = \frac {1}{T} \int_ {0} ^ {T} \lambda (t) \cdot \mathbb {E} _ {\boldsymbol {x} \sim p _ {t}} \left[ \| \boldsymbol {s} _ {t, \boldsymbol {\theta}} (\boldsymbol {x}) - \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}}} (\boldsymbol {x}) \| _ {2} ^ {2} \right] d t \\ \lesssim \left(C _ {T, \delta} + C _ {T, e}\right) ^ {2} \frac {\log (1 / \delta)}{m} d \cdot \frac {1}{T} \int_ {0} ^ {T} \lambda (t) d t \leq \frac {\log^ {2} (1 / \delta^ {2})}{m} d. \\ \end{array}
$$
Obviously,  $\frac{1}{m}\| A\| _F^2 = \frac{1}{m}\sum_{i = 1}^{m}\| a(\pmb {w}_i,\pmb {u}_i)\| _2^2 = \| s_{t,\pmb{\theta}}\|_{\mathcal{H}_{k_{\rho_0}}}^2 < \infty$  . The proof is completed.
Now we are ready to prove the main theorem.
Proof of Theorem 1. Based on Lemma 2, we have
$$
D _ {\mathrm {K L}} \left(p _ {0} \| p _ {0, \hat {\boldsymbol {\theta}} _ {n} (\tau)}\right) \leq \tilde {\mathcal {L}} \left(\hat {\boldsymbol {\theta}} _ {n} (\tau); g ^ {2} (\cdot)\right) + D _ {\mathrm {K L}} \left(p _ {T} \| \pi\right).
$$
To bound the term  $\tilde{\mathcal{L}} (\hat{\pmb{\theta}}_n(\tau);g^2 (\cdot))$ , we use the following decomposition:
$$
\begin{array}{l} \tilde {\mathcal {L}} (\hat {\boldsymbol {\theta}} _ {n} (\tau)) = \left[ \tilde {\mathcal {L}} (\hat {\boldsymbol {\theta}} _ {n} (\tau)) - \tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau)) \right] + \tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau)) \\ \lesssim \left[ \tilde {\mathcal {L}} \left(\hat {\boldsymbol {\theta}} _ {n} (\tau)\right) - \tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau)) \right] + \tilde {\bar {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) + \text {M o n t e C a r l o} \triangleq I _ {1} + I _ {2} + I _ {3}. \\ \end{array}
$$
According to Lemma 4, we obtain
$$
I _ {2} \lesssim \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} ^ {*}) + \frac {1}{\tau} \left(\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}\right).
$$
The term  $I_{3}$  can be further divided into
$$
\begin{array}{l} I _ {3} := \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} (\tau)} (\boldsymbol {x} (t)) - \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}} (\tau)} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right] \\ \lesssim \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \left[ \lambda (t) \cdot \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \| \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}} (\tau)} (\boldsymbol {x} (t)) - \bar {\boldsymbol {s}} _ {t, \bar {\boldsymbol {\theta}} ^ {*}} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right] \\ + \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] \\ + \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} (\tau)} (\boldsymbol {x} (t)) - \boldsymbol {s} _ {t, \boldsymbol {\theta} ^ {*}} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right] \\ =: I _ {3, 1} + I _ {3, 2} + I _ {3, 3}, \\ \end{array}
$$
where  $\theta^{*}$  is the Monte Carlo estimator of  $\bar{\theta}^{*}$ . By Lemma 6, for any  $\delta >0$ ,  $\delta \ll 1$ , with the probability of at least  $1 - \delta$ , it holds that
$$
I _ {3, 2} = \operatorname {E r r} _ {\mathrm {M C}} \left(\boldsymbol {\theta} ^ {*}, \bar {\boldsymbol {\theta}} ^ {*}; T, \lambda (\cdot)\right) \lesssim \frac {\log^ {2} \left(1 / \delta^ {2}\right)}{m} d, \tag {34}
$$
while Lemma 4 gives
$$
I _ {3, 1} \lesssim \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} (\tau)) + \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} ^ {*}) = I _ {2} + \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} ^ {*}) \lesssim \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} ^ {*}) + \frac {1}{\tau} \left(\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}\right),
$$
$$
I _ {3, 3} \lesssim \tilde {\mathcal {L}} (\boldsymbol {\theta} ^ {*}) + \frac {1}{\tau} \left(\| \boldsymbol {s} _ {0, \boldsymbol {\theta} _ {0}} \| _ {\mathcal {H}} ^ {2} + \| \boldsymbol {s} _ {0, \boldsymbol {\theta} ^ {*}} \| _ {\mathcal {H}} ^ {2}\right).
$$
Hence, for any  $\delta > 0$ ,  $\delta \ll 1$ , with the probability of at least  $1 - \delta$ , it holds that
$$
I _ {3} \lesssim \frac {\log^ {2} (1 / \delta^ {2})}{m} d + \bar {\tilde {\mathcal {L}}} (\bar {\boldsymbol {\theta}} ^ {*}) + \tilde {\mathcal {L}} (\boldsymbol {\theta} ^ {*}) + \frac {1}{\tau} \left(\left\| \bar {\mathbf {s}} _ {0, \bar {\boldsymbol {\theta}} _ {0}} \right\| _ {\mathcal {H}} ^ {2} + \left\| \bar {\mathbf {s}} _ {0, \bar {\boldsymbol {\theta}} ^ {*}} \right\| _ {\mathcal {H}} ^ {2} + \| \mathbf {s} _ {0, \boldsymbol {\theta} _ {0}} \| _ {\mathcal {H}} ^ {2} + \| \mathbf {s} _ {0, \boldsymbol {\theta} ^ {*}} \| _ {\mathcal {H}} ^ {2}\right).
$$
For the term  $I_{1}$ , we have
$$
\begin{array}{l} \sqrt {\tilde {\mathcal {L}} (\hat {\boldsymbol {\theta}} _ {n} (\tau))} - \sqrt {\tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau))} \\ = \left\{\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \left[ \lambda (t) \cdot \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \| \boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right] \right\} ^ {\frac {1}{2}} \\ \left. \right. - \left\{\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} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right]\right]\right\} ^ {\frac {1}{2}} \\ = \left\{\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \| \boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right\} ^ {\frac {1}{2}} \\ \left. - \left\{\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \| \boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right\} ^ {\frac {1}{2}} \right. \\ \leq \left\{\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \| \boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau)} (\boldsymbol {x} (t)) - \boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) \| _ {2} ^ {2} \right] \right\} ^ {\frac {1}{2}} \\ = \left\{\mathbb {E} _ {t \sim \mathcal {U} (0, T)} \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \lambda (t) \left\| \frac {1}{m} \sum_ {i = 1} ^ {m} \hat {\boldsymbol {a}} _ {i} (\tau) \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) - \frac {1}{m} \sum_ {i = 1} ^ {m} \boldsymbol {a} _ {i} (\tau) \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) \right\| _ {2} ^ {2} \right] \right\} ^ {\frac {1}{2}}. \\ \end{array}
$$
$$
\begin{array}{l} \left\| \frac {1}{m} \sum_ {i = 1} ^ {m} (\hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau)) \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) \right\| _ {2} ^ {2} \\ \leq \frac {1}{m ^ {2}} \left(\sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} | \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) |\right) ^ {2} \\ \leq \frac {1}{m ^ {2}} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \sum_ {i = 1} ^ {m} | \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) | ^ {2} \\ \lesssim \frac {1}{m ^ {2}} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \sum_ {i = 1} ^ {m} \left(\left| \boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) \right| ^ {2} + \left| \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t) \right| ^ {2}\right) \\ \leq \frac {1}{m ^ {2}} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \sum_ {i = 1} ^ {m} \left(\| \boldsymbol {w} _ {i} \| _ {1} ^ {2} \| \boldsymbol {x} (t) \| _ {\infty} ^ {2} + \| \boldsymbol {u} _ {i} \| _ {1} ^ {2} \| \boldsymbol {e} (t) \| _ {\infty} ^ {2}\right) \\ \lesssim \frac {1}{m ^ {2}} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \sum_ {i = 1} ^ {m} \left(C _ {T, \delta} ^ {2} + C _ {T, \boldsymbol {e}} ^ {2}\right), \\ \end{array}
$$
$$
\begin{array}{l} \sqrt {\tilde {\mathcal {L}} (\hat {\pmb {\theta}} _ {n} (\tau))} - \sqrt {\tilde {\mathcal {L}} (\pmb {\theta} (\tau))} \lesssim \left\{\frac {1}{m} \sum_ {i = 1} ^ {m} \| \hat {\pmb {a}} _ {i} (\tau) - \pmb {a} _ {i} (\tau) \| _ {2} ^ {2} \left(C _ {T, \delta} ^ {2} + C _ {T, e} ^ {2}\right) \right\} ^ {\frac {1}{2}} \\ \leq \left(C _ {T, \delta} + C _ {T, e}\right) \left\{\frac {1}{m} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \right\} ^ {\frac {1}{2}}. \\ \end{array}
$$
Here, the triangle inequality, Cauchy-Schwartz inequality, the fact that  $|\sigma(y)| = |\mathrm{ReLU}(y)| \leq |y|$  for any  $y \in \mathbb{R}$ , Hölder's inequality, the positive homogeneity property of the ReLU activation, and
the boundness of the input data, embedding function  $e(t)$  and weighting function  $\lambda(t)$ . Thus, we get
$$
\tilde {\mathcal {L}} (\hat {\boldsymbol {\theta}} _ {n} (\tau)) - \tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau))
$$
$$
\lesssim \frac {1}{m} \left(C _ {T, \delta} ^ {2} + C _ {T, e} ^ {2}\right) \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} + \sqrt {\tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau))} \left(C _ {T, \delta} + C _ {T, e}\right) \left\{\frac {1}{m} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \right\} ^ {\frac {1}{2}}
$$
$$
\begin{array}{l} \lesssim \sqrt {\tilde {\mathcal {L}} (\boldsymbol {\theta} ^ {*}) + \frac {1}{\tau} \left(\| \boldsymbol {s} _ {0 , \boldsymbol {\theta} _ {0}} \| _ {\mathcal {H}} ^ {2} + \| \boldsymbol {s} _ {0 , \boldsymbol {\theta} ^ {*}} \| _ {\mathcal {H}} ^ {2}\right)} \left(C _ {T, \delta} + C _ {T, e}\right) \left\{\frac {1}{m} \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2} \right\} ^ {\frac {1}{2}} \\ + \frac {1}{m} \left(C _ {T, \delta} ^ {2} + C _ {T, e} ^ {2}\right) \sum_ {i = 1} ^ {m} \| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \| _ {2} ^ {2}, \\ \end{array}
$$
where the last inequality follows from Lemma 4. We further deduce that
$$
\begin{array}{l} \frac {1}{m} \sum_ {i = 1} ^ {m} \left\| \hat {\boldsymbol {a}} _ {i} (\tau) - \boldsymbol {a} _ {i} (\tau) \right\| _ {2} ^ {2} \\ = \frac {1}{m} \sum_ {i = 1} ^ {m} \left\| \int_ {0} ^ {\tau} \frac {d}{d \tau_ {0}} \left(\hat {\boldsymbol {a}} _ {i} (\tau_ {0}) - \boldsymbol {a} _ {i} (\tau_ {0})\right) d \tau_ {0} \right\| _ {2} ^ {2} \\ = \frac {1}{m} \sum_ {i = 1} ^ {m} \left\| \int_ {0} ^ {\tau} \left(\nabla_ {\boldsymbol {\theta} _ {i} (\tau_ {0})} \tilde {\mathcal {L}} (\boldsymbol {\theta} (\tau_ {0})) - \nabla_ {\hat {\boldsymbol {\theta}} _ {n, i} (\tau_ {0})} \hat {\tilde {\mathcal {L}}} _ {n} (\hat {\boldsymbol {\theta}} _ {n} (\tau_ {0}))\right) d \tau_ {0} \right\| _ {2} ^ {2} \\ = \frac {1}{m ^ {2}} \sum_ {i = 1} ^ {m} \left\| \int_ {0} ^ {\tau} \left(2 \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \left[ \lambda (t) \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \left(\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau_ {0})} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right. \right. \right. \\ \left. - 2 \mathbb {E} _ {t \sim \mathcal {U} (0, T)} \left[ \lambda (t) \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t} ^ {(n)}} \left[\left(\boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau_ {0})} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) \right]\right]\right) d \tau_ {0} \left. \right\| _ {2} ^ {2}, \\ \end{array}
$$
where  $p_t^{(n)}$  denotes the empirical distribution of  $p_t$ . Note that
$$
\left\| \boldsymbol {\theta} \right\| _ {2} ^ {2} = \left\| \operatorname {v e c} (\boldsymbol {A}) \right\| _ {2} ^ {2} / m = \left\| \boldsymbol {A} \right\| _ {F} ^ {2} / m = \left\| s _ {0, \boldsymbol {\theta}} \right\| _ {\mathcal {H}} ^ {2}, \tag {35}
$$
by Lemma 5 we get  $\| \pmb{\theta}(\tau)\| _2 = \left\| s_{0,\pmb{\theta}(\tau)}\right\|_{\mathcal{H}}\lesssim \left\| s_{0,\pmb{\theta}_0}\right\|_{\mathcal{H}} + \sqrt{\tau / m}$ . For any  $t\in [0,T]$ , define the function space
$$
\mathcal {F} _ {t} := \left\{\boldsymbol {f} _ {1} (\boldsymbol {x} (t); \boldsymbol {\theta} _ {1} (\tau)) f _ {2} (\boldsymbol {x} (t); \boldsymbol {\theta} _ {2}): \boldsymbol {f} _ {1} \in \mathcal {F} _ {1, t}, f _ {2} \in \mathcal {F} _ {2, t} \right\},
$$
$$
\mathcal {F} _ {1, t} := \left\{\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t)): \| \boldsymbol {\theta} (\tau) \| _ {2} \lesssim \left\| \boldsymbol {s} _ {0, \boldsymbol {\theta} _ {0}} \right\| _ {\mathcal {H}} + \sqrt {\tau / m} \right\},
$$
$$
\mathcal {F} _ {2, t} := \left\{\sigma (\boldsymbol {w} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} ^ {\top} \boldsymbol {e} (t)): \| \boldsymbol {w} \| _ {1} + \| \boldsymbol {u} \| _ {1} \leq 1 \right\}.
$$
Then, according to Theorem A.5 in [61], for any  $\delta \in (0,1)$ , with the probability at least  $1 - \delta$  over the choice of the dataset  $\mathcal{D}_{\boldsymbol{x}} = \{\pmb{x}_i\}_{i=1}^n$ , it holds that
$$
\begin{array}{l} \left| \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \left(\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right. \\ \left. - \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t} ^ {(n)}} \left[ \left(\boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right| \\ \leq \left| \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t}} \left[ \left(\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma (\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)) \right] \right. \\ \left. - \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t} ^ {(n)}} \left[ \left(\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right| \\ + \left| \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t} ^ {(n)}} \left[ \left(\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right. \\ \left. - \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t} ^ {(n)}} \left[ \left(\boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau)} (\boldsymbol {x} (t)) - \nabla_ {\boldsymbol {x} (t)} \log p _ {t} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right| \\ \end{array}
$$
$$
\begin{array}{l} \lesssim \widehat {\operatorname {R a d}} _ {n} (\mathcal {F} _ {t}) + \sup  _ {\boldsymbol {f} \in \mathcal {F} _ {t}, \boldsymbol {x} (t) \in \left[ - C _ {T, \delta}, C _ {T, \delta} \right] ^ {d}} | \boldsymbol {f} (\boldsymbol {x} (t)) | \sqrt {\frac {\log (2 / \delta)}{n}} \\ + \left| \mathbb {E} _ {\boldsymbol {x} (t) \sim p _ {t} ^ {(n)}} \left[ \left(\boldsymbol {s} _ {t, \boldsymbol {\theta} (\tau)} (\boldsymbol {x} (t)) - \boldsymbol {s} _ {t, \hat {\boldsymbol {\theta}} _ {n} (\tau)} (\boldsymbol {x} (t))\right) \sigma \left(\boldsymbol {w} _ {i} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} _ {i} ^ {\top} \boldsymbol {e} (t)\right) \right] \right| =: J _ {1} + J _ {2} + J _ {3}, \\ \end{array}
$$
where  $\widehat{\mathrm{Rad}}_n$  denotes the (empirical) Rademacher complexity of  $\mathcal{F}_t$  on  $\mathcal{D}_{\pmb{x}} = \{\pmb{x}_i\}_{i=1}^n$ , and all the inequalities hold in the element-wise sense.
(i) For  $J_{1}$ , according to Lemma A.6 in [61], we have
$$
\begin{array}{l} \widehat {\operatorname {R a d}} _ {n} (\mathcal {F} _ {t}) \leq \left(\sup  _ {\boldsymbol {f} _ {1} \in \mathcal {F} _ {1, t}, \boldsymbol {x} (t) \in [ - C _ {T, \delta}, C _ {T, \delta} ] ^ {d}} | \boldsymbol {f} _ {1} (\boldsymbol {x} (t)) | + \sup  _ {f _ {2} \in \mathcal {F} _ {2, t}, \boldsymbol {x} (t) \in [ - C _ {T, \delta}, C _ {T, \delta} ] ^ {d}} | f _ {2} (\boldsymbol {x} (t)) |\right) \\ \cdot \left(\widehat {\operatorname {R a d}} _ {n} (\mathcal {F} _ {1, t}) + \widehat {\operatorname {R a d}} _ {n} (\mathcal {F} _ {2, t})\right). \\ \end{array}
$$
$$
\begin{array}{l} \left| \sigma (\boldsymbol {w} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} ^ {\top} \boldsymbol {e} (t)) \right| \leq \left| \boldsymbol {w} ^ {\top} \boldsymbol {x} (t) + \boldsymbol {u} ^ {\top} \boldsymbol {e} (t) \right| \\ \lesssim \| \boldsymbol {w} \| _ {1} \| \boldsymbol {x} (t) \| _ {\infty} + \| \boldsymbol {u} \| _ {1} \| \boldsymbol {e} (t) \| _ {\infty} \\ \leq C _ {T, \delta} + C _ {T, e}, \tag {36} \\ \end{array}
$$