SlowGuess/NeurIPS_2023_OCR huggingface dataset:
$$
\begin{array}{l} \sum_ {i = n + 1} ^ {M} \prod_ {j = i} ^ {M - 1} A _ {j} \leq \sum_ {i = n + 1} ^ {M} \frac {\alpha_ {t _ {M}} \mathbf {A} \left(\lambda_ {t _ {i}} , \lambda_ {t _ {M}}\right)}{\alpha_ {t _ {i}}} \left(1 + \frac {\sigma}{M}\right) ^ {M - i} \\ \leq \beta \sum_ {i = 0} ^ {M - n - 1} \left(1 + \frac {\sigma}{M}\right) ^ {i} \tag {40} \\ = \frac {\beta M}{\sigma} \left[ \left(1 + \frac {\sigma}{M}\right) ^ {M - n} - 1 \right] \\ \leq C _ {6} \left(e ^ {\sigma} - 1\right) M \\ \end{array}
$$
Then we substitute Eq. (39) and Eq. (40) into Eq. (38). Note that  $M = \mathcal{O}(1 / h)$  by Assumption B.5, and  $f_{n} = \mathcal{O}(h^{n + 1})$  by Assumption B.6, finally we conclude that  $|\Delta_M| \leq f_M = \mathcal{O}(h^{n + 1})$ .
B.3 Pseudo-Order Solver
First, we provide a lemma that gives the explicit solution to Eq. (15).
Lemma B.8. The solution to Eq. (15) is
$$
\frac {\hat {\mathbf {g}} _ {\mathrm {s}} ^ {(k)}}{k !} = \sum_ {p = 1} ^ {k} \frac {\mathbf {g} _ {i _ {p}} - \mathbf {g} _ {i _ {0}}}{\prod_ {q = 0 , q \neq p} ^ {k} \left(\delta_ {p} - \delta_ {q}\right)} \tag {41}
$$
$$
\boldsymbol {R} _ {k} = \left( \begin{array}{c c c c} \delta_ {1} & \delta_ {1} ^ {2} & \dots & \delta_ {1} ^ {k} \\ \delta_ {2} & \delta_ {2} ^ {2} & \dots & \delta_ {2} ^ {k} \\ \vdots & \vdots & \ddots & \vdots \\ \delta_ {k} & \delta_ {k} ^ {2} & \dots & \delta_ {k} ^ {k} \end{array} \right) \tag {42}
$$
Then the solution to Eq. (15) can be expressed as
$$
\frac {\hat {\boldsymbol {g}} _ {s} ^ {(k)}}{k !} = \sum_ {p = 1} ^ {k} \left(\boldsymbol {R} _ {k} ^ {- 1}\right) _ {k p} \left(\boldsymbol {g} _ {i _ {p}} - \boldsymbol {g} _ {i _ {0}}\right) \tag {43}
$$
where  $(\pmb{R}_k^{-1})_{kp}$  is the element of  $\pmb{R}_k^{-1}$  at the  $k$ -th row and the  $p$ -th column. From previous studies of the inversion of the Vandermonde matrix [12], we know
$$
\left(\boldsymbol {R} _ {k} ^ {- 1}\right) _ {k p} = \frac {1}{\delta_ {p} \prod_ {q = 1 , q \neq p} ^ {k} \left(\delta_ {p} - \delta_ {q}\right)} = \frac {1}{\left(\delta_ {p} - \delta_ {0}\right) \prod_ {q = 1 , q \neq p} ^ {k} \left(\delta_ {p} - \delta_ {q}\right)} \tag {44}
$$
Substituting Eq. (44) into Eq. (43), we finish the proof.
Then we prove Theorem 3.4 below:
Proof. (Proof of Theorem 3.4) First, we use mathematical induction to prove that
$$
D _ {l} ^ {(k)} = \tilde {D} _ {l} ^ {(k)} := \sum_ {p = 1} ^ {k} \frac {\boldsymbol {g} _ {i _ {l + p}} - \boldsymbol {g} _ {i _ {l}}}{\prod_ {q = 0 , q \neq p} ^ {k} \left(\delta_ {l + p} - \delta_ {l + q}\right)}, \quad 1 \leq k \leq n, 0 \leq l \leq n - k \tag {45}
$$
For  $k = 1$ , Eq. (45) holds by the definition of  $D_l^{(k)}$ . Suppose the equation holds for  $k$ , we then prove it holds for  $k + 1$ .
Define the Lagrange polynomial which passes  $(\delta_{l + p},\pmb{g}_{i_{l + p}} - \pmb{g}_{i_l})$  for  $0\leq p\leq k$  ..
$$
P _ {l} ^ {(k)} (x) := \sum_ {p = 1} ^ {k} \left(\boldsymbol {g} _ {i _ {l + p}} - \boldsymbol {g} _ {i _ {l}}\right) \prod_ {q = 0, q \neq p} ^ {k} \frac {x - \delta_ {l + q}}{\delta_ {l + p} - \delta_ {l + q}}, \quad 1 \leq k \leq n, 0 \leq l \leq n - k \tag {46}
$$
Then  $\tilde{D}_l^{(k)} = P_l^{(k)}(x)[x^k]$  is the coefficients before the highest-order term  $x^k$  in  $P_{l}^{(k)}(x)$ . We then prove that  $P_{l}^{(k)}(x)$  satisfies the following recurrence relation:
$$
P _ {l} ^ {(k)} (x) = \tilde {P} _ {l} ^ {(k)} (x) := \frac {\left(x - \delta_ {l}\right) P _ {l + 1} ^ {(k - 1)} (x) - \left(x - \delta_ {l + k}\right) P _ {l} ^ {(k - 1)} (x)}{\delta_ {l + k} - \delta_ {l}} \tag {47}
$$
By definition,  $P_{l+1}^{(k-1)}(x)$  is the  $(k-1)$ -th order polynomial which passes  $(\delta_{l+p}, \pmb{g}_{i_{l+p}} - \pmb{g}_{i_l})$  for  $1 \leq p \leq k$ , and  $P_l^{(k-1)}(x)$  is the  $(k-1)$ -th order polynomial which passes  $(\delta_{l+p}, \pmb{g}_{i_{l+p}} - \pmb{g}_{i_l})$  for  $0 \leq p \leq k-1$ .
Thus, for  $1 \leq p \leq k - 1$ , we have
$$
\tilde {P} _ {l} ^ {(k)} \left(\delta_ {l + p}\right) = \frac {\left(\delta_ {l + p} - \delta_ {l}\right) P _ {l + 1} ^ {(k - 1)} \left(\delta_ {l + p}\right) - \left(\delta_ {l + p} - \delta_ {l + k}\right) P _ {l} ^ {(k - 1)} \left(\delta_ {l + p}\right)}{\delta_ {l + k} - \delta_ {l}} = \boldsymbol {g} _ {i _ {l + p}} - \boldsymbol {g} _ {i _ {l}} \tag {48}
$$
For  $p = 0$ , we have
$$
\tilde {P} _ {l} ^ {(k)} (\delta_ {l}) = \frac {\left(\delta_ {l} - \delta_ {l}\right) P _ {l + 1} ^ {(k - 1)} (\delta_ {l}) - \left(\delta_ {l} - \delta_ {l + k}\right) P _ {l} ^ {(k - 1)} (\delta_ {l})}{\delta_ {l + k} - \delta_ {l}} = \boldsymbol {g} _ {i _ {l}} - \boldsymbol {g} _ {i _ {l}} \tag {49}
$$
for  $p = k$ , we have
$$
\tilde {P} _ {l} ^ {(k)} \left(\delta_ {l + k}\right) = \frac {\left(\delta_ {l + k} - \delta_ {l}\right) P _ {l + 1} ^ {(k - 1)} \left(\delta_ {l + k}\right) - \left(\delta_ {l + k} - \delta_ {l + k}\right) P _ {l} ^ {(k - 1)} \left(\delta_ {l + k}\right)}{\delta_ {l + k} - \delta_ {l}} = \boldsymbol {g} _ {i _ {l + k}} - \boldsymbol {g} _ {i _ {l}} \tag {50}
$$
Therefore,  $\tilde{P}_l^{(k)}(x)$  is the  $k$ -th order polynomial which passes  $k + 1$  distinct points  $(\delta_{l + p},\pmb{g}_{i_{l + p}} - \pmb{g}_{i_l})$  for  $0\leq p\leq k$ . Due to the uniqueness of the Lagrange polynomial, we can conclude that  $P_{l}^{(k)}(x) = \tilde{P}_{l}^{(k)}(x)$ . By taking the coefficients of the highest-order term, we obtain
$$
\tilde {D} _ {l} ^ {(k)} = \frac {\tilde {D} _ {l + 1} ^ {(k - 1)} - \tilde {D} _ {l} ^ {(k - 1)}}{\delta_ {l + k} - \delta_ {l}} \tag {51}
$$
where by the induction hypothesis we have  $D_{l + 1}^{(k - 1)} = \tilde{D}_{l + 1}^{(k - 1)}, D_l^{(k - 1)} = \tilde{D}_l^{(k - 1)}$ . Comparing Eq. (51) with the recurrence relation of  $D_{l}^{(k)}$  in Eq. (16), it follows that  $D_{l}^{(k)} = \tilde{D}_{l}^{(k)}$ , which completes the mathematical induction.
Finally, by comparing the expression for  $\tilde{D}_l^{(k)}$  in Eq. (45) and the expression for  $\hat{\pmb{g}}_s^{(k)}$  in Lemma B.8, we can conclude that  $\hat{\pmb{g}}_s^{(k)} = k! D_0^{(k)}$ .
B.4 Local Unbiasedness
Proof. (Proof of Theorem 3.2) Subtracting the local exact solution in Eq. (9) from the  $(n + 1)$ -th order local approximation in Eq. (14), we have the local truncation error
$$
\begin{array}{l} \hat {\boldsymbol {x}} _ {t} - \boldsymbol {x} _ {t} = \alpha_ {t} \boldsymbol {A} (\lambda_ {s}, \lambda_ {t}) \left(\int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \boldsymbol {g} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) \mathrm {d} \lambda - \sum_ {k = 0} ^ {n} \hat {\boldsymbol {g}} _ {\theta} ^ {(k)} (\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s}) \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \frac {\left(\lambda - \lambda_ {s}\right) ^ {k}}{k !} \mathrm {d} \lambda\right) \\ = \alpha_ {t} \boldsymbol {A} (\lambda_ {s}, \lambda_ {t}) \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \left(\boldsymbol {g} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {g} _ {\theta} (\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s})\right) d \lambda \\ - \alpha_ {t} \boldsymbol {A} (\lambda_ {s}, \lambda_ {t}) \sum_ {k = 1} ^ {n} \hat {\boldsymbol {g}} _ {\theta} ^ {(k)} (\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s}) \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \frac {(\lambda - \lambda_ {s}) ^ {k}}{k !} \mathrm {d} \lambda \\ = \alpha_ {t} \boldsymbol {A} (\lambda_ {s}, \lambda_ {t}) \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \left(\boldsymbol {g} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {g} _ {\theta} (\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s})\right) d \lambda \\ \left. - \alpha_ {t} \boldsymbol {A} \left(\lambda_ {s}, \lambda_ {t}\right) \sum_ {k = 1} ^ {n} \left(\sum_ {l = 1} ^ {n} \left(\boldsymbol {R} _ {n} ^ {- 1}\right) _ {k l} \left(\boldsymbol {g} _ {\theta} \left(\boldsymbol {x} _ {\lambda_ {i _ {l}}}, \lambda_ {i _ {l}}\right) - \boldsymbol {g} _ {\theta} \left(\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s}\right)\right)\right) \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \frac {\left(\lambda - \lambda_ {s}\right) ^ {k}}{k !} d \lambda \right. \tag {52} \\ \end{array}
$$
where  $\pmb{x}_{\lambda}$  is on the ground-truth ODE trajectory passing  $\pmb{x}_{\lambda_s}$ , and  $(R_{n}^{-1})_{kl}$  is the element of the inverse matrix  $R_{n}^{-1}$  at the  $k$ -th row and the  $l$ -th column, as discussed in the proof of Lemma B.8. By Newton-Leibniz theorem, we have
$$
\boldsymbol {g} _ {\theta} \left(\boldsymbol {x} _ {\lambda}, \lambda\right) - \boldsymbol {g} _ {\theta} \left(\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s}\right) = \int_ {\lambda_ {s}} ^ {\lambda} \boldsymbol {g} _ {\theta} ^ {(1)} \left(\boldsymbol {x} _ {\tau}, \tau\right) \mathrm {d} \tau \tag {53}
$$
Also, since  $\pmb{x}_{\lambda_{ij}}, l = 1, \dots, n$  are on the ground-truth ODE trajectory passing  $\pmb{x}_{\lambda_s}$ , we have
$$
\boldsymbol {g} _ {\theta} \left(\boldsymbol {x} _ {\lambda_ {i _ {l}}}, \lambda_ {i _ {l}}\right) - \boldsymbol {g} _ {\theta} \left(\boldsymbol {x} _ {\lambda_ {s}}, \lambda_ {s}\right) = \int_ {\lambda_ {s}} ^ {\lambda_ {i _ {l}}} \boldsymbol {g} _ {\theta} ^ {(1)} \left(\boldsymbol {x} _ {\tau}, \tau\right) \mathrm {d} \tau \tag {54}
$$
$$
\boldsymbol {g} _ {\theta} ^ {(1)} \left(\boldsymbol {x} _ {\tau}, \tau\right) = e ^ {- \int_ {\lambda_ {s}} ^ {\tau} \boldsymbol {s} _ {r} \mathrm {d} r} \left(\boldsymbol {f} _ {\theta} ^ {(1)} \left(\boldsymbol {x} _ {\tau}, \tau\right) - \boldsymbol {s} _ {\tau} \boldsymbol {f} _ {\theta} \left(\boldsymbol {x} _ {\tau}, \tau\right) - \boldsymbol {b} _ {\tau}\right) \tag {55}
$$
Note that  $s_{\lambda}, l_{\lambda}$  are the solution to the least square problem in Eq. (11), which makes sure  $\mathbb{E}_{p_{\lambda_s}^{\theta}(\boldsymbol{x}_{\lambda_s})}\left[\boldsymbol{f}_{\theta}^{(1)}(\boldsymbol{x}_{\tau}, \tau) - \boldsymbol{s}_{\tau}\boldsymbol{f}_{\theta}(\boldsymbol{x}_{\tau}, \tau) - \boldsymbol{b}_{\tau}\right] = 0$ . It follows that  $\mathbb{E}_{p_{\lambda_s}^{\theta}(\boldsymbol{x}_{\lambda_s})}\left[\boldsymbol{f}_{\theta}^{(1)}(\boldsymbol{x}_{\tau}, \tau) - \boldsymbol{s}_{\tau}\boldsymbol{f}_{\theta}(\boldsymbol{x}_{\tau}, \tau) - \boldsymbol{b}_{\tau}\right] = 0$ , since  $\boldsymbol{x}_{\tau}$  is on the ground-truth ODE trajectory passing  $\boldsymbol{x}_{\lambda_s}$ . Therefore, we have  $\mathbb{E}_{p_{\lambda_s}^{\theta}(\boldsymbol{x}_{\lambda_s})}\left[\boldsymbol{g}_{\theta}(\boldsymbol{x}_{\lambda}, \lambda) - \boldsymbol{g}_{\theta}(\boldsymbol{x}_{\lambda_s}, \lambda_s)\right] = 0$  and  $\mathbb{E}_{p_{\lambda_s}^{\theta}(\boldsymbol{x}_{\lambda_s})}\left[\boldsymbol{g}_{\theta}(\boldsymbol{x}_{\lambda_{i_l}}, \lambda_{i_l}) - \boldsymbol{g}_{\theta}(\boldsymbol{x}_{\lambda_s}, \lambda_s)\right] = 0$ . Substitute them into Eq. (52), we conclude that  $\mathbb{E}_{p_{\lambda_s}^{\theta}(\boldsymbol{x}_{\lambda_s})}\left[\hat{\boldsymbol{x}}_t - \boldsymbol{x}_t\right] = 0$ .
C.1 Computing the EMS and Related Integrals in the ODE Formulation
The ODE formulation and local approximation require computing some complex integrals involving  $l_{\lambda}, s_{\lambda}, b_{\lambda}$ . In this section, we'll give details about how to estimate  $l_{\lambda}^{*}, s_{\lambda}^{*}, b_{\lambda}^{*}$  on a few datapoints, and how to use them to compute the integrals efficiently.
C.1.1 Computing the EMS
First for the computing of  $l_{\lambda}^{*}$  in Eq. (5), note that
$$
\nabla_ {\boldsymbol {x}} \boldsymbol {N} _ {\theta} \left(\boldsymbol {x} _ {\lambda}, \lambda\right) = \sigma_ {\lambda} \nabla_ {\boldsymbol {x}} \epsilon \left(\boldsymbol {x} _ {\lambda}, \lambda\right) - \operatorname {d i a g} \left(l _ {\lambda}\right) \tag {56}
$$
Since  $\mathrm{diag}(l_{\lambda})$  is a diagonal matrix, minimizing  $\mathbb{E}_{p_{\lambda}^{\theta}(\boldsymbol{x}_{\lambda})}\left[\|\nabla_{\boldsymbol{x}}\boldsymbol{N}_{\theta}(\boldsymbol{x}_{\lambda}, \lambda)\|_F^2\right]$  is equivalent to minimizing  $\mathbb{E}_{p_{\lambda}^{\theta}(\boldsymbol{x}_{\lambda})}\left[\|\mathrm{diag}^{-1}(\nabla_{\boldsymbol{x}}\boldsymbol{N}_{\theta}(\boldsymbol{x}_{\lambda}, \lambda))\|_2^2\right] = \mathbb{E}_{p_{\lambda}^{\theta}(\boldsymbol{x}_{\lambda})}\left[\|\mathrm{diag}^{-1}(\sigma_{\lambda}\nabla_{\boldsymbol{x}}\boldsymbol{\epsilon}(\boldsymbol{x}_{\lambda}, \lambda)) - \boldsymbol{l}_{\lambda}\|_2^2\right]$ , where
$\mathrm{diag}^{-1}$  denotes the operator that takes the diagonal of a matrix as a vector. Thus we have  $\boldsymbol{l}_{\lambda}^{*} = \mathbb{E}_{p_{\lambda}^{\theta}(\boldsymbol{x}_{\lambda})}\left[\mathrm{diag}^{-1}(\sigma_{\lambda}\nabla_{\boldsymbol{x}}\epsilon (\boldsymbol{x}_{\lambda},\lambda))\right]$ .
However, this formula for  $l_{\lambda}^{*}$  requires computing the diagonal of the full Jacobian of the noise prediction model, which typically has  $\mathcal{O}(d^2)$  time complexity for  $d$ -dimensional data and is unacceptable when  $d$  is large. Fortunately, the cost can be reduced to  $\mathcal{O}(d)$  by utilizing stochastic diagonal estimators and employing the efficient Jacobian-vector-product operator provided by forward-mode automatic differentiation in deep learning frameworks.
For a  $d$ -by- $d$  matrix  $\pmb{D}$ , its diagonal can be unbiasedly estimated by [4]
$$
\operatorname {d i a g} ^ {- 1} (\boldsymbol {D}) = \left[ \sum_ {k = 1} ^ {s} \left(\boldsymbol {D} \boldsymbol {v} _ {k}\right) \odot \boldsymbol {v} _ {k} \right] \oslash \left[ \sum_ {k = 1} ^ {s} \boldsymbol {v} _ {k} \odot \boldsymbol {v} _ {k} \right] \tag {57}
$$
where  $\pmb{v}_k \sim p(\pmb{v})$  are  $d$ -dimensional i.i.d. samples with zero mean,  $\odot$  is the element-wise multiplication i.e., Hadamard product, and  $\varnothing$  is the element-wise division. The stochastic diagonal estimator is analogous to the famous Hutchinson's trace estimator [20]. By taking  $p(\pmb{v})$  as the Rademacher distribution, we have  $\pmb{v}_k \odot \pmb{v}_k = \mathbf{1}$ , and the denominator can be omitted. For simplicity, we use regular multiplication and division symbols, assuming they are element-wise between vectors. Then  $l_{\lambda}^{*}$  can be expressed as:
$$
\boldsymbol {l} _ {\lambda} ^ {*} = \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda}) p (\boldsymbol {v})} \left[ \left(\sigma_ {\lambda} \nabla_ {\boldsymbol {x}} \epsilon_ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) \boldsymbol {v}\right) \boldsymbol {v} \right] \tag {58}
$$
which is an unbiased estimation when we replace the expectation with mean on finite samples  $\pmb{x}_{\lambda} \sim p_{\lambda}^{\theta}(\pmb{x}_{\lambda}), \pmb{v} \sim p(\pmb{v})$ . The process for estimating  $l_{\lambda}^{*}$  can easily be paralleled on multiple devices by computing  $\sum (\sigma_{\lambda} \nabla_{\pmb{x}} \epsilon_{\theta}(\pmb{x}_{\lambda}, \lambda) \pmb{v}) \pmb{v}$  on separate datapoints and gather them in the end.
Next, for the computing of  $s_{\lambda}^{*}, b_{\lambda}^{*}$  in Eq. (11), note that it's a simple least square problem. By taking partial derivatives w.r.t.  $s_{\lambda}, b_{\lambda}$  and set them to 0, we have
$$
\left\{ \begin{array}{l} \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} \left[ \left(\boldsymbol {f} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {s} _ {\lambda} ^ {*} \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {b} _ {\lambda} ^ {*}\right) \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) \right] = 0 \\ \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} \left[ \boldsymbol {f} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {s} _ {\lambda} ^ {*} \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {b} _ {\lambda} ^ {*} \right] = 0 \end{array} \right. \tag {59}
$$
And we obtain the explicit formula for  $s_{\lambda}^{*}, b_{\lambda}^{*}$
$$
\boldsymbol {s} _ {\lambda} ^ {*} = \frac {\mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} \left[ \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda) \boldsymbol {f} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda} , \lambda) \right] - \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} \left[ \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda) \right] \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} \left[ \boldsymbol {f} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda} , \lambda) \right]}{\mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} [ \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda) \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda) ] - \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} [ \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda) ] \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} [ \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda) ]} \tag {60}
$$
$$
\boldsymbol {b} _ {\lambda} ^ {*} = \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} [ \boldsymbol {f} ^ {(1)} (\boldsymbol {x} _ {\lambda}, \lambda) ] - \boldsymbol {s} _ {\lambda} ^ {*} \mathbb {E} _ {p _ {\lambda} ^ {\theta} (\boldsymbol {x} _ {\lambda})} [ \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) ] \tag {61}
$$
which are unbiased least square estimators when we replace the expectation with mean on finite samples  $\boldsymbol{x}_{\lambda} \sim p_{\lambda}^{\theta}(\boldsymbol{x}_{\lambda})$ . Also, the process for estimating  $s_{\lambda}^{*}, b_{\lambda}^{*}$  can be paralleled on multiple devices by computing  $\sum f_{\theta}, \sum f_{\theta}^{(1)}, \sum f_{\theta} f_{\theta}, \sum f_{\theta} f_{\theta}^{(1)}$  on separate datapoints and gather them in the end. Thus, the estimation of  $s_{\lambda}^{*}, b_{\lambda}^{*}$  involving evaluating  $f_{\theta}$  and  $f_{\theta}^{(1)}$  on  $\boldsymbol{x}_{\lambda}$ .  $f_{\theta}$  is a direct transformation of  $\epsilon_{\theta}$  and requires a single forward pass. For  $f_{\theta}^{(1)}$ , we have
$$
\begin{array}{l} \boldsymbol {f} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda}, \lambda) = \frac {\partial \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda} , \lambda)}{\partial \lambda} + \nabla_ {\boldsymbol {x}} \boldsymbol {f} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) \frac {\mathrm {d} \boldsymbol {x} _ {\lambda}}{\mathrm {d} \lambda} \\ = e ^ {- \lambda} \left(\boldsymbol {\epsilon} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda}, \lambda) - \boldsymbol {\epsilon} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda)\right) - \frac {\dot {\boldsymbol {l}} _ {\lambda} \alpha_ {\lambda} - \dot {\alpha} _ {\lambda} \boldsymbol {l} _ {\lambda}}{\alpha_ {\lambda} ^ {2}} \boldsymbol {x} _ {\lambda} - \frac {\boldsymbol {l} _ {\lambda}}{\alpha_ {\lambda}} \left(\frac {\dot {\alpha} _ {\lambda}}{\alpha_ {\lambda}} \boldsymbol {x} _ {\lambda} - \sigma_ {\lambda} \boldsymbol {\epsilon} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda)\right) \\ = e ^ {- \lambda} \left(\left(\boldsymbol {l} _ {\lambda} - 1\right) \boldsymbol {\epsilon} _ {\theta} \left(\boldsymbol {x} _ {\lambda}, \lambda\right) + \boldsymbol {\epsilon} _ {\theta} ^ {(1)} \left(\boldsymbol {x} _ {\lambda}, \lambda\right)\right) - \frac {\dot {\boldsymbol {l}} _ {\lambda} \boldsymbol {x} _ {\lambda}}{\alpha_ {\lambda}} \tag {62} \\ \end{array}
$$
After we obtain  $l_{\lambda}, \dot{l}_{\lambda}$  can be estimated by finite difference. To compute  $\pmb{\epsilon}_{\theta}^{(1)}(\pmb{x}_{\lambda}, \lambda)$ , we have
$$
\begin{array}{l} \boldsymbol {\epsilon} _ {\theta} ^ {(1)} (\boldsymbol {x} _ {\lambda}, \lambda) = \partial_ {\lambda} \boldsymbol {\epsilon} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) + \nabla_ {\boldsymbol {x}} \boldsymbol {\epsilon} _ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) \frac {\mathrm {d} \boldsymbol {x} _ {\lambda}}{\mathrm {d} \lambda} \tag {63} \\ = \partial_ {\lambda} \epsilon_ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) + \nabla_ {\boldsymbol {x}} \epsilon_ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda) \left(\frac {\dot {\alpha} _ {\lambda}}{\alpha_ {\lambda}} \boldsymbol {x} _ {\lambda} - \sigma_ {\lambda} \epsilon_ {\theta} (\boldsymbol {x} _ {\lambda}, \lambda)\right) \\ \end{array}
$$
which can also be computed with the Jacobian-vector-product operator.
In conclusion, for any  $\lambda, l_{\lambda}^{*}, s_{\lambda}^{*}, b_{\lambda}^{*}$  can be efficiently and unbiasedly estimated by sampling a few datapoints  $x_{\lambda} \sim p_{\lambda}^{\theta}(x_{\lambda})$  and using the Jacobian-vector-product.
C.1.2 Integral Precomputing
In the local approximation in Eq. (14), there are three integrals involving the EMS, which are  $\mathbf{A}(\lambda_s,\lambda_t),\int_{\lambda_s}^{\lambda_t}\mathbf{E}_{\lambda_s}(\lambda)\mathbf{B}_{\lambda_s}(\lambda)\mathrm{d}\lambda ,\int_{\lambda_s}^{\lambda_t}\mathbf{E}_{\lambda_s}(\lambda)\frac{(\lambda - \lambda_s)^k}{k!}\mathrm{d}\lambda$ . Define the following terms, which are also evaluated at  $\lambda_{j_0},\lambda_{j_1},\ldots ,\lambda_{j_N}$  and can be precomputed in  $\mathcal{O}(N)$  time:
$$
\boldsymbol {L} _ {\lambda} = \int_ {\lambda_ {T}} ^ {\lambda} \boldsymbol {l} _ {\tau} \mathrm {d} \tau
$$
$$
\boldsymbol {S} _ {\lambda} = \int_ {\lambda_ {T}} ^ {\lambda} \boldsymbol {s} _ {\tau} \mathrm {d} \tau
$$
$$
\boldsymbol {B} _ {\lambda} = \int_ {\lambda_ {T}} ^ {\lambda} e ^ {- \int_ {\lambda_ {T}} ^ {r} \boldsymbol {s} _ {\tau} \mathrm {d} \tau} \boldsymbol {b} _ {r} \mathrm {d} r = \int_ {\lambda_ {T}} ^ {\lambda} e ^ {- \boldsymbol {S} _ {r}} \boldsymbol {b} _ {r} \mathrm {d} r \tag {64}
$$
$$
\boldsymbol {C} _ {\lambda} = \int_ {\lambda_ {T}} ^ {\lambda} \left(e ^ {\int_ {\lambda_ {T}} ^ {u} \left(\boldsymbol {l} _ {\tau} + \boldsymbol {s} _ {\tau}\right) \mathrm {d} \tau} \int_ {\lambda_ {T}} ^ {u} e ^ {- \int_ {\lambda_ {T}} ^ {r} \boldsymbol {s} _ {\tau} \mathrm {d} \tau} \boldsymbol {b} _ {r} \mathrm {d} r\right) \mathrm {d} u = \int_ {\lambda_ {T}} ^ {\lambda} e ^ {\boldsymbol {L} _ {u} + \boldsymbol {S} _ {u}} \boldsymbol {B} _ {u} \mathrm {d} u
$$
$$
\boldsymbol {I} _ {\lambda} = \int_ {\lambda_ {T}} ^ {\lambda} e ^ {\int_ {\lambda_ {T}} ^ {r} \left(\boldsymbol {l} _ {\tau} + \boldsymbol {s} _ {\tau}\right) \mathrm {d} \tau} \mathrm {d} r = \int_ {\lambda_ {T}} ^ {\lambda} e ^ {\boldsymbol {L} _ {r} + \boldsymbol {S} _ {r}} \mathrm {d} r
$$
Then for any  $\lambda_s, \lambda_t$ , we can verify that the first two integrals can be expressed as
$$
\boldsymbol {A} (\lambda_ {s}, \lambda_ {t}) = e ^ {\boldsymbol {L} _ {\lambda_ {s}} - \boldsymbol {L} _ {\lambda_ {t}}}
$$
$$
\int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \boldsymbol {B} _ {\lambda_ {s}} (\lambda) \mathrm {d} \lambda = e ^ {- \boldsymbol {L} _ {\lambda_ {s}}} \left(\boldsymbol {C} _ {\lambda_ {t}} - \boldsymbol {C} _ {\lambda_ {s}} - \boldsymbol {B} _ {\lambda_ {s}} \left(\boldsymbol {I} _ {\lambda_ {t}} - \boldsymbol {I} _ {\lambda_ {s}}\right)\right) \tag {65}
$$
which can be computed in  $\mathcal{O}(1)$  time. For the third and last integral, denote it as  $E_{\lambda_s,\lambda_t}^{(k)}$ , i.e.,
$$
\boldsymbol {E} _ {\lambda_ {s}, \lambda_ {t}} ^ {(k)} = \int_ {\lambda_ {s}} ^ {\lambda_ {t}} \boldsymbol {E} _ {\lambda_ {s}} (\lambda) \frac {(\lambda - \lambda_ {s}) ^ {k}}{k !} \mathrm {d} \lambda \tag {66}
$$
We need to compute it for  $0 \leq k \leq n$  and for every local transition time pair  $(\lambda_s, \lambda_t)$  in the sampling process. For  $k = 0$ , we have
$$
\boldsymbol {E} _ {\lambda_ {s}, \lambda_ {t}} ^ {(0)} = e ^ {- \boldsymbol {L} _ {\lambda_ {s}} - \boldsymbol {S} _ {\lambda_ {s}}} \left(\boldsymbol {I} _ {\lambda_ {t}} - \boldsymbol {I} _ {\lambda_ {s}}\right) \tag {67}
$$