Diffusion Transformers for Imputation: Statistical Efficiency and Uncertainty Quantification
Authors: Zeqi Ye
Minshuo Chen
Venue: NeurIPS 2025 poster ()
Keywords: Time Series Imputation
Uncertainty Quantification
Statistical Sample Complexity
Diffusion Transformer
Primary area: deep_learning

Abstract: Imputation methods play a critical role in enhancing the quality of practical time-series data, which often suffer from pervasive missing values. Recently, diffusion-based generative imputation methods have demonstrated remarkable success compared to autoregressive and conventional statistical approaches. Despite their empirical success, the theoretical understanding of how well diffusion-based models capture complex spatial and temporal dependencies between the missing values and observed ones remains limited. Our work addresses this gap by investigating the statistical efficiency of conditional diffusion transformers for imputation and quantifying the uncertainty in missing values. Specifically, we derive statistical sample complexity bounds based on a novel approximation theory for conditional score functions using transformers, and, through this, construct tight confidence regions for missing values. Our findings also reveal that the efficiency and accuracy of imputation are significantly influenced by the missing patterns. Furthermore, we validate these theoretical insights through simulation and propose a mixed-masking training strategy to enhance the imputation performance.

Diffusion Transformers for Imputation: Statistical Efficiency and Uncertainty Quantification
Zeqi Ye
Minshuo Chen
Northwestern University Evanston, IL, USA begiyy22290u..ooobweseeer
Northwestern University Evanston, IL, USA
minshuo. htt.aoo.cccennooteeseeere..eee
Abstract
Imputation methods play a critical role in enhancing the quality of practical time- series data, which often suffer from pervasive missing values. Recently, diffusion- based generative imputation methods have demonstrated remarkable success com- pared to autoregressive and conventional statistical approaches. Despite their empirical success, the theoretical understanding of how well diffusion-based mod- els capture complex spatial and temporal dependencies between missing values and observed ones remains limited. Our work addresses this gap by investigating the statistical efficiency of conditional diffusion transformers for imputation and quantifying the uncertainty in missing values. Specifically, derive statistical sample complexity bounds based on a novel approximation theory for conditional score functions using transformers, and, through this, construct tight confidence regions for missing values. Our findings also reveal that the efficiency and accuracy of imputation are significantly influenced by the missing patterns. Furthermore, we validate these theoretical insights through simulation and propose a mixed-masking training strategy to enhance the imputation performance.
Introduction
Sequential data are ubiquitous real-world applications such as finance [John et al., 2019, Chen al., 2016]. healthcare [Tonekaboni al., 2021, Kazijevs and Samad, 2023], transportation [Li et al., 2020, Tedjopurnomo et al., 2020], and meteorology [Yozgatligil et al., 2013]. However, these datasets often suffer from missing values due to factors such sensor malfunctions, data transmission errors, and human oversight [Greco et al., 2012, Yi et al., 2016). Missing data can significantly degrade the performance of downstream tasks [Ribeiro and Castro, 2022, Alwateer et al., 2024], making accurate and robust imputation a critical challenge.
One of the carliest imputation methods dates back to Allan and Wishart [1930]. which provided formulas for estimating single missing observations. Over the past century, this foundational idea of imputation has been extended to broader application domains. Statistical imputation methods have gained sustained attention due their computational efficiency and case of implementation. These approaches range from simple techniques, such as imputation using the mean or median of observations, interpolation-base based methods [Tukey, 1952], and more sophisticated model- based techniques, including Kalman filters and autoregressive models [Gómez and Maravall, 1994, Shumway et al., 2000J. However, these methods often rely strong assumptions such linearity and stationarity, which may not hold complex real-world scenarios, thereby limiting their applicability and accuracy [Fuller, 2009].
To address the limitations of statistical methods, recent research has increasingly turned machine learning approaches for imputation. These methods are capable capturing complex spatio-temporal patterns and nonlinear dependencies without requiring strict assumptions [Fang and Wang, 2020]. Typical examples include training neural networks such as recurrent neural networks and transformer
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

architectures for inferring missing values [Wang et al., 2024]. In parallel, generative models such as Variational AutoEncoders (VAEs) and Generative Adversarial Networks (GANs) have shown promise by introducing uncertainty-awart imputations [Fortuin et al., 2020, Miao et al., 2021]. However, these generative models often spell limitations in expressiveness or training stability. More recently, diffusion-basee generative models have emerged as a powerful alternative, offering robust imputations and strong empirical performance across diverse and high-dimensional time series datasets [Tashiro et al., 2021, Zhou et al., 2024].
Despite their widespread empirical success, diffusion-based imputation methods exhibit two key challenges. First, their performance is highly sensitive to dataset characteristics, often displaying sub- stantial variability across benchmarks [Zhang et al., 2024, Zheng and Charoenphakdee, 2022, Tashiro et al., 2021]. Second, they are significantly affected by missing patterns, leading to inconsistencies in imputation quality [Zhang et al., 2024, Ouyang et al., 2023, Zhou et al., 2024). These observations motivate the following fundamental questions:
How well can diffusion models caplure the underlying conditional distribution of missing values?
How does the missing pattern affect the imputation performance?
In this paper, we answer the two questions from a statistical learning perspective. Our analysis centers on Diffusion Transformers (DiT, Peebles and Xie [2022]) applied to imputation tasks with Gaussian process (GP) data. Despite their conceptual simplicity, GPs exhibit rich spatio-tempont dependencies and long- -orizon dependencies that pose challenges for modeling and imputation. On the other hand, GPs are powerful statistical tools widely used in regression, classification, and forecasting tasks [Seeger, 2004, Banerjee al., 2013, Borovitskiy al., 2021].
We establish sample complexity bounds for DiT's in learning the underlying conditional distribution of missing values given observed ones. The obtained bounds demonstrate the role of missing patterns in imputation performance, highlighting how the condition number of the covariance matrix for the missing values and distribution shifts contribute to variability in accuracy. Furthermore, we derive confidence intervals for imputed values and show the coverage probability of them converging the desired level. We summarize our contributions as follows.
Statistical Efficiency. We show that DìTs capture the conditional distribution of missing values effectively. The sample complexity in Theorem 2 scales at a rate (H /Vn). where n denotes 1/2 -1/2 convergence rate with a mild polynomial dependence on the sequence length H. In addition, is the condition number induced by the missing patterns. To establish Theorem 2, we develop a novel score representation theory (Theorem 1) for DiT', where we utilize an algorithm unrolling technique.
the training sample size. We obtain a n
Uncertainty Quantification. Leveraging the generative power of trained DiTs, construct confidence regions (intervals) from massive generated missing values. This approach possesses its natural appeal and enjoys strong coverage guarantees (Corollary 1). We show that the coverage probability converges the desired level at a ((n-1/2) rate. Meanwhile, the missing patterns influences the convergence.
Mixed-Masking Training Strategy. Motivated by our theoretical results, we propose a training strategy blending different masking schemes to cover diverse missing patterns. The performance of our method on synthetic datasets validates our findings and outperforms benchmark methods.
Notations We use bold lowercase letters to denote vectors and bold uppercase letters to denote matrices. For a vector v. /|vl|2 denotes its Euclidean norm. For matrix A, ||A ll2 and ||F denote its spectral norm and Frobenius norm, respectively, and ||A = max,, (Ayll- When matrix A is positive definite, condition number is k(A)= = Amax (A) Amin (A). We denote f g if there exists a constant C' > 0 such that f < Cg. Notation (-)) suppresses constants, while ((-) further hides logarithmic factors.
denote )max(A) (A) and Auni (A) as its largest and smallest eigenvalues; its
Due to space limit, the related work section is deferred to Appendix A.
2 Imputation in Gaussian Processes via Conditional Diffusion Models
In this section, we formalize the imputation task as a conditional distribution estimation problem. When the data are sampled from Gaussian process, we identify rich structures in the conditional distribution. We then utilize a DiT to learn the distribution of missing values. Lastly, we summarize diffusion- based imputation method in Algorithm 1.

---

2.1 Imputation for Gaussian Process Data
Imputation refers to the task of inferring missing values given the observed ones. Denote by I = {l..... H} the set of all time indices. For multivariate sequence X = [x....,,X. X,rl R of length H, we consider a block-missing setting, where certain time frames are entirely unobserved. The subset of observed indices is denoted by Inbs {i1,.... ... ([Zaii) J. where | denotes the cardinality. Correspondingly, Timiss I Tabs denotes the time indices of missing frames. To avoid degenerate cases, we assume O < [T < H. In the sequel, we focus on the Missing Completely at Random case [Little, 1988], where each index in Tmiss is independently sampled from some underlying distribution.
We represent the vectorized observed partial sequence as Xobs [3 X [Zoball LJ E Rd and
the vectorized missing part as Xuiss [x]...... X Fmmnia IJ E [dd/Tmmeel We estimate the missing
values by learning the conditional distribution P(x P(xuisss Xobs). Notably, learning the conditional distribution goes beyond point estimates of the missing values, but provides easy access to confidence regions. We slightly abuse the notation by using x to simultaneously refer random vectors.
Throughout our theoretical analysis, we focus d-dimensional Gaussian process data. To uniquely distinguish a Gaussian process, it suffices to specify its mean and covariance functions. In particular, we denote the mean as μı; E[x;] and we parameterize the covariance matrix by Cov Cov(x, x,,
((,, j) A, where A Var xx, E Rdxd for any h and 7 is a kernel function. It is worth mentioning that A captures the spatial dependencies and function 7 represents temporal correlation. The kernel function y dictates the strength and decay of the temporal dependencies among different data frames. The joint distribution of sequence vec(X)
[x] . xx1 is Gaussian A M(,, P & A), where
F1:
71,1)A
71, A)
1
: H
and
(H,
(H,
Here = j). and @ is the matrix Kronecker product. We impose the following assumption for characterizing the temporal dependencies.
=
Assumption 1. There exists di-diiensional embedding fc, E ER such that [|e:ll2 for constant r. Moreover, for any i, it holds that llee - ejll2 = -(ji- jili-j], and I (új - -iil # lig - jzl.
I(li: jil) + f((62 f2l) Kernel function {i, only depends on lle ezllg- Furthermore, we assume T and are positive definite.
Assumption 1 ensures that the pairwise distances in the embedding uniquely identifies positional gaps. We do not specify a particular form of the kernel function, which encodes many commonly ones such as Gaussian Radial Basis Function (RBF), Ornstein-Uhlenbeck kernels, and Matérn kernels [Rasmussen and Williams, 2006]. As concrete example, sinusoidal embedding is widely used in transformer networks [Vaswani et al., 2017]. Consider a two-dimensional embedding defined as where r > 0 is a fixed radius and C > 0 is a scaling constant. The Euclidean distance between any two embedding is lle e, ||2 2r sin (T - /)/C),, I, which is strictly positive for i * j. and approximately linear in i when C' is sufficiently large.
e, Ir sin ( r cos(2Ti//))
Under the Gaussian process setting, the conditional distribution of x Xohs is still Gaussian [Bishop and Nasrabadi, 2006]. The conditional mean and covariance are given by
Meond (Xons) Xobs) = Mmiss + DJ
oor obs (Nobs Mobs), Σxond Emiss
Σ D-11 Dcor:
where we denote Hobs E[xobes] (the same holds for pmiss). Ecor Cov [Nobs, Xmiss] and Eobs (resp. Emiss) as the covariance of Xobe (resp. Xmiss). See Figure for graphical demonstration. We check that Eobs = Tobe 8 with Tobs E RR X Zaball capturing correlation among index set Tobs-
2.2 Training Diffusion Transformers for Imputation
We estimate the conditional distribution P(xmiss Xohs) using diffusion transformers. A diffusion model consists of two coupled processes-a forward and backward process. We adopt a continuous- time description. In the forward process, we gradually corrupt data by
(1) and w, is Wiener process. The forward process terminates at sufficiently large time T and we denote the distribution of x, as Pr(-Ixohs) with density pe( p+(-\xobx). Note that we only corrupt the missing values by Gaussian noise, but keep the observed partial sequence Xoba unchanged.
dx; 2 xtdt + dw, with Xg P(- Xobs):
3

---

Corresponding to the forward process, the backward process simulates the reverse evolution of the forward process. As a result, it generates new samples by progressively removing noise:
dv. 2 Vv. logpr-:(V. Xohs) d/ dw.
with Vn Pr( Xohs),
where W, is another Wiener process and Vv, log Pv Xobs) is the conditional score function. In the remaining of the paper, we drop the subscript v, in the score function for simplicity. Unfortunately, V log pT- Xobs is typically unknown and must be estimated using neural network. We denote the estimated score function by S(v Xobs: t). Consequently, the sample generation process follows alternative backward SDE:
dv. V. + S(V,, Xohs: 1) dt + dw, with Vo ~ (O Ldj 4/mmn)))
Here, we also replace the unknown Pr by a standard Gaussian distribution.
When training the score estimator S, we assume access to fully observed sequences. To simulate partially observed sequence, sample masking sequence (Tı,..., TH) {0, 1) where 0 denotes missing the observation and keeping the observation. 'hen Xobs is extracted according to the masking sequence. later context, we will investigate how choose masking strategies. We summarize the diffusion-based method for sequence imputation in Algorithm 1.
Algorithm Ceetuusion-Basee Sequence Imputation
1: Module I: Training
2: Input: Fully observed sequences D = {x.)__ a masking strategy.
3: Simulate (x X i=1 pairs via the masking strategy, and train conditional diffusion model.
(i)
)a
4: Output: A well-trained conditional diffusion model.
5: Module II: Imputation
6: Input: Conditional diffusion model from Module I, new partial sequence X repetition time Z, and confidence level a.
7: Conditioned on Xibs independently generate B missing sequences x(1) for 1,..., Z.
8: * Point estimate: Mean X E s(()
miss (or median of the generated sequences).
9: * Confidence region: ČR {Xmiss [|Xmiss- - & ll2 S Dr
where Di-a is the upper quantile of || X ll for 2 = 1,..., Z.
10: Return: X² miss and 1
For the rest of the paper, we parameterize the conditional score function using transformer network. transformer [Vaswani et al., 2017], comprises a series of blocks and each block encompasses a multi-head attention layer and feedforward layer. Let Y = , Y ] stacking matrix of H patches. In a transformer block, the multi-head attention layer computes
R DDD be the (column)
AYtn(Y) = Y m= V or ((Q (TTT))) K"Y),
(4)
where V Q" K are weight matrices corresponding sizes in the m-th attention head, and or is an activation function. The attention layer is followed by feedforward layer, which computes
FFN(Y) Y + W1 ReLU(W2 + b²1") bb1
Here, WT, W2 are weight matrices, by and b2 are offset vectors, denotes a vector of ones, and the ReLU activation function is applied entry-wise. This feedforward layer performs linear transformation to the output of the attention module with more flexibility. For our study, the raw input to transformer's is H patches of d-dimensional vectors and time in the backward process. We refer to T(D, L, M, B, R) as a transformer architecture defined by
T(D, I, M, B. R = { f= fout o (FFN, o Attng) o ... (FFN, o Attni) o sin-
Attn; uses entrywise ReLU activation for = 1..... number of heads in each Atta is bounded by M, the Frobenius norm each weight matrix is bounded by B, the output range || is bounded by R}.
(5)

---

Conditional Score Approximation via Algorithm Unrolling
Suggested by the sample generation process (3), the key is to learn the conditional score function.
This section devotes to establishing a novel score approximation theory of transformers based on algorithm unrolling.
Since Xmiss Xobs
Xobs is Gaussian, the forward process (1) yields the following closed-form score function:
V log Pe (Vx|Xobs) = - (a, Zcond + o (v: Oit//oond (Xobs));
where 0; = = and 0; = v1- e-c. The matrix inverse poses a challenge in representing the score by transformer, as it may deteriorate structures in Lcond. Therefore, reformulate the conditional score function as the optimal solution of a quadratic optimization problem:
1
V logggu (v Xoba) arg min (s) == 2 (a. + oi + (v, ac (Xobs)) (7)
It suffices obtain approximate optimal solution of (7) using gradient descent algorithm. At the k-th iteration, with a step size T/2 we have
(k1-1) =g()) - Tl (of1+ at ms + at D Dabe - s + (Ve Cts Moond (Xobs)) ))
(8)
(L.(s(()
for k- = 0. O..... 1. Unfortunately, we encounter another matrix inverse in Enth Ecors(t), Analo- gous (7), consider auxiliary quadratic optimization problem:
obs Ycors (8) = arg min C(k) (u) := 2 Eots Eoor (A)
(9)
Via gradient descent algorithm with step size 0, the update reads
(Maxx+1) u (Kaux) - OV LVV (k) aux (u)
(Eohau- Ecors((t)
(10)
where iteration index kaux O,..., Kaux
We substitute the last iterate u (oo)) into the right-hand side of (8) obtain (s()s as an approximation (s(A)) We summarize the nested gradient descent algorithm for calculating the conditional score function in Algorithm 2.
Algorithm 2 Nested Gradient Descent for Representing Score Function
1: Input: Observation Xobs current state Vt. time f, step sizes T): 6, iteration counts Kaux, K. (Major) Gradient Descent:
2: Initialize ( = =0
3: for k 0, 1, 1,... K 1 do
Auxiliary Gradient Descent:
4: Initialize (0) 0.
5: for Kaux = 0, 0,1,... Kaux 1 do
6: u (Maux u (KAux) OV Laux (kan))
(k))
7:
Calculate (s)) using u (Kau)).
8:
(R++1) (()) (s(H))
9:
Return: (K))
With sufficiently large Kaux and K, the representation error of Algorithm 2 can be well-controlled. Lemma (Representation error of Algorithm 2). Suppose Assumption I holds. For an arbitrarily fixed time E (0, T), given an error tolerance E (0, 1), choose K, Kaux as
K = o (Erond) log
dx Econa) )il (Eooa))) ))
Kaux = =0 -(k(Eott)) log
HHds( (Eois))))
Tzc
Then, given 8 > O, for any Xoba and Va in a compact region Cs, there exist step sizes q and 0 such that running Algorithm 2 gives rise to
[| - log Xoba)ll2 o0
5

---

Detailed proof of Lemma is provided in Appendix B. The compact region C. truncates the norm of Xobs and V,, which is plausible due the Gaussian tail; see a precise definition of C2 in Appendix Equation (13). Lemma suggests that the computational complexity of Algorithm 2 for approximating the score function governed by the condition numbers of Econd and Eobs. A large condition number on cond implies that the variability of missing values among different directions changes significantly. Equivalently, with a large condition number, given Xobs. the missing values exhibit strong anistropic uncertainty that complicates the imputation.
Representing the conditional score function by a nested gradient descent algorithm enables an effective transformer network approximation. We show that transformers can realize each gradient descent iteration using constant number of attention blocks. We provide the following score approximation theory using transformers.
Theorem 1. Suppose Assumption I holds. Given an early stopping time to [0, and an error level e E (0, 1), for any Xobs, E Cā, there exists transformer architecture T(D, L, M, B, R) such that, with proper weight parameters, it yields an approximation satisfying
[[[(v, Xobs, 1) V logge(Vv Xobs) ll S o, for all L E [lo T).
The configuration of the transformer architecture satisfies
L 0=0 |mmax Evond )max (ooa) log" Hddimax (Eooud Kmax [ooss)
D = O(d+ de),
( (V (r² + Kmax Eoba) (oo1) R = 0 Hddkkmax Eobe)).
M = 4H, B B-O
where we define kima SuPI.rr x(-).
The proof is provided Appendix C. Figure 1 depicts the transformer architecture in our constructive proof, which unrolls Algorithm 2 efficiently. To obtain the approximation error bound, develop careful analysis of the error propagation in the auxiliary gradient descent for calculating C. Theorem also reinforces the insights from Lemma 1, where we observe that the size of the transformer network scales with the worst-case condition number. We will further discuss the relation between missing patterns and the condition number in Theorem 2.
g(K) as arg min, (s) V Pilyp Xata)
Attention Heads:
Attention Heads: Correlation Part
Attention Heads: Observed Purt
Missing Part
K
Emise
E
Approximating major GD iteration:
Compute Ma Dagg (axa
Transformer
En
K K
Approximating auxiliary GD iteration Compute E.s (hed (
Capture Different Parts of Sputial and Temporal Correlations Sepurately
Transtomne
Perform Matrix- Vector Multiplications Approximate GD
V (u) = Eows Eeu (+)))) x (af of E E + at E. I (Kuz)) + (v: asstoond (Xatx))
Samples & Current States Time Emboddings
Figure I: Constructed transformer architecture: Within each transformer block, attention heads focus on capturing information of different covariance components (Zobs. Ecors Emiss) separately, and approximate corresponding matrix-vect multiplications. A total of K block groups perform major GD steps, with Kaux inner blocks in each group dedicated solving the auxiliary problem.
4 Capturing Conditional Distribution and Uncertainty Quantification
Given properly chosen transformer architecture, establish guarantees for learning the conditional distribution of missing values and uncertainty quantification. We consider an estimated score network I obtained by minimizing the following empirical score matching loss (a detailed derivation is deferred to Appendix E):
I E arg minseT Ż(s) En (( X 5),
(6)
(1I)

---

where
(T E
(6) {(x m[xmss>> Xbsi is
(0) [|s(ve, xobs xobe: t) -vv - (vv-oovo)oofl dt.
v,
(12)
to
mina
Substituting the learned score S into the backward SDE (3) yields generated distribution Vta ~ Pio Xobs). We introduce an early-stoppine time to to stabilize the training and sample generation [Song et al., 2020]. We now present a convergence guarantee of P.o to the true conditional distribution. Theorem Referring to the training procedure in Algorithm 1, by choosing the transformer architecture as in Theorem with e c= n 2, terminal time T (logn)) n), and earty-stopping time to OO (oond)) 3). it holds that Exxra
(m) dditt
(D(())
(VV((()xxbb)) (- Öly H& (Econd) {> Vn).
The proof of Theorem 2 is provided in Appendix D. This result establishes that DiT can efficiently learn the true conditional distribution of missing values. The sample complexity mildly depends on the sequence length. More importantly, the bound highlights that the estimation error depends on the condition numbers of Econd and ooba. reflecting the discussion after Lemma 1.
We provide example to demonstrate that different missing patterns can lead to distinet condition numbers. Consider data of length H = 96 with time correlation modeled by a Laplace kernel ((.5) exp(-Iles eyll2/ /128), and missing length |mmiss| = 16. Clustered missingness- consecutive missing entries at the tail pields a large condition number x( ((Coona) = 415.40, making the task challenging. In contrast, dispersed missing patterns, 16 randomly placed missing entries, result in much smaller x( (Zoond) = 3.00, making estimation easier. We provide numerical results on this example in Section 5.
Confidence Region Construction Given the learned conditional distribution Pt, and new ob- served sequence Xass we deploy the model to generate samples and form point estimates and confidence regions as in Algorithm 1. Since Xebs may not be seen in the training samples, we encounter a distribution shift, meaning that we need to transfer the knowledge in the learned model to the new testing instance. The subtlety here is how to quantify the knowledge transfer rate. Our proposal the following class-dependent distribution shift coefficient.
Definition 1.
The distribution shift between two probability distributions P1 and P2 with respect to BymP, [g(y)]
function class G is defined as DS(P1, P2; = suP by [8y]
In our analysis, we specialize G to a function class induced by the transformer network:
g [g(y) Ex..... [[(xiine y;s] s T(D, M, B, R)}.
Since G might be insensitive to certain distinctions, it introduces some smoothing effect capture the difference between P1 and Pz. We consider P1 and P2 as the marginal training distribution Xohs and the point mass of the testing distribution 1 x }. denoted as Px and Px;,, respectively. The following corollary provides a guarantee for the coverage probability of the constructed CR.
Corollary 1. Under the setting of T'hoorem 2, given X Algorithm yields CRI-a satisfying
Epin)) [(xxinss E cRRiaa) 2 (1 a) Edist
(n)
DS(Px:.. Px... : n (Xxns))
where (xxabs) is independent of n and proportional to ||x and x( ((Eoou)).
Detailed proof is provided in Appendix D. Corollary 1 says that the coverage probability of the constructed CR converges to the desired level at the same rate of the conditional distribution estimation. More importantly, the distribution shift coefficient directly influences the coverage probability. We present detailed discussion in the following remark.
Remark 1. There two factors controlling the distribution shift coefficient: 1) the observed values in Ab and 2) the missing pattern. From our thcoretical analysis, we identify a profound impact of the missingness patterns the learning efficiency and the choice of transformer architectures. (n)
Indeed, when the masking strategy in Algorithm relatively easy, ddist small. However, x can deviate significantly from the training samples, causing a large distribution shift. On the contrary,
7

---

including harder masks can effectively reduce the distribution shift, but elevates learning difficulty. As result, there is a trade-off between the masking strategy and the reliability of the trained diffusion transformer for imputation. In Section 5, we introduce a mixed-masking training strategy to enhance the performance of diffusion transformers, where diverse masking patterns are randomly sampled. This reduces distribution shift and improves robustness to varying imputation difficulty.
Experiments
We evaluate the performance of DiT through simulation to validate our theoretical results on im- putation efficiency, uncertainty quantification, and the effectiveness of the mixed-masking train- ing strategy. Experiments are conducted on Gaussian processes and, additionally, on more com- plex latent Gaussian processes to assess generalization beyond our theoretical scope. The DiT implementation builds on the DiT codebase [Peebles and Xie, 2022]. Further experimental de- tails and real-world dataset experiments are provided in Appendix F, Our code is available at https:// ggt.. com hom/liang http:/laammyga/aa /DiT cime_series imputation.
5.1 Gaussian Processes
96, dimension 8, and define the missing segment length as [Tmiss| = 16. In addition to applying Algorithm to construct 95% confidence regions (CRs), we sample from the true conditional distribution to evaluate CR coverage the proportion of true values that fall within the estimated CR for comparison.
generate Gaussian process data with sequence length
P1: k=415.4
P2: K=29.57
P3: K=9.47
P4: k=3.0
Figure 2: Visualization of the four missing patterns for sequence of length 96. Each horizontal line shows the positions of missing values (highlighted in blue, orange, green and red for Patterns 1-4), and annotations on the right indicate the pattern number and its condition number ki( <((oond))
We first vary two factors: training sample size E (103, 103.5, 10 10 108), }. and missing patterns 1-4 (denoted as PI-P4) as shown in Figure 2. As discussed in Theorem 2, x( <<oonnd) acts as key varying parameter. To mitigate distribution shift, we apply the same missing patterns to both training and test data. Results in Figure 3 show that small training sets (n = 10 10a.3) result in low variability and poor distribution estimation. As sample size increases, DiT yields CRs that significantly better match the true distribution. We further vary sequence length (H) and report the results in Table 1. The results suggest that CR coverage rate decreases sequence length increases, which supports our theoretical findings. Regarding missing patterns, those with lower condition numbers reduce the sample complexity needed for effective estimation. These findings are consistent with our theory, suggesting that the conditional covariance condition number serves practical measure of estimation difficulty. Patterns with lower condition numbers retain richer temporal correlations, enabling accurate estimation with fewer samples.
Table I: Sequence length vs CR coverage rates (% (t).
P1
CR Coverage 195)
P2
(2..7))))
(=201)
(2.14()1..0)
64)
(=1.77
P4
Table 2: CR coverage rates (%) (t) of models trained using different strategies different missing patterns.
Pl
P4
34.58 (a1.22)
58.46 (21189)
72.42 (a)
50.25 (a1.64)
10'
101
Sample Size
204
1044
107
S2
66.22 (+3.86)
83.71 (+2.86)
74.59
74.04 (+1.90)
(+2.11
83,09 (x1
57.27 (15.34)
79.00 (12.42)
74.38 (23.00)
82.74 (22.40)
Only 8x2
36.74 (+1.31)
60.51 (+1.65)
71.24 (+1 (21
80.46 (+2.01)
Figure 3: Percentage of real data samples that fall within the DiT -eneratedd 95% CR.
Only 4sc4
34.15 (+1.16)
59.23
73.08
Only Ix16
32.68 (11.50)
$4,23 (11.76)
69.16 (11.53)
76.72 (13.20)

---

Sixxed-Masking training strategy. Based on our insights from our distribution shift analysis, we introduce mixed-masking training strategy. Remark 1 highlights that discrepancies between training and test distributions can impair CR estimation, especially in real-world settings with limited training data. common practice to train on fully random masks, which tend to have lower condition numbers and thus pose easier estimation tasks. However, this intensifies the mismatch with test cases featuring more challenging, clustered missing patterns, limiting model adaptability. To address this, we propose mixed-masking training strategy. Using the same n. 10 training samples, we evaluate the four test patterns in Figure 2. During training, we define four different mixed-masking strategies (each with 16 missing entries):
S1: 100% random missing pattern (16x1, sixteen randomly placed missing entries).
S2: 50% random (16x1)+ + 50% weakly grouped (8x2, eight randomly placed blocks of two consecutive missing entries).
S3: 33.3% random (16x1) 33.3% weakly grouped (8x2) + 33.3% moderately grouped (4x4, four randomly placed blocks of four consecutive missing entries).
S4: 25% random (16x1)+ 1)+ + 25% weakly grouped (8x2) + 25% moderately grouped (4x4) + 25% strongly grouped (1x16, one randomly placed block of sixteen consecutive missing entries).
Results in Table 2 show that models trained with mixed masking consistently outperform the baseline trained with completely random placed masks (S1). We also evaluate the strategies only containing individual patterns (8x2, 4x4, and 1x16 separately), and the results suggest that they yield inferior imputation performance compared to appropriately mixing different patterns. This supports our proposed mixed-masking strategies and aligns well with our theoretical insights. Yet determining optimal mixing ratios is instance based and remains an open question for future work.
Regarding how these strategies relate to our theoretical results, intuitively, different missing patterns during training lead to different training distributions Px resulting in varying condition numbers and consequently different DS values. Training with diverse missing sntterss-ranging from easy to hard helps the model adapt to imputation tasks with varying levels of difficulty by effectively covering more scenarios. As for a more concrete example, let us denote the training distributions
corresponding to S1 and S4 as Rota and p(4) respectively. Consider test sample x following the strongly grouped missing pattern Pl (consecutive missing entries). Intuitively, the resulting
(())
distribution Px: is closer to ((1)
Roba than to P(l... which implies the distribution shift coefficient
Xobe Empirically, we calculate the average ratio across all test samples with missing pattern P1 and find that:
of p(4) is smaller than the one of (1)) Xabx
(S Phhy
47.93.
DS(Px; plh.
This clearly indicates that the mixed-masking training strategy (S4) yields significantly smaller distribution-shift coefficients compared to purely random missingness (S1). According to Corollary 1, this provides strong theoretical support for the superior empirical performance achieved by our mixed-maad maseking strategy.
5.2 Latent Gaussian Processes
We conduct additional experiments to assess whether our findings generalize beyond the theo- retical soppgggissplllll......oo whether different missing patterns affect imputation and uncertainty quantification performance, and whether the mixed-masking training strategy improves them. For X drawn from the Gaussian process in Section 5.1, we consider a corresponding latent Gaus- sian process: Y ((X) + with vec(e) (0, 0.1 LaH), where the non-linear transform σ(x) = r + exp(- =2)) + 2 sin(z) is applied entry-wise. We adopt a training sample size of n = 10 This introduces nonlinearity and noise, increasing the difficulty of distribution estimation.
We evaluate DiT on this transformed dataset using the same four missing patterns and four training strategies from Section 5. 1. For comparison, we implement two representative generative imputation models-CST [Tashroo et al., 2021] and GPVAE [Fortuin et al., 2020], ensuring all models have comparable numbers of trainable parameters. We report Mean Squared Error (MSE) against the true conditional mean and CR coverage rates, following the setup in Section 5. Results are shown in

---

Table 3: MSE (L) on latent Gaussian process data.
Table 4: CR coverage rates (%) (1).
DiT
CSDI
GPVAE
DiT
CSDI
S1
0.70 (#0.03)
0.75 (#0.03)
5.24 (+0.75)
SI
36.46 (#1.62) 54.75 (+1.89)
S2
0.68 (#0.02)
0.69 (+0.02)
5,45 (+1.05)
S2
53.68 (+3.26) 56.68 (+2.75)
P1
P1
S3
0.67 (#0.03)
0.70 (+0.03)
5.13 (40.49)
S3
54.26 (+2.79) 58.64 (#3.11)
S4
0.67 (#0.02)
0.68 (+0.02)
5.28 (+0.68)
S4
56.43 (+3.76) 55.67 (+4.03)
S1
0.64 (#0.03)
0.66 (#0.03)
5.09 (#0.70)
S1
55.81 (#1.55) 63.67 (+1.77)
S2
0.62 (#0.02)
0.63 (#0.03)
5.01 (40.62)
S2
65.77 (+2.87) 64.89 (+3.43)
P2
P2
S3
0.60 (#0.03)
0,62 (#0.02)
4,94 (+0.56)
S3
66.24 (#3.22) 63.13 (+2.95)
S4
0.62 (#0.03)
0.63 (+0.03)
4.84 (40.60)
S4
63.95 (#4.38) 65.97 (+3.59)
S1
0.62 (=0.02)
0.65 (+0.02)
4.63 (40.58)
S1
63.53 (+1.72) 61.35 (#1.49)
S2
0.60 (#0.03)
0.64 (+0.03)
5.12 (+1.0
S2
71.29 (+2.99) 65.69 (+2.79)
P3
P3
S3
0.58
(#0.02)
0.63 (#0.03)
4.50 (40.52)
S3
70.89 (#2.45) 63.48 (+2.90)
S4
0.58 (=0.03)
0.61 (+0.02)
4.59 (+0.54)
S4
73.36 (+4.37)
67.17 (+3.93)
S1
0.56 (=0.01)
0.59 (+0.03)
4.89 (#0.69)
S1
76.46 (+1.33) 68.60 (+1.74)
S2
0.53 (+0.03)
0.60 (#0.02)
4.79 (40.61)
S2
78.63 (+2.62) 70.48 (+2.34)
P4
P4
S3
0.53 (=0.01)
0.58 (+0.03)
4.39 (#0.49)
S3
78.79 (+2.67) 73.46 (+2.53)
S4
0.53 (#0.02)
0.58 (=0.02)
4.45 (40.54)
S4
80.64 (#3.72) 72.89 (+3.78)
Tables 3 and 4. Since GPVAE performs poorly in point estimation, omit its CR coverage. DiT consistently outperforms in both MSE and CR coverage, indicating transformers may better suit this task than CSDI's convolutional design. Moreover, mixed-masking training improves performance not only for DiT but also for other models, demonstrating its broader benefit. These findings reinforce conclusions from Gaussian process experiments and support the generalization of our theory and training methodology to more complex, nonlinear settings.
6 Conclusion and Discussion
Our work addresses a critical gap the theoretical understanding of diffusion- basond time series impu- tation and uncertainty quantification by investigating the statistical efficiency of diffusion transformers on Gaussian process data. This result enables efficient and accurate imputation and confidence region construction. Motivated by the theory, we propose a mixed-maskin diverse missing patterns during training, rather than relying solely on completely random masks. Our experiments validate the theoretical findings and further demonstrate that the proposed strategy performs well and generalizes to more complex data beyond our analytical scope.
training strategy that introduces
Looking ahead, investigating the behavior of diffusion transformers on heavy-tailed time series (e.g.. financial data) would further clarify their limitations and guide practical design choices. Moreover, more detailed analysis of optimal mixed-masking training strategies essecially those leveraging prior knowledee-coube significantly improve the performance of imputation models.
10

---

References
Juan Miguel Lopez Alcaraz and Nils Strodthoff. Diffusion-based time series imputation and forecast- ing with structured state space models. arXiv preprint (tt:::222200999999 2022.
FE Allan and John Wishart. A method of estimating the yield of a missing plot in field experimental work. The Journal of Agricultural Science, 20(3):399-406, 1930.
Majed Alwateer, El-Sayed Atlam, Mahmoud Mohammed Abd El-Raouf, Osama A Ghoneim, and Ibrahim Gad. Missing data imputation: comprehensive review. Journal of Computer and Communications, 12(11 ((11))53-75, 2024.
Theodore W Anderson. The statistical analysis of time series. John Wiley & Sons, 2011.
Anjishnu Banerjee, David B Dunson, and Surya T Tokdar. Efficient gaussian process regression for large datasets. Biometrika, 100(1):75-89, 2013.
Parikshit Bansal, Prathamesh Deshpande, and Sunita Sarawagi. Missing value imputation on multidi- mensional time series. arXiv preprint arXiv:2103.0160) 2021.
Joe Benton, Valentin De Bortoli, Arnaud Doucet, and George Deligiannidis. Nearly d-linear conver- gence bounds for diffusion models via stochastic localization. arXiv preprint arXiv:2308. 03686, 2023.
Christopher Bishop and Nasser Nasrabadi. Puttern recognition and muchine learning, volume Springer, 2006.
Viacheslav Borovitskiy, Iskander Azangulov, Alexander Terenin, Peter Mostowsky, Mare Deisenroth, and Nicolas Durrande. Matérn gaussian processes on graphs. In International Conference Artificial Intelligence and Statistics, pages 2593-2601. PMLR, 2021.
Sébastien Bubeck et al. Convex optimization: Algorithms and complexity. Foundations and Trends in Machine Learning, 8(3-4):231-357, 2015.
Clément L Canonne. A short note on an inequality between kl and tv. arXiv preprint 17X.9.2202.07198, 2022.
Defu Cao, Wen Ye, Yizhou Zhang, and Yan Liu. Timedit: General-purpose diffusion transformers for time series foundation model. arXiv preprint arXiv:2409.00 02322, 2024.
Wei Cao. Dong Wang, Jian Li, Hao Zhou, Lei Li, and Yitan Li. Brits: Bidirectional recurrent imputation for time series. Advances in neural information processing systems, 31, 2018.
Zhengping Che, Sanjay Purushotham, Kyunghyun Cho, David Sontag, and Yan Liu. Recurrent neural networks for multivariate time series with missing values. Scientific reports, 8(1):6085, 2018
Jou-Fan Chen, Wni-Lun Chen, Chun-Ping Huang, Szu-Hao Huang, and An-Pin Chen. Financial time-series data analysis using deep convolutional neural networks. In 2016 7th International conference on cloud computing and big data (CCBD), pages 87-92. IEEE, 2016.
Minshuo Chen, Kaixuan Huang. Tuo Zhao, and Mengdi Wang. Score approximation, estimation and distribution recovery of diffusion models on low-dimensional data. In International Conference on Machine Learning, pages 4672-4712 PMLR, 2023.
Minshuo Chen, Song Mei, Jianqing Fan, and Mengdi Wang. Opportunities and challenges of diffusion models for generative ai. National Science Review, 11(12)nwae3448 2024.
Sitan Chen, Sinho Chewi, Jerry Li, Yuanzhi Li, Adil Salim, and Anru R Zhang. Sampling is as easy as learning the score: theory for diffusion models with minimal data assumptions. arXiv preprint arXiv:2209. 11215, 2022.
Andrea Cini, Ivan Marisca, and Cesare Alippi. Filling the g_ap a__s_ Multivariate time series imputation by graph neural networks. arXiv preprint arXiv:2108.0022 2021.

---

David R Cox, Gudmundur Gudmundsson, Georg Lindgren, Lennart Bondesson, Erik Harsaae, Petter Laake, Katarina Juselius, and Steffen L Lauritzen. Statistical analysis of time series: Some recent developments [with discussion and reply]. Scandinavian Journal of Statistics, pages 93-115, 1981.
Wenjie Du PyPOTS: Python toolbox for data mining on Partially-Oserreed Time Series. arXiv preprint 2rX.v:2305.188// 2023.
Wenjie Du, David Côté, and Yan Liu. Saits: Self-attention-bass imputation for time series. Expert Systems with Applications, 219:1 19619, 2023.
Wenjie Du Jun Wang, Linglong Qian, Yiyuan Yang, Zina Ibrahim, Fanxing Liu, Zepu Wang. Haoxin Liu, Zhiyuan Zhao, Yingjie Zhou, et al. Tsi-bench: Benchmarking time series imputation. arXiv preprint arXiv:2406. 12747, 2024.
Chenguang Fang and Chen Wang. Time series data imputation: A survey on deep learning approaches. arXiv preprini urXiv:201 1347, 2020.
Vincent Fortuin, Dmitry Baranchuk, Gunnar Rätsch, and Stephan Mandt. Gp-vae: Deep probabilistic time series imputation. In International conference on artificial intelligence and statisties, pages 1651-1661. PMLR, 2020.
Chun Fu, Matias Quintana, Zoltan Nagy, and Clayton Miller. Filling time-series gaps using image techniques: Multidimensional context autoencoder approach for building energy data imputation. Applied Thermal Engineering. 236:121545, 2024a.
Hengyu Fu, Zehao Dou, Jiawei Guo, Mengdi Wang, and Minshuo Chen. Diffusion transformer captures spatial-temporal dependencies: A theory for gaussian process data. arXiv preprint arXiv:2407.16134 2024b.
Hengyu Fu, Zhuoran Yang, Mengdi Wang, and Minshuo Chen. Unveil conditional diffusion models with classifier-free guidance: A sharp statistical theory. arXiv preprint arXiv:2403. 1968, 2024c.
Wayne A Fuller. Introduction to statistical time series. John Wiley & Sons, 2009.
Víctor Gómez and Agustín Maravall. Estimation, prediction, and interpolation for nonstationary series with the kalman filter. Journal of the American Statistical Association, 89(42)):611-22 6-1624 1994.
Sergio Greco, Cristian Molinaro, and Francesca Spezzano. Incomplete data and data dependencies in relational databases, volume 29. Morgan Claypool Publishers, 2012.
José M Jerez, Ignacio Molina, Pedro Garcfa-Laeneina, Emilio Alba, Nuria Ribelles, Miguel Martín, and Leonardo Franco. Missing data imputation using statistical and machine learning methods in real breast cancer problem. Artificial intelligence in medicine, 1052)1 2010.
Chisimkwuo John, Emmanuel Ekpenyong, and Charles C Nworu. Imputation of missing values in economic and financial time series data using five principal component analysis approaches. CBN Journal of Applied Statistics (JAS), 10(1):3, 2019.
Maksims Kazijevs and Manar D Samad. Deep imputation of missing values in time series health data: A review with benchmarking. Journal of biomedical informatics, page 104440, 2023.
SeungHyun Kim, Hyunsu Kim, Eunggu Yun, Hwangrae Lee, Jaehun Lee, and Juho Lee. Probabilistic imputation for time-series classification with missing data. In International Conference on Machine Learning, pages 16654-16667. PMLR, 2023.
Gen Li, Yu Huang, Timofey Efimov, Yuting Wei, Yuejie Chi, and Yuxin Chen. Accelerating convergence of score-based diffusion models, provably. arXiv preprint arXiv:2403.0388 2024.
Huiping Li, Meng Li, Xi Lin, Fang He, and Yinhai Wang. spatiotemporal approach for traffic data imputation with complicated missing patterns. Transportation research part C: emerging technologies, 109:10 2020.
Roderick Little. A test of missing completely at random for multivariate data with missing values. Journal of the American statistical Association, 83(404):1198-1200 1988.
12

---

Mingzhe Liu, Han Huang, Hao Feng, Leilei Sun, Bowen Du, and Yanjie Fu. Pristi: A conditional diffusion framework for spatiotemporal imputation. In 2023 IEEE 39th International Conference on Data Engineering (ICDE), pages 1927-1939. IEEE, 2023.
Yonghong Luo, Xiangrui Cai, Ying Zhang, Jun Xu, al. Multivariate time series imputation with generative adversarial networks. Advances in neural information processing systems, 31, 2018.
Pierre- Alexandee Mattei and Jes Frellsen. Miwae: Deep generative modelling and imputation of incomplete data sets. In International conference on machine learning, pages 4413-4423, PMLR, 2019.
Song Mei and Yuchen Wu. Deep networks denoising algorithms: Sample-efficient learning of diffusion models in high-dimensional graphical models. IEEE Transactions on Information Theory, 2025.
Xiaoye Miao, Yangyang Wu, Jun Wang, Yunjun Gao, Xudong Mao, and Jianwei Yin. Generative semi-supervised learning for multivariate time series imputation. In Proceedings of the AAAI conference on artificial intelligence, volume 35, pages 8983-8991, 2021.
Ahmad Wisnu Mulyadi, Eunji Jun, and Heung-Il Suk. Uncertainty-aware variational-recurrentt imputation network for clinical time series. IEEE' Transactions Cybernetics, 52(9))9684 9694, 2021.
Kazusato Oko, Shunta Akiyama, and Taiji Suzuki. Diffusion models are minimax optimal distribution estimators. In International Conference on Machine Learning, pages 26517-26582. PMLR, 2023.
Yidong Ouyang, Liyan Xie, Chongxuan Li, and Guang Cheng. Missdiff: Training diffusion models on tabular data with missing values. arXiv preprint arXiv:2307.00467 2023.
Leandro Pardo. Statistical inference based on divergence measures. Chapman and Ha
William Peebles and Saining Xie. Scalable diffusion models with transformers. arXiv preprint arXiv:2212.09748, 2022.
Ignacio Peis, Chao Ma, and José Miguel Hemándezz-obato. Missing data imputation and acquisition with deep hierarchical models and hamiltonian monte carlo. Advances Neural Information Processing Systems, 15:35899-35855 2022.
Carl Edward Rasmussen and Christopher K. I. Williams. Gaussian Processes for Machine Learning. MIT Press, 2006.
S Mara Ribeiro and CL Castro. Missing data in time series: A review of imputation methods and case study, Learning and Nonlinear Models, 20(1):31-46, 2022.
Matthias Seeger. Gaussian processes for machine learning. International journal of neural systems, 14(02):69-106, 2004.
Robert H Shumway, David S Stoffer, and David S Stoffer. Time series analysis and its applications, volume 3. Springer, 2000.
Yang Song, Jascha Sohl-Dickstein, Diederik P Kingma, Abhishek Kumar, Stefano Ermon, and Ben Poole. Score-based generative modeling through stochastic differential equations. arXiv preprint 1.Xiv:20111.13456 2020.
Wenpin Tang and Hanyang Zhao. Score-based diffuusion models via stochastic differential equations-a technical tutorial. arXiv preprint arXiv:2402.07487, 2024.
Yusuke Tashiro, Jiaming Song, Yang Song, and Stefano Ermon. Csdi: Conditional score-based diffu- sion models for probabilistic time series imputation. Advances in Neural Information Processing Systems, 34:24804-24816. 2021.
David Alexander Tedjopurnomo, Zhifeng Bao, Baihua Zheng, Farhana Murtaza Choudhury, and Alex Kai Qin. A survey on modern deep neural network for traffic prediction: Trends, methods and challenges. IEEE Transactions Knowledge and Data Engineering, 34(4)15544-1561, 2020.
13

---

Sana Tonekaboni, Danny Eytan, and Anna Goldenberg. Unsupervised representation learning for time series with temporal neighborhood coding. arXiv preprint vrXiv:2106.00750, 2021.
John W Tukey, The extrapolation, interpolation and smoothing of stationary time series with engineering applications, 1952.
Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N Gomez, Fukasz Kaiser, and Illia Polosukhin. Attention all you need. Advances in neural information processing systems, 30, 2017.
Peerapon Vateekul and Kanoksri Sarinnapakom. Tree-based approach to missing data imputation. In 2009 IEEE International Conference Data Mining Workshops, pages 70-75. IEEE, 2009.
Pascal Vincent. A connection between score matching and denoising autoencoders. Neural computa- tion, 23(7):1661-1674 2011.
Jun Wang, Wenjie Du, Wei Cao, Keli Zhang, Wenjia Wang, Yuxuan Liang, and Qingsong Wen. Deep learning for multivariate time series imputation: A survey, arXiv preprint 240i::2200.005599 2024.
Xu Wang, Hongbo Zhang, Pengkun Wang, Yudong Zhang, Binwu Wang, Zhengyang Zhou, and Yang Wang. An observed value consistent diffusion model for imputing missing values in multivariate time series. In Proceedings of the 29th ACM SIGKDD Conference on Knowledge Discovery and Data Mining, pages 2409-2418, 2023.
Haixu Wu, Tengge Hu, Yong Liu, Hang Zhou, Jianmin Wang, and Mingsheng Long, Timesnet: Temporal 2d-variation modeling for general time series analysis. arXiv preprint arXiv:2210.02186, 2022.
Shin-Fu Fu Wu, Chia- Yung Chang, and Shie-Jue Lee. Time series forecasting with missing values. In 2015 1st International Conference on Industrial Networks and Intelligent Systems (INISCom), pages 151-156. IEEE, 2015.
Jun-He Yang. Ching-Hsue Cheng, and Chia-Pan Chan, time-series water level forecasting model based on imputation and variable selection method. Computational intelligence and neuroscience, (7:34114 2017.
Xiuwen Yi, Yu Zheng, Junbo Zhang, and Tianrui Li. St-mvl: Filling missing values in geo-sensory time series data. In Proceedings of the 25th international joint conference on artificial intelligence, 2016.
Jinsung Yoon, James Jordon, and Mihaela Schaar. Gain: Missing data imputation using generative adversarial nets. In International conference on machine learning, pages 5689-5698. PMLR, 2018a.
Jinsung Yoon, William R Zame, and Mihaela van der Schaar. Estimating missing data in temporal data streams using multi-directional recurrent neural networks. IEEE Transactions on Biomedical Engineering, 167541147, 2018b.
Ceylan Yozgatligil, Sipan Aslan, Cem lyigun, and Inci Batmaz. Comparison of missing value imputation methods in time series: the case of turkish meteorological data. Theoretical and applied climatology, 112:143-677 2013.
Hengrui Zhang, Liancheng Fang, and Philip S Yu. Unleashing the potential of diffusion models for incomplete data imputation. arXiv preprint arXiv:2405. 20690, 2024.
Shuyi Zhang, Bin Guo, Anlan Dong, Jing He, Ziping Xu, and Song Xi Chen. Cautionary tales on air-quality improvement in beijing. Proceedings the Royal Society A: Mathematical, Physical and Engineering Sciences, 473(2205):20104557, 2017.
Shuhan Zheng and Nontawat Charoenphakdee. Diffusson models for missing value imputation in tabular data. arXiv preprint arXiv:2210.17128 2022.
Jianping Zhou, Junhao Li, Guanjie Zheng, Xinbing Wang, and Chenghu Zhou. Mtsci: A conditional diffusion model for multivariate time series consistent imputation. In Proceedings the the 33rd ACM International Conference Information and Knowledge Management, pages 3474-3483 2024.
14

---

NeurIPS Paper Checklist
1. Claims
Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Question:
Answer: [Yes]
Justification:
Section 3 and 4 provide theoretical results, while Section 5 provides numerical results.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation as long as is clear that these goals attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 6 discuss possible limitations.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors encouraged to create separate "Limitations" section in their paper.
The paper should point any strong assumptions and how robust the results are violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and prool's
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [Yes]
15

---

Proper assumptions are provided in Section 2, and proofs are provided in the Appendix.
Justification:
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proof's can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide a short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should properly referenced.
4. Experimental result reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification:
Experiment settings are provided in Section 5 and Appendix F.
Guidelines:
The answer NA means that the paper does include experiments.
If the paper includes experiments, No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how reproduce that algorithm. (b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) If the contribution is a new model (e.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
16

---

Answer: [Yes]
Justification:
Code is provided in the supplementary materials.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.. cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, so "No" an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https:
/ //nips. wttpccccpppccccceeececceccccccct /ttpa//gguiieee/eeeuuuubiiiiiiii
/CodeSubmie sswonnoolll. onPolicy))) for more details.
The authors should provide instructions on data access and preparation, including how access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information possible in supplemental material (appended to the paper) recommended, but including URLs to data and code permitted.
6. Experimental setting/details
Question: Does the paper specify all the training and test details (e.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: Simulation details are provided in Section 5 and Appendix F. Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to level of detail that is necessary to appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, or as supplemental material.
7. Experiment statistical significance
Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Question:
Answer: [Yes]
Justification: Standard deviations of the experimental results are reported.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, at least for the experiments that support main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call a library function, bootstrap, etc.)
The assumptions made should be given (e.g.. Normally distributed errors).
should be clear whether the error bar is the standard deviation the standard error of the mean.
17

---

It is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of enrors is not verified.
For asymmetric distributions, the authors should be careful to show in tables or figures symmetric error bars that would yield results that are out of range (e.g. negative eтror rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of exccution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Details are provided Appendix F.
Guidelines:
The answer NA means that the paper does include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount compute required for each of the individual experimental runs well estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurIPS Code of Ethics https: htt////eeur....eee cc cc/public//eeeeccccccccccececeee www.ethhhiieeeeeneeeeeeeeee
Answer: [Yes]
Justification: The research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics.
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make to preserve anonymity (e.g.. if there is special consid- cration due to laws or regulations their jurisdiction..
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification:
Our theoretical research is not tied to particular negative societal impacts.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
18

---

The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition attacks, mechanisms for monitoring misuse, mechanisms to monitor how system learns from feedback over time, improving the efficiency and accessibility ML).
11. Safeguards
Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, scraped datasets)?
Question:
Answer: [NA]
Justification: Our work poses no such risks.
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: The creators and original owners of assets used the paper are properly credited.
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset is used and, if possible, include URL.
The name of the license (c.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
19

---

If this information is not available online, the authors are encouraged to reach out to the asset's creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [NA]
Justification: Our work does release assets.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/mmoel as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent obtained from people whose asset is used.
At submission time, remember to anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsoureing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well details about compensation (if any)?
Answer: [NA]
Justification: Our research does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsoureing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included the main paper.
According the NeurIPS Code of Ethics, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question:
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or an equivalent approval/review based on the requirements of your country or institution) were obtained?
Answer: [NA]
Justification:
Our research does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such the institution conducting the review.
16. Declaration of LLM usage
20

---

Question: Does the paper describe the usage of LLMs if it is important, original, or non-standard component of the core methods this research? Note that if the LLM used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, or originality the research, declaration is not required.
Answer: [NA]
Justification: The core method development in this research does not involve LLMs any important, original, or non-standard components.
Guidelines:
The answer NA
that the method development in this research does involve LLMs as any important, original, or non-standard components.
Please refer to our LLM policy (https : h/tp://ies...//eeeeeeeeee http///rrs..c//ooeeeeeeeeeeeeeee / for what should or should be described.
21

---

A Related Work
In the early stages of time series imputation, statisticians developed wide range of traditional statistical methods aimed at both imputation (point estimation) and quantifying uncertainty, often by leveraging well-established statistical tools to construct confidence intervals [Cox et al., 1981, Shumway et al., 2000]. Initial techniques were relatively simple, such as imputing missing values using the mean median of observed entries. These were later followed by more advanced interpola- tion approaches based on regression models, including linear regression and splines [Shumway et al., 2000]. To better exploit the spatio-temporal structure inherent in time series data, model-based meth- ods emerged, such as ARIMA, GARCH, Kalman filters, and Bayesian inference frameworks [Fuller, 2009]. These approaches knowledge, and support for formal statistical testing. Moreover, many of them naturally allow for uncertainty quantification through predictive intervals or posterior distributions. However, these methods come with notable limitations: they typically rely on strong assumptions about stationarity, linearity, noise distributions, making them less effective for complex real-world data with nonlinear or -igensimensional apatio-tempora dependencies [Anderson, 201 1J. Additionally, their computa- tional cost often scales poorly with data dimensionality, posing challenges for modern large-scale applications.
advantageous for their interpretability. ability to incorporate domain
To address the limitations of statistical approaches, machine- based imputation methods have be- come increasingly popular in recent years. Early approaches include classical machine learning models [Jerez et al., 2010] such as support vector machines [Wu et al., 2015] and tree-based meth- ods (including bagging and boosting techniques) [Vateekul and Sarinnapakorn, 2009, Yang et al., 2017]. With the advancement of model architectures and increasing computational power, deep learning-based models have gained prominence for their ability to capture complex temporal de- pendencies [Fang and Wang, 2020, Wang et al., 2024, Du et al., 2024]. Predictive models such as RNNs [Che et al., 2018, Yoon et al., 2018b, Cao al., 2018], CNNs [Wu al., 2022, Fu et al., 2024a]. GNNs [Cini et al., 2021]. and transformer-basee networks [Bansal et al., 2021, Du et al., 2023] directly estimate missing values using well-designed architectures. Generative imputation methods model the distribution of missing data and perform better in quantifying uncertainty; representative techniques include GAN-based methods [Luo et al., 2018, Yoon al., 2018a, Miao et al., 2021], VAE-based approaches [Mattei and Frellsen, 2019, Fortuin et al., 2020, Mulyadi et al., 2021, Peis et al., 2022. Kim al., 2023], and diffusion models. Among diffusion approaches, CSDI [Tassiro et al., 2021] introduced conditional diffusion for time series imputation, and subsequent work [Alcaraz and Strodthoff, 2022, Wang al., 2023, Liu et al., 2023, Zhou et al., 2024] improved conditioning strategies and computational efficiency. DiT [Peebles and Xie, 2022, Cao et al., 2024] extends this line by integrating a transformer backbone into the diffusion framework, achieving better imputation accuracy and uncertainty quantification. These methods resolve certain issues and perform well empirically, however, are still limited by lacks of uncertainty quantification in many methods and theoretical understanding.
Our work also contributes towards the theoretical foundations of diffusion models [Chen et al., 2024, Tang and Zhao, 2024]. Some prior works have established sample efficiency and learning guarantees for diffusion models when modeling the original data distribution. Chen et al. [2022], Benton et al. [2023], Li et al. [2024] show that the generated distribution remains close to the target distribution, assuming access to relatively accurate score function. By incorporating score approximation procedures and corresponding theoretical analysis, Chen et al. [2023], Oko et al. [2023]. Mci and Wu [2025] provide end-to-end guarantees, covering various types of data including manifold data and graphical models. In the case of conditional diffusion models, sharp statistical bounds of distribution estimation have been derived in Fu et al. [2024c]. Additionally, Fu et al. [2024b] explores the theoretical regime of modeling spatio-temporal dependencies in sequential data. However, these results do not directly apply to more concrete and complex scenarios, such how conditional DiT models can learn intricate dependencies to accomplish time series imputation tasks.
B Proof of Lemma 1
We provide the detailed proof of Lemma in this section.
To simplify our analysis, we begin by making some assumptions. Firstly, without loss of generality, we assume the mean of the Gaussian process data = O. Large norms in x and v, often lead to training instability, making practical to perform clipping. Inspired by this, leveraging the Gaussian and light-tailed nature of x and Vz, we truncate the domain of the data and diffused samples by
22

---

defining an event that occurs with high probability 1 - 8:
Cs [1|x||2 S Castar l|val12 S Cdstal
(13)
where Cå (d) is threshold depending 8. Our score approximation analysis of Lemma and Theorem 1 is conducted under the condition of event C. (ensuring the conclusions hold with high probability 1 ö), which significantly simplifies the process. The relationship between the truncation range Caata and high probability å is deferred to Lemma 11. Outside event Cj (i.e., on C the unbounded range complicates obtaining a meaningful score approximation in the second-norm sense. However, Cg occurs with a small probability, we can still achieve reliable results in distribution estimation, where evaluation is based on expectation.
Some Useful Results In this part, present some key results regarding the eigenvalues and condition numbers of covariance matrices, which will be instrumental in our analysis.
We first define:
Amax (a Doond + II
Kg Amtn (o E cond + oīT)
Using the positive definiteness of r and A, we obtain:
Amax (A) = ||A||2 > O, Amiu (A) = || > 0.
Furthermore, by the properties of the Kronecker product, we derive:
Amax Eoba) = Amax Am Amax
Amin (Dobs) Amin Amin) (A),
Amax Ammax (Tohs) = K k(Tobs))(
Kc <(Ebb))
Amin (A) Amii (Tobs)
Finally, we assume:
Amax (obs), Amax (T miss), Amax -
B.1 Key Steps for Proving Lemma 1
In Lemma 1, we aim to show that the gradient-basee Algorithm provides good approximation of the conditional score function.
The algorithm employs gradient descent to solve two types of optimization problems: the major GD problem (7), and the auxiliary GD problem (9), which is solved within each update step of the major GD. It critical to note that the major GD updates are inherently noisy due to various reasons, such as the auxiliary GD approximating certain quantities at each step, and later using transformers to approximate each step. Therefore, establish the result in Lemma 1, our proof consists of two key steps:
Step 1.
We demonstrate that, with a sufficient number of auxiliary iterations Kaux. the approximation error of the auxiliary GD loop's result can be controlled below specified threshold.
Step 2. We then show that, by controlling the perturbation level in each major GD update step, the score approximation error (i.e., the gap between the output of the major GD and the ground truth score function) can also be bounded, provided there are enough major iterations N.
In the following, we elaborate each step by providing precise statements and subsequently use them to prove Lemma 1. All supporting results are deferred later sections.
B.2
Detailed Statements in Steps 1-2 and Proof of Lemma
Now we present formal statements in Step 1-2 and use them to prove Lemma 1.
B.2.1 Formal Statements in Steps 1-2
This section contains the statements of Lemma and Lemma
Lemma 2. For arbitrarily fixed time I. E (0, T] and given an error tolerance eo E (0,1), let b == (ko)s(A) E [RdlIli running the auxiliary gradient descent in (10) with suitable step size 0 = 2/(miin (Eolsa) + Amax(E for
xs( (oobs)
|
Kaux
log Xmmin (Totas Arinn DEn
2
iterations produces a solution u (Kau)) that satisfies
u (Kaux)
-ull2 5 ED.
23

---

Here, we introduce to distinguish the noise arising from the auxiliary GD loop approximation from the error level stated in Lemma This distinction provides additional flexibility adjust 6o in subsequent proofs.
Next, we establish a lemma for the convergence of the major GD. In each major GD step (referring to (8)), incorporate error term and represent gradient update as:
(+++) s N/ (
g(().
(14)
where (A)) represents the error term in each perturbed major GD step. Explicitly accounting for the noise present each perturbed gradient step, we can establish:
Lemma 3. For arbitrarily fixed time E (0, and given an tolerance E (0, 1), suppose [| ||2 < then running the major gradient descent in (14) and suitable step size = 2/ (2min (ai Beond + o + Amax (aixx(ajeeonn Bcond + +o(1) o
== (1, (gg log
Hdk((A) obs)
OrE
iterations produces a solution (K) that satisfies
[ks(k) - s112 S
lit
1)
2
where == k(aa Eoodd + oi
With the convergence of both the auxiliary and major GD established, we are ready to prove Lemma
B.2.2 Proof of Lemma
Proof. By the statement in Lemma 3, we need to control the noise level in each major GD step, i.e.
ensure (kk) S... We analyze this error as
112 - ||a T
cor obs soo aj a2200 S So1. Ecors (
ΣT D-1 Xoba 112
Xobs
< 0 Eoo
Ecors (N) E --1
Erors (8)
ohs
obss
oor
ll E Xohs
Sobs Xohn 2|
Here, the latter term arises from approximating Meond (Xobe), and Σ To (Xoba Mobs) represents the Kaux-iteration auxiliary GD approximation of the matrix vector product. We provide useful lemma to help control the error above.
Lemma 4, For an arbitrarily fixed time E (0, we have
[Foor (Toorl2x((A)
1Ive Meoud (lobs)l12
[++
Amiu (Tobe)
data
and
|s/112 S o, 2| -eoond (nona)XXb))lll
Invoking Lemma 2 and Lemma 4, letting ég = ((aî + 1) Ecorlle) -1 can bound the required auxiliary iteration steps by:
e, to ensure that [|5(() S c.
16 E(Eobs) log
VVHd(Clau + [Is|/oo) Ecor
Kaux
2
Amin (T (obss) Amia Aoain(A (A)e
=0 c0(mtt
Hdd((A)s(Taa
(((A)) ((AA))((aaa))) log
Ote
Lastly, we invoke Lemma 3, substitute with e =
met2 e, we have
K = 0 20(m) Ki log
Hdrri((Tbba)l ))
24

---

Finally, notice that
Amax (a Σ Lcond + +0²1)
16g
Amin (a, Eooud + +,,1) S 16 (cond)
and leveraging
k( [ohh)) = xi(I H(TTobs)s(AA)
we have
K O (x(x(2 ((oondd) log
dxx Zcond )is( (Dote)) )) Kaux =O(((Eobs) log
HHdsss (nnbs)
OzE
This completes the proof Lemma
B.3 Proofs of Lemma and Lemma 3
To prove the lemmas, we first state standard result in convex optimization.
Lemma 5 (Theorem 3.12 in [Bubeck et al., 2015)). Let f be 3-smooth and a--strongy conver on R and be the global minimizer. Then gradient descent with v) a+3 satisfies
|xxlht)) x" $11
K
|x(00
x
k = 0,
where x (k+1)) x(() nV (x((()) is the outcome at the (Ic + 1)-th iteration of gradient descent, and K a'
Equipped with this lemma, the proof process is straightforward.
B.3.1 Proof of Lemma 2
Proof. Referring to (10) (expression of auxiliary GD step), the update steps are
u u Cinner (u)
u - 0 (Zoba - b)
We should notice that Laux is Amax Eobs) )abbe))sm)) -smooth and Amin (Eobe) -stronglyy convex. Then by Lemma 5, we have
Ki Eobs) -
Iu (1)x - u/12 x
|/u (Maux) u112
Ki Eohs) +
(1
xi( Dobs) + y
(lu - u1/2
- 2KKaxx + +1 )
< exp
(Eobs) + 1 ||u||2-
We also have
|ull || bll2 < Eoba 11 ||bll S
Amin (T (A)-11|bll2.
With preset error eo >> 0, taking number of iterations Kaux obtain
N ([Eab)) +11 2
log( Amat obu Ammin (A)co )
blI
I (Kaux) -u//2 I cu-
B.3.2 Proof of Lemma? 3
Proof. In each step, we incorporate an error term and represent our gradient update as in (14):
s (k+1) (k) - (s))) + (k))
We should also notice that L, Amax (a I Doond + 0, -mmooh and Amin (a Doond + o, -strongly
convex. Then with [k{()) ||2 5 by Lemma 5,
25

---

[|(k+1]]
s//2 2 ||s(k) (s)))((
112 + (1 2|
||s(A) L, 1122 +
Nt
[s( - 5|/2 tE.
Kip +
Then we have
[|s(t+]) -8|12
Kt
Kt 1
u+1
||s - -8112
Kr
2 {)
2
Kit + 1
Similar to the proof of Lemma 2, we obtain
|s(K) ss|12- Kt c cs exp
Ki + 1f ||s||2-
2N
Ki
2 c)
exp
-2N Kic + }
(1|s||2).
By invoking Lemma 4, we also have
[1+
m((A)
||||| I1(a; oi (1(v: acMooad)ll S ot²²
Cdata
Amin (Tobs)
Lastly, taking K 0 (Kt log (dd(())e(Tatn) )) we obtain
[| s 1) E.
This finishes the proof of Lemma 3.
26

---

C Proof of Theorem 1
We provide the detailed proof of Theorem in this section by explicitly construct transformer architecture to unroll Algorithm 2. Firstly, we assume that the mean function Mobs: Pmiss can be constructed by an additional preprocessing network. Thus, the analysis in this section can also be conducted under the condition of event Cs and still assuming = 0 as stated in Appendix B.
C.1
Key Steps for Proving Theorem
The proof of Theorem is presented in a constructive framework. Revisiting the architecture in (5), we observe that comprises the encoder fin. which transforms the original input into a form compatible with the unrolling of Algorithm 2; the raw transformer blocks, which perform the algorithm unrolling; and the decoder four. which extracts and truncates the output to provide the final score approximation.
We define the major GD step with major GI steps. Similarly, we categorize the auxiliary GD steps. Notably, the first major GI step is relatively simpler, while the later major GD steps are analogous to it. Accordingly, we separate our analysis into the first and later major GD steps. To establish Theorem 1, the proof proceeds through the following steps:
the first major GD step and those with k >> 1 as the later
Step Construct the encoder, decoder, and essential components that critical for constructing the subsequent raw transformer architectures.
Step 2. Construct the raw transformer architecture for the first major GD step.
Step Construct the raw transformer architecture for the later major GD steps analogously.
Step Analyze the error and configuration the raw transformer architectures constructed in the previous steps.
Step Summarize the constructions and analyses to establish the result in Theorem 1.
C.2
Constructing Encoder, Decoder and Some Crucial Transformer Components
For sake of simplicity, given time step E (to, T). we denote (vr), = X; E Rd in the following analysis. Additionally, define each (FFN, Attni) transformer block. The architecture composed solely of transformer blocks, excluding the encoder fin and decoder four. is referred to as the transformer, denoted by Traw (D, L, B).
The encoder need is mapping our input higher dimensions in attepmpt to include some useful values (e.g. time embeddings) and also some buffer spaces to finish the gradient descent process. For simplicity, specific time t, suppose the encoder converts the initial input into Y fin([x1, X2,..., XH; ()) = [yI,... ......... NJ] E R DR which satisfies
Encoder
x e o()) T 1,0,1,x, OIl
yj [x] e, (o() 0 1,0,1,0
where (t) [m, 0t of a E Rdd with d, 4. Specifically, we use different subscriptions for observed indices and missing indices, i... i E Iobs. Imiss For simplicity, omitted the subscript here, and o og 0, serve as the buffer space for storing the components necessary for unrolling the algorithm.
Decoder Suppose the output tokens from the transformer blocks has produced conditional score approximator in matrix shape, and the stability for the computation inside the network, we design the decoder fout = faorm fiieerr- Where flnnesr Dx flattens the input into a vector that aligns with the dimension of the conditional score function, and fnorm= RddiZuiml + didii controls the output range of the network by the upper bound of the
diditli extracts and
score function. By Lemma 4, denote o, (1+ Aaiia (Toke) Ca as Re, we can set
Jnorm (s) {(s, s,
if |||||| I Rr, otherwise.
||=||a
fmuit Module At the end of this part, we also provide the construction of the multiplication module, which approximates the product between scalars and vectors. This is crucial component in constructing JGD later. We introduce a lemma, which a modified version of Corollary [Fu et al., 2024b]:
27

---

Lemma 6. Suppose input to be Y = 1y1, .Hll E Dx H with = {x Ogd We, z where
X; E [-B, BJ wa E 1- B) and Z; E Rd,, Given any Emult > 0, there exists (FFN-only) transformer architecture such that
fmah FFN,, FFN o... FFN,
with L (log(B/ Emult)) layers that approximately multiply each component x, with the weight w; and put it into a buffer, keeping other dimensions the same. This can be formally written as
X1,
XH
fmult ( x1),
mullc (WH XH)
Emult(Y
=
O2d,
O2d
where Il x.) - w; / oc I Cmult.
01
21.
Z.
The number of nonzero coefficients in each weight matrices or bias vectors is at most (d), and the norm the matrices and bias all bounded by ((d)).
C.3 First Major GD Step
In this section, we construct the transformer architecture unrolling the first step of major GD procedure, and the result can be summarized as:
Lemma 7 (Construct first major GD step). There exists transformer architecture JGD,1 Tram (D, M, B), which can construct an approximate first step major GD result 5(1) from the output of the encoder.
Given an error level e and learning rate qi > O, the approximated first step GD result satisfies
5()) -(1) =
( where (1 Se
where s( is the groundtruth gradient step. The configuration of the transformer architecture satisfies
Hdd(A)) (Enba))
D 12d de de + +,,11 O
(((A)x(T obbs)) log?
c
M 4H, B = (H (( We defer the analysis of transformer configuration to C.5. Following the encoder network construction above, we obtain the following input:
X;
X;
e;
e
((t)
(())
Ond
Osd
Y
0
1
X;
O₂
Oaa
O4a
Revisiting the major GD gradient step in (8), our goal form the following major GD first step:
(0) + Monnd TJ,8 0 (ΣT corr abso (Smuit (Nrar, Xohs))); -fmute (1mu x,). (15)
Wea gooo
Initially, we apply multiplication module construct:
fmult (n,, x4)
fmalt Smule(Tre x,)
e;
((t)
((t)
funtt (niOe, X2
pi.obs
Osd
Osd 1
fautt Y)
O
0
1
X;
Oa
Ozd
Old
This encapsulates the fundamental operations required for constructing the first major GD iteration.
28

---

C.3.1 Auxiliary GD for First Major GD Step
In this section, we approximate the term E (7)os))bss) using an iterative auxiliary GD procedure.
First Auxiliary GD Step Starting from the initialization u, (0) Oa for i E Tobs, want the first auxiliary GD iteration finish the update:
(1) u
(0)
+ (nio)Xobs)).
Using fmult, we can easily obtain
Smnte xi) Bj
Smllt (1te xz)
((1)
(1)
fmule( fmule (ner,: xx)) )) u ())
Od
mmll fmult( Y)-
0
Osd
0
X,
Og
O4d
Oad
This finished the first step of auxiliary GD.
Auxiliary
GD Later Steps For subsequent iterations, the updated rule for the auxiliary GD
becomes:
un (1) (kaux) + +(11)()Xohss)
Σ
Tu..Au.
KEIaba H-1
(Roux) + ((T)C((Xohm)
Σ
Σ
7'm Il {1i kl mmmmAAu Aug.
(16)
m= kEZoba
for kaux 1,
Kaux - 1.
Similar to the first auxiliary GD step performed above, we firstly use fmuli to obtain
fmute x,)
Smult (M, x
Ci
<(t)
(tt)
(kaus)
u, Os
Oa
Oa
fmnltt (Oro02, x3)
Od
mmllt 0 fmult mullt (Y)
fraute (0, u.
(kaux)
0₂
O2d
O2d
1
1
0
0
X;
Og
Oad Then, by constructing a 4H -had attention block as described in C.6. we obtain
O4d
fmult (Th x;)
fmult (1) x,)
(tt)
(tt)
u (kaux)
0a
E#1-1 Eketom
Im (/i k| m) A Smale (0, up
(kaux)
0.
Smalt (Bm,0r, x:)
0₂
f fmute (8, u
(kaux))
0,
O2d
Ozd
1
0
0
X;
0a
O4d
O4d
29

---

Lastly, after combining an linear transformation FFN block with the attention block above to build up the basic transformer block T abs. we will have
fuult (MG, x;)
fault (Tc, x;)
ej
((t)
(())
(kaux+1)
u,
Oa
TBobs o fmult fmuh fmuit( (Y)
0dd
Osd
1
1
0
O
X;
Og
O4d
Old
where
H-1
(hanx +1)
(Raux)
Σ
Σ
Ym 1 - = mm) fmult (6, u (khox)) ) ) Smult(mnaaa (Bnaa, X; - Mi.obs).
u.
U;
m== TToo
This completes later step auxiliary GD update.
We denote the iterative blocks (TBohs o fmmlt o fmmat ) Kaux as finner- The result of the auxiliary GD for approximating (E (neare (Xobs Mobs)) is expressed as:
Final result for Auxiliary GD
u - munle (ne (Kobs)));
After completing Kaux auxiliary GD iterations, we obtain at the following transformation:
muut (M,, x,)
fmutt (Ne. x,)
(f)
(t)
ü,
Og
Oad
Oad
finner mmlte (Y)
1
1
I
0
0
1
X;
0g
O4d
0ad
which incorporates the iterative updates. The resulting output includes ü;; for each observation entry, alongside the original data. We therefore finish the auxiliary GD procedure for the first major GD step.
C.3.2 Matrix Multiplication
After Kaux steps of auxiliary GD iterations, we proceed with an additional matrix multiplication step compute Σ II
The multiplication can be expressed
T ū) Σ
Σ
Ym Aii - m) A.A.
m=0 ieZ.he
Referring to the construction in C.6.2, this computation can be implemented using 4H-head attention block:
funute (me. x,) e; ((t)
fmult (7,,))
e;
@(()
U
0.
Oa
Em= Σ Eiezar mm 1(li jl = Aü, nia (P
Oad
Ozd
1
1
0
0
X;
0,1
O4d
04d
30

---

Combining the attention block described above with linear transformation through FFN block, which we denote as TB we obtain:
mult (7), (1)
fmllt (nt,,)) xz)
ej
((t)
((t)
U;
0g
0g
(1)
O4d
04d
fGD,1 TBeortt o finner Emuti (Y)
1
1
0
1
Og
Og
(2 2 obe fuute {obs))]
Osd
O3d
We now define 5/0) as:
5/ fumtt (nprr Mj,miss)
Σ
(Eote (n, Xaba))
Smult (Mt, (T))))))
Comparing this result with (15), we observe that the first step of the major gradient descent is now complete. We represent this step fGD,1-
For simplicity, introduce the notation:
(Door (mu Xobe))
Th TioiifM..cond
C.4 Major GD Later Steps
In this section, we construct the transformer architecture unrolling the later steps of major GD procedure, and the result can be summarized as
Lemma 8 (Construct major GD later steps). There exists raw transformer architecture fan
Traw (D, L, M, B) (i.e. without encoder and decoder), which can construct a new approximate later step GD result s+ from the output of the latest step of major GD. Given error level (0, 1) and learning rate Th > 0, the approximated first step GD result satisfee
g(++1) -g =t ,where | Se,
where (K+1) is the groundtrush gradient step. Furthermore, the configuration of the raw transformer architecture satisfies
H A )os(( obs))
))
= 12d + de + dg + 3, =
(K(K(A) (T)o))))))
OeC
M = 4H, B B= O (v (rd +(A)ss(TTt))))
We defer the analysis of transformer configuration to C.5.
In the later steps of the major GD, we need to compute the following update:
Li (k+1) -(()) - -0 T E-1 Ecors Otepoond (Xonss)) + az Emisos (k) + (1) (k) + vi
(k)) -T 5-1 S-1E Fenrfmuit an ((1)) + mi Mr (Xons) + Σ Emiss fmutt (n s s(k)) obss
ITT
- fmuut (noo, s(x) fmule Vz).
(17)
In the following proof, for the sake of simplicity, respectively.
st , abbreviation for (k+1)) s(k)
In each new major GD step, the input the iteration is the output the most recent GD step. For simplicity, continue to represent this input using Y.
31

---

Similar to the construction in the first step, we first apply a multiplication module to obtain:
[mule (T/,, x.)
Smatt (Mte xj)
(())
(())
0g
Sj
03d
O3d mmilt (m sq)
0,d
muli (Y)
Og
fmntt (no) sy)
1
1
0
X;
Og
Og
03d
03d
Next, proceed to the matrix multiplication step, followed by the auxiliary GD process. C.4.1 Matrix Multiplication
To begin, we compute Ecors, which can be expressed as:
T cor & A A)s.
For each i corresponding to the observations, this can be further rewritten as:
#-1
Σ
Σ
7miiii- As,.
m== J.T
To perform this computation, similar the construction in the first step major GD, employ 4H -had attention block combined with an identical FFN block to form 7 Beo::
TBoor YY
funle (M, x;)
fmult (T,, Xz) ej
C;
o(t)
d(t)
0g
Sj
H E Emmm Mm 1 {li
5l
m) Afmule Ammute(1oo?ss
Oa Oza Smute (na 8y) 2
Oza
0
0₂
fmatt (na),
1
0
X;
0a T7 Oig Hi.oMo Ogd
0g
0gd
C.4.2 Auxiliary GD
In this step, we compute Σ-1
abs Ecor8.
Following the auxiliary GD procedure described in C.3.1, we employ similar iterative approach using finner to approximate the multiplication between matrix inverse and vectors:
32

---

Smute (Th, xi)
Smnte (Tle xj)
<(t)
(())
Og
5j
( E cor Smatt (nea s));
obs
Oad
Old Smatte (na 85)
0g
finner Teor((Y)
0g
Smult (nrof, 8z)
1
0
X;
0g
03d
03d
C.4.3 Matrix Multiplication
In this step, compute Σ. ubs Eca E.
Following similar procedure as in C.3.2, we construct transformer block TB This block similarly employs 4H-heead attention mechanism combined with identity FFN to perform the matrix multiplication:
TBeort finner o Tooor (Y)
finte (1t,
fmute (n (Tt,,))
(tt) (n.)).s)
((t)
02
8j
{ ohs E cor fmut
for
0₂
0₂
Σ S obs Ecor mule (n s));
02
Oa
funte (n sy)
S (n 8j)
0 1
0
X;
02 Th Ti Ozd
0a
Ozd
C.4.4 Matrix Multiplication
In this step, we compute Emiss..
Following a similar procedure in C.6.1, we construct transformer block TBiise. This attention layer utilizes 4H4 -head attention block to obtain:
33

---

TB TBoot o flnner TBoor (Y)
(It xs)
Smutt (Tis xj)
((1)
((t)
Oa
S
(E obs
10-10 D smntt (na s));
0a
4
Σ Σ cor obs Ecor (m))ss))
misss Funlt (noo Smute on Sj)
0g
Smate (m Sq
0 1
0
X;
Od
0a
O3d
Odd
Then, with a linear transformation FFN layer, we get the final output
[mutt (Yh, X,)
Smute (mh x,)
((t)
0g
S
Ond
0sd
TBuisa o TBoo finner o T TBoo
1
1
1
0
0 X;
1
0a
0a
T/ Tl 03d
O3d
where
5] =S + Do (na :s) + 170 (Xobs)
;
+ f (naa, s)), - fmult on s,) fmult (T) xz).
Comparing this expression with (17), we conclude that one major GD update has been completed. We represent the later major GD steps compactly
fGD = T Buiss Beort finner T TB
C.5 Error Analysis and Transformer Configurations
In this section, we analyze the error induced by using transformer architectures to unroll the gradient descent procedure, and derive the corresponding transformer configurations to formally establish the result in Lemma and 8. Lastly, we combine these results to finish the proof of Theorem 1.
Above all, we should notice that, by our construction above, leveraging transformer to approximate each step of Auxiliary GD also induces noise. Consequently, similar to (14), in each auxiliary GD step, we incorporate an error term and represent the update as:
u (Maux
u (Kamx) OV au, (u (Huuz)) 50
(koux)
(18)
where cf (kaux) represents the approximation error term in each auxiliary GD step. We can also state corresponding Lemma that sharing the same proof strategy with counterpart in major GD (Lemma 3):
Lemma 9. For arbitrarily fixed time (0, and given an error tolerance 6o E (0, 1), if
can control || (haus) || S Co. then running the auxiliary GD in (14) with suitable step size for
Ki( (obs))
+
|b||
K Kanx
log
2
Amnin (Tobe) Amin (A)
iterations gives:
u//2 5 x( M(Dors)) +3 Co.
l (Kar))
2
34

---

C.5.1 First Major GD Step
In this part, we analyze the noise introduced by the construction in C.3, and corresponding transformer architecture configuration (i.e. D, L, M, B).
Bounding Approximation Error Lets(() denote the exact major GD update, and 5(1) represent our approximation. According to (15), in first major GD step, transformer blocks are utilized to compute:
8 E Eouhsmute ( Xobs)); Liuite (Ne, Xj):
for j // as an approximation to:
s TJt (as (Demsss
Σ cor
Σ-1 obs
Xobs vz))
In first major GD step, auxiliary GD is responsible for computing E (1e (lbbs))) From (16), each auxiliary GD step updates as (here we only analyze the later auxiliary GD step, and we use ut for (Kau +1), ñ (Haan) for the sake of simplicity):
H-1
= u, + ((iaa(oob)) - Σ
7m 1 kl mm
m=u TETu
and we approximate using:
i Σ
Σ Yi f1i k- m) Afmlt (0, us) + fmult (Omior, xx).
keZoba
|1Eoll2 S co. Lemma indicates that setting To ensure control over the error ||ū u+ u+112
3/ H suffices. This requires Luuutt O (log [AN]]AllE) co
)) itera-
aux,t,
tions.
With each step noise controlled, the entire auxiliary GD procedure, combined with the subsequent Moba). According to Lemma 9, setting
matrix product blocks, yields Ecor 5-1 E.11
obs mull (Te Xobs
Dohs) +
17. Xobxl //2
Kanx
log Amin Tobes Amin (A) Co
2
ensures that:
I T cor Eo (J),k,, Xobs
Mobs) Σ
E cor
obs Nobelll S Ecorl1z
kll
(obb)) + 3
(19)
2
This provides an approximation I Sobe with controlled error bounds.
Finally, the fmuit module approximates scalar and vector multiplications complete the first step of gradient descent. The overall error for each j is computed as:
|3 S 1l2 S fmute (Maa,, (M).mss) 04
+ ETT Eoo fmult (n, Xobs) - ΣT Eote naay (x
cor
+ mmutt (7),
77x 1|
16 (obs)) + 3
< Eoor 1|2
co + 2V ddduult.
2
By setting ep step count of
(1 Eeor (x= + +3)/H
which leads to auxiliary gradient descent
[100 Ecor ( (oha) + 3) VddC5 data
K Kaux,1
(As (oob)) + log
Amin (T Pobs Amiul
and by Lemma 6, setting Emult
(8VdN)
which leads to
Imultt.1 = 0 log d.N
successfully control the error (12
([5(1)
5(1) || S _ <
35

---

Configuration of Transformer Architecture for Approximating the First Major GD Step We finally summarize our construction by characterizing the configuration the architecture:
The input to the transformer is dimension D X H with D 12d + de + de +
In each auxiliary GD step, we use transformer block to form the matrix product and some fmult modules, requiring a total of Lmult,aux transformer blocks. We need perform Pi auxiliary GD steps. After completing the auxiliary GD, additional fault modules are used to compute scalar and vector products, which require 1ult,, blocks. Thus, the number of the transformer blocks is bounded by
= Kaux, (1 + Luxut,,ux.1) + 1mult,1
=0 n (AA))(TTan) (ht) log? [dx(A)rt ((d((A))(Eots)) (ooba))) ))
The number of transformer blocks is bounded by M = 4H.
With the constructions above, referring to Lemma 6, the norm of the multiplication module is bounded by (d(( [x][ sc + ||s|| co)l );, and Lemma helps us bound X and || Referring to the attention module constructed in C.6.2 and C. the norm of the attention matrices bounded by ((d(r2 + Amax (A ))); and considering the weight matrices in the FFN, since they only do linear transformations and only have most (cd) nonzero weights, their norm are bounded by Odd). To sum up, we have the norm the transformer parameters bounded
(V (r² + K(A)R((T obb
And this finishes the proof of Lemma C.5.2 Major GD Later Steps
In this part, we analyze the noise introduced by the construction in C.4, and corresponding transformer architecture configuration.
Bounding Approximation Error In later steps, according to 17, we use
g+ = S ET, obe Ecor ( 8) Me
+ Emiss (na s) - finlt (ni fmull (71, Vt)
to approximate
=s mel--aa - ooh Scor OltMooond + a 2 Eiiss + 0758 + VrJ.
E We first consider the auxiliary gradient descent which computes E-1 E mult an s). Compared to the first step analysis, we simply replace Xobs with So we can completely follow the procedure in (19). To control [ u |[]|| S Eo, we the corresponding inside multiplication module error as
60
Emalt,aux
(H3/2||A1I 11/23/Vd
which requires
Ecor Leorllalls| dN
Lmult, = log
Combining the auxiliary GD output with the following matrix multiplication blocks, we obtain
ΣT obs Emort (ma.. By Lemma 9, with Kaux ne <<Eobbr )+1 log( Xmia Taba Dama (A)eg )7, we
near Eose |r
have
16 k(Eoba)) +
Σ E E-11 - Deor
(m) (
n E.n E Eobe
Slors112 S Σ Ecorlie
cu.
2
Next, decompose the overall error term. Recall that
CiCiddocoo fmult (1) O'ty j.miss)) +
Eor obss
Σ
(-11
f-ult fiult (1)C0e, (xobb))))
36

---

similar to the analysis in the first iteration, we can derive the error bound for approximating s as:
Σ-1 | s7 2 - Σ oor obs Scor (na Ma T D
cor obs ccor8|12
e (obs) 1/, Oroondd (xobs) 1|2
+ miss mult (1, , 1/ misss
+ mmllt (neo,
+ mutt (Te, xj) - 1):2j] 1l2
16 (obs)
+
6o + 2V de + c + Emissll2 demult.
coo 11
2
2
where the second term, Monnd bounded by 5
Xobs), (Xobs))
computed in the first iteration bound and is thus
By setting co (2) Eor ( +3)/H) e, which leads to the auxiliary gradient descent
step count:
(had Door
+ 3) H ||s||2
K K (k(Cobs) + log
Amin ( Amin (A)e
and setting Emult (8 Emiss ||F + +2) e, which leads to
Lmult,+ = =0(10 [dH([ss|loc + Cduta)) E
Esmisollo
successfully control the error ||5 ||s - st Sc.
Size of Transformer Architecture for Approximating the Later Steps Major GD We finally summarize our construction by characterizing the size of the architecture:
The input to transformer is dimension D H with D 12d + de + d+ +33
In the later major GD steps,, GD step, and a total of Kaux.+ auxiliary GD steps are required. After completing the auxiliary GD, we perform additional matrix multiplications (e.g., multiplying vectors by Ecor. ΣT oorr and miss)), which require transformer blocks for attention. Subsequently, fmuht modules are used complete the major GD, requiring mmult, blocks. Thus, the total number of transformer blocks required for each subsequent major GD step is bounded by:
+ aux,+ transformer blocks for each auxiliary
L = Kaux,+ + (1 Luultt,uxx++) + 3 Lmult,+
(; log? (ddx(A)(Tobbs
OtE
The number of transformer blocks is bounded by M = 4H.
Same as the analysis in the first step of major GD, we have the norm of the transformer parameters bounded by
(VHd (((++(AA))))) ((A) (Tob)o))
And this finishes the proof of Lemma 8. C.5.3 Proof of Theorem
Proof. We formally construct the conditional score approximation transformer as follows: 5(V,, Xohs) Sout SCD o... fCD ofGD, Jin(Ve, Xobe)-
(K (K-1)xfoo
Recalling that Kr ><aa"oond Doond + o²)), by the major GD convergence result in Lemma 3, to ensure [s $112 < o, te, the total major GD iteration number required upper bounded by K 0 Ki log (()(((A))((Too))) )). which obtained by substituting with o, krt2
37

---

Utilizing Lemma and and substituting with of ( mc+2 e, the following transformer configura-
tion can control the error in each major GD step:
L= (1O(m(((AA)) ((( lob)
(A)(TT)) )),
D 12d de de +3
M 4H,
B B=O O (VHP ((A)a((T (Tott)))))
where l is computed by K times the transformer block required in each major GD step.
Finally, by substituting o; Nt with 0io Net. and considering the truncation range which induced by the decoder (R O " (Hd((A)) ( ob)))). we obtain a uniform bound for any t E [to. T7. Taking supremum over all admissible Tabs C I, and leverage the relationship that n( (<Eond)) = K(A A (o(Tobs). finish the proof of Theorem 1.
C.6 Construction of Attention Layers In this section, we construct the attention layers used in the transformer architectures built up in C.3 and C.4.
We utilize the added
x;
Xj
to construct different types of interaction between different types of samples. (i.e. When we want attention exclusively among observed samples or missing samples, we use the construction method as described in TBobs below; When we want attention between observed samples and missing samples, while setting all other interactions to zero, we use the construction method as described in Boot below.
The intuition of (16) suggests construction of a multi-head attention layer. Formally, for an arbitrary value of m. construct four attention heads with ReLU activation. The indicator function 1[ - m} can be realized by calculating the auxiliary product of e, of time embedding. To see this, we observe
1 (2r² (212 f (0-1111 j1)
e,
2 (2r² -le, - -,,1
Therefore, it holds that
1 { = (m)
a ji m -n)==1
since f Assumption ensures that the time embeddings uniquely identify discrete time gaps through their pairwise distances. Directly approximating an indicator function using a ReLU network can be difficult. Yet we note that - jl can only take integer values. Therefore, we can slightly widen the decision band for the indicator function. Specifically, we denote a minimum gap f2()). Thus, we deduce
min:
{ fote, p2
1 ((() 2
1 A, ,² 4
1 P()) 4 }
1 - ýl m
We can use four ReLU functions to approximate the right-hand side of the last display, and simultane- ously take different type of interaction types into account. We use another indicator function (which can be realized by the Os and 1s added above) to represent what types of interaction we want this specific transformer block.
38

---

We construct trapezoid function as follows:
1 j| m) ReLU
8
e 6j p² + / (m) + 1 )
T
1
A
8
e 6 ,² 1 (mm) +
1
ReLU
8 ) 1 5/))
A
8
1
ReLU
e; 6j
, + 2 (m))
A
8
e p² ( }
1 2
ReLU
C.6.1
Construction of attention matrices related to the observed part We construct the attention matrices for TB here. For particular m, we utilize
(Q') T diag (fouxca Ld., OcteOexxd ((6)) 0, -r- 2 (m) + A 0, 0 (c0onaxcaa)) 1)
1
Oc4d+do + (2d
(4dddo+ 8 Am
(4d4 +.de) x
V
Odx(d+d++dy 0 (6d+ (2dd)))(22+d++4))
(Gd+3)) (6d+1))(6dd3))
(66+1)xx
and
((o.rit Idc, Od2 xd. O (6d)x O, -r2+ (m) A ((5))))) 1):
(Q9) TK? -diag
2
8
V -
and
(Daxci Lec Ode (66))x 0, 0r- 2 (m) A, ((4)) 1).
(Qa) K -diag
33 --
and
(oosee It,, Xd,x ((d))x(6)) 0, -p2 2
1 A ((5)))(50)) 1):
T
T -diag
v V
It is casy to verify that
4 Σ (TT(Q) K K"y, = 1[ji - = mi, E Tons].
We can claim that 4/ attention heads and identity FFN enough for constructing this block.
39

---

C.6.2 Construction of attention matrices related to the correlation part We only need to do some small changes to Toobs. For particular m, let
(Q') '
(Descaiuu, Inl, Ode (6d)) -6
1
+ - A. 2. 1
( (15d) )
= diag
2
((2d+de+d)) ((ddd++)))) (dx(d+++
(2d4d. +de)xd
0 (2d44
V
3 Am
(6d++33
(dd
(8d+1)xd
O 18d+
and
(Q²)
(Doord Ida, Odexde ((d)) (od), - + {
1
1
1
(650)((50) )
= diag
(m)
+
-
A.
A.
0
2
8
2
2
V2 =-=
and
(Q²)
= diag (Dooas Id., Od,xa, xd. (6d))
1
1 A,
((5))(6a)) 1):
-p² + (m))
-
A
2
8
2
V3 -v
and
(Q") K
(Dooat Ide, Od,xd (6d) -p2 + 2 P()m
1
1
1 A,
0(60)(((0) 1)
= =diag
A A,
4 2 2
v-= V
It is easy to verify that
Σ (T)) 1iji -- = -m]lltie Iobas E Imiss] U {i Imiss, E Ioha]}.
We can thus state that 4H attention heads and identity FFN are enough for constructing this block.
40

---

D Proofs of Theorem 2 and Corollary 1
In this section, we provide the detailed proof of Theorem and Corollary 1.
Firstly, we introduce some notations specifically for this part for sake of simplicity, We denote our training set with i.i.d. samples as
D = {x (8) ji=1
miss: {(x
(6)
(bo)) = ((x
(0) y
(0) )
We introduce the corollary below which will act as an significant role in our later proof:
Corollary 2. By choosing the transformer architecture T(D, L, M, B, R) as in Theorem 1, the early-stopping time to 1 and the terminal time T (log n). it holds that
Hd kig R² (A) (Tab))
E [(x(9) Va [R(S)]
log (H (A)(( oobbba
where Kto = x(a Eoodd + I;
The proof of Corollary 2 is deferred to Appendix E.
D.1 Proof of Theorem 2
Although our assumption on Gaussian processes does not ensure the Novikov's 's condition to hold, according to [Chen et al., 2022], long have bounded the second moment for the score estimation error and finite KL divergence w.r.t the standard Gaussian, we could still adopt Girsanov's Theorem and bound the KL divergence between the two distribution. We restate the Lemma as follows:
Lemma 10 (Corollary D.1 in [Oko et al., 2023], see also Theorem [Chen et al., 2022]). Let Po be probability distribution, and let Y {Yi)eJO and Y (EJOT, be two stochastic processes that satisfy the following SDEs:
dY. ()d, dW,, Yo Po: dY; s'(Y,, l)dt + dW,, Yg Po-
We further define the distributions Yi and Y by Pt and pt.. Suppose that
f p(()))((s- (a
(20)
for any t E [0, TY. Then we have
KL (pr)) I [01 2 p((=)]((s sx) t)]1 dt.
Equipped with Corollary 2 and Lemma 10, we are ready to prove Theorem 2.
Proof of Theorem 2, Firstly, following the proof Lemma 12, we can easily verify that for any SE L,, B,
J. pi(v.. y)), - V log (ogg,((v 1 $ o7
Thus, the condition (10) holds for all / E [fo, TJ. which means that we could apply Girsanov's theorem in this time range.
To further distinguish the SDE defined in (1), (2), and (3), we denote the distribution of x,, Vt. Ve Pt, ,, PPt. respectively. Additionally, need to introduce another intermediate backward process between Pt Pt follows
dv
Vlogpr- (y) dt dw, with v² NN L Fymmiim
and we denote the marginal distribution of vi- (conditioned on y) (-,)).
41

---

Equipped with these notations, we can decompose the total variation between P and Pro as By [V P - By [V(P, + TV 11 + (V P% + TV PT) = By TV(P, + TV(P. Po TV(P% .
(21)
We denote {5: ( Eanical as the eigenvalues of Eoond, and can do eigenvalue decompositions to Boond 2oond QEQ', where (E) ii E
Considering the second last term,, by Data Processing Inequality and Pinsker's Inequality (see e.g. Lemma 2 in [Canonne, 2022]), we have
[TV Pto ))
By (KL to PI (Pinsker's 's Inequality)
By [KL(PP]NN (O, Im
(Data Processing Inequality)
2 Ey (KL( Idj exp(-I)
Fyy log( Exond) + tr Egond) + y TT E-1 Evor ET obs
cor Soboy) - d) Imissl exp(-77
V log( Econd tr Econd) Ex-AV(o,1/2 /z E 1/2
obe cor DT 2 obs z d Imissl exp(- -T)
log( Econd) tr Econd + E E obso Ecor) d Imiss exp( (-T)
log( Econd) + tr Econd + ΣT E- obs Ecor) d[Imiasl exp(-T)
Hd (lloi(Smin) + tr [misss) - d|Imiss exp(-T)
$ v Hd cxp(-T),
(22)
where we leverage the close-form solution of the KL-divergence between two gaussian distributions in the first equality.
Regarding the first term, by Pinsker's Inequality and the close-form solution of the KL-divergence between Gaussian distributions [Pardo, 2018), we have
By (-l() Peo (-1x))
Ey (TV ((bb))
$ Ey VKL((Po ((Hx))
Ey /KL(N(a DTer E, a Econd o IdjI WE y, Econd))
s SVF, Ey [KT,(N ET Dibay aio Econd + 0% I dj I[ EJLy. Σcand)
Leveraging the close-form solution of the KL-divergence between two gaussian distributions, we further have
Ey (L ET, DoLy. y, 0to Econd + [ cor Eobyy, Econd)) ]]
||?? Σ Fond + tg
- log
2 tr -I- cond (a Econd + o I))
Doond
x-1 T + (1 - Oito) 2 [y -n-1 Ecor cond cor obsy y - -[[mmissl
obs
Considering term A, we have
|0 Ecoud 0² I Eeouel
10 I oz E Dold
d[Iminal
Ecouel
(a + o
Em1
42

---

Then we obtain
A= Σ
log(o + + +
Regarding term B, we have
= tr( E11
cond (02 Doond + T2 = 3+ 1Q (Q=
= =tr( =+ 1 1)=1
d/Lmiel
Σ
(a + o,
S d Tmissl + dl Imiss S
Considering term C, we have
C (1 (it)) " [y En E oor cond
D-1 D D obsy
(1 0ity )2 2 10. Sotal E 1 cor I-1
cond IJ Σ
obs
ohs
= (1 ont)) tr tr(E T
cor Dobs Deor Doold):
Thus, with at. = 2 o = vI - e0 we can take to = OSEmin 2). and
P.o Ey (( Sn (Hd))
(23)
Combining (21), (23), (22), and invoking Lemma 10, have:
By (PPP, z n (Hd) + Hd + *(A)) exp(-
[ to 2 / pv(v.)))(S(vv....) log (Bvv,d (24)
+ Ey
An 3 (Hd) + Hdlgmm + +(A)) exp( (T) + RR()).
Plugging in the result in Corollary 2 (taking T Olog and to O (( +)), we finally obtain
< H dr. (T (Tcond)ri )o obx)) log²
ED By TV(P, Pr))
]]
Ō Hd Dcond 222 //tx)/ Vn))
where I Teond
r misss
cor ohs r ceor- We complete proof.
D.2 Proof of Corollary 1
With Theorem established, the conclusion in Corollary 1 goes straightforward.
Proof. Firstly, by the definition of total variation distance, we have the relationship
Pro (xinte E CRR(aa) - P(x CRR-aWl (VP. (.(xaba), P(: (x)b))).
Following the decomposition in (21), we can obtain
TV(P,, P(. ((-))))))) (P (CV(PP(-x))))))) (-1x))))) Poo (-1x))))) TV (P.. (-1x))m) Pf (-1x"ba))
+ TV(P. (-1xa Pio (- (-1x)
43

---

Regarding the right hand side, following the derivation in (24), we can bound each term similarly by taking to = (n and T = n).
For the second term, we leverage the close-form solution of the KL-divergence between two gaussian distributions:
TV (Pro ( (+/2bbb)) Plo (- Xubu))
log( Doond )) + tr (cond) + oohs T "obb Coor Door 2 Xoba d mmisol exp(-I)
log( Doondl 1) + Econd) + xass D abs
cor D 2hs Xobs
miss
24
Hdlog(Amin Doond) ) Xobs "Eobe Ecor E E
obs abs
1 Vobs "Eobe Ecor ET obs Kobs
$ Siif
For the last term, recalling the definition DS Pxors Prise (immi)) 9),, we have
Diii
TV(P (1/xbbe> Pro (-1Xoba)))
V 1n
T 1
2 J./. (e(ve(x)hm) I[S(vv), Xibs 1) 1o
X Epis)
1/0 (VVV t) V loog PV
< B
T
E x [e(y)] Egg E Ey [e(y)]
< E
| Evs.yr [1|((Ve.,,)) log (v
sup
S EEDON [R(S)] V DS( DS(Pxion P9)
$ dit
V DS((x Pii i).
Let
((xxba) = max V Amiin (cond)xx
X T Ecor Ecor cond
- E-1 ET 5 obs
obs?
Tx-1 - Door ΣT --1 obs
xabs
obs
For the first term, leveraging the result in (23), and the decomposition of term C the Theorem 2, we have
proof
(V(P(- Pto (-1x)baa)) S n - 2 (Hd)+ + n (xxx)) S (n)
Caur + n (x(x
Finally, combine all the bounds above to obtain
E PLr (ximis E CRi P(x CRR-aall
S EDD (TV Xibba) (-o(-1xabee) + (V.. (:)x)ts)) Pto (-1x))1s)
+ TV(P, (-1x (-x(((x)aa))))
S n 3/b (x(()))))) + diiff
/DS(Px Pxin.>> :9)
(n) S Caff
/DS(Px Pini Pxino () + (xxx)))
and the corollary follows.
44

---

E Proof of Corollary 2
In this section, we provide the detailed proof of Corollary 2.
Training Loss During training, given state V. + 02 O 2 NO0 Tmmioo vo we aim minimize the ideal risk function:
Xmiss
/ Ev, [[]s(V Xobs, 4) Vv, log Pe (V. (Ve)Xobs)
R(s)=
[2] dl.
(25)
However, in practice, the objective (25) not directly accessible. According to Lemma C.3 in Vincent [2011], an equivalent objective function (s)), which differs from R(s) only by a constant, can be used for optimization:
((s)
E (Xmiis Xohne) [vv [[]s(Ve, Xobs, t) Vv log te (v.l dt.
(26)
Here, 0e is the Gaussian transition kernel of the forward process, satisfying V Vlog ói = (V.-au aavvo
ot
Thus, we can leverage the corresponding empirical loss (11):
$ arg min L(s), where L(s) Σ (xx()) ( SET
where the loss function is defined in (12)
((x y (1) (5)
Ev.vo=xx [[ls(v..y,t) Vv, log vv dt.
E.1
Steps for Proving Corollary 2
E.1.1 Risk Decomposition
The proof procedure is analogous to the proof of Theorem 4.1 in [Fu et al., 2024c], provided in Appendix D of the same work. Our goal derive a bound on [(xx0, [R()]] We denote
the ground truth score function as s* and set R(s*) O.
Following the setup, the risk can be decomposed as:
R(S) R()) R(s") L()) L(s*),
where I is the score function trained on dataset D(V using the empirical risk. By creating n i.i.d. ghost samples
(D')(m) (xx
( ~ Pxmiss Noba!
the population risk of can be rewritten as:
((x (0:)) ((x(1) , (0:))))
R(S) R(s*)= (
To bound the rewritten population risk, we can further decompose it by analyzing its behavior in truncated area (aligning with our score approximation analysis in Theorem I, we analyze the error conditioning on the event C, [||x||2, ||v||2 I C Caatal). and the error induced by truncation.
The truncated loss function is defined as
Arunc (x,y:
Ev. [ls(v - Vv, log de (velvo) dt.
(27)
Accordingly, we denote the truncated domain the score functiinn by X' = C- Ca C5 CantaldiZmmal [-C Ce and the truncated loss function class defined as
X
S(Cåata)
{ drmme (., s) R|s T
(28)
Define the following intermediate terms (s depends on D(V):
C
[(x y (0)
f(x (1) y (0:)
45

---

Airune
Arune (x -(0) - prune (x(11 (y(0))8))
i=1
((xx()))((e) - ((x110) y ;,)))) (t))
L2
i=1
and
Ctrune
grnnn (x (0)) y'() - finine (x , -(()))
71
The decomposition for the expected empirical risk over D(V then becomes:
E
[R($)]
EDOw (DrjO [C2 L
[ L1]
Dto
R
(Dry))))
L.
Ept [G
(29)
The terms A, B, and C respectively represent the emror incurred due truncation, approximation in truncation, and the in-sample empirical risk expectation.
E.1.2 Bound of Each Component We first bound the data range with high probability. The proof of the lemmas stated in this section deferred to E.3.
Lemma 11 (Range of the data). Given a sufficiently large data truncation range Co Canta > O, we have
IP U 11, ||v l12 S Canta)) 2 1-ddd
where dn 2n exp { -C 8(Hd- 11) }, and C2 the absolute constant defined Lemma 16. We also have P[E] 2 5d. where 1 8n.d
In the following analysis, for the sake of simplicity, we denote Cs: = + [[Feo))22<(A)
(Toba) which origins
from Lemma 4. Then state Lemma 12 to bound term A in (29), which is the counterpart of (D.12) in [Fu et al., 2024c].
Lemma 12. For any s T.
(r+t)
Ex.y - & (os(CCa (x,
e"rum
+ Hd]
is straightforward to conclude that
A S (a + Hd
(T+1)
T
(30)
to
Then we proceed the term C in (29). For any E T, we have the following relationship
C Expc
[6]]
n
E
(1) (1xx00,,(0003 ((003) ((x!)) (0;)
n
i=1
(((x(() (r())))),((0)0)))) (r((0)s)))))))))))))))))))));;) (1x11) y0(0)ss))
< Ep:,)
$ R(s),
the inequality holds due to 8 minimizes Ć. Taking minimum w.r.t. T. we have
C < min R(s) min xET
Ev,y [Ils(v.,,,)) V., (v. dt. (31)
sET
46

---

Lemma 13. Given error level e (0,1),
- to (1))
min R(s) S x² ( log SET
Vool((Cr (Ittaaataa* Hd]
+++ t0
T
Finally, proceed to the bound of term B. In attempt to providing the bound, we first need to calculate the covering number of the loss function class S(Caata): and correspondingly, the covering number of our transformer architecture function class. The covering number defined as follows: Definition 2. We denote F,, II to be the 5 -covering number of any function class F w.r.f the norm 1|, i.e.,
F, II) arg min (3(1) F. VS E F. Fi (J) II; fl 53..
A modified version of Lemma 23 in [Fu et al., 2024b] provides the following result on transformer covering numbers:
Lemma 14. Consider the entire transformer architecture F = -,, B, R) (ie. with encoder and decoder). If the input to the transformer satisfy || 11 | S C the time embedding e and the diffusion time time-step embedding Φ(t) satisfy ||e||2
r, llo(t)ll2 < Caut and r, Cain 3O(HHd)
then the log-cove rog-ooveinng number of the transformer architecture is bounded by
log. FF, (1F,,0) & D2M (17
(I.² (T2 log (BMNRC data C2) log
BM
data
data C3 )) (
6,
Then, we can leverage the following lemma to calculate the covering number of the corresponding truncated loss function class.
Lemma 15. Suppose ((1) ((2) E ET(D, L, M, B, R) such that ([S( (ve,y; - -(( s 8,
for any 2 yll,, |x||| S Canta and 2 lo. then we have
Jinnn (5(1) prrane (S (2) S 46, T log (1))) (CrC data VHd).
Equipped with this lemma, it is straight forward to derive that
BMTHdCd Cr:
log (6;; S(Caata) 1,11 (lx) DPM
L2 log (BMN RCCtta data Cx.) + log
date
6
and 8 satisfies
T log )
61 46,
(C2Ca data + VHd).
to
Invoking the bound provided in (D.16) [Fu et al., 2024c], we have
1 log, (br; ll (loo) /" o, "d + 78
T
BSC+AL
"
(T to
BMLH data Cx:
CC+A+
D'M
L log (BMNRC) data C2) log
data
n
+88,
T log
(
(Cr VHd).
(32)
Combining the bound of A, B and C ((30),32),, (31)), we can leverage the empirical risk decomposi- tion (29) finalize the proof of Corollary
E.2 Proof of Corollary 2
Proof of Corollary 2. By (29), (30).(3)) and (31), we have E EDD [R()) A + B C
22 5a (CDC&
Hd] (T+
(T' + 1.) D2M [.2log (BMN
BM LH
daaa Cs:
data Cr:) log
8y
) (Cr data + VHd) +
1
+ +88,
T log to
T log
to
47

---

Plugging in the configuration of our transformer architecture in Theorem 1, and take
1 vni 0, nC data Cg
=
I (log(n)))
the inequality above gives rise
tg Hd 6² (A)))22TTa))) log( (H (Aoaa(A)ssii ():(TT))a))))))
p(n) R(()] S
E.3 Proof' of Supporting Lemmas in E.1.2
Proof of Lemma 11. We first state the polynomial concentration lemma for Gaussian random vari- ables.
Lemma 16 (Lemma 24 in [Fu al., 2024b]). Let be a polynomial of degree p and NN Ia). Then there exists an absolute positive constant Cp. depending only on p, such that for any < 1,
P [1(xx) E[g(+)]] (VVar(g(())
s exp (-C,82/0)
For random variable (O0 Eg), consider g(t) = || -1 have
B[g(u)] tr(EE) Eig [ s 3| Eoli
Applying Lemma 16, conclude that with high probability at least 1 - 2exp(-C20),
||ا|ا E| s 5. Var(llg()) 13) S V35l|Eol1F.
Considering Vr. we have E) = (T (T miss & A) + Ii and for x, we have Eg = r Therefore.
[| 2| 5 V (r(Ei) 2
Ir(T mis)) (r(A) d Imiss! + misso T A| 1/ + d[Imiss)) 1)
A| d + V |/ ||]] || + dl d(Imiss))
s V(Hd+ A|e + 1) + v28),
the last inequality holds for 5 < 1. Similar inequalities hold for ||x(+) 1|2
Consider Ca 2 2 (Hd+ [)[[ 1) and let 8 8Hd+ D Ft1)*
(CA )²
We can then obtain
union bound. With probability at least 1 2n exp 8[H A| }
max |xx 10 |||
|2) S Chata
-C
We finish the proof by setting 5n.d 2n exp
48

---

Proof of Lemma 12. For any s E T (s can depend on x, y),
Exy [[((x)y:)) frume (x, s;
(1SV - v Flogds oo y)dxdyd
x.y
JSg Evelvo=x
[Is(..,,01B+
Vi VvVo
2
1cs
p(x, ydxdyyd
of
- [ R(C)) + to
0,z
of
Co
dt
f(r of dr +
2
I 6 )2
s 6dl(CCCattal }2 Tt a dt + voa 1" or
Hddt
10
S Vv data))
12 + Hd]
or dl
20
/T
2
= -oall[Ca Catta)
)² + Hd]
dl
to
et-
s Voll(CsCCaaa) + Hd]
T
to
where R2 is the truncation range of the decoder, and we apply triangular inequality in the third line.
Proof of Lemma 13. Since & E T, we can invoke Theorem and triangular inequality:
min R(s) = min SET
J. Ev,y [Iss(ve.,, - Vv, pvg dt
SET
T
Evay [B((),))) Vv, logg(v..l))11.... di
+2
Ev,y [[]s(V.. v. log (v dt.
ta
oi 2 20d I" y.
S
2
to
to
22 1' [T
E1/2 [VV (v de
/" [[V(.. 12dt
< 2
o, 2d2 + 28,
30
20
T /'
+ +2V5a Ev..y [V log lo dt,
where we apply Cauchy-Schae -Shwwaz inequality in the second step, Jensen's inequality in the last step. For the second last term, similar to the proof of Lemma 12, we have
S" EV. [Is((v.. y, +)]2] de (Cá Cx) I' o, dt.
10
49

---

For the last term, we have
1( Ev.y Vv, log dt
$
o, 'Broy [llve an Deor Sobu y12) dt
2
-BBxya-N(0.. [||x Ecor D ooo
11 dt
to
T
2
"EunsV(.
|| + o Hd
t
{ dr
2 [o] Hd tr E miss D Sho Ecor
2 of Hd + tr([ mis) tr(A))
o dt,
"td
10
where we utilize the positive definiteness of Σ DeEr in the last inequality. Combining all the terms together, we have
oor obs
/" or + 2/888 [" EvEyy (1 + Vv. log (v de
min R(s) I 2
#ET
tu
S 2 T + log (2) to (1))
+ vōd (Cssaatta) + Hd] ( to
Proof of Lemma 15. We have
E vo==x
-
(ve, - Vlog - V log d. (v.
E Ao..
|B
-33 ,1 + - 11
(150) -
((()) (g())
( (ax
2z
Ee 1
01
||(00) -((2) 0:2
2z
E2 Noo
(8()) $ (axx ,,
0+
||(300 + ((r)) +Zz,,;
2z
26, Ex
0
$ 28, 22|6||2 + 20, EE [||]||]]
S 48,0, 3 (Cx HHd).
Therefore, we obtain
( grrun (1(2)) < (CrCC VHdd I' o, 2t
to
< 40 T +looog log (t)) (C&Ch
data v VHd).
50

---

C564
Ground Truch
Mika ng Data
22
Figure 4: Comparison of imputation methods on the Electricity dataset, with 95% CR.
F Experiment Details
For our numerical experiments, we trained the models using a batch size 64. Our adapted DiT model architecture used hidden size of 256, 12 transformer layers, and 16 attention heads per layer. We utilized the PyPOTS [Du, 2023] framework to implement and handle hyperparameter tuning for the baseline methods CSDI and GP-VAE This tuning process aimed find the best settings and ensure the models had a comparable number of trainable parameters. Experiments were conducted on hardware consisting of an NVIDIA RTX A6000 GPU (48GB) and Intel(R) Xeon(R) Gold 6242R CPU @ 3. 10GGz. We report all the results as the average runs. Our implementation of DiT for imputation is attached in supplementary materials.
F1 Real World Datasets
Dataset Descriptions. We utilize two real-world datasets, Beijing [Zhang et al., 2017] and ETT ml, benchmark the imputation performance of DiT. The BeijingAir dataset comprises hourly measurements of six air pollutants and metcorological variables collected from 12 monitoring sites in Beijing. The ETT ml dataset, part of the Electricity Transformer Temperature benchmark, records clients' electricity consumption data, including power load and oil temperature. Detailed statistics for both datasets are provided in Table 5.
Dataset
Total Samples
Sequence Length
Time Interval
Number of Variables
Air Quality 1168
30
1H
132
Electricity
2321
48
15min
Table 5: 80% of the data is used for training, and 20% for testing.
Results. We report the Mean Absolute Error (MAE) in Table 6, the Mean Squared Error (MSE) in Table 7 and the Mean Relative Error (MRE) in Table 8. Results are shown across different missing data rates (10%, 20%, and 50%) for both datasets. The experimental results indicate that DiT consistently outperforms the baseline methods on both datasets, demonstrating its effectiveness, and our mixed-masking strategy can also enhance DiT's performance on real-world datasets.
Figure 4 presents comparison of imputation results the ETT ml dataset, where randomly select samples from a 50% missing data scenario. From the plots, it is evident that although both DiT and CSDI generate CRs that largely encompass the true data points, DiT achieves tighter bandwidth, leading to improved uncertainty quantification performance.
51

---

Model
ETTm. 1 (Missing %)
Beijing Air (Missing %)
10%
20%
50%
10%
20%
50%
[Tsshiro al, 2021]
0.1448
(+00114)
(410007)
(+00))
(+0:01
19)
GP VAE [Fortuin et al., 2020]
0.2786 (+0,0077)
0.3267 (00,)
0.666 (003)
0,4152
(+0.0088)
0.4401 (+0.0080)
0.5268 (#0.0054)
DiT
0.1269 (:0.0076)
0.1377 (20,0995)
0.1543 (20.0102)
0,1753 (10.0094)
0.1815 (a0.0208)
0,2057 (20.0145)
Table Time Series Imputation MAE Results
Model
ETT ml (Missing %)
Beijing Air (Missing %)
10%
20%
50%
10%
209
50%
CSDI [Thchiro al., 2021]
0.0155
0.0688 (=04006)
(40.0106)
(40. 1726)
9926
(-0U039))
GP-VAE [Poctuin al. 2020)
(=00044)
(=00067)
(#0.0127)
0.4096
0.4777
(#00179) (7770
0.7017 (+00189)
DiT
0.0534 (400003))
0,0606 (=0.006))
0.0684
(20.0070)
0.3683
(40.035)
0.4025
(200424)
0.4255 (=0,0670)
ITT w/ mixedd-makking strategy
0.0502 (+0.0055)
OLASNN (=(L00NT)
8,7711 (+0.0092)
0.342N (=0.0275)
0.3864 (x0.0403)
0.4229 (5100539)
Table 7:
Time Series Imputation MSE Results
Model
RTT. ml (Missing %)
BeijingAir (Missing %)
10%
20%
509
10%
20%
50%
[Tsbiro al.,
(10.0174)
GP-VAE [Fortuin al., 2020)
0.3285 (a0.0091)
0,3882 (10.0052)
0.5478
0085)
0.5595
(20,0118)
0.5917 (10.0107)
0,7042 (110 (072)
DIT
0.1592 (*100084)
0.1701 (a0.0102)
0.1825 (=10,094)
0.2154 (a1L0125)
0.2578 (*0.0375)
0.9075 (=0.0241)
Table 8: Time Scries Imputation MRE Results
52


Reviews:
[{"review_id": "yHUiWKN9dd", "reviewer": "Reviewer_i349", "summary": "This paper addresses the problem of theoretically understanding the effectiveness of score-based diffusion models for imputation. Under a Gaussian data assumption, the authors propose a nested optimization algorithm to approximate conditional scores and show that a Transformer can emulate this process via a tailored architecture. A key contribution is the explicit modeling of how missingness patterns affect score estimation. The paper introduces tools such as distribution shift measures and confidence regions to quantify and mitigate this effect. Theoretical guarantees are provided, with non-asymptotic error bounds that scale with the condition number and missing data structure.", "questions": "- The conclusion in Remark 1 is not immediately obvious. Could the authors provide further clarification or intuition behind this result?", "limitations": "While the authors acknowledge the assumption of Gaussianity in their analysis, the implications of this modeling choice merit further discussion. The assumption enables theoretical tractability but may limit applicability to real-world data distributions, which are often multimodal or heavy-tailed.\n\nNo ethical or broader societal risks are apparent from the method itself.", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "## Strengths\n- The paper shows that Transformers can approximate conditional score functions through unrolled optimization, supported by a constructive proof (Theorem 1).\n- The proposed score approximation under Gaussian assumptions is elegant, despite the simplified distributional setting.\n- The method explicitly accounts for the impact of missing patterns and introduces tools like confidence regions and shift metrics. The method explicitly accounts for the impact of missing patterns and introduces tools like confidence regions and shift metrics.\n- Theoretical analysis provides non-asymptotic error bounds, which scale with the condition number and model parameters.\n\n## Weaknesses\n- While not a major weakness, the paper overlooks several relevant works on missing data imputation using GANs and VAEs. The claim that these models “often spell limitations in expressiveness or training stability” is overly broad and insufficiently supported. Recent methods [1–3] demonstrate that VAE- and GAN-based approaches can be both efficient and accurate. A more balanced discussion would improve the positioning of this work.\n- Clarity and flow could be improved, particularly in Sections 3 and 4, where the theoretical framework and algorithmic insights are harder to follow.\n\n## References\n[1] Yoon, Jinsung, James Jordon, and Mihaela Schaar. \"Gain: Missing data imputation using generative adversarial nets.\" International conference on machine learning. PMLR, 2018.\n\n[2] Mattei, Pierre-Alexandre, and Jes Frellsen. \"MIWAE: Deep generative modelling and imputation of incomplete data sets.\" International conference on machine learning. PMLR, 2019.\n\n[3] Peis, Ignacio, Chao Ma, and José Miguel Hernández-Lobato. \"Missing data imputation and acquisition with deep hierarchical models and hamiltonian monte carlo.\" Advances in Neural Information Processing Systems 35 (2022): 35839-35851.", "was_revised": true, "final_justification": "After reading the authors’ detailed rebuttal and considering the clarifications provided, I am satisfied that my initial concerns have been fully addressed. The authors offered further intuition behind Remark 1, clarified the treatment of GP data, and acknowledged the modeling limitations associated with Gaussian assumptions. These clarifications, along with a more nuanced discussion of related work, have increased my confidence in the paper. While the theoretical exposition in Sections 3 and 4 could still benefit from improved clarity, the paper makes a solid theoretical and algorithmic contribution to the field of imputation with diffusion models. I maintain my rating and recommend acceptance.", "rebuttal": "Thank you for the valuable and constructive comments. We respond to the specific questions, weaknesses, and limitations you have raised in the following.\n\n> **Q1**: Discussions on GAN- and VAE-based methods.\n\n**A1**: While our main focus is on diffusion-based imputation, we appreciate the references to recent GAN- and VAE-based methods [1–3]. We will cite and discuss these works in the Related Work section to present a more balanced view of expressiveness and training stability. We also note that we have already compared our method with a VAE-based approach (GP-VAE) in our experiments. Through our experiments on real-world and synthetic datasets, we observe that diffusion-based methods outperform GP-VAE. Additionally, thank you for highlighting clarity and flow issues in Sections 3 and 4; we will revise these sections to ensure the theoretical framework and algorithmic insights are presented more smoothly.\n\n\n> **Q2**: Assumptions on Gaussian process data, and heavy-tail data extension.\n\n**A2**: We appreciate the limitations you pointed out. Indeed, our theoretical analysis assumes Gaussian process (GP) data. However, as discussed in the introduction, GPs naturally exhibit rich spatio-temporal dependencies and long-horizon correlations, which pose central challenges in time series tasks [4] and are widely adopted in practical modeling scenarios such as climate [5] and financial forecasting. Additionally, as noted in lines 119–121, we allow for a wide range of kernel functions for GPs, ensuring our theoretical results remain generalizable across various settings. Empirical validation in Section 5.2 (latent GP data) and the Appendix (real-world datasets) further supports the applicability of our theoretical findings beyond the synthetic GP settings. It is also worth noting that although our theory is developed under GP assumptions, the insights gained from the theory motivate the development of our mixed-masking strategy. This strategy has demonstrated strong generalizability to real-world applications, clearly indicating that the theoretical assumptions do not limit the practical applicability of our method.\n\nYour insightful point regarding heavy-tailed data distributions common in real-world scenarios is particularly valuable and promising. Yet we suspect that diffusion-based models are relatively weak in capturing heavy-tailed data distributions accurately without careful modifications. Therefore, we leave it as an interesting and meaningful direction for future research.\n\n\n\n\n> **Q3**:  The conclusion in Remark 1 is not immediately obvious. Could the authors provide further clarification or intuition behind this result?\n\n**A3**:  Remark 1 builds directly on Corollary 1 by leveraging the distribution-shift coefficient $\\mathsf{DS}$ (defined in Definition 1, measuring how much the test sample deviates from the training distribution). Corollary 1 shows that a large $\\mathsf{DS}$ can lead to poor imputation outcomes. Remark 1 explains under **what conditions** $\\mathsf{DS}$ becomes large (or small) and **how** we can benefit from this insight:\n\n- **Under what conditions $\\mathsf{DS}$ is large:**  \n  If $\\mathbf{x}^*  _ {\\mathrm{obs}}$ lies in the region where the training distribution has a weak coverage (e.g., hardly seen similar examples in the training data), then $\\mathsf{DS}$ becomes large, making imputation difficult.\n\n- **How we leverage this insight:**  \n  Different missing patterns during training lead to different training distributions $P _ {\\mathbf{x} _ {\\mathrm{obs}}}$, resulting in varying condition numbers and consequently different $\\mathsf{DS}$ values. Training with diverse missing patterns—ranging from easy to hard—helps the model adapt to imputation tasks with varying levels of difficulty by effectively covering more scenarios.\n\nTo illustrate this clearly, consider the concrete example in our numerical simulation in Section 5. Missing patterns P1–P4 have varying difficulties (estimated by condition numbers), with P1 (most clustered missingness) yielding the highest condition number and P4 (completely random missingness) having the lowest. Specifically, our training strategies S1 and S4 differ as follows:\n\n- **S1** uses only random missing patterns during training (P4-like, low condition number).\n- **S4** includes various missing patterns (both P1-like and P4-like, high and low condition numbers).\n\nDenoting the training distributions corresponding to S1 and S4 as $P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$ and $P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$, respectively, if we consider a new sample $\\mathbf{x}^*  _ {\\mathrm{obs}}$ with missing pattern P1, it is intuitive that the resulting distribution $P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}$ is closer to $P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$. Hence, we have:\n\n$$\n\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G}) < \\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G}).\n$$\nEmpirically, in Section 5.1, we approximately computed the average ratio over all test samples $\\mathbf{x}^*  _ {\\mathrm{obs}}$ and obtained:\n\n$$\n\\frac{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})}{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})} \\approx 47.93.\n$$\nThis clearly demonstrates that the mixed-masking training strategy (S4) yields significantly smaller distribution-shift coefficients compared to purely random missingness (S1). According to Corollary 1, this strongly supports the superior imputation performance achieved by our mixed-masking strategy. We also provide another example with an even larger difference in distribution-shift coefficients (with the ratio approximately 72), which correspondingly leads to a larger performance gap. Please refer to **A4** of reviewer iceu for further details.\n\n\n\n## References\n[1] Yoon, Jinsung, James Jordon, and Mihaela Schaar. \"Gain: Missing data imputation using generative adversarial nets.\" International conference on machine learning. PMLR, 2018.\n\n[2] Mattei, Pierre-Alexandre, and Jes Frellsen. \"MIWAE: Deep generative modelling and imputation of incomplete data sets.\" International conference on machine learning. PMLR, 2019.\n\n[3] Peis, Ignacio, Chao Ma, and José Miguel Hernández-Lobato. \"Missing data imputation and acquisition with deep hierarchical models and hamiltonian monte carlo.\" Advances in Neural Information Processing Systems 35 (2022): 35839-35851.\n\n[4] Roberts, Stephen, Michael Osborne, Mark Ebden, Steven Reece, Neale Gibson, and Suzanne Aigrain. \"Gaussian processes for time-series modelling.\" *Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences* 371, no. 1984 (2013): 20110550.\n\n[5] Camps-Valls, Gustau, Jochem Verrelst, Jordi Munoz-Mari, Valero Laparra, Fernando Mateo-Jimenez, and Jose Gomez-Dans. \"A survey on Gaussian processes for earth-observation data analysis: A comprehensive investigation.\" *IEEE Geoscience and Remote Sensing Magazine* 4, no. 2 (2016): 58-78."}, {"review_id": "JnuawEQP2q", "reviewer": "Reviewer_PUbh", "summary": "This paper addresses the problem of imputing missing values in time-series data using diffusion-based generative models, specifically focusing on the diffusion transformers (DiT) method. This work contributes a theoretical analysis of their statistical properties under Gaussian Process data. The paper provides sample complexity bounds for the imputation of the DiT architecture. Based on the analysis, confidence regions for quantifying uncertainty in the imputed values can be derived. The analysis also reveals that missingness patterns critically affect both the efficiency and accuracy of imputation. In addition to the theory, the paper proposes a mixed-masking training strategy to improve imputation quality and supports the theoretical findings with an evaluation using synthetic data.", "questions": "1. Clarification on Benchmark Experiments\nAppendix F presents experimental results on benchmark datasets, which seem to largely confirm prior findings that DiT outperforms CSDI. Could the authors clarify what new insights are gained from this particular experiment in the context of the paper’s theoretical contributions? Additionally, have the authors considered applying the mixed-masking strategy to these benchmark experiments as well? It would be interesting to see whether the proposed strategy generalizes beyond the synthetic setting.\n2. Purpose of Algorithm 2\nIf I understand correctly, Algorithm 2 is introduced solely for analytical purposes and is not used in practice. Is this correct? \n3. Connection Between Condition Numbers and Mixed-Masking\nIt is insightful that the paper identifies the difficulty of imputation as being influenced by the condition numbers of Σcond and Σobs. However, the link between this and the effectiveness of the mixed-masking strategy remains unclear. While it is intuitive that training on a variety of missing patterns might enhance robustness to unseen test patterns, I did not see a formal theoretical justification showing that mixing patterns with different condition numbers leads to more effective training. Could the authors clarify whether such a result exists or whether this remains an empirical hypothesis?\n4. Coverage of Masking Pattern Combinations\nIn the synthetic experiments, not all combinations of S1-S4 were explored. While it is understandable that exhaustive evaluation would be too costly, to support the claim that combining different masking patterns is beneficial, it would be helpful to also report results from training on individual subsets, that is, using only S2/S3/S4. This would make the benefits of combining patterns more evident and quantifiable.", "limitations": "Not so many limitations were discussed. It might be useful to include more discussion, for example, the limitations of the theoretical analysis compared with the real use cases, and the need to rely on assuming Gaussian process data.", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. The paper addresses a highly relevant problem in time-series imputation using deep learning, an area where theoretical understanding remains limited.\n2. The authors present a range of theoretical results, which yield practical insights. For example, based on the condition number of the covariance matrix, we can interpret why one missing pattern is more difficult than the other. Convergence rate is also derived to provide a good theoretical grounding of DiT and also its uncertainty quantification.\n\nWeaknesses\n1. The experiments are primarily conducted on a Gaussian Process dataset that satisfies the model’s assumptions. While Appendix F includes benchmark experiments demonstrating the superiority of DiT over CSDI, these results do not directly support the effectiveness of the proposed mixed-masking training strategy, as DiT itself is not introduced in this paper.\n2. While it is encouraging that the results support the use of transformer architectures, the theoretical argument does not appear strong enough to imply that CSDI (Conditional Score-based Diffusion Imputer) would be inherently less effective unless the paper explicitly shows structural limitations of CSDI compared to transformers.\n3. The work does not propose a novel algorithm for diffusion-based imputation. It focuses solely on the theoretical analysis of an existing DiT model under GP assumptions.\n4. The paper lacks a related work section in the main body, making it overly dependent on the supplementary material and limiting its self-contained readability.\n\nWhile I appreciate the authors' effort in deriving several theoretical results for DiT, I have several concerns that currently lead me to assign a borderline (slightly negative) rating. Please refer to the weaknesses and questions sections for more information.", "was_revised": true, "final_justification": "Additional results in the rebuttal clarified my concern on the effectiveness of mixed masking strategy and clarified my misunderstanding on the definition of S1-S4 of the training strategies. I increased the score to 4.", "rebuttal": "Thank you for your insightful comments and constructive suggestions. We now address the specific questions, weaknesses, and limitations you have raised.\n\n> **Q1**: Assumptions on Gaussian process data.\n\n**A1**: We impose Gaussian process (GP) assumption only for theoretical analysis. The insights gained from the analysis motivate our mixed-masking strategy, which applies to real-world data well beyond the GP assumption.\n\nAs discussed in the introduction, GPs naturally exhibit rich spatio-temporal dependencies and long-horizon correlations, which pose central challenges in time series tasks [1] and are widely adopted in practical modeling scenarios such as climate [2] and financial forecasting.  Additionally, as noted in lines 119–121, we intentionally do not restrict the kernel function for GPs, ensuring our theoretical results remain generalizable across various settings. \n\nEmpirical validation in Section 5.2 (latent GP data) and the Appendix (real-world datasets) further supports the applicability of our findings beyond strictly synthetic settings. Also worth noting is that although our theory is developed under GP assumptions, the insights gained from this theory directly informed the development of our mixed-masking strategy. This strategy has demonstrated strong generalizability to real-world applications, clearly indicating that the theoretical assumptions do not limit the practical applicability of our method.\n\n> **Q2**:  The theoretical argument does not appear strong enough to imply that CSDI would be inherently less effective.\n\n**A2**: Thank you for raising this point. Theorem 1 demonstrates how transformers capture temporal correlations effectively through the attention mechanism (Theorem 1). While convolutional architectures used by CSDI may also model temporal correlations, theoretical analysis of their efficiency is elusive. Recent study [3] suggests that transformers are inherently suitable for sequential modeling, though rigorously proving DiT’s superiority over CSDI remains open and challenging.\n\n>**Q3**: What are the new insights gained from the real-world data experiment in Appendix F?  Additionally, have the authors considered applying the mixed-masking strategy to these benchmark experiments?\n\n**A3**: Thank you for this valuable suggestion. The original Appendix F experiments primarily confirmed DiT’s superior imputation performance, supporting the theoretical result from Theorem 2 . Following your advice, we conducted additional experiments to specifically validate our mixed-masking strategy on real-world benchmark datasets. In these experiments, we designed mixed-masking strategies inspired by those described in Section 5 (patterns tend to have more grouped missing values rather than purely random). The updated Mean Squared Error results are shown below:\n\n| Model| **ETT _ m1 (Missing %)** ||| **BeijingAir (Missing %)** |||\n|-|-|-|-|-|-|-|\n| | 10% | 20%| 50%| 10%  | 20%| 50% |\n| CSDI                              | 0.0615 (±0.0097)       | 0.0698 (±0.0106)     | 0.0797 (±0.0106)     | 0.4196 (±0.1726)           | 0.3926 (±0.0790)     | 0.4534 (±0.0379)     |\n| GP-VAE                            | 0.1567 (±0.0094)       | 0.2138 (±0.0067)     | 0.4249 (±0.0127)     | 0.4096 (±0.0202)           | 0.4777 (±0.0179)     | 0.7017 (±0.0189)     |\n| **DiT**                           | 0.0534 (±0.0063)       | 0.0606 (±0.0076)     | **0.0684 (±0.0070)** | 0.3683 (±0.0351)           | 0.4025 (±0.0424)     | 0.4255 (±0.0670)     |\n| **DiT w/ mixed-masking strategy** | **0.0502 (±0.0055)**   | **0.0588 (±0.0081)** | 0.0711 (±0.0092)      | **0.3428 (±0.0275)**       | **0.3864 (±0.0403)** | **0.4229 (±0.0539)** |\n\nAs can be seen, incorporating our proposed mixed-masking strategy further improves DiT's performance on real-world datasets with varying missing rates, supporting our theoretical and methodological contributions.\n\n> **Q4**: Is Algorithm 2 introduced solely for analytical purposes and is not used in practice?\n\n**A4**: Yes, Algorithm 2 serves for analytical purposes to understand how transformers unroll an algorithm to represent the score function. Our theoretical guarantees in Lemma 1, Theorem 1 and Theorem 2 are built upon the insights from this algorithm.\n\n> **Q5**: Connection between condition numbers and mixed-masking strategies.\n\n**A5**: Thank you for highlighting this important aspect.  First, we clarify the definitions of mixed-masking strategies S1–S4 used in Section 5:\n\n- **S1**: 100% random (16×1, sixteen fully random placed missing entries)\n- **S2**: 50% random (16×1) + 50% weakly grouped (8×2, eight randomly placed blocks of two consecutive entries)\n- **S3**: 33.3% random (16×1) + 33.3% weakly grouped (8×2) + 33.3% moderately grouped (4×4, four randomly placed blocks of four consecutive entries)\n- **S4**: 25% random (16×1) + 25% weakly grouped (8×2) + 25% moderately grouped (4×4) + 25% strongly grouped (1×16, one randomly placed blocks of sixteen consecutive entries)\n\nOur theoretical analysis indeed strongly supports the effectiveness of the mixed-masking strategy by linking it directly to the distribution-shift coefficient $\\mathsf{DS}$ (Corollary 1). Using mixed-masking strategies reduces $\\mathsf{DS}$, directly improving imputation performance.\n\nTo illustrate this clearly, consider the concrete example in our numerical simulations (Section 5). Missing patterns P1–P4 have varying difficulties , with P1 (strongly grouped missing) yielding the highest condition number and P4 (purely random missingness) having the lowest. Specifically, our training strategies S1 and S4 differ as follows:\n\n- **S1** uses only random missing patterns during training (P4-like, low condition number).\n- **S4** includes various missing patterns (both P1-like and P4-like, high and low condition numbers).\n\nDenoting the training distributions corresponding to S1 and S4 as $P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$ and $P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$, respectively, if we consider a new sample $\\mathbf{x}^*  _ {\\mathrm{obs}}$ with missing pattern P1 (strongly grouped missing), it is intuitive that the resulting distribution $P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}$ is closer to $P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$. Correspondingly, we have:\n\n$$\n\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G}) < \\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G}).\n$$\nEmpirically, in Section 5.1, we approximately computed the average ratio over all test samples $\\mathbf{x}^*  _ {\\mathrm{obs}}$ and obtained:\n\n$$\n\\frac{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})}{\\mathsf{DS}(P _ {\\mathbf{x}^* _ {\\mathrm{obs}}}, P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})} \\approx 47.93.\n$$\nThis clearly demonstrates that the mixed-masking training strategy (S4) yields significantly smaller distribution-shift coefficients compared to purely random missingness (S1). According to Corollary 1, this strongly supports the superior imputation performance achieved by our mixed-masking strategy. We also provide another example with an even larger difference in distribution-shift coefficients (with the ratio approximately 72), which correspondingly leads to a larger performance gap. Please refer to **A4** of reviewer iceu for further details.\n\n> **Q6**: Not all combinations of masking patterns were explored; it would be helpful to also report results from training on individual subsets, that is, using only S2/S3/S4.\n\n**A6**: Thank you for suggesting additional masking-pattern evaluations.\n\nFor the definition of the mixed-masking strategies, please refer to **A5**. We agree that evaluating the individual patterns (8×2, 4×4, and 1×16 separately) would provide useful insights. Following your suggestion, we performed these additional experiments under the Gaussian process data setting. The updated confidence region coverage (%) results are shown below:\n\n|   | S1 | S2 | S3 | S4 | Only 8×2 | Only 4×4 | Only 1×16 |\n|-|-|-|-|-|-|-|-|\n|P1  |34.58(±1.22)|**66.22(±3.86)**|56.04(±6.48)   |57.27(±5.34)|36.74(±1.31)|34.15(±1.16)|32.68(±1.50)|\n|P2  |58.46(±1.89)|**83.71(±2.86)**|81.05(±2.09)   |79.00(±2.42)|60.51(±1.65)|59.23(±1.88)|54.23(±1.76)|\n|P3  |72.42(±1.66)|74.04(±1.90)   |**74.59(±1.27)**|74.38(±3.00)|71.24(±1.52)|73.08(±1.10)|69.46(±1.53)|\n|P4  |80.25(±1.64)|81.50(±2.12)   |**83.09(±1.48)**|82.74(±2.40)|80.46(±2.01)|79.83(±1.84)|76.72(±2.20)|\n\nThese results clearly illustrate that training exclusively on single masking patterns, regardless of their complexity, yields inferior imputation performance compared to appropriately mixing different patterns. This supports our proposed mixed-masking strategies and aligns well with our theoretical insights.\n\n\n\n## References\n\n[1] Roberts, Stephen, Michael Osborne, Mark Ebden, Steven Reece, Neale Gibson, and Suzanne Aigrain. \"Gaussian processes for time-series modelling.\" *Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences* 371, no. 1984 (2013): 20110550.\n\n[2] Camps-Valls, Gustau, Jochem Verrelst, Jordi Munoz-Mari, Valero Laparra, Fernando Mateo-Jimenez, and Jose Gomez-Dans. \"A survey on Gaussian processes for earth-observation data analysis: A comprehensive investigation.\" *IEEE Geoscience and Remote Sensing Magazine* 4, no. 2 (2016): 58-78.\n\n[3] Wen, Qingsong, Tian Zhou, Chaoli Zhang, Weiqi Chen, Ziqing Ma, Junchi Yan, and Liang Sun. \"Transformers in time series: A survey.\" *arXiv preprint arXiv:2202.07125* (2022)."}, {"review_id": "BfIowjK6U0", "reviewer": "Reviewer_iceu", "summary": "This paper takes a highly theoretical approach to the problem of imputing missing values in time series using diffusion transformers (DiT). The authors focus on Gaussian process data and derive detailed statistical sample complexity bounds. They also construct confidence regions for the imputations based on theoretical guarantees. Synthetic experiments with varied missing patterns are used to validate the theory and match intuitive expectations. The paper further explores how to mitigate distribution shift between training and testing, improving robustness. Extensive real-world results are included in the supplement, but the core contribution remains largely theoretical.", "questions": "1. In Section 5.1, the missing rate is about 17%, and the real data experiments use 10%, 20%, and 50% missingness. Can DiT handle much sparser settings, such as 90% missing?\n\n2. Can DiT be directly applied (or easily modified) to incorporate time-invariant covariates?\n\n3. Are all simulation settings run with the same transformer architecture, regardless of the condition number $\\kappa$?\n\n4. In Table 1, under P1 (first row), there is a noticeable gap in coverage between imputation strategies S1 and S2. While this may be due to empirical variation, is there any theoretical intuition or insight that might help explain this discrepancy?", "limitations": "yes", "rating": 5, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strength:\n1. Figure 1 is well-designed and helps clarify the structure of DiT.\n2. The paper is theory-heavy, and I’m impressed by the depth and completeness of the analysis.\n3. Theorem 1 provides a detailed characterization of score approximation error in terms of transformer architecture, with explicit dependence on the missing pattern via the condition number. This is much more realistic than assuming missing completely at random.\n4. The background section (Appendix A) offers a solid overview of time series imputation, which is helpful for non-experts.\n\nWeaknesses:\nNo major weaknesses observed.\n\nSuggestion:\nIn Figure 4, consider plotting CSDI first and then DiT, so that DiT appears on top. Since CSDI’s confidence intervals are wider, they visually overshadow DiT, making it hard to see DiT’s advantage in the current version.", "was_revised": true, "final_justification": "I believe the author has addressed all my questions with high clarity, and the additional numerical results further highlight the significance of this research. As a result, I will maintain my score of 5.", "rebuttal": "We sincerely thank you for your positive evaluation of our work, and we greatly appreciate your helpful suggestion on improving Figure 4 to better highlight our results. We now address the specific questions you raised.\n\n> **Q1**: In Section 5.1, the missing rate is about 17%, and the real data experiments use 10%, 20%, and 50% missingness. Can DiT handle much sparser settings, such as 90% missing?\n\n**A1**: Yes. For Section 5.1 (simulated Gaussian process dataset), we intentionally selected a relatively challenging setup—featuring small model size, complex kernels, long sequence length, and high dimensionality—to clearly differentiate the effects of different mixed-masking strategies. With a larger model, more training samples, or simplified kernels and lower dimensionality, we are confident that DiT can handle much higher missing rates, including up to 90%.\n\nRegarding real-world data experiments, we followed your suggestion and conducted additional experiments for both DiT and CSDI under 90% missingness. Below, we report the Mean Absolute Error (MAE) and Mean Squared Error (MSE) results:\n\n**MAE Results (90% Missing Rate)**\n\n| Model              | ETT _ m1                  | BeijingAir              |\n|--------------------|--------------------------|--------------------------|\n| CSDI               | 0.2654 (±0.0297)         | 0.3340 (±0.0410)         |\n| DiT                | **0.2445 (±0.0244)**     | **0.3012 (±0.0247)**     |\n\n**MSE Results (90% Missing Rate)**\n\n| Model              | ETT _ m1                  | BeijingAir              |\n|--------------------|--------------------------|--------------------------|\n| CSDI               | 0.2167 (±0.0286)         | 0.3733 (±0.0552)         |\n| DiT                | **0.1991 (±0.0163)**     | **0.3178 (±0.0326)**     |\n\nThese results demonstrate that DiT continues to perform robustly even in extremely sparse settings, achieving better results than CSDI under 90% missingness.\n\n\n> **Q2**: Can DiT be directly applied (or easily modified) to incorporate time-invariant covariates?\n\n**A2**: Yes, we address this question from both a theoretical and a practical perspective. \n\nTime-invariant covariates indeed play a critical role in time series analysis. From a theoretical perspective, time-invariant covariates can be modeled as conditioning variables in the diffusion process, allowing them to be treated similarly to the observed data. As such, only minimal extensions are needed to incorporate them into our current analytical framework. \n\nOn the application side, while this extension is not the main focus of our paper, relevant studies such as TimeDiT [1] have demonstrated practical approaches for integrating time-invariant covariates (as conditions) into DiT for time series forecasting and imputation. DiT can be readily adapted to follow similar strategies.\n\n\n> **Q3**: Are all simulation settings run with the same transformer architecture, regardless of the condition number?\n\n**A3**: Yes, as detailed in Appendix F (Experiment details), we consistently employed the same DiT architecture across all conditions, specifically with a hidden size of 256, 12 transformer layers, and 16 attention heads per layer. This ensured comparability across all experiments.\n\n> **Q4**: In Table 1, under P1 (first row), there is a noticeable gap in coverage between imputation strategies S1 and S2. While this may be due to empirical variation, is there any theoretical intuition or insight that might help explain this discrepancy?\n\n**A4**: Thank you for pointing out this intriguing observation. First, we briefly restate the definitions of the mixed-masking strategies S1 and S2:\n\n- **S1**: 100% random (16×1, sixteen fully random placed missing entries)\n- **S2**: 50% random (16×1) + 50% weakly grouped (8×2, eight randomly placed blocks of two consecutive entries)\n\nIn row P1 (missing 16 consecutive entries at the end of the sequence, corresponding to a very challenging condition number of 415.4), strategy S2 significantly outperforms S1. Intuitively, training with S2 exposes the model to a combination of simpler (random) and more challenging (weakly grouped) missing patterns, enhancing the model’s ability to handle difficult imputation scenarios. In contrast, the model trained only on simpler random patterns (S1) is less robust against difficult test patterns like P1.\n\nFrom a theoretical perspective, denoting training distributions corresponding to strategies S1 and S2 as $P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$ and $P^{(2)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$, respectively, we find that the test distribution $P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}$ (with pattern P1) is closer to $P^{(2)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$. Empirically, we approximately computed the distribution-shift coefficient ratio over all test samples with pattern P1 and obtained:\n\n$$\n\\frac{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})}{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(2)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})} \\approx 72.46.\n$$\n\nThus, the training distribution under strategy S2 yields a substantially lower distribution-shift coefficient compared to S1 for the test pattern P1. According to Corollary 1, this theoretically confirms our empirical observation that S2 significantly improves imputation performance over S1 for challenging missing patterns.\n\n## References\n\n[1] Cao, Defu, Wen Ye, Yizhou Zhang, and Yan Liu. \"TimeDiT: General-purpose diffusion transformers for time series foundation model.\" *arXiv preprint arXiv:2409.02322* (2024)."}, {"review_id": "n4mCCgIA17", "reviewer": "Reviewer_F7up", "summary": "This paper studies the statistical efficiency of imputation using conditional diffusion models. It provides a theoretical analysis on the sample complexity required for accurate imputation. The main result characterizes sample complexity as a function of the number of training samples, sequence length, and the condition number of the data covariance.  In addition, it proposes a mixed masking procedure that improves the imputation performance empirically.", "questions": "1. In Line 146, the paper states that the training procedure assumes fully observed training data. However, in real-world scenarios, training data often contains missing values. How would the proposed framework and theoretical analysis adapt under this more realistic setting? \n\n2. The definitions of S1–S4 in the mixed-reamsking strategy are not clearly provided. Could the authors explicitly define what each of the S1–S4 masking strategies entails, and how they are used during training?", "limitations": "yes", "rating": 4, "confidence": 2, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 2, "clarity": 2, "significance": 2, "originality": 2, "strengths_and_weaknesses": "**Strengths**\n\n1. The paper is well-written and presents its technical contributions in a clear and accessible manner.\n\n2. It addresses an important and underexplored problem: understanding the statistical efficiency of imputation using diffusion models, which is increasingly relevant for structured data applications.\n\n3. The theoretical derivations appear technically sound and novel, offering valuable insights into how sample complexity scales with factors like sequence length and data covariance.\n\n**Weaknesses**\n\n1. **Strong Assumptions Undermine the Practical Utility of Theoretical Results**  \n   Theorem 2 relies on multiple strong assumptions, including a Gaussian process prior and Assumption 1. More critically, it presumes access to a well-trained diffusion transformer that approximates the score function accurately (per Theorem 1). However, training conditional diffusion models for score estimation is a key open challenge. By assuming this is solved, the theory sidesteps a central difficulty, reducing its real-world applicability.\n\n2. **Limited Empirical Validation of Theoretical Claims**  \n   While the experiments confirm that imputation performance improves with more training samples, they do not test the predicted \\( O(1/\\sqrt{n}) \\) scaling. Additionally, dependencies on sequence length \\( H \\) are not empirically explored. This weakens the connection between the theoretical results and observed behavior.\n\n3. **Clarity Issues in Mixed-Masking Strategy**  \n   The mixed-masking training approach is not clearly described, and the paper provides little intuition for why it should improve imputation performance. A more thorough explanation or comparison with existing masking strategies would make this contribution clearer and more convincing.", "was_revised": true, "final_justification": "I recommend a borderline acceptance of this paper.\nThe authors have addressed most of my concerns regarding the practicality of the assumptions made in the theorem. However, I still believe that a more general study would be beneficial. Therefore, I recommend a borderline acceptance of this paper.", "rebuttal": "Thank you very much for your insightful comments and suggestions. We now address the specific questions, weaknesses, and limitations you have raised.\n\n> **Q1**: Gaussian process data assumption undermines practical utility.\n\n**A1**: We respectfully disagree with the view that the Gaussian process (GP) assumption undermines the practical utility of our theoretical contributions and developed methodologies. We impose GP assumption only for theoretical analysis. The insights gained from the analysis motivate our mixed-masking strategy, which applied to real-world data well beyond the GP assumption.\n\nAs discussed in the introduction, GPs are not only mathematically convenient but also widely used in real-world time series modeling, including in domains such as climate science [2] and finance. Their ability to capture rich spatio-temporal and long-horizon dependencies makes them highly representative of the core challenges in time series imputation [1]. Moreover, our theoretical analysis intentionally avoids specifying a particular kernel class (lines 119–121), thereby preserving generality across a wide range of covariance structures. This flexibility allows our theoretical findings to cover many practical scenarios beyond the synthetic GP case. \n\nEmpirically, we further validate this point through experiments on latent GP data (Section 5.2) and real-world datasets (Appendix), both of which confirm the alignment between our theoretical insights and observed model performance. Also worth noting is that although our theory is developed under GP assumptions, the insights gained from this theory directly informed the development of our mixed-masking strategy. This strategy has demonstrated strong generalizability to real-world applications, clearly indicating that the theoretical assumptions do not limit the practical applicability of our method. This strategy has shown robust generalization across real-world datasets and varying missing patterns. \n\n\n> **Q2**: Experiments lack explicit testing of performance scaling with $\\mathcal{O}(1/\\sqrt{n})$ and dependencies on sequence length $H$.\n\n**A2**: Thank you for highlighting this. Regarding sample-size scaling ($n$), we summarize the results obtained using training strategy S1 (vanilla strategy with purely random missing) and missing pattern P4 (purely random missing) with $10^5$ training samples in Section 5.1 as follows:\n\n| Sample Size (n)      | $10^3$ | $10^{3.5}$ | $10^4$ | $10^{4.5}$ | $10^5$ |\n|----------------------|--------|------------|--------|------------|--------|\n| CR Coverage Rate (%) |  0.19  |   0.64     | 20.28  |   73.41    | 80.25  |\n\nApproximately, we observe that the difference between the achieved coverage rate and the target rate scales up to constant factors, supporting the theoretical prediction. The Python code snippet below plots the difference versus $1/\\sqrt{n}$, and we also perform a linear regression, yielding $R^2 \\approx 0.7$, which indicates a reasonably strong linear trend supporting the theoretical prediction.\n\n```python\nimport numpy as np\nfrom sklearn.linear_model import LinearRegression\nfrom sklearn.metrics import r2_score\nimport matplotlib.pyplot as plt\n\nn = np.array([1e3, 10**3.5, 1e4, 10**4.5, 1e5])\ncoverage = np.array([0.19, 0.64, 20.28, 73.41, 80.25])\ndiff = 95 - coverage\n\nx_reg = 1 / np.sqrt(n)\nx_reg = x_reg.reshape(-1, 1)\ny_reg = diff\n\nreg_model = LinearRegression()\nreg_model.fit(x_reg, y_reg)\ny_pred_reg = reg_model.predict(x_reg)\n\nr2 = r2_score(y_reg, y_pred_reg)\n\nplt.plot(x_reg, y_reg, marker='o')\nplt.xlabel(r'$1/\\sqrt{n}$')\nplt.ylabel('Difference from Target Coverage (%)')\nplt.title('Coverage Gap vs $1/\\sqrt{n}$')\nplt.tight_layout()\nplt.show()\n```\n\nAdditionally, following your suggestion, we conducted further experiments to study the impact of sequence length ($H$). Keeping the model architecture fixed (dimension $d=8$) and maintaining a consistent missing rate ($16/96$), we varied sequence lengths $H=16, 32, 64, 96, 128$. Due to computational limitations, we again adopted training strategy S1 and missing pattern P4, with $10^5$ training samples. The results are as follows:\n\n| Sequence Length (H)  | 16     | 32     | 64     | 96     | 128    |\n|----------------------|--------|--------|--------|--------|--------|\n| CR Coverage Rate (%) | 92.67  | 88.63  | 82.14  | 80.25  | 77.81  |\n\nThese empirical findings also support our theoretical claims in Theorem 2 and Corollary 1, demonstrating how imputation performance scales with both sample size and sequence length.\n\n\n> **Q3**: Definitions of strategies S1–S4, and why they can improve imputation performance.\n\n**A3**:Thank you for pointing out the need for clarification. Here, we explicitly restate the definitions of the mixed-masking strategies S1–S4 used in Section 5:\n\n- **S1**: 100% random missing pattern (16×1, sixteen randomly placed missing entries).\n- **S2**: 50% random (16×1) + 50% weakly grouped (8×2, eight randomly placed blocks of two consecutive missing entries).\n- **S3**: 33.3% random (16×1) + 33.3% weakly grouped (8×2) + 33.3% moderately grouped (4×4, four randomly placed blocks of four consecutive missing entries).\n- **S4**: 25% random (16×1) + 25% weakly grouped (8×2) + 25% moderately grouped (4×4) + 25% strongly grouped (1×16, one randomly placed block of sixteen consecutive missing entries).\n\nRegarding how these strategies relate to our theoretical results, let us denote the training distributions corresponding to S1 and S4 as $P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$ and $P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$, respectively. Consider a test sample $\\mathbf{x}^*  _ {\\mathrm{obs}}$ following the strongly grouped missing pattern P1 (consecutive missing entries). Intuitively, the resulting distribution $P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}$ is closer to $P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$ than to $P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}$. Formally, this implies:\n\n$$\n\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G}) < \\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G}).\n$$\n\nEmpirically, in Section 5.1, we calculated the average ratio across all test samples and found that:\n\n$$\n\\frac{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(1)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})}{\\mathsf{DS}(P _ {\\mathbf{x}^*  _ {\\mathrm{obs}}}, P^{(4)} _ {\\mathbf{x} _ {\\mathrm{obs}}}, \\mathcal{G})} \\approx 47.93.\n$$\n\nThis clearly indicates that the mixed-masking training strategy (S4) yields significantly smaller distribution-shift coefficients compared to purely random missingness (S1). According to Corollary 1, this provides strong theoretical support for the superior empirical performance achieved by our mixed-masking strategy. We also provide another example with an even larger difference in distribution-shift coefficients (with the ratio approximately 72), which correspondingly leads to a larger performance gap. Please refer to **A4** of reviewer iceu for further details.\n\n> **Q4**: In Line 146, the paper states that the training procedure assumes fully observed training data. However, in real-world scenarios, training data often contains missing values. How would the proposed framework and theoretical analysis adapt to this more realistic setting?\n\n**A4**: Thank you for raising this important point. Practically, TimeDiT [3] proposes effective strategies to handle naturally occurring missing values when applying DiT to time series tasks. Alternatively, methods like those presented in CSDI [4], which fill missing entries with dummy values (e.g., zeros) and treat them as observed, can also be combined with our framework.\n\nFrom a theoretical perspective, your insight is highly relevant. In fact, our theoretical framework can readily adapt to scenarios involving training data with missing values, provided the missingness follows a \"Missing Completely at Random\" scheme (i.e., missing patterns do not depend on observed or missing values themselves). Consider Gaussian-process data with Bernoulli masks applied to each entry: since Bernoulli random variables are also sub-Gaussian (light-tailed), our theoretical analysis naturally extends to this setting with minimal adjustments. \n\n\nHowever, handling cases where missingness depends on observed or even missing values represents a more challenging and promising direction for theoretical advancement. We currently leave this as an open research direction.\n\n\n## References\n\n[1] Roberts, Stephen, Michael Osborne, Mark Ebden, Steven Reece, Neale Gibson, and Suzanne Aigrain. \"Gaussian processes for time-series modelling.\" *Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences* 371, no. 1984 (2013): 20110550.\n\n[2] Camps-Valls, Gustau, Jochem Verrelst, Jordi Munoz-Mari, Valero Laparra, Fernando Mateo-Jimenez, and Jose Gomez-Dans. \"A survey on Gaussian processes for earth-observation data analysis: A comprehensive investigation.\" *IEEE Geoscience and Remote Sensing Magazine* 4, no. 2 (2016): 58–78.\n\n[3] Cao, Defu, Wen Ye, Yizhou Zhang, and Yan Liu. \"TimeDiT: General-purpose diffusion transformers for time series foundation model.\" *arXiv preprint arXiv:2409.02322* (2024).\n\n[4] Tashiro, Yusuke, Jiaming Song, Yang Song, and Stefano Ermon. \"CSDI: Conditional score-based diffusion models for probabilistic time series imputation.\" *Advances in Neural Information Processing Systems* 34 (2021): 24804–24816."}]

Author rebuttal:

Meta-review: The paper develops a statistical framework for capturing statistical efficiency and uncertainty of diffusion-based time-series imputation with conditional diffusion transformers (DiT). All reviewers (and area chair) found the paper investigates an important problem in time-series imputation with diffusion models that can lack theoretical guarantees, where the paper provides theory and the architectural insight with clear intuition, and the experimental results highlight the significance of the proposed work. 

During rebuttal, the authors provided clarifications on intuition for Remark 1 and its link to the shift coefficient, and a more balanced discussion of VAE/GAN work and existing GP-VAE comparisons. The authors provided additional real-world experiments with high (i.e., 90%) missingness, and mixed-masking strategies, to support DiT’s robustness and exhibited empirical trends. These additions strengthen the paper and resolve the main concerns by reviewers. The reviewers also left some important notes on improving parts of the paper, including theoretical exposition in Sections 3 and 4. We highly encourage the authors to more clearly discuss those points and include the additional analysis in revised paper.  

The potential limitation pointed by reviewer is the GP assumption bounds’ formal generality, though classified by authors it is well-motivated and clearly acknowledged. We hope the authors clearly describe the potential limitations of their assumptions, while we still agree the experiments, algorithmic and structure design well support the contribution of this paper to related domains.
Decision: Accept (poster)