ReDiTT: Retrieval Augmented Conditional
Diffusion Transformers for Asynchronous Time Series
July 14, 2026
We present a diffusion based model for asynchronous time series prediction, where the goal is to predict the next inter event time and event type. To address the inherent uncertainty of future events, we introduce ReDiTT, a retrieval augmented conditional diffusion transformer that operates in latent space. ReDiTT retrieves structurally similar latent sequences from a memory bank during both training and inference and incorporates them as reference conditions through cross attention. This retrieval based conditioning allows the model to attend to relevant temporal dynamics and provides global structural guidance for generation. As a result, ReDiTT stabilizes long horizon forecasting and improves sample diversity. Experiments on seven real world datasets demonstrate state of the art performance on next event prediction and long horizon forecasting. Our code is available at https://github.com/BorealisAI/ReDiTT.
Asynchronous time series (a.k.a. continuous-time event sequence) prediction arises in a wide range of real world applications, including event driven systems [1], healthcare monitoring [2], [3], finance [4], [5], and user behavior modeling [6]–[8], where observations occur at irregular time intervals rather than on a fixed grid. Unlike regularly sampled time series, asynchronous data encodes information jointly in both event values and inter-event times, leading to complex temporal dynamics that are highly stochastic and nonstationary [9]. Accurately modeling such data is crucial for downstream tasks such as forecasting, simulation, and decision making, yet remains challenging due to the sparsity, irregularity, and long-range temporal dependencies inherent in these processes [10], [11].
Asynchronous time series prediction is further complicated by the need to model uncertainty and multimodality over future events, particularly in long-horizon forecasting. Classical autoregressive and likelihood-based temporal point process models often rely on strong parametric assumptions or Markovian dynamics [12]–[16], limiting their ability to generalize to more complex asynchronous time series, particularly those exhibiting rich global structure, long-range dependencies, or mixed continuous and discrete observations.
Recent progress has shown that generative modeling in latent space via variational autoencoders (VAEs) [17] combined with diffusion models [18], can effectively capture the stochastic structure of asynchronous time series. In particular, [19] demonstrates that a VAE can learn a compact and expressive latent representation that supports both accurate reconstruction and diffusion-based next-event and long-horizon forecasting. While diffusion models are powerful generative models, their application to asynchronous time series remains underexplored, despite their ability to generate sequences holistically and mitigate error accumulation compared to autoregressive approaches.
Despite recent progress, unconditional or weakly conditional diffusion models still face significant challenges in long-horizon time series prediction [20]. Unlike image diffusion models, time series data usually do not come with explicit semantic labels or other strong sources of supervision, which limits the availability of informative global conditioning signals during generation. Without such guidance, the model has less ability to preserve the specific dynamics of an individual trajectory over extended horizons. As a result, long-term forecasts often drift toward generic, high-probability patterns from the training distribution, rather than maintaining trajectory-specific temporal structure and fine-grained event evolution.
To address these limitations, we propose ReDiTT : Retrieval Augmented Conditional Diffusion Transformers for Asynchronous Time Series. During training, each sequence retrieves its top \(k\) nearest neighbors from a token memory bank, which are then used as conditions. These retrieved sequences share the same latent format as the input and are incorporated through cross-attention modules within the diffusion transformer blocks, allowing the model to explicitly attend to structurally similar temporal dynamics. At inference time, we retrieve the top \(k\) references from the token memory bank built from training set to guide generation. An overview of ReDiTT is illustrated in 1. This retrieval-augmented conditioning provides global structural guidance that stabilizes long-horizon forecasting, and improves sample diversity by anchoring generation to concrete examples of temporal dynamics rather than relying purely on learned parameters. Our main contributions are:
(char) 1;] We introduce ReDiTT, the first retrieval-based diffusion framework for asynchronous time series prediction that conditions on top-\(k\) retrieved latent priors from a pre-constructed latent token bank.
(char) 2;] We propose a novel conditioning approach for retrieve-based diffusion transformer and show it effectively integrates prior information and provides reasonable guidance.
(char) 3;] We demonstrate that ReDiTT significantly improves next-event and long-horizon prediction with state-of-the-art performances by experiments on seven real world datasets with a comprehensive analysis.
Marked TPPs are widely adopted as a standard approach for modeling asynchronous time series. A TPP is a stochastic process that generates sequences of discrete events over time. A sequence of \(n\) events can be represented as \({\mathbf{x}}=\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_n\}\), where each event \({\mathbf{x}}_i=(t_i,e_i)\) consists of the index \(i\) indicating the chronological order of events, the inter-event time \(t_i\) since the previous event and the corresponding event type \(e_i\). Another widely used variant is to represent events by their inter-event intervals \(\tau_i=t_i-t_{i-1}\) rather than the absolute timestamps \(t_i\). These two parameterizations are essentially equivalent, so the literature often switches between them without loss of generality. Classical TPP models [13] parameterize conditional intensity functions for next-event prediction and are trained by maximizing the log-likelihood of observed event sequences. Prior work on temporal point processes has largely focused on neural architectures that extend likelihood-based intensity modeling, beginning with RNN-based approaches [13], [21] and later incorporating more expressive designs to better capture uncertainty [22], [23]. More recent Transformer-based TPPs leverage attention mechanisms to improve long range dependency modeling [14]–[16], but still rely on autoregressive intensity formulations. However, the goal has extended beyond forecasting the immediate next event to generating the full future event trajectory, and classical and neural TPPs can be unsatisfactory since errors introduced at early steps accumulate as the autoregressive rollout proceeds. On the optimization side, likelihood-based training typically requires computing distributional quantities implied by the learned intensity, which can become expensive.
The shortcomings of classic neural TPP approaches motivate diffusion and flow matching approaches. These methods treat forecasting as conditional generation, transforming simple base noise into future event sequences given the history. By modeling the joint continuation rather than repeatedly sampling one event at a time, diffusion based approaches can better support long-horizon generation and produces diverse futures for uncertainty quantification, offering a strong non-autoregressive alternative to intensity-based neural TPPs. Early diffusion work in time series focused on learning conditional distributions for tasks like imputation [24] and probabilistic forecasting [25], demonstrating that iterative denoising can capture uncertainty without being constrained to a single greedy rollout. For temporal point processes specifically, Add and Thin [23] formulates diffusion over full marked event sequences and outperforms autoregressive TPP forecasters. [26] embed event sequences into a vector space and diffuse to forecast the full future sequence in one go. [27] extended diffusion to spatio-temporal setting. [28] coupled two diffusions, one for inter-arrival times and one for event types, to model the joint distribution. More recently, ADiff4TPP [19] proposed an asynchronous noise schedule in a VAE latent space to strengthen conditioning via earlier event history when predicting further into the future. While promising, these approaches yield only limited empirical gains, and diffusion for asynchronous event modeling is still relatively understudied. We propose retrieval based diffusion transformers to tackle these challenges and better predict long-horizon future.
Retrieval has become a practical way to enhance generative models with non-parametric memory. For text, \(k\)NN language models [29] augment a base LM by retrieving semantically similar contexts from a data store at inference time. For image, retrieval-augmented diffusion models [30] condition denoising on retrieved reference images to better capture specific visual structure and long-tail concepts. In time series forecasting, several recent methods [31]–[34] leverage retrieval by selecting historically similar segments and using their continuations as additional context. RATD [20] further combines retrieval with diffusion by injecting retrieved references into the denoising process. However, these approaches are primarily designed for regularly sampled time series segments and do not directly apply to asynchronous marked event streams without nontrivial changes. In contrast, our method is tailored to event prediction: we perform retrieval in a VAE latent space that preserves event wise structure, and condition a latent diffusion transformer on retrieved references represented in the same event-sequence format. This design enables trajectory-specific guidance for coherent long horizon event generation, rather than serving only as a segment level short range forecasting booster.
Asynchronous time series forecasting is typically evaluated using two tasks. Suppose we observe an event history \(\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_i\}\). Next event prediction requires a model to forecast the immediate next event \(\hat{\mathbf{x}}_{i+1} = (\hat{t}_{i+1}, \hat{e}_{i+1}\)), including both its occurrence time or inter-arrival time and its type or mark, conditioned on the observed history. Long horizon prediction extends this setting by requiring the model to generate a sequence of future events \(\{\hat{\mathbf{x}}_{i+1},\cdots,\hat{\mathbf{x}}_{i+m}\}\) over a prediction window of length \(m\). This task evaluates the model’s ability to capture how uncertainty accumulates as the forecasting horizon increases.
Flow matching (FM) [35]–[38] provides a diffusion style way to train continuous time generative models by directly regressing a velocity field instead of learning scores. Concretely, FM views the generation as solving an ODE that transports a simple base distribution (typically Gaussian noise) to the data distribution via a learned vector field \(\frac{d{\mathbf{z}}_s}{ds}=v_\theta({\mathbf{z}}_s,s)\), using an explicit interpolation path \({\mathbf{z}}_s=a_s{\mathbf{z}}+b_s\boldsymbol{\epsilon}\), where \(s\) denotes a random time \(s\sim \mathcal{U}(0,1)\) (we use \(s\) to denote the timestep in diffusion model, and \(t\) to represent the time value for asynchronous time series data), \({\mathbf{z}}\) is the input clean sample, and \(\boldsymbol{\epsilon}\sim \mathcal{N}(0,1)\) is the Gaussian noise. Rectified flow [36], [37] defines a simpler interpolation \({\mathbf{z}}_s=(1-s){\mathbf{z}}+s\boldsymbol{\epsilon}\). Conditional Flow Matching (CFM) makes training simulation-free by deriving a direct closed-form conditional vector field \(u_s({\mathbf{z}}_s|\boldsymbol{\epsilon})\) for this path, and learning \(v_\theta\) by a simple regression loss \(\mathbb{E}_{s,\boldsymbol{\epsilon},{\mathbf{z}}}\big[\|v_\theta({\mathbf{z}}_s,s)-u_s({\mathbf{z}}_s|\boldsymbol{\epsilon})\|^2\big]\).
To facilitate efficient retrieval and ease the diffusion training, each event \({\mathbf{x}}_i\) is encoded as a compact latent representation \({\mathbf{z}}_i\in\mathbb{R}^d\) using a pretrained VAE \(E_\phi\). Let \({\mathbf{x}}\in\mathcal{X}\) be an event sequence of length \(n\). We pad \({\mathbf{x}}\) to a fixed length \(N\), where \(N\) is the maximum sequence length in \(\mathcal{X}\). The encoder then maps the padded sequence to a latent representation \({\mathbf{z}}=E_\phi({\mathbf{x}})\in\mathbb{R}^{N\times d}\), which serves as a single training sample for latent flow matching.
To make flow matching respect the causal and unevenly spaced nature of event sequences, asynchronous matrix valued interpolation [19] defines: \[{\mathbf{z}}_s={\boldsymbol{A}}(s){\mathbf{z}}+({\boldsymbol{I}}-{\boldsymbol{A}}(s))\boldsymbol{\epsilon}, \label{eqn:interpolation}\tag{1}\] where \(\epsilon=\{\epsilon_1,\cdots,\epsilon_N\}\) is a Gaussian noise of same dimension \(N\times d\) as \({\mathbf{z}}\). And \({\boldsymbol{A}}(s)\in\mathbb{R}^{N\times N}\) is a diagonal matrix whose per event schedule is designed so that later events are injected with noise earlier (and therefore are trained to be denoised earlier) than earlier events, that being said, later events are corrupted earlier and the model learns to denoise and forecast the tail under stronger noise.
This yields a generative ODE that incorporates the schedule derivative with the chain rule: \[\label{eqn:vectorfield} \frac{d{\mathbf{z}}_s}{ds}={\boldsymbol{A}}'(s)v_\theta({\mathbf{z}}_s,{\boldsymbol{A}}(s)).\tag{2}\] In each training iteration, we randomly select a time \(s\in[0,1]\) and generate the corresponding intermediate state \({\mathbf{z}}_s\) using 1 . Training follows a CFM objective in this asynchronous setting by weighting the regression with \({\boldsymbol{A}}'(s)\): \[\mathcal{L}_{\text{CFM}}(\theta)=\mathbb{E}_{s,\boldsymbol{\epsilon},{\mathbf{z}}}\big[\,\,\|{\boldsymbol{A}}'(s)\big(\,v_\theta({\mathbf{z}}_s,{\boldsymbol{A}}(s))-u_s({\mathbf{z}}_s|\boldsymbol{\epsilon})\,\big)\|^2\,\,\big].\]
During inference, given a observed history \({\mathbf{x}}_1,\cdots,{\mathbf{x}}_n\) and a target prediction horizon size \(m\). We can form a initial condition \(\{{\mathbf{z}}_1,\cdots,{\mathbf{z}}_n,\epsilon_{n+1},\cdots,\epsilon_{n+m}\}\) and we can incorporate the history by filling the corresponding entries of \(\frac{d{\mathbf{z}}_s}{ds}\) in 2 with the known latents: \[\frac{d{\mathbf{z}}_s}{ds} = \begin{cases} {\boldsymbol{A}}'(s)_{ii}({\mathbf{z}}_i-\epsilon_i) & \text{if } 1\leq i\leq n , \\ {\boldsymbol{A}}'(s)_{ii}v_\theta({\mathbf{z}}_s,{\boldsymbol{A}}(s)) & \text{if } n<i\leq n+m. \end{cases}\] Standard ODE solvers can solve this without introducing numerical error, which leads us to the predictions \(\{\hat{\mathbf{z}}_{n+1},\cdots,\hat{\mathbf{z}}_{n+m}\}\). The decoder then maps them back to the original space to obtain \(\{\hat{\mathbf{x}}_n,\cdots,\hat{\mathbf{x}}_{n+m}\}\).
To better ground diffusion-based generation in trajectory-specific dynamics, we propose ReDiTT, a retrieval augmented latent Diffusion Transformer for temporal point processes, as illustrated in 2. Given an observed event history \(\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_i\}\) as input, our model generates a future event sequence \(\{\hat{\mathbf{x}}_{i+1},\cdots,\hat{\mathbf{x}}_{i+m}\}\), supporting both next event prediction (\(m=1\)) and long horizon prediction (\(m>1\)). ReDiTT couples latent space embedding retrieval with conditional generation: we retrieve the top-\(k\) most similar historical trajectories from a reference bank and use them as additional conditioning to guide the diffusion model towards context consistent futures. We first describe our latent space retrieval strategy in 4.1. We then introduce the reference guided conditional flow matching objective in 4.2. Finally, we detail the conditional DiT architecture that integrates retrieved references via cross-attention in 4.3.
Directly retrieving asynchronous sequences in the original observation space requires a carefully designed similarity function: each sequence mixes continuous inter-event times with discrete event types, has variable length with padding, and may admit multiple plausible alignments, such as shifted temporal patterns or local event reorderings. Consequently, naive raw-space distances, such as Euclidean distance over concatenated padded features, can be dominated by scale mismatch, sparsity, or padding artifacts, and may not reflect trajectory-level similarity. We therefore perform retrieval in the latent space of the pretrained VAE \(E_\phi\), where each event \({\mathbf{x}}_i=(t_i,e_i)\) is mapped to a compact representation \({\mathbf{z}}_i\in\mathbb{R}^d\) that is optimized to preserve the information needed for sequence reconstruction. This choice is primarily motivated by representation alignment with the downstream generative model. Since ReDiTT performs diffusion and conditioning over event-sequence latents, latent-space retrieval allows the model to consume retrieved references in the same space as the generation target, without requiring an additional hand-designed metric over mixed continuous-discrete observations. Our raw-space retrieval ablation in 10.1 shows that carefully normalized raw similarities can also provide useful neighbors, indicating that retrieval-augmented conditioning is robust to the retrieval representation. Nevertheless, latent retrieval offers a simple and model-aligned way to obtain reference trajectories for coherent multi-event forecasting.
Consider an asynchronous sequence \({\mathbf{x}}=\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_N\}\) with a padding mask \({\mathbf{m}}\in\{0,1\}^N\) (this mask is common in data processing to keep a constant total event length among all the sequence), we use \(E_\phi\) to produce per-event latent tokens \({\mathbf{z}}=E_\phi({\mathbf{x}})\in\mathbb{R}^{N\times d}\). In addition, we precompute a latent token bank \(\mathcal{B}=\{({\mathbf{z}}=E_\phi({\mathbf{x}}), {\mathbf{m}})| \,{\mathbf{x}}\in \mathcal{X}_{\text{train}} \}\) for the training database. We then define a masked cosine similarity between a query \((\boldsymbol{\em q}, {\mathbf{m}}^q)\in E_\phi(\mathcal{X})\) and a reference candidate \(({\mathbf{r}},{\mathbf{m}}^r)\in\mathcal{B}\) as a token-wise cosine averaged over valid time steps: \[\texttt{sim}(\boldsymbol{\em q},{\mathbf{r}})=\frac{1}{\sum_i\mathbf{1}[{\mathbf{m}}_i^{q} \wedge {\mathbf{m}}_i^r]}\sum_{i:{\mathbf{m}}_i^q\wedge {\mathbf{m}}_i^r}\frac{\langle \boldsymbol{\em q}_i,{\mathbf{r}}_i\rangle}{\|{\boldsymbol{\em q}}_i\|\|{\mathbf{r}}_i\|}, \label{eqn:sim}\tag{3}\] where \(\frac{\boldsymbol{\em q}}{\|\boldsymbol{\em q}\|_2},\frac{{\mathbf{r}}}{\|{\mathbf{r}}\|_2}\) are \(\ell_2\) normalized tokens. Finally, we retrieve the top \(k\) references for \(\boldsymbol{\em q}\): \[R(\boldsymbol{\em q})=\{{\mathbf{r}}\,|\,\text{TopK}_{r\in\mathcal{B}}\,\texttt{sim}(\boldsymbol{\em q},{\mathbf{r}}), {\mathbf{r}}\neq\boldsymbol{\em q}\},\] which will enter the diffusion process as extra condition. During training, the query sequence is part of the training reference bank, so its nearest neighbor is typically the query sample \(\boldsymbol{\em q}\) itself and would dominate the conditioning signal. Including this exact match would create a degenerate conditioning signal, allowing the model to rely on a duplicate of the input rather than learning to use informative neighboring trajectories. We therefore exclude the closest reference during training and retrieve from the second closest neighbor onward. At inference time, the query comes from the validation or test set while the reference bank contains only training sequences, so the top retrieved neighbor is not an exact duplicate of the query and can be safely used for conditioning.
During evaluation, retrieval follows the natural forecasting setup used for all baselines: only the observed prefix is available when predicting future events. Accordingly, before computing retrieval similarity, we mask all target/future positions in the query sequence and perform retrieval using only the observed prefix. The retrieved sequence itself may contain its full continuation, but it is selected without access to the query’s future events. To further examine the effect of this prefix-only retrieval protocol, we include an additional retrieval-only analysis in 9.2. The results show that prefix-only retrieval remains competitive with full-sequence retrieval on most datasets, indicating that the observed prefix is generally sufficient to retrieve useful neighboring trajectories.
Let \(R\) denote the external reference condition associated with a target latent event sequence \({\mathbf{z}}\). Conditioning does not alter the flow matching construction and 1 still holds, but we now learn a conditional vector field \(v_\theta(\,\cdot\,|\,R)\) that transports samples along this path while being guided by \(R\). The conditional flow matching objective becomes: \[\mathcal{L}_{\text{CFM}}(\theta)=\mathbb{E}_{s,\boldsymbol{\epsilon},{\mathbf{z}}}\big[\,\,\|{\boldsymbol{A}}'(s)\big(\,v_\theta({\mathbf{z}}_s,{\boldsymbol{A}}(s)\,|\,R)-u_s({\mathbf{z}}_s|\boldsymbol{\epsilon})\,\big)\|^2\,\,\big].\] Intuitively, \(R\) does not modify the geometry of the bridge from \(\boldsymbol{\epsilon}\) to \({\mathbf{z}}\), but it provides additional information that biases the learned dynamics toward trajectory-consistent solutions, which is especially helpful for long-horizon generation.
During training, we retrieve \(k\) references \(R\) from the same training bank \(\mathcal{B}\) and feed the latent tokens as an addition condition in every DiT block, encouraging the conditional flow to model a set of plausible futures consistent with similar historical trajectories. During inference, we retrieve from the same training bank again , so the sequences that model has already encountered can serve as an in-distribution prototype that stabilizes long-horizon generation and reduces drift when extrapolating far beyond the local context.
We modify the Transformer based DiT blocks because they have the following issues when applying to asynchronous time series: conditioning is often injected through adaptive normalization (e.g., AdaLN) where a global conditioning vector (typically built from the diffusion timestep and a sparse label) modulates each block. For asynchronous sequences, this kind of conditioning is intrinsically limited: the timestep carries no instance-specific semantics, and event “labels” (if any) are usually coarse or unavailable, so the conditioning signal is weak and globally broadcast to all tokens in the same way.
After retrieving \(k\) nearest-neighbor reference sequences in latent space, we therefore construct a sequence-shaped condition that mirrors the input format. Say we have retrieve \(R=\{{\mathbf{r}}^j\}_{j=1}^k\) for input latent \({\mathbf{z}}\), we first obtain a reference sequence by weighted pooling: \[{\mathbf{r}}=\sum_{j=1}^kw_j{\mathbf{r}}^j\in\mathbb{R}^{N\times d}, \,\,w_j=\texttt{softmax}\big(\texttt{sim}({\mathbf{z}},{\mathbf{r}}^j)\big).\]
Before entering the DiT blocks, both current latent input and retrieved reference are expanded to hidden size \(D\) by channel repetition, for \({\mathbf{r}}\), we also broadcast the mask
\({\mathbf{m}}^{{\mathbf{z}}}\) across the channel dimension. Then \({\mathbf{z}}\) and \({\mathbf{r}}\) are augmented with fixed positional embeddings
pos to get \({\mathbf{z}}=\texttt{Repeat}({\mathbf{z}})+\texttt{pos}_{{\mathbf{z}}}\in\mathbb{R}^{N\times D}\) and \({\mathbf{r}}=\texttt{Repeat}({\mathbf{r}})+\texttt{pos}_{{\mathbf{r}}}\in\mathbb{R}^{N\times D}\). Considering the reference has the same format as the input, it is naturally to inject \({\mathbf{r}}\) via cross-attention after self attention in each DiT block as illustrated in 2. After the LayerNorm, the module forms: \[{\boldsymbol{Q}}={\boldsymbol{W}}_q(\boldsymbol{LN}({\mathbf{z}})),\,\,{\boldsymbol{K}}={\boldsymbol{W}}_k(\boldsymbol{LN}({\mathbf{r}})),\,\,{\boldsymbol{V}}={\boldsymbol{W}}_v(\boldsymbol{LN}({\mathbf{r}})),\] splits into
\(\alpha\) heads, and computes masked attention weights over the flattened memory positions: \[\boldsymbol{Attn}=\texttt{MaskedSoftmax}(\frac{{\boldsymbol{Q}}{\boldsymbol{K}}^\intercal}{\sqrt{D/\alpha}};{\mathbf{m}}^r).\]
The attention output is \({\boldsymbol{Y}}=\boldsymbol{Attn}{\boldsymbol{V}}\), followed by an output projection \({\boldsymbol{W}}_{\text{out}}\) and a learned scalar sigmoid gate \(c\in\mathbb{R}\): \[{\mathbf{z}}={\mathbf{z}}+\sigma(c){\boldsymbol{W}}_{\text{out}}{\boldsymbol{Y}}.\]
Our cross-attention conditioning is effective as it injects the retrieved reference as a token level external memory, letting each latent token selectively attend to the most relevant reference positions rather than relying on a coarse global conditioning vector. The sigmoid gated residual makes this guidance stable and adaptive, i.e. starting weak and strengthening only when it improves generation, so the model can leverage retrieval without overwhelming the base dynamics. More details are discussed in Appendix 8.3.
Datasets. Following [9], we run experiments on five standard temporal point process datasets: Amazon [39] with 16 event types; Retweet [40] with 3 event types; StackOverflow [41] with 22 event types; Taobao [42] with 20 event types; Taxi [43] with 10 event types. We also evaluate ReDiTT on two action datasets: Breakfast [44] with 177 action classes, we scale the timestamp by dividing by 100 to avoid VAE training loss explosion; Multithumos [45] with 65 action classes, we take the 212 sequences with length \(\leq\) 100 out of total 289 sequences to avoid model ran out of memory. More details are in Appendix 8.1 and ¿tbl:tab:hyper?.
Baselines. We compare ReDiTT with four kinds of state-of-the-art TPP models for asynchronous time series: RNN-based models, including RMTPP [21] and NHP [13]; Attention-based models, including SAHP [14], THP [15], AttNHP [16], and DTPP [46]; Diffusion-based models, including Add&Thin [23] and ADiff4TPP [19]; Other popular models including IFTPP [47] and HYPRO [42].
Metrics. Following [9], we examine our models on next event prediction. We evaluate the time prediction by root mean square error (RMSE) between predicted time and true time, and the event type prediction by accuracy between predicted type and true type. Note to compute the accumulated RMSE and accuracy, when predicting \(\hat{\mathbf{x}}_{i+2}\), the model uses the true history \(\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_{i},{\mathbf{x}}_{i+1}\}\) rather than a previously predicted event \(\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_{i},\hat{\mathbf{x}}_{i+1}\}\). We also examine ReDiTT on long-horizon prediction by computing the optimal transport distance (OTD), a measurement of edit distance between the predicted \(\{\hat{\mathbf{x}}_{i+1},\cdots,\hat{\mathbf{x}}_{i+m}\}\) and the ground truth \(\{{\mathbf{x}}_{i+1},\cdots,{\mathbf{x}}_{i+m}\}\). We follow the implementation of OTD by [48] with dynamic programming to efficiently find alignment and compute distance between predictions and true events. Following [19], we compute the OTD with horizon window size \(m=5,10,20,30\) respectively.
To facilitate reproducibility, we provide full implementation details and hyperparameter settings in Appendix 8.
width=
In this section, we evaluate our model against existing methods on next-event prediction and long-horizon prediction tasks. Overall, ReDiTT outperforms existing baselines on both tasks and is particularly effective in regimes with long event sequences and limited training samples, where purely parametric models struggle to capture rare transitions and long-range temporal dependencies. By retrieving trajectory-level neighbors as explicit context, ReDiTT provides strong guidance for coherent long-horizon generation under sparse supervision. Beyond retrieval itself, the results demonstrate that the manner in which retrieved information is injected into the model is critical. The proposed conditioning strategy, based on cross attention in diffusion transformers, enables the denoising network to selectively attend to the most relevant retrieved context, resulting in more accurate and coherent predictions.
¿tbl:tab:main? summarizes the main results on seven benchmark datasets. Performance is evaluated on next-event prediction, using RMSE for time prediction and accuracy for type prediction. Overall, ReDiTT exhibits the strongest and most consistent performance across a broad range of dynamics. It achieves the best RMSE and the highest type accuracy on all datasets. ReDiTT substantially outperforms both classic neural TPP baselines and recent diffusion-based or flow-based forecasters. In addition, ReDiTT delivers large gains in average performance across datasets compared to the strongest baseline, ADiff4TPP [19]. The average RMSE is reduced by \(22\%\), decreasing from \(3.887\) for ADiff4TPP to \(3.008\) for ReDiTT. At the same time, the average type accuracy increases by 12 percentage points, from \(42.4\%\) to \(55.2\%\). The advantages of ReDiTT are especially pronounced on Breakfast and MultiThumos, which feature large event vocabularies but substantially fewer training sequences. In this low-data, high-cardinality regime, purely parametric models are prone to sparse supervision for rare event types. Retrieval provides valuable exemplar-based guidance by exposing the model to structurally similar trajectories and richer type co-occurrence patterns. On Breakfast, which contains 177 event types, ReDiTT improves type accuracy from \(8.2\%\) to \(15.1\%\). A consistent trend is observed on MultiThumos with 65 event types, where ReDiTT yields a large accuracy gain from \(16.3\%\) to \(25.0\%\). These results indicate that retrieval is particularly effective when the label space is large and training coverage per class is limited.
For long-horizon prediction, ReDiTT consistently reduces OTD at horizons \(m = 5, 10, 20, 30\) on six datasets compared to the strongest diffusion baseline [19]. This result demonstrates that retrieval guidance substantially improves sequence-level coherence beyond single-step prediction. The gains are most pronounced on Amazon, where ReDiTT lowers OTD from 12.4 to 10.7 at horizon 10 and from 32.9 to 25.8 at horizon 30, indicating that the advantages of retrieval-augmented conditioning grow as the prediction horizon increases. Similar patterns are observed on Retweet and StackOverflow, where ReDiTT consistently achieves the best OTD across horizons, highlighting the robustness of the approach under diverse temporal dynamics. These long-horizon improvements are driven by the retrieval-conditional design, which injects trajectory-relevant training examples as explicit context. This design enables the model to generate multiple future events coherently while preserving temporal ordering and capturing long-range sequential structure. On Taxi, ReDiTT achieves the best RMSE and accuracy, while OTD is optimized at smaller retrieval sizes. This behavior aligns with the dataset characteristics, which include short sequences and a limited event vocabulary. In this setting, retrieval provides smaller incremental benefits because sequence-level alignment is already well constrained. Overall, these findings support the core motivation of retrieval-augmented conditioning. By injecting trajectory-relevant training examples as explicit context, the model is more effectively guided toward plausible future dynamics.
Different Conditioning Architecture. As discussed in [18], adaptive layer norm often outperforms cross-attention conditioning in diffusion transformers for class-labeled image generation. We therefore conduct an ablation on five datasets with fixed \(k=1\) to analyze which conditioning architecture better suits asynchronous event sequences. Following the design in [18], we implement adaLN conditioning on \({\mathbf{r}}\) and also consider Concat\((s,{\mathbf{r}})\), which concatenates the diffusion timestamp \(s\) with \({\mathbf{r}}\) to encode temporal information, i.e. three paths for \(s\), \({\mathbf{r}}\) and the fused Concat\((s,{\mathbf{r}})\), each is controlled by a scale hyperparameter. More details can be found in Appendix 8.3. As reported in ¿tbl:tab:adaln?, cross-attention consistently yields stronger long-horizon behavior, achieving substantially lower OTD, whereas adaLN only improves event-type accuracy on three of the five datasets and often degrades sequence-level alignment. StackOverflow illustrates the trade-off clearly. adaLN achieves better next-event metrics, but its long-horizon alignment is much worse. In contrast, cross-attention produces more coherent multi-step forecasts (OTD 10.8 at horizon 10 and 27.4 at horizon 30), even though its next event prediction accuracy is lower and RMSE is higher. For the other three datasets except for Taxi, we observe a similar pattern: any gains adaLN provides in next-event RMSE or type accuracy are relatively modest, while its long-horizon OTD deteriorates substantially. This indicates that adaLN’s improvements are largely limited to short-term, local prediction, whereas cross-attention is far more reliable for preserving sequence-level alignment over extended horizons.
We attribute this difference to the mismatch between the inductive bias of adaLN and the structure of conditioning signals in asynchronous time series. In class-labeled image generation, the condition is typically a single global discrete cue shared across the entire sample; in this setting, adaLN’s global, layer-wise modulation can effectively inject the label signal and amplify class conditional priors. In contrast, for asynchronous time series, retrieved input-like references are high-entropy and event-dependent, with information that varies across events and time. Cross-attention therefore provides a more suitable mechanism by enabling state-dependent access to context and selectively integrating temporal patterns as the generated history evolves. Overall, the observed trade-off suggests that adaLN primarily strengthens categorical discrimination, whereas cross-attention better supports precise, context-aligned modeling of event times. Moreover, adaLN is also more expensive in our implementation, resulting in a 294M parameter model, whereas the cross-attention variant is substantially lighter with only 209M parameters. Cross-attention also converges faster in practice, reaching strong performance in fewer training iterations. To obtain both predictive quality and computational efficiency, we adopt cross-attention conditioning in our final model.
width=.5
Different Aggregation of Retrieved References. We further compare two strategies for incorporating the top \(k\) retrieved references during conditioning: average pooling and concatenation. We perform the experiments on StackOverflow, Breakfast, and MultiThumos, and the results are illustrated in ¿tbl:tab:concat?. Under a fixed retrieval budget \(k\), we observe that average pooling of the retrieved references consistently outperforms concatenation across three metrics. Moreover, as \(k\) increases, the performance gap widens, mostly notably on event-type prediction accuracy. For concatenation, the accuracy at \(k=7\) is lower than at \(k=5\) for StackOverflow and Breakfast. We hypothesize that pooling becomes increasingly beneficial with larger retrieved sets because it acts as a robust aggregation operator, that being said, it emphasizes signals that are consistent across references with high similarity weights while dampening irrelevant information. In contrast, concatenation scales the conditioning length linearly with \(k\), which increases exposure to irrelevant retrievals and makes it easier for the model to overfit to noisy reference. The results showed that this effect becomes more obvious as \(k\) grows, leading to a larger degradation in categorical prediction. We therefore adopt average pooling in our conditioning implementation.
width=
Different \(k\) for Retrieval Mechanism. We also conducted an ablation study on the number of retrieved references \(k\) to further investigate how the amount of retrieved context influences conditioning effectiveness while keeping all other settings fixed, specifically, whether a small set of references provides sufficient guidance for generation, or whether increasing \(k\) introduces irrelevant matches that act as noise and degrade performance. As listed in ¿tbl:tab:concat?, larger \(k\) does not improves the results for concatenation conditioning much. However, for average pooling conditioning, increasing \(k\) consistently improves both RMSE and event type accuracy, whereas OTD remains largely stable across different values of \(k\). This suggests that retrieving more neighbors can provide additional contextual signal that helps the model refine local predictions, particularly for event type inference, while preserving the global sequence level structure captured by OTD. Notably, for StackOverflow, event type accuracy reaches 53% at \(k=7\), representing a 10.9 point improvement over \(k=1\). This indicates that multi-reference conditioning is beneficial, and a moderate retrieval size can yield significant gains, especially for categorical prediction, while preserving long-horizon sequence alignment.
Effect of Time and Type Conditioning Components. We further ablate the conditioning signal by using time-only or event-type-only guidance. Because retrieval is performed in the latent space, we first retrieve reference sequences via latent similarity, and then mask out the undesired modality by replacing the corresponding latent block (time or type) with a zero latent (i.e., the latent obtained by encoding an all-zero input in the original space), while keeping the other block unchanged. This design isolates the contribution of each modality without altering the retrieval set.
We evaluate this ablation on Taobao (20 event types), Breakfast (177 event types), and Multithumos (64 event types) with fixed \(k=1\). As shown in ¿tbl:tab:cond?, each modality contributes useful but incomplete information. Time-only conditioning can improve metrics that are more sensitive to temporal dynamics, while event-type-only conditioning can perform better on metrics that depend more on categorical structure. However, neither single component consistently dominates across all evaluation criteria, indicating that time and type provide complementary guidance. When both conditioning blocks are used together, the model achieves strong and balanced performance across all metrics, suggesting that jointly conditioning on temporal and event-type information enables the retrieval guidance to capture both when events occur and what events occur. This supports the design choice of using the full time-and-type conditioning signal rather than relying on either modality alone.
width=
We propose ReDiTT, a retrieval augmented conditional Diffusion Transformer for asynchronous event sequence forecasting that models the joint uncertainty of the next inter event time and event type. ReDiTT operates in latent space and uses a memory bank to retrieve structurally similar latent sequences as reference conditions during training and inference. The retrieved references are injected through cross attention in DiT blocks, enabling the model to selectively use trajectory specific temporal patterns that are hard to capture with purely parametric conditioning. Across seven real world datasets, ReDiTT achieves state-of-the-art results on both next event prediction and long horizon forecasting. Ablations confirm the value of retrieval based conditioning and cross attention integration, and show that structured guidance is critical for stable long horizon generation in temporal point processes.
Limitations. Although ReDiTT achieves substantial empirical gains, it introduces additional system overhead by requiring a reference bank and performing retrieval at inference time, which increases both memory usage and latency. Another limitation is that gains can be limited on small-scale datasets, especially when the sequence length is relatively short, the model might have limited improvement on capturing long range behaviors.
Future Work. Future directions include learning retrieval representations that are better aligned with forecasting objectives, and making conditioning more robust via reference weighting or retrieval dropout. On the efficiency side, we can reduce cost with compressed or hierarchical reference banks and caching. It is also promising to incorporate richer context such as covariates or text, and to explore training objectives that more directly target long-horizon metrics.
Broader Impact Statement. This paper presents a diffusion-based model for predicting future events from historical data. Such models may support improved planning and decision-making by enabling probabilistic forecasts that capture uncertainty and multiple plausible futures. This can be beneficial in domains where anticipating diverse outcomes is important. However, predictions of future events can be misinterpreted or misused, particularly if treated as deterministic or deployed in high-stakes settings without appropriate oversight. The model may also inherit biases present in training data and does not address issues of fairness or causality. Additionally, diffusion-based methods can be computationally demanding, and their environmental and efficiency costs should be considered. This work is intended as a methodological contribution, and we encourage responsible use and further research on interpretability, fairness, and efficiency.
Following [9], we pad Amazon, Retweet, StackOverflow, Taobao, and Taxi to the maximum sequence length within each dataset. For Breakfast, we rescale time values by dividing by 100 to stabilize VAE training. For the two text-based action datasets, Breakfast and MultiThumos, we encode event types as integer indices in the same way as the other datasets, and split the data into train/validation/test sets with a 70/10/20 ratio. Additional dataset details are provided in dataset settings of ¿tbl:tab:hyper?.
For baseline models, we use the implementations from EasyTPP [9] for RMTPP [21], NHP [13], SAHP [14], THP [15], AttNHP [16], and IFTPP [47].
For other baselines, we adapt DTPP [46] to predict inter-event times in their original scale rather than in log space, following ADiff4TPP [19]. For Add&Thin [23], the method outputs only inter-event times, so we do not report event-type accuracy or OTD.
Our VAE implementation for asynchronous time series follows ADiff4TPP [19]. For each event \({\mathbf{x}}=(t,e)\), we train a \(\beta\)-VAE consisting of an encoder \(E_\phi\) that maps \({\mathbf{x}}\) to a latent variable \({\mathbf{z}}\), modeled as a Gaussian distribution parameterized by an inferred mean (and variance). A decoder \(D_\phi\) then reconstructs the event from \({\mathbf{z}}\), producing \(\tilde{\mathbf{x}}=(\tilde{t},\tilde{e})=D_\phi({\mathbf{z}})\). The whole \(\beta\)-VAE is optimized with the following objective: \[\mathcal{L}=(t-\tilde{t})^2+\text{CE}(e,\tilde{e})+\beta\mathcal{L}_{KL},\]
\(\beta\) controls the weight of the KL regularization term, which penalizes deviations of the latent distribution \(\mathcal{X}\) from a standard Gaussian prior. As shown in the VAE settings of ¿tbl:tab:hyper?, we use the same \(\beta\) choices reported in the appendix of [19].
Adapting DiT for Asynchronous Time Series. Following [18], we instantiate \(v_\theta(\cdot,{\boldsymbol{A}}(s))\) with a Diffusion Transformer (DiT) backbone, since its Transformer blocks provide a flexible way to model long-range dependencies while remaining compatible with diffusion or flow-style training objectives. We only make minimal modifications to adapt DiT from images to asynchronous event sequences. In particular, we remove the image-specific patchify and patch-embedding modules and instead feed the model with latent event tokens produced by the encoder \(E_\phi\). This replacement preserves the core DiT computation (stacked Transformer blocks with conditioning) while ensuring that the model operates directly on a sequence of event-level representations, which is the natural input format for asynchronous time series forecasting.
When initializing the model, we specify a hyperparameter \(N\) that sets the maximum sequence length the model can represent and generate. Concretely, sequences shorter than \(N\) are padded and masked, allowing us to use a fixed token length for efficient batched training and inference while still supporting variable-length histories and horizons. Finally, to inject the conditioning matrix \({\boldsymbol{A}}(s)\in\mathbb{R}^{N\times N}\), we extend the original DiT conditioning design by constructing an embedding that reflects both temporal progression and structural interactions encoded by \({\boldsymbol{A}}(s)\). We achieve this by broadcasting the entries of \({\boldsymbol{A}}(s)\) through sinusoidal frequency features, producing a dense embedding that provides the Transformer blocks with a smooth, scale-aware representation of the matrix. This design allows the network to exploit the temporal and structural dynamics captured by \({\boldsymbol{A}}(s)\) without introducing additional architectural complexity.
Timestep Embedding. We set the maximum period to \(T_{\max}=10000\) and use an embedding horizon of \(f=128\) sinusodial frequencies. Given the diagonal elements of \({\boldsymbol{A}}(s)\in\mathbb{R}^{N\times N}\), we construct a frequency-scaled argument matrix \({\boldsymbol{B}}\in\mathbb{R}^{N\times h}\) to encode each event position with multiple time scales. Concretely, for \(i\in\{1,\cdots,N\}\) and \(j\in\{1,\cdots,h\}\), we define: \[{\boldsymbol{B}}_{ij}={\boldsymbol{A}}(s)_{ii}\cdot T_{\max}^{-\frac{j-1}{f}}\]
so that \(j\) indexes a geometric progression of frequencies spanning from coarse to fine resolutions. We then form the final timestep embedding \(\texttt{emb}_s\in\mathbb{R}^{N\times 2f}\) by concatenating sine and cosine features: \[\texttt{emb}_s=[\cos({\boldsymbol{B}}),\sin({\boldsymbol{B}})].\]
Then we map this to the model hidden size \(D\) via an MLP: \[\texttt{emb}_s={\boldsymbol{W}}_2\boldsymbol{SiLU}({\boldsymbol{W}}_1\texttt{emb}_s+b_1)+b_2\in\mathbb{R}^{N\times D}\]
Following the choice of [19], we define the schedule function \({\boldsymbol{A}}\) as: \[{\boldsymbol{A}}(s)_{ij} = \begin{cases} 0 & \text{if } i\neq j , \\ \max(0,\min(\frac{2N-i-s(2N-1)}{N},1)) & \text{if } i=j, \end{cases}\]
which is a diagonal matrix that assigns a separate noise schedule to each event position. The schedule is constructed so that tokens corresponding to later events receive noise at earlier diffusion times than tokens for earlier events. Equivalently, later events are corrupted sooner, and the model is trained to denoise them earlier during the reverse process, forcing it to learn forecasting of the sequence tail under higher noise levels.
This sinusoidal construction provides a smooth multi-scale representation of the time conditioning signal \({\boldsymbol{A}}(s)\), allowing the Transformer blocks to access both low frequency (global) and high frequency (local) temporal variations through a fixed dimension embedding.
Positional Embedding. Before entering the DiT blocks, both current latent input and retrieved reference are expanded to hidden size \(D\) by channel repetition, for \({\mathbf{r}}\), we also broadcast the mask \({\mathbf{m}}^z\) across the channel dimension. Then \({\mathbf{z}}\) and \({\mathbf{r}}\) are augmented with fixed positional embeddings pos to get: \[{\mathbf{z}}=\text{Repeat}({\mathbf{z}})+\texttt{pos}_{{\mathbf{z}}}\in\mathbb{R}^{N\times
D},\,\,{\mathbf{r}}=\text{Repeat}({\mathbf{r}})+\texttt{pos}_{{\mathbf{r}}}\in\mathbb{R}^{N\times D},\]
where pos is defined as follows. For input sequence \({\mathbf{z}}\) with \(N\) being \(\texttt{num\_rows}\), we assign each position an
integer “row index" \(p\in\{0,1,\cdots,N-1\}\). For embedding dimension \(D\), define frequency coefficients: \[\gamma_i=10000^{-\frac{i}{D/2}},i=0,\cdots,\frac{D}{2}-1.\] The fixed positional embedding for row \(p\) is: \[\texttt{pos}(p)=[\sin(p\gamma),\cos(p\gamma)]\in\mathbb{R}^D.\]
Stacking all rows gives: \[\texttt{pos}\in\mathbb{R}^{N\times D}\]
Since our retrieved reference \({\mathbf{r}}\) and input \({\mathbf{z}}\) has the same dimension, so we apply the same pos embedding to both of them.
Mask Mechanism. When training or sampling sequences whose length \(n\) is smaller than the maximum token budget \(N\), we apply an attention mask within multi-head self-attention and the multi-head cross attention. so that padded positions do not participate in the computation. The similar strategy is designed in the retrieval process to mask the padded positions. This masking enforces that the Transformer operates only on valid event tokens, preserving the natural ordering of the sequence and preventing information leakage through padding. Combined with our asynchronous noise schedule, the masked attention mechanism aligns the model’s computation with the underlying temporal structure of the data. It also makes the framework compatible with variable-length inputs and flexible forecasting horizons, since the same architecture can condition on any observed prefix and generate different prediction window sizes without changing the model.
AdaLN Conditioning Implementation. We also include the details of our adaLN conditioning for retrieved asynchronous time series. Inside the DiT block, let \({\mathbf{z}}\in\mathbb{R}^{N\times D}\) be the token features entering the block, we have three conditioning streams:
time embedding emb\(_s\),
reference embedding \({\mathbf{r}}\),
fused conditioning Concat\((s,{\mathbf{r}})\).
We first apply separate LayerNorms and project the concatenation to the desired dimension: \[s=\boldsymbol{LN}_s(s),\,{\mathbf{r}}=\boldsymbol{LN}_{\mathbf{r}}({\mathbf{r}}),\,\text{cond}=\texttt{Proj}(\texttt{Concat}(s,{\mathbf{r}}))={\boldsymbol{W}}_{\text{cond}}[s;{\mathbf{r}}]+b_{\text{cond}}\in\mathbb{R}^{N\times D},\] where \([\,\cdot\,;\,\cdot\,]\) concatenates along channels and \({\boldsymbol{W}}_{\text{cond}}\in\mathbb{R}^{D\times 2D}\). Each stream produces a modulation vector via an MLP: \[{\boldsymbol{M}}_s=f_s(s)\in\mathbb{R}^{N\times 6D},\,{\boldsymbol{M}}_{\mathbf{r}}=f_{\mathbf{r}}({\mathbf{r}})\in\mathbb{R}^{N\times 6D},\,\,{\boldsymbol{M}}_\text{cond}=f_{\text{cond}}(\text{cond})\in\mathbb{R}^{N\times 6D}.\]
Then we scale each modulator by a learned scalar passed through a sigmoid: \[{\boldsymbol{M}}_s=\sigma(\lambda_s){\boldsymbol{M}}_s,\,{\boldsymbol{M}}_{\mathbf{r}}=\sigma(\lambda_{\mathbf{r}}){\boldsymbol{M}}_{\mathbf{r}},\,{\boldsymbol{M}}_{\text{cond}}=\sigma(\lambda_{\text{cond}}){\boldsymbol{M}}_{\text{cond}},\] with trainable \(\lambda_s,\lambda_{\mathbf{r}},\lambda_{\text{cond}}\in\mathbb{R}\).
The combined modulation is: \[{\boldsymbol{M}}={\boldsymbol{M}}_s+{\boldsymbol{M}}_{\mathbf{r}}+a\cdot{\boldsymbol{M}}_{\text{cond}}\in\mathbb{R}^{N\times 6D},\,\text{where }a\in\{0,1\}\text{ is a flag hyperparameter.}\]
Then \({\boldsymbol{M}}\) is split into six \(D\)-dimension tensors and enters the layers as vanilla adaLN in DiT.
Multi-reference Cross-attention by Concatenating \(k\) References into Length \(kN\) memory. The main query is \({\mathbf{z}}\in\mathbb{R}^{N\times D}\), and in the average pooling case \({\mathbf{r}}\in\mathbb{R}^{N\times D}\), we just need to adapt to \({\mathbf{r}}\in\mathbb{R}^{N\times k\times D}\) with per reference masks \({\mathbf{m}}\in\{0,1\}^{k\times N}\). Similarly to 4.3, the only changes are \({\boldsymbol{K}}\in\mathbb{R}^{N\times k\times D},{\boldsymbol{V}}\in\mathbb{R}^{N\times k\times D}\), then split them into \(\alpha\) heads with \({\boldsymbol{K}}\in\mathbb{R}^{\alpha\times (kN)\times \frac{D}{\alpha}},{\boldsymbol{V}}\in\mathbb{R}^{\alpha\times (kN)\times \frac{D}{\alpha}}\).
We report the dataset specifications, model hyperparameters, and computation costs in ¿tbl:tab:hyper?.
The long time reported for next-event prediction mainly comes from the autoregressive evaluation protocol over the full sequence length \(N\). Following prior work [9], [19], we compute this metric by iterating \(i=1,\cdots,N-1\) and predicting the next event \(\hat{\mathbf{x}}_{i+1}\) conditioned on the prefix \(\{{\mathbf{x}}_1,\cdots,{\mathbf{x}}_i\}\). Consequently, obtaining the full set of one-step-ahead predictions \(\hat{\mathbf{x}}_2,\cdots,\hat{\mathbf{x}}_{N}\) requires \(N-1\) separate forward passes, after which RMSE and accuracy are evaluated between \(\hat{\mathbf{x}}\) and \({\mathbf{x}}\). The detailed algorithm is provided in Appendix E of [19]. We use the implementation from [19], and this implementation is not optimized for inference; therefore, the reported time could be reduced with further implementation improvements. For a single sequence, ReDiTT takes 2.1s for one-step next-event prediction, compared to 2.0s for the unconditional baseline of [19]. This indicates that the retrieval-conditioned design adds only a reasonable runtime overhead. Because next-event prediction is relatively slow under the standard autoregressive evaluation protocol, this also motivates our long-horizon prediction setting: instead of repeatedly running one-step inference many times, we generate an entire future segment in a single run, enabling more efficient forecasting of long sequences.
We also include the iteration numbers we used in ablation tables in ¿tbl:tab:abl95iter? and ¿tbl:tab:abl95iter952?. Training iteration is indeed one of the main hyperparameters we tuned, and we will clarify the checkpoint-selection protocol in the revision. For each configuration, we trained the model for up to 500k iterations, saved checkpoints every 50k iterations, and evaluated each checkpoint using all three reported criteria: next-event time error, next-event type accuracy, and long-horizon OTD. We then reported the checkpoint that gave strong and balanced performance across these metrics, rather than selecting using only a single metric. Table 1 shows an example for ReDiTT on MultiThumos with \(k=7\). We select the 350k checkpoint because it provides the best overall tradeoff across the three objectives. Although the 300k checkpoint gives the lowest RMSE and later checkpoints slightly improve OTD at some horizons, the 350k checkpoint substantially improves type accuracy while keeping RMSE and OTD competitive.
| Checkpoint | RMSE \(\downarrow\) | Accuracy \(\uparrow\) | OTD@5 \(\downarrow\) | OTD@10 \(\downarrow\) | OTD@20 \(\downarrow\) | OTD@30 \(\downarrow\) |
|---|---|---|---|---|---|---|
| 50k | 5.4777 | 0.2548 | 8.3745 | 14.4535 | 21.3398 | 21.3327 |
| 100k | 5.0321 | 0.2355 | 8.3314 | 14.6365 | 21.0014 | 19.8137 |
| 150k | 4.8918 | 0.2382 | 8.6275 | 14.2741 | 20.6826 | 20.2415 |
| 200k | 4.8550 | 0.2371 | 8.4456 | 14.2449 | 20.9167 | 20.0028 |
| 250k | 4.8380 | 0.2355 | 8.5000 | 14.3425 | 20.6929 | 19.9026 |
| 300k | 4.7628 | 0.2376 | 8.3457 | 14.1603 | 20.9749 | 19.7040 |
| 350k | 4.7974 | 0.2500 | 8.5046 | 14.2710 | 20.8755 | 19.8915 |
| 400k | 4.8566 | 0.2425 | 8.4571 | 14.0243 | 20.5657 | 19.8726 |
| 450k | 4.8451 | 0.2446 | 8.5062 | 14.0904 | 20.5492 | 19.7714 |
width=
width=0.5
width=0.8
We further quantify the computational and memory overhead introduced by retrieval. ReDiTT stores a latent reference bank containing the VAE latent tokens of all training sequences, together with their padding masks. For a dataset with \(N_{\mathrm{train}}\) training sequences, padded sequence length \(N\), latent dimension \(d\), and 32-bit floating point storage, the dominant memory cost is \[N_{\mathrm{train}} \times N \times d \times 4\] bytes, plus a small mask cost of approximately \[N_{\mathrm{train}} \times N\] bytes. Thus, the reference-bank memory scales linearly in the number of training sequences, sequence length, and latent dimension.
To measure retrieval latency, we isolate the exact nearest-neighbor retrieval step from the rest of inference. For each dataset, we use a test batch of up to 128 sequences and compute the same masked token-wise cosine similarity used by ReDiTT between the query prefix latents and all training reference latents. The measured operation includes latent-token normalization, the batched similarity computation, masking over valid prefix/reference positions, averaging over valid positions, and the top-\(k\) selection with \(k=7\). It does not include the subsequent diffusion ODE solve, so it measures the overhead of the retrieval step itself. All measurements are performed on a single GPU using the full training reference bank.
| Dataset | \(N_{\mathrm{train}}\) | \(N\) | Memory (MB) | Retrieval latency (ms) |
|---|---|---|---|---|
| Amazon | 6454 | 94 | 222.75 | 3.33 |
| Retweet | 9000 | 97 | 320.54 | 5.14 |
| StackOverflow | 1401 | 101 | 51.95 | 0.67 |
| Taobao | 1300 | 64 | 30.55 | 0.36 |
| Taxi | 1400 | 38 | 19.53 | 0.28 |
| Breakfast | 179 | 131 | 8.61 | 0.22 |
| MultiTHUMOS | 209 | 100 | 7.67 | 0.14 |
Table 2 shows that the reference-bank memory is modest for the benchmark-scale datasets. The largest bank is Retweet, requiring approximately \(320.5\) MB, followed by Amazon at \(222.8\) MB. The remaining datasets require substantially less memory, ranging from \(7.7\) MB to \(52.0\) MB. Exact retrieval latency is also small relative to the diffusion-based forecasting cost: full-bank top-\(k\) retrieval takes about \(5.1\) ms on Retweet and \(3.3\) ms on Amazon, and less than \(1\) ms on the smaller datasets.
The main deployment consideration is scaling to much larger reference banks. Exact retrieval scales linearly with bank size because each query prefix is compared against all stored training sequences: \[O(B\,N_{\mathrm{train}}\,N\,d),\] where \(B\) is the query batch size. This exact search is practical for the datasets considered here, but for larger deployments the same framework can use approximate nearest-neighbor search, reference-bank subsampling, clustering/prototype selection, or cached retrieval results. Since retrieval is used only to provide conditioning references, these approximations are natural ways to trade a small amount of retrieval precision for lower memory and latency.
Retrieval-augmented forecasting relies on an analog-forecasting assumption: if two observed histories are similar, then their future continuations should also be similar. This assumption is natural for repeated or locally stationary event dynamics, but it may fail when the process is highly stochastic, non-stationary, or when similar prefixes can lead to different futures. We therefore add a diagnostic analysis to evaluate whether the references retrieved by ReDiTT are meaningful future templates.
We use the same retrieval similarity as in the main ReDiTT model. Consider a padded asynchronous sequence \({\mathbf{x}}=\{{\mathbf{x}}_1,\ldots,{\mathbf{x}}_N\}\) with padding mask \({\mathbf{m}}\in\{0,1\}^N\), where \(N\) is the maximum padded length. The pretrained VAE encoder \(E_\phi\) maps the sequence to per-event latent tokens \({\mathbf{z}}= E_\phi({\mathbf{x}}) \in \mathbb{R}^{N\times d}\). We precompute a latent reference bank \(\mathcal{B}=\{({\mathbf{z}}^r,{\mathbf{m}}^r): {\mathbf{z}}^r=E_\phi({\mathbf{x}}^r),\,{\mathbf{x}}^r\in\mathcal{X}_{\mathrm{train}}\}\). For a test sequence \({\mathbf{x}}^q\) with unpadded length \(L\), prediction horizon \(m\), and observed prefix length \(n=L-m\), we construct a prefix mask \({\mathbf{m}}^{q,n}\) such that \[{\mathbf{m}}^{q,n}_i = 1[i \leq n],\] with padded positions also masked out. Let \(\boldsymbol{\em q}=E_\phi({\mathbf{x}}^q)\) denote the latent tokens of the test sequence. For each candidate reference \(({\mathbf{r}},{\mathbf{m}}^r)\in\mathcal{B}\), we compute the same masked token-wise cosine similarity used in ReDiTT as in 3 : \[\texttt{sim}(\boldsymbol{\em q},{\mathbf{r}}) = \frac{1}{\sum_{i=1}^{N}\mathbf{1}[{\mathbf{m}}_i^{q,n}\wedge {\mathbf{m}}_i^r]} \sum_{i:{\mathbf{m}}_i^{q,n}\wedge {\mathbf{m}}_i^r} \left\langle \hat{\boldsymbol{\em q}}_i, \hat{{\mathbf{r}}}_i \right\rangle,\] where \(\hat{\boldsymbol{\em q}}_i=\boldsymbol{\em q}_i/\|\boldsymbol{\em q}_i\|_2\) and \(\hat{{\mathbf{r}}}_i={\mathbf{r}}_i/\|{\mathbf{r}}_i\|_2\) are \(\ell_2\)-normalized latent tokens. We then retrieve the top-\(k\) references according to this similarity: \[R(\boldsymbol{\em q}) = \operatorname{TopK}_{({\mathbf{r}},{\mathbf{m}}^r)\in\mathcal{B}} \texttt{sim}(\boldsymbol{\em q},{\mathbf{r}}).\] The only difference from standard ReDiTT inference is that, for this diagnostic, the query mask is explicitly set by the horizon \(m\) so that retrieval uses only the observed prefix \({\mathbf{x}}^q_{1:n}\), while the held-out suffix \({\mathbf{x}}^q_{n+1:n+m}\) is reserved for evaluating whether the retrieved reference has a similar future.
Let \({\mathbf{r}}^{\star}\) be the top-1 retrieved reference according to this similarity: \[{\mathbf{r}}^\star = \arg\max_{({\mathbf{r}},{\mathbf{m}}^r)\in\mathcal{B}} \texttt{sim}(\boldsymbol{\em q},{\mathbf{r}}).\]
We then evaluate whether the future of the retrieved reference is close to the true future of the query. For test example \(j\), with prefix length \(n_j=L_j-m\), we compute \[D_{\mathrm{ret}}^{(j)}(m) = \mathrm{OTD} \left( {\mathbf{x}}^{q,j}_{n_j+1:n_j+m}, {\mathbf{x}}^{r^\star,j}_{n_j+1:n_j+m} \right).\] As a random baseline, we sample a training reference \({\mathbf{x}}^{r_{\mathrm{rand}},j}\) uniformly at random and compute \[D_{\mathrm{rand}}^{(j)}(m) = \mathrm{OTD} \left( {\mathbf{x}}^{q,j}_{n_j+1:n_j+m}, {\mathbf{x}}^{r_{\mathrm{rand}},j}_{n_j+1:n_j+m} \right).\] Finally, we report the empirical retrieved/random future distance ratio for each dataset: \[\rho(m) = \frac{ \frac{1}{N_m}\sum_{j=1}^{N_m}D_{\mathrm{ret}}^{(j)}(m) }{ \frac{1}{N_m}\sum_{j=1}^{N_m}D_{\mathrm{rand}}^{(j)}(m) },\] where \[N_m = \left|\{j:L_j>m\}\right|\] is the number of test sequences with at least \(m\) real future events. A value \(\rho(m)<1\) indicates that retrieved futures are closer to the true futures than random training futures. We also report the correlation between the prefix similarity \(\texttt{sim}(\boldsymbol{\em q},{\mathbf{r}}^\star)\) and \(D_{\mathrm{ret}}(m)\); a negative correlation indicates that more similar prefixes tend to have closer futures.
| Dataset | \(\rho(5)\) | \(\rho(10)\) | \(\rho(20)\) | \(\rho(30)\) | Corr. at \(m=30\) |
|---|---|---|---|---|---|
| Amazon | 0.968 | 0.976 | 0.947 | 0.934 | -0.260 |
| Retweet | 1.012 | 1.008 | 0.971 | 0.937 | -0.304 |
| StackOverflow | 0.875 | 0.809 | 0.751 | 0.763 | -0.408 |
| Taobao | 0.727 | 0.770 | 0.772 | 0.773 | -0.463 |
| Taxi | 0.890 | 0.890 | 0.844 | 0.917 | -0.103 |
| Breakfast | 0.937 | 0.914 | 0.888 | 0.829 | -0.459 |
| MultiTHUMOS | 0.968 | 1.010 | 1.016 | 1.002 | 0.003 |
Table 3 shows that the analog-forecasting assumption holds strongly on some datasets, but not uniformly across all domains. Taobao, StackOverflow, and Breakfast provide the clearest evidence. On Taobao, retrieved futures are much closer than random futures across all horizons, with \(\rho(m)\approx 0.73\)–\(0.77\), and prefix similarity is strongly negatively correlated with future OTD. StackOverflow shows a similar pattern, with \(\rho(m)\) reaching approximately \(0.75\) at longer horizons, indicating that retrieved references provide useful future templates. Breakfast also supports the assumption, with retrieved futures improving over random futures by roughly \(6\%\)–\(17\%\) depending on the horizon.
Amazon and Retweet show weaker but still partially positive evidence. On Amazon, retrieved futures are only modestly closer than random futures, with \(\rho(m)\approx 0.93\)–\(0.98\), suggesting that retrieval provides useful but relatively weak future information. Retweet is mixed at short horizons, where retrieved futures are comparable to random futures, but becomes more favorable at longer horizons. Taxi exhibits a different behavior: retrieved futures are better than random futures, with \(\rho(m)\approx 0.84\)–\(0.92\), but the correlation between prefix similarity and future OTD is close to zero. This suggests that retrieval is useful on average, but the similarity score does not reliably rank which references have the closest futures. MultiTHUMOS shows the weakest analog structure under this diagnostic, with \(\rho(m)\) close to or slightly above \(1\) for most horizons, indicating that nearest latent neighbors are often no better future templates than random references.
These results clarify when retrieval is expected to help in asynchronous event forecasting. Retrieval is most reliable when event dynamics contain repeated, locally predictable patterns, so that similar observed prefixes provide informative future templates. This explains the strong diagnostic results on Taobao, StackOverflow, and Breakfast. In contrast, when the prefix-future relationship is noisy, ambiguous, or weakly captured by the retrieval metric, retrieved references may be imperfect or only weakly informative, as observed for MultiTHUMOS and partly for Taxi.
MultiTHUMOS is therefore best viewed as a stress test for the analog-forecasting assumption. Its weak retrieved/random OTD ratios suggest that visually or temporally similar prefixes do not always imply similar future action continuations, possibly due to greater variability in activity progression or multiple plausible futures. This highlights why ReDiTT should not be interpreted as simply copying retrieved futures. Instead, retrieved references are used as conditioning signals inside a diffusion model. When references are strong analogs, gated cross-attention allows the model to exploit them directly; when references are noisy or imperfect, diffusion-based refinement can adapt the conditioning signal to the query. Thus, retrieval provides useful non-parametric evidence when similar histories imply similar futures, while the diffusion model improves robustness when this relationship is weaker or dataset-dependent.
A potential concern for retrieval-augmented forecasting is whether the retrieval query uses information that is unavailable at inference time. In the forecasting setting, the model observes only a prefix \(x_{1:n}\) and predicts future events \(x_{n+1:N}\). Therefore, a retrieved training sequence may contain its full continuation, but it should be selected using only the observed prefix. To study the practical impact of this issue without retraining all models, we perform a retrieval-only diagnostic comparing two retrieval queries: Prefix retrieval: neighbors are selected using only the observed prefix \(x_{1:n}\); future target tokens are masked before computing cosine similarity. Full retrieval: neighbors are selected using the full sequence \(x_{1:N}\), which acts as an oracle retrieval query.
After retrieval, we use the retrieved training sequence as a simple retrieval-only forecast and evaluate the predicted suffix with OTD. This diagnostic isolates the retrieval step: if full-sequence retrieval is much better than prefix retrieval, then future target tokens materially change neighbor selection. Conversely, if prefix retrieval is close to full retrieval, then the observed prefix already contains enough information to retrieve useful analogs, which suggests that the train-test retrieval mismatch has limited practical effect.
| Dataset | Horizon 5 | Horizon 10 | Horizon 20 | Horizon 30 | Avg. |
|---|---|---|---|---|---|
| Amazon | 0.957 | 0.974 | 1.028 | 1.053 | 1.003 |
| Retweet | 0.990 | 1.002 | 1.036 | 1.052 | 1.020 |
| StackOverflow | 0.949 | 0.993 | 0.999 | 0.994 | 0.984 |
| Taobao | 0.918 | 0.893 | 0.855 | 0.809 | 0.869 |
| Taxi | 0.754 | 0.747 | 0.758 | 0.764 | 0.756 |
| Breakfast | 0.988 | 0.984 | 1.013 | 1.002 | 0.997 |
| MultiTHUMOS | 0.981 | 0.969 | 1.003 | 0.995 | 0.987 |
Table 4 reports the OTD ratio \(\mathrm{OTD}_{\mathrm{full}} / \mathrm{OTD}_{\mathrm{prefix}}\) for prediction horizons 5, 10, 20, and 30. Lower OTD is better. Thus, a ratio below 1 indicates that the oracle full-sequence query improves retrieval, while a ratio near 1 indicates that prefix retrieval performs similarly to full retrieval. The results show that for five of the seven datasets, prefix retrieval is very close to full-sequence retrieval. On Amazon, Breakfast, MultiTHUMOS, Retweet, and StackOverflow, the average ratio lies between 0.984 and 1.020. In these datasets, using the future suffix in the retrieval query does not substantially improve the retrieval-only forecast. This supports the claim that the observed prefix is already sufficient to identify useful neighbors.
Taobao and Taxi are different: full-sequence retrieval gives noticeably lower OTD, especially for longer horizons. This indicates that the future suffix contains additional information that can improve oracle neighbor selection. Therefore, full-sequence retrieval should not be used as the forecasting-time query. This also motivates our evaluation protocol, where future target tokens are masked before retrieval.
Overall, this analysis helps contextualize the train-test mismatch concern. Since prefix retrieval is competitive with full-sequence retrieval on most datasets, ReDiTT can still perform well when retrieval is restricted to the information available at forecast time. This suggests that, in many event sequence datasets, the observed prefix already provides enough signal to identify useful training analogs. The stronger gap on Taobao and Taxi also indicates that the retrieval-visible token set should be stated explicitly, and that prefix-masked retrieval during training is an important direction for future work to further align training with the forecasting protocol. With further exploration on prefix based training, the Taobao and Taxi results might be also improved.
To disentangle the contribution of retrieval from the diffusion-based conditioning architecture, we add a retrieval-only baseline alongside the diffusion-only model and full ReDiTT. The retrieval-only baseline uses the same reference bank and top-\(k\) retrieval procedure as ReDiTT, but removes the diffusion model and gated cross-attention layers: after retrieving the nearest reference sequences, it directly aggregates their continuations in latent space and decodes the result with the pretrained VAE. This provides a non-parametric analog forecaster that measures how much predictive signal is already contained in the retrieved futures.
As empirically shown in ¿tbl:tab:retrieval-only?, diffusion-only and retrieval-only provide complementary but comparable baselines: neither dominates the other consistently, with diffusion-only performing better on some datasets and retrieval-only performing better on others. In contrast, full ReDiTT achieves the best performance on nearly all metrics across the seven datasets, with two exceptions: retrieval-only obtains the best next-event prediction accuracy on StackOverflow, and diffusion-only achieves the best long-horizon result on Taxi. These results suggest that retrieval alone is not uniformly sufficient, and diffusion alone also does not fully explain the gains. Rather, ReDiTT’s main advantage comes from combining the two: retrieved references provide useful analog futures, while the diffusion model with gated cross-attention can adapt those references to the query instead of directly copying or averaging them. The exceptions further indicate that the relative value of retrieval and generative refinement depends on the dataset, but the combined ReDiTT model is the most robust overall.
width=
We further evaluate a raw-space retrieval ablation to test whether ReDiTT’s performance depends specifically on retrieving neighbors in the VAE latent space. In this variant, we keep the trained ReDiTT model, VAE, reference bank size, and diffusion-based conditioning architecture fixed, and change only the neighbor-selection rule. Instead of computing similarity between latent event representations, we retrieve references using normalized raw event tuples, combining log-scaled inter-event time distance with event-type matching over the observed prefix. The retrieved references are still provided to the diffusion model as latent sequences, so the ablation isolates the effect of the retrieval space while preserving the rest of the pipeline. As listed in ¿tbl:tab:raw-space?, latent-space retrieval performs better on most datasets and metrics, but the gap is generally small: raw-space retrieval remains competitive and often produces similar results. This suggests that retrieval-augmented conditioning is itself a robust source of improvement, while latent-space retrieval provides a modest but consistent advantage and is better aligned with ReDiTT’s latent diffusion target. It also avoids the need to hand-design raw-space similarities over mixed continuous-discrete, padded asynchronous sequences.
width=
Corresponding author. Work done during an internship at RBC Borealis.↩︎