SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Zheng Wang*
Kahlert School of Computing
Salt Lake City, UT 84112
u1208847@utah.edu
Shikai Fang*
shikai.fang@utah.edu
Shibo Li
shibo@cs.utah.edu
Shandian Zhe†
zhe@cs.utah.edu
Tensor decomposition is an important tool for multiway data analysis. In practice, the data is often sparse yet associated with rich temporal information. Existing methods, however, often under-use the time information and ignore the structural knowledge within the sparsely observed tensor entries. To overcome these limitations and to better capture the underlying temporal structure, we propose Dynamic EMbedIngs fOr dynamic Tensor dEcomposition (DEMOTE). We develop a neural diffusion-reaction process to estimate dynamic embeddings for the entities in each tensor mode. Specifically, based on the observed tensor entries, we build a multi-partite graph to encode the correlation between the entities. We construct a graph diffusion process to co-evolve the embedding trajectories of the correlated entities and use a neural network to construct a reaction process for each individual entity. In this way, our model can capture both the commonalities and personalities during the evolution of the embeddings for different entities. We then use a neural network to model the entry value as a nonlinear function of the embedding trajectories. For model estimation, we combine ODE solvers to develop a stochastic mini-batch learning algorithm. We propose a stratified sampling method to balance the cost of processing each mini-batch so as to improve the overall efficiency. We show the advantage of our approach in both simulation study and real-world applications. The code is available at https://github.com/wzhut/Dynamic-Tensor-Decomposition via-Neural-Diffusion-Reaction-Processes.
Multiway data is common in real-world applications and naturally represented by tensors. For example, online shopping and promotion activities can be expressed as a three-mode tensor (customer, commodity, online merchant). Tensor decomposition is an important tool for multiway data analysis. It estimates embeddings for the entities in each tensor mode, with which to recover the observed
entry values. The embeddings can reflect the underlying structures within the entities and can be used as predictive features, such as for recommendation and ads auction.
In practice, tensor data is often very sparse. That is, the observed entries only take a tiny portion of all possible entries, say,  $0.01\%$ . In addition, the data often includes timestamps for the observed entry values, which imply rich, complex temporal variation patterns. Current tensor decomposition approaches often ignore the structure knowledge within the sparsely observed entries and under-use the temporal information, e.g., simply binning the timestamps into crude time steps (Xiong et al., 2010; Rogers et al., 2013; Zhe et al., 2016a, 2015; Du et al., 2018). More important, standard tensor decomposition estimates a static embedding for each entity. However, as the representation of entities, these embeddings summarize the underlying properties of the entities, which can naturally evolve with time, such as customer interests, user income, product popularity, and fashion. Learning static embeddings can miss capturing these interesting, important temporal knowledge. While the most recent work (Wang et al., 2022) has proposed the first decomposition method to estimate embedding trajectories, it never considers the structural knowledge within the data.
To overcome these limitations, we propose DEMOTE, a dynamic embedding approach for dynamic tensor decomposition. We construct a nonlinear diffusion-reaction process in an Ordinary Differential Equation (ODE) framework to estimate embedding trajectories for tensor entities. The ODE framework is known to be flexible and convenient to handle irregularly sampled timestamps and sparsely observed data (Rubanova et al., 2019). In addition, since ODE models focus on learning the dynamics (i.e., time derivatives) of the target function, they have promising potential for providing robust, accurate long-term predictions (via integration with the dynamics). Specifically, to leverage the structural knowledge within the data, we first build a multi-partite graph based on the observed entries. The graph encodes the correlations between entities at different modes in terms of their interactions. We then construct a graph diffusion process in the ODE to co-evolve the embedding trajectories of correlated entities. Next, we use a neural network to construct a reaction process to model the individual-specific evolution for each entity. In this way, our neural diffusion-reaction process captures both the commonalities and personalities of the entities in learning their dynamic embeddings. Given the embedding trajectories, we model the entry value as a latent function of the associated entities' trajectories. We use another neural network to flexibly estimate the function and to capture the complex relationships of the entities. For efficient training, we base on ODE solvers to develop a stochastic mini-batch learning algorithm. We develop a stratified sampling scheme, which can balance the cost of executing the ODE solvers in each mini-batch so as to improve the efficiency.
We evaluated our method in both simulation and real-world applications. The simulation experiments show that DEMOTE can successfully capture the underlying dynamics of the entities from their temporal interactions and recover the hidden clustering structures within the trajectories. Then in three real-world applications, we tested the accuracy in predicting the tensor entry values at different time points. DEMOTE consistently outperforms the state-of-the-art decomposition methods that incorporate temporal information, often by a large margin. We also demonstrated that both the diffusion and reaction processes contribute to the learning performance. Finally, we investigated the learned embedding trajectories and found interesting evolution paths and hidden structures.
2 Notations and Background
Suppose we have collected data for a  $K$ -mode tensor. Each mode  $k$  includes  $d_{k}$  entities, which we index by  $1, \ldots, d_{k}$ . We then index each tensor entry by a tuple  $\ell = (l_{1}, \ldots, l_{K})$  where for each  $k$ , we have  $1 \leq l_{k} \leq d_{k}$ . Suppose we observed  $N$  tensor entry values and timestamps. The dataset is denoted by  $\mathcal{D} = \{(\ell_{1}, t_{1}, y_{1}), \ldots, (\ell_{N}, t_{N}, y_{N})\}$  where  $\{t_{n}\}$  and  $\{y_{n}\}$  are the timestamps and entry values, respectively. Our goal is for each entity  $j$  of mode  $k$ , to estimate a dynamic embedding  $\mathbf{u}_{j}^{k}(t): \mathbb{R}_{+} \to \mathbb{R}^{R}$ . That is, the embedding is a time function (trajectory) with an  $R$ -dimensional output. Standard tensor decomposition introduces a static embedding representation for each entity, namely,  $\mathbf{u}_{j}^{k}$  is considered as time invariant. Tensor decomposition aims to estimate the embeddings (or factors) to reconstruct the tensor. For example, the classical Tucker decomposition (Tucker, 1966) employs a multilinear factorization model,  $\mathcal{M} = \mathcal{W} \times_{1} \mathbf{U}^{1} \times_{2} \ldots \times_{K} \mathbf{U}^{K}$ , where  $\mathcal{M} \in \mathbb{R}^{d_{1} \times \ldots \times d_{k}}$  is the entire tensor,  $\mathcal{W} \in \mathbb{R}^{R_{1} \times \ldots \times R_{K}}$  is the tensor-core parameter,  $\mathbf{U}^{k}$  comprises all the embeddings of the entities in mode  $k$ , and  $\times_{k}$  is the tensor-matrix product at mode  $k$  (Kolda, 2006). The popular CANDECOMP/PARAFAC (CP) decomposition (Harshman, 1970) can be viewed as a simplified version of Tucker decomposition, where we set  $R_{1} = \ldots = R_{K} = R$  and the tensor-core  $\mathcal{W}$  to be
Figure 1: The illustration of the embedding model in DEMOTE.
diagonal. Hence, each entry value is factorized as  $m_{\ell} = (\mathbf{u}_{l_1}^1 \circ \ldots \circ \mathbf{u}_{l_K}^K)^\top \boldsymbol{\lambda}$ , where  $\circ$  is the Hadamard (element-wise) product, and  $\boldsymbol{\lambda}$  corresponds to  $\mathrm{diag}(\mathcal{W})$ . While CP and Tucker decomposition are popular, their multilinear modeling can be oversimplistic for complex applications. To estimate nonlinear relationships of the entities, Xu et al. (2012); Zhe et al. (2015, 2016a) used a Gaussian process (GP) (Rasmussen and Williams, 2006) to model the entry value as a random function of the embeddings,  $m_{\ell} = g(\mathbf{u}_{l_1}^1, \dots, \mathbf{u}_{l_K}^K)$ , where  $g \sim \mathcal{GP}(0, \kappa(\mathbf{x}_{\ell}, \mathbf{x}_{\ell'})$ ,  $\mathbf{x}_{\ell} = [\mathbf{u}_{l_1}^1; \dots; \mathbf{u}_{l_K}^K]$  and  $\mathbf{x}_{\ell'} = [\mathbf{u}_{l_1'}^1; \dots; \mathbf{u}_{l_K'}^K]$  are the embeddings of the entities in entry  $\ell$  and  $\ell'$ , respectively, and  $\kappa(\cdot, \cdot)$  is the covariance (kernel) function. Given the GP prior, any finite set of  $N$  entry values follow a multi-variate Gaussian distribution,  $\mathbf{m} \sim \mathcal{N}(\mathbf{0}, \mathbf{K})$ , where  $\mathbf{m} = [m_{\ell_1}, \dots, m_{\ell_N}]$ ,  $\mathbf{K}$  is the  $N \times N$  kernel matrix, and each  $[\mathbf{K}]_{i,j} = \kappa(\mathbf{x}_{\ell}, \mathbf{x}_{\ell'})$ . Suppose we have collected continuous observations for the  $N$  entries  $\mathbf{y} = [y_1, \dots, y_N]$ . We can use a Gaussian noise model:  $y_n = m_{\ell_n} + \epsilon_n$  where  $\epsilon_n \sim \mathcal{N}(0, \sigma^2)$ . The marginal likelihood is  $p(\mathbf{y}) = \mathcal{N}(\mathbf{y} | \mathbf{0}, \mathbf{K} + \sigma^2 \mathbf{I})$ , which we can maximize to estimate the embeddings and the model parameters.
Practical data often includes temporal information, i.e., the timestamp when each observed entry value is generated. To leverage this information, existing methods often bin the timestamps into a series of steps, say, by weeks or months (Xiong et al., 2010; Rogers et al., 2013; Zhe et al., 2016a; Song et al., 2017). The tensor is then expanded with an additional time-step mode, and one can apply any decomposition algorithm to estimate embeddings for both the entities and time steps. To capture the temporal dependency, a conditional model is often used (Xiong et al., 2010), say,  $p(\mathbf{t}_{j+1} | \mathbf{t}_j) = \mathcal{N}(\mathbf{t}_{j+1} | \mathbf{t}_j, \tau \mathbf{I})$  where  $\mathbf{t}_j$  is the embedding of  $j$ -th step. To leverage the continuous time information, Zhang et al. (2021) recently developed continuous CP decomposition, where the coefficients  $\lambda$  are modeled as a time function with polynomial splines.
Standard tensor decomposition assumes the embeddings are static and time-invariant. However, the embeddings summarize and extract the properties of entities, which can evolve with time, such as customer interests, health status, and product popularity. Therefore, only estimating static embeddings can miss important temporal variations of the entities' properties, resulting in poor representations and predictive performance. In addition, practical tensor data are typical sparse, and only a small portion of entries actually have data. Within these entries can be valuable structural knowledge. Current methods, however, are rarely aware of such knowledge. To overcome these limitations, we propose DEMOTE, a novel dynamic embedding approach.
Specifically, we propose an ODE model to learn the embedding trajectories  $\{\mathbf{u}_j^k (t)|1\leq k\leq$ $K,1\le j\le d_k\}$ . The ODE framework is known to be amenable for irregularly sampled, sparsely observed data. More important, ODE models concentrate on learning the time derivative  $\mathrm{d}\mathbf{u}_j^k /\mathrm{d}t$  (i.e., dynamics), rather than the trajectory function itself. Therefore, they have a promising potential to give reliable, long-term trajectory prediction (via numerical integration) even at time points far away from the training timestamps, provided the time derivative is well captured. We construct a joint ODE model for all the embedding trajectories. The ODE consists of a diffusion process and a
reaction process. The diffusion process leverages the structural knowledge in data to co-evolve the embeddings of correlated entities, so as to better overcome the data sparsity. The reaction process models the entity-specific evolution so that it can capture the individual differences in the embedding evolution. The ODE model synergizes the two processes to capture both the commonalities and personalities of these embedding trajectories.
Diffusion Process on Multi-Partite Graphs. First, we construct a graph-based diffusion process to exploit the entity correlations reflected in data. Intuitively, if an observed entry involves entity A (e.g., customer A) and B (e.g., commodity B), the two entities are likely correlated. Thus, we can draw an edge between A and B to express the correlation. We then generalize this intuition to create a  $K$ -partite undirected graph  $\mathbb{G}(E,V)$ , to encode such correlations across all the entities in the  $K$  tensor modes. Each vertex represents a particular entity, and the entire collection of the entities is partitioned into  $K$  groups,  $V = V^{1} \cup \ldots \cup V^{K}$ , where group  $V^{k} = \{v_{1}^{k}, \ldots, v_{d_{k}}^{k}\}$  represents the entities of mode  $k$ . Two entities (at different modes) are connected if they were observed to interact, namely,  $(v_{j}^{k}, v_{j'}^{k'}) \in \mathbb{E}$  if  $\exists \ell_{n} \in \mathcal{D}$  such that  $\ell_{n} = (\ldots, j, \ldots, j', \ldots)$  where  $j$  and  $j'$  are indices at mode  $k$  and  $k'$ , respectively. See Fig. 1 for an illustration. This graph naturally implies underlying information diffusion across the entities within their interactions. For example, if customer A connects to products B and C, it might mean that A distributes their interests/willingness/budgets to purchase B and C. The edges between one merchant A and a list of products  $\{\mathrm{B}, \mathrm{C}, \ldots\}$  might indicate the diffusion of willingness to increase the inventory of these products.
To flexibly estimate the diffusion rate, we introduce a weight  $w_{j,j'}^{k,k'}$  for each edge  $(v_j^k, v_{j'}^{k'}) \in E$ . We then arrange these weights into  $K(K - 1)$  adjacent matrices,  $\mathcal{W} = \{\mathbf{W}^{k,k'} | 1 \leq k, k' \leq K, k \neq k'\}$  where  $\mathbf{W}^{k,k'} = (\mathbf{W}^{k',k})^\top$ . Each  $\mathbf{W}^{k,k'}$  is a sparse  $d_k \times d_{k'}$  matrix that represents the edges and edge weights between  $V^k$  and  $V^{k'}$ , i.e.,  $[\mathbf{W}^{k,k'}]_{j,j'} = w_{j,j'}^{k,k'}$  if  $(v_j^k, v_{j'}^{k'}) \in E$  and 0 otherwise. We now construct a diffusion process based on the  $K$ -partite graph. We view the embedding trajectory as a kind of concentration. For each entity  $j$  at mode  $k$ , the change rate of its concentration (embedding)  $\mathbf{u}_j^k(t)$  is determined by the difference from the concentrations of its neighbors. Since the neighbors come from entities of all the other  $K - 1$  modes, we have
$$
\frac {\mathrm {d} \mathbf {u} _ {j} ^ {k}}{\mathrm {d} t} = \sum_ {s \in \{1, \dots , K \} \backslash k} \sum_ {j ^ {\prime} = 1} ^ {d _ {s}} [ \mathbf {W} ^ {k, s} ] _ {j, j ^ {\prime}} \left(\mathbf {u} _ {j ^ {\prime}} ^ {s} (t) - \mathbf {u} _ {j} ^ {k} (t)\right) = \sum_ {s \in \{1, \dots , K \} \backslash k} \left(\mathbf {w} _ {j} ^ {k, s} \mathbf {U} ^ {s} (t)\right) ^ {\top} - a _ {j} ^ {k, s} \mathbf {u} _ {j} ^ {k},
$$
where  $\mathbf{w}_j^{k,s}$  is the  $j$ -th row of  $\mathbf{W}^{k,s}$ ,  $\mathbf{U}^s (t) = [\mathbf{u}_1^s (t),\dots ,\mathbf{u}_{d_s}^s (t)]^\top$  is the embeddings of all the entities at mode  $s$ , of size  $d_{s}\times R$ , and  $a_j^{k,s} = \sum_{j' = 1}^{d_s}[\mathbf{W}^{k,s}]_{j,j'}$  is the degree of vertex  $j$  in  $\mathbf{W}^{k,s}$ . We can see that the evolution of the embeddings for different modes are coupled. Hence, it is natural to formulate the diffusion process jointly for all the embeddings,  $\frac{\mathrm{d}\mathcal{U}(t)}{\partial t} = \mathrm{d}\left(\mathbf{U}^1 (t),\ldots ,\mathbf{U}^K (t)\right) / \mathrm{d}t = \mathcal{W}\mathcal{U}(t) - \mathcal{A}\mathcal{U}(t) = (\mathcal{W} - \mathcal{A})\mathcal{U}(t)$  where
$$
\mathcal {W} = \left( \begin{array}{c c c c} \mathbf {0} & \mathbf {W} ^ {1, 2} & \ldots & \mathbf {W} ^ {1, K} \\ \mathbf {W} ^ {2, 1} & \mathbf {0} & \ldots & \vdots \\ \vdots & & \ddots & \mathbf {W} ^ {K - 1, K} \\ \mathbf {W} ^ {K, 1} & \ldots & \mathbf {W} ^ {K, K - 1} & \mathbf {0} \end{array} \right),
$$
$\mathcal{A} = \mathrm{diag}\left(\sum_{s\in \{1\dots K\} \setminus 1}\mathbf{A}^{1,s},\dots ,\sum_{s\in \{1\dots K\} \setminus K}\mathbf{A}^{K,s}\right)$  and each  $\mathbf{A}^{k,s} = \mathrm{diag}(a_1^{k,s},\ldots ,a_{d_k}^{k,s})$  is the degree matrix of  $\mathbf{W}^{k,s}$ .
Reaction Process of Individual Entities. Next, to capture the individual difference of each entity in evolving their embeddings, we model a local reaction process for each entity,  $\mathbf{f}_{\theta_k}(\mathbf{u}_j^k (t),t)$ , where  $\mathbf{f}(\cdot)$  is a neural network (NN), and  $\theta_{k}$  are the NN (reaction) parameters for mode-  $k$  entities. The metaphor from the chemical physics is as follows. While substances are being diffused across different sites, at each site a chemical reaction process happens concurrently, which varies the concentration locally. We extend the model as a joint diffusion-reaction process,
$$
\frac {\partial \mathcal {U} (t)}{\partial t} = (\mathcal {W} - \mathcal {A}) \mathcal {U} (t) + \mathcal {F} (\mathcal {U}, t), \quad \mathcal {U} (0) = \mathcal {U} _ {0}, \tag {1}
$$
where  $\mathcal{F}(\mathcal{U},t) = [\mathbf{f}_{\pmb{\theta}_1}(\mathbf{u}_1^1,t),\dots ,\mathbf{f}_{\pmb{\theta}_1}(\mathbf{u}_{d_1}^1,t),\dots ,\mathbf{f}_{\pmb{\theta}_K}(\mathbf{u}_1^K,t),\dots ,\mathbf{f}_{\pmb{\theta}_K}(\mathbf{u}_{d_K}^K,t)]^\top$
Entry Value Generation. Given the embedding trajectories, to obtain the tensor entry value  $m_{\ell}$  at arbitrary time  $t$ , we model  $m_{\ell}(t)$  as a function of the relevant embeddings at time  $t$ ,
$$
m _ {\ell} (t) = g \left(\mathbf {u} _ {l _ {1}} ^ {1} (t), \dots , \mathbf {u} _ {l _ {K}} ^ {K} (t)\right). \tag {2}
$$
While one can follow (Xu et al., 2012; Zhe et al., 2016b) to assign a GP prior over  $g(\cdot)$ , the GP model needs to compute a giant kernel matrix over all the observed entry values (see Sec. 2). It is computationally too expensive or infeasible when the number of observations is large. Hence one has to seek for complex low-rank approximations. To avoid this problem, we model  $g$  with another neural network, which is not only as flexible as GP, but is more scalable and convenient for computation. Since now, the input to  $g(\cdot)$  consists of the trajectory values, which vary with time, our NN model for  $g$  can flexibly capture the complex temporal relationship between the entities. We finally sample the observed entry values with a Gaussian noise model,  $p(\mathbf{y}|\mathbf{m}) = \mathcal{N}(\mathbf{y}|\mathbf{m},\sigma^2\mathbf{I})$  where  $\mathbf{y} = [y_1,\dots,y_N]^\top$  and  $\mathbf{m} = [m_{\ell_1}(t_1),\dots,m_{\ell_N}(t_N)]^\top$ . We focus on real-valued data in this paper. However, it is straightforward to extend our approach to other types of data.
4 Model Estimation
Given data  $\mathcal{D} = \{(\ell_1, t_1, y_1), \ldots, (\ell_N, t_N, y_N)\}$ , the joint probability of our model is
$$
p (\boldsymbol {\beta}, \{\boldsymbol {\theta} _ {k} \}, \mathbf {y}) = p (\boldsymbol {\beta}) \cdot \prod_ {k = 1} ^ {K} p (\boldsymbol {\theta} _ {k}) \cdot \prod_ {n = 1} ^ {N} \mathcal {N} \left(y _ {n} \mid g \left(\mathbf {u} _ {l _ {n 1}} ^ {1} (t _ {n}), \dots , \mathbf {u} _ {l _ {n K}} ^ {K} (t _ {n})\right), \sigma^ {2} \mathbf {I}\right), \tag {3}
$$
where  $\beta$  is the NN parameters for  $g$ , each  $\pmb{\theta}_k$  is the NN reaction parameters for mode- $k$  entities,  $p(\beta)$  and  $p(\pmb{\theta}_k)$  are element-wise standard Gaussian, and  $\mathbf{y} = (y_1,\dots,y_N)^\top$ . To obtain the trajectory values in the Gaussian likelihood of each  $y_{n}$ , we need to solve the ODE in (1) to time  $t_n$ ,
$$
\mathcal {U} \left(t _ {n}\right) = \operatorname {O D E S o l v e} \left(\mathcal {U} _ {0}, 0, t _ {n}, \Theta\right) \tag {4}
$$
where  $\Theta = \{\mathcal{W},\pmb {\theta}_1,\dots ,\pmb {\theta}_K\}$  consists of the ODE parameters. Our goal is to estimate  $\Theta$  , the initial state  $\mathcal{U}_0$  , the NN parameters  $\beta$  , and the noise variance  $\sigma^2$
Stratified Mini-Batch Sampling. We use stochastic mini-batch optimization to maximize the log joint probability so as to estimate all the required parameters,
$$
\mathcal {L} = \log p (\pmb {\beta}, \{\pmb {\theta} _ {k} \}, \mathbf {y}) = \log (\mathrm {P r i o r}) - \sum_ {n = 1} ^ {N} \log \mathcal {N} \left(y _ {n} | g \left(\mathbf {x} _ {n}\right), \sigma^ {2} \mathbf {I}\right)
$$
where  $\log (\mathrm{Prior}) = \log p(\pmb {\beta}) + \sum_{k = 1}^{K}\log p(\pmb{\theta}_k)$ , and  $\mathbf{x}_n = \left(\mathbf{u}_{l_{n1}}^1 (t_n),\dots ,\mathbf{u}_{l_{nK}}^K (t_n)\right)$ . Each time, we sample a mini-batch of observations  $\mathcal{B}$ , and obtain an unbiased stochastic estimate of the log probability,  $\widehat{\mathcal{L}} = \log (\mathrm{Prior}) - \frac{N}{B}\sum_{n\in \mathcal{B}}\left[\log \mathcal{N}(y_n|g(\mathbf{x}_n),\sigma^2)\right]$ . We compute  $\nabla \widehat{\mathcal{L}}$  as the stochastic gradient to update all the parameters.
For each data point  $n$  in the mini-batch, we need to run ODE solving (4) to obtain  $\mathbf{x}_n = \left(\mathbf{u}_{l_{n1}}^1(t_n),\dots,\mathbf{u}_{l_{nK}}^K(t_n)\right)$ . To back-propagate the gradient so as to compute the gradient w.r.t the ODE parameters  $\Theta$  and initial state  $\mathcal{U}_0$ , we can either construct a computational graph during the running of the solver (e.g., the Runge-Kutta method (Dormand and Prince, 1980)), or use the adjoint state method (Pontryagin, 1987; Chen et al., 2018) that solves an adjoint backward ODE to compute the gradient. In whichever case, we need to sort the time points in the mini-batch and solve the ODE sequentially for these time points. As a result, the number of unique time points in the mini-batch greatly influences the speed of processing the mini-batch. The standard mini-batch sampling (based on the training example indices) can result in an uneven allocation of the computational cost across the mini-batches — some mini-batch is fast and some including more unique time points is much slower. To address this issue, we use a simple stratified sampling approach.
- We collect the unique time points in the whole dataset,  $\mathcal{T} = \{\tau_1,\tau_2,\dots \}$  at the beginning.
- To conduct each stochastic update, we first sample  $B$  unique time points  $\mathcal{C}$  from  $\mathcal{T}$ , then for each time point  $\tau_{j} \in \mathcal{C}$ , we look at all the observed entry values produced at  $\tau_{j}$ , namely  $\mathcal{D}_{\tau_{i}} = \{(\ell_{n}, t_{n}, y_{n}) \in \mathcal{D} | t_{n} = \tau_{j}\}$ .
- We randomly sample one example from each  $\mathcal{D}_{\tau_j}$  to collect the mini-batch  $\mathcal{B}$ .
In this way, we ensure the cost of running ODE solvers and related gradient computation in each mini-batch is identical. There are no fluctuations in cost/running time when processing different min-batches. Empirically, we found that the overall speed of our method is much faster than vanilla stochastic mini-batch optimization (see Table 2 in Appendix).