CSDI: Conditional Score-based Diffusion Models for Probabilistic Time Series Imputation
Year: 2021

Abstract:
The imputation of missing values in time series has many applications in healthcare and finance. While autoregressive models are natural candidates for time series imputation, score-based diffusion models have recently outperformed existing counterparts including autoregressive models in many tasks such as image generation and audio synthesis, and would be promising for time series imputation. In this paper, we propose Conditional Score-based Diffusion models for Imputation (CSDI), a novel time series imputation method that utilizes score-based diffusion models conditioned on observed data. Unlike existing score-based approaches, the conditional diffusion model is explicitly trained for imputation and can exploit correlations between observed values. On healthcare and environmental data, CSDI improves by 40-65% over existing probabilistic imputation methods on popular performance metrics. In addition, deterministic imputation by CSDI reduces the error by 5-20% compared to the state-of-the-art deterministic imputation methods. Furthermore, CSDI can also be applied to time series interpolation and probabilistic forecasting, and is competitive with existing baselines. The code is available at https://github.com/ermongroup/CSDI.

heading: 1 Introduction
text: Multivariate time series are abundant in real world applications such as finance, meteorology and healthcare. These time series data often contain missing values due to various reasons, including device failures and human errors [1, 2, 3]. Since missing values can hamper the interpretation of a time series, many studies have addressed the task of imputing missing values using machine learning techniques [4, 5, 6]. In the past few years, imputation methods based on deep neural networks have shown great success for both deterministic imputation [7, 8, 9] and probabilistic imputation [10]. These imputation methods typically utilize autoregressive models to deal with time series.
Score-based diffusion models – a class of deep generative models and generate samples by gradually converting noise into a plausible data sample through denoising – have recently achieved state-ofthe-art sample quality in many tasks such as image generation [11, 12] and audio synthesis [13, 14], outperforming counterparts including autoregressive models. Diffusion models can also be used to impute missing values by approximating the scores of the posterior distribution obtained from the prior by conditioning on the observed values [12, 15, 16]. While these approximations may work well in practice, they do not correspond to the exact conditional distribution.
In this paper, we propose CSDI, a novel probabilistic imputation method that directly learns the conditional distribution with conditional score-based diffusion models. Unlike existing score-based approaches, the conditional diffusion model is designed for imputation and can exploit useful information in observed values. We illustrate the procedure of time series imputation with CSDI in
35th Conference on Neural Information Processing Systems (NeurIPS 2021).
Figure 1. We start imputation from random noise on the left of the figure and gradually convert the noise into plausible time series through the reverse process pθ of the conditional diffusion model. At each step t, the reverse process removes noise from the output of the previous step (t+ 1). Unlike existing score-based diffusion models, the reverse process can take observations (on the top left of the figure) as a conditional input, allowing the model to exploit information in the observations for denoising. We utilize an attention mechanism to capture the temporal and feature dependencies of time series.
For training the conditional diffusion model, we need observed values (i.e., conditional information) and ground-truth missing values (i.e., imputation targets). However, in practice we do not know the ground-truth missing values, or training data may not contain missing values at all. Then, inspired by masked language modeling, we develop a self-supervised training method that separates observed values into conditional information and imputation targets. We note that CSDI is formulated for general imputation tasks, and is not restricted to time series imputation.
Our main contributions are as follows:
• We propose conditional score-based diffusion models for probabilistic imputation (CSDI), and implement CSDI for time series imputation. To train the conditional diffusion model, we develop a self-supervised training method.
• We empirically show that CSDI improves the continuous ranked probability score (CRPS) by 40-65% over existing probabilistic methods on healthcare and environmental data. Moreover, deterministic imputation with CSDI decreases the mean absolute error (MAE) by 5-20% compared to the state-of-the-art methods developed for deterministic imputation.
• We demonstrate that CSDI can also be applied to time series interpolations and probabilistic forecasting, and is competitive with existing baselines designed for these tasks.

heading: 2 Related works
text: Time series imputations with deep learning Previous studies have shown deep learning models can capture the temporal dependency of time series and give more accurate imputation than statistical methods. A popular approach using deep learning is to use RNNs, including LSTMs and GRUs, for sequence modeling [17, 8, 7]. Subsequent studies combined RNNs with other methods to improve imputation performance, such as GANs [9, 18, 19] and self-training [20]. Among them, the combination of RNNs with attention mechanisms is particularly successful for imputation and interpolation of time series [21, 22]. While these methods focused on deterministic imputation, GP-VAE [10] has been recently developed as a probabilistic imputation method.
Score-based generative models Score-based generative models, including score matching with Langevin dynamics [23] and denoising diffusion probabilistic models [11], have outperformed existing methods with other deep generative models in many domains, such as images [23, 11],
audio [13, 14], and graphs [24]. Most recently, TimeGrad [25] utilized diffusion probabilistic models for probabilistic time series forecasting. While the method has shown state-of-the-art performance, it cannot be applied to time series imputation due to the use of RNNs to handle past time series.

heading: 3 Background

heading: 3.1 Multivariate time series imputation
text: We consider N multivariate time series with missing values. Let us denote the values of each time series as X = {x1:K,1:L} ∈ RK×L where K is the number of features and L is the length of time series. While the length L can be different for each time series, we treat the length of all time series as the same for simplicity, unless otherwise stated. We also denote an observation mask as M = {m1:K,1:L} ∈ {0, 1}K×L where mk,l = 0 if xk,l is missing, and mk,l = 1 if xk,l is observed. We assume time intervals between two consecutive data entries can be different, and define the timestamps of the time series as s = {s1:L} ∈ RL. In summary, each time series is expressed as {X,M, s}. Probabilistic time series imputation is the task of estimating the distribution of the missing values of X by exploiting the observed values of X. We note that this definition of imputation includes other related tasks, such as interpolation, which imputes all features at target time points, and forecasting, which imputes all features at future time points.

heading: 3.2 Denoising diffusion probabilistic models
text: Let us consider learning a model distribution pθ(x0) that approximates a data distribution q(x0). Let xt for t = 1, . . . , T be a sequence of latent variables in the same sample space as x0, which is denoted as X . Diffusion probabilistic models [26] are latent variable models that are composed of two processes: the forward process and the reverse process. The forward process is defined by the following Markov chain:
q(x1:T | x0) := T∏ t=1 q(xt | xt−1) where q(xt | xt−1) := N (√ 1− βtxt−1, βtI )
(1)
and βt is a small positive constant that represents a noise level. Sampling of xt has the closed-form written as q(xt | x0) = N (xt; √ αtx0, (1 − αt)I) where α̂t := 1 − βt and αt := ∏t i=1 α̂i. Then,
xt can be expressed as xt = √ αtx0 + (1− αt) where ∼ N (0, I). On the other hand, the reverse process denoises xt to recover x0, and is defined by the following Markov chain:
pθ(x0:T ) := p(xT ) T∏ t=1 pθ(xt−1 | xt), xT ∼ N (0, I),
pθ(xt−1 | xt) := N (xt−1;µθ(xt, t), σθ(xt, t)I).
(2)
Ho et al. [11] has recently proposed denoising diffusion probabilistic models (DDPM), which considers the following specific parameterization of pθ(xt−1 | xt):
µθ(xt, t) = 1
αt
( xt −
βt√ 1− αt θ(xt, t)
) , σθ(xt, t) = β̃ 1/2 t where β̃t = { 1−αt−1 1−αt βt t > 1
β1 t = 1
(3)
where θ is a trainable denoising function. We denote µθ(xt, t) and σθ(xt, t) in Eq. (3) as µDDPM(xt, t, θ(xt, t)) and σDDPM(xt, t), respectively. The denoising function in Eq. (3) also corresponds to a rescaled score model for score-based generative models [23]. Under this parameterization, Ho et al. [11] have shown that the reverse process can be trained by solving the following optimization problem:
min θ L(θ) := min θ Ex0∼q(x0), ∼N (0,I),t|| − θ(xt, t)|| 2 2 where xt =
√ αtx0 + (1− αt) . (4)
The denoising function θ estimates the noise vector that was added to its noisy input xt. This training objective also be viewed as a weighted combination of denoising score matching used for training score-based generative models [23, 27, 12]. Once trained, we can sample x0 from Eq. (2). We provide the details of DDPM in Appendix A.

heading: 3.3 Imputation with diffusion models
text: Here, we focus on general imputation tasks that are not restricted to time series imputation. Let us consider the following imputation problem: given a sample x0 which contains missing values, we generate imputation targets xta0 ∈ X ta by exploiting conditional observations xco0 ∈ X co, where X ta and X co are a part of the sample space X and vary per sample. Then, the goal of probabilistic imputation is to estimate the true conditional data distribution q(xta0 | xco0 ) with a model distribution pθ(x ta 0 | xco0 ). We typically impute all missing values using all observed values, and set all observed values as xco0 and all missing values as x ta 0 , respectively. Note that time series imputation in Section 3.1 can be considered as a special case of this task.
Let us consider modeling pθ(xta0 | xco0 ) with a diffusion model. In the unconditional case, the reverse process pθ(x0:T ) is used to define the final data model pθ(x0). Then, a natural approach is to extend the reverse process in Eq. (2) to a conditional one:
pθ(x ta 0:T | xco0 ) := p(xtaT ) T∏ t=1 pθ(x ta t−1 | xtat ,xco0 ), xtaT ∼ N (0, I),
pθ(x ta t−1 | xtat ,xco0 ) := N (xtat−1;µθ(xtat , t | xco0 ), σθ(xtat , t | xco0 )I).
(5)
However, existing diffusion models are generally designed for data generation and do not take conditional observations xco0 as inputs. To utilize diffusion models for imputation, previous studies [12, 15, 16] approximated the conditional reverse process pθ(xtat−1 | xtat ,xco0 ) with the reverse process in Eq. (2). With this approximation, in the reverse process they add noise to both the target and the conditional observations xco0 . While this approach can impute missing values, the added noise can harm useful information in the observations. This suggests that modeling pθ(xtat−1 | xtat ,xco0 ) without approximations can improve the imputation quality. Hereafter, we call the model defined in Section 3.2 as the unconditional diffusion model.

heading: 4 Conditional score-based diffusion model for imputation (CSDI)
text: In this section, we propose CSDI, a novel imputation method based on a conditional score-based diffusion model. The conditional diffusion model allows us to exploit useful information in observed values for accurate imputation. We provide the reverse process of the conditional diffusion model, and then develop a self-supervised training method. We note that CSDI is not restricted to time series.

heading: 4.1 Imputation with CSDI
text: We focus on the conditional diffusion model with the reverse process in Eq. (5) and aim to model the conditional distribution p(xtat−1 | xtat ,xco0 ) without approximations. Specifically, we extend the parameterization of DDPM in Eq. (3) to the conditional case. We define a conditional denoising function θ : (X ta ×R | X co)→ X ta, which takes conditional observations xco0 as inputs. Then, we consider the following parameterization with θ:
µθ(x ta t , t | xco0 ) = µDDPM(xtat , t, θ(xtat , t | xco0 )), σθ(xtat , t | xco0 ) = σDDPM(xtat , t) (6)
where µDDPM and σDDPM are the functions defined in Section 3.2. Given the function θ and data x0, we can sample xta0 using the reverse process in Eq. (5) and Eq. (6). For the sampling, we set all observed values of x0 as conditional observations xco0 and all missing values as imputation targets x ta 0 . Note that the conditional model is reduced to the unconditional one under no conditional observations and can also be used for data generation.

heading: 4.2 Training of CSDI
text: Since Eq. (6) uses the same parameterization as Eq. (3) and the difference between Eq. (3) and Eq. (6) is only the form of θ, we can follow the training procedure for the unconditional model in Section 3.2. Namely, given conditional observations xco0 and imputation targets x ta 0 , we sample noisy
targets xtat = √ αtx ta 0 + (1− αt) , and train θ by minimizing the following loss function:
min θ L(θ) := min θ Ex0∼q(x0), ∼N (0,I),t||( − θ(x ta t , t | xco0 ))||22 (7)
where the dimension of corresponds to that of the imputation targets xta0 .
However, this training procedure has an issue. Since we do not know the ground-truth missing values in practice, it is not clear how to select xco0 and x ta 0 from a training sample x0. To address this issue, we develop a self-supervised learning method inspired by masked language modeling [28]. We illustrate the training procedure in Figure 2. Given a sample x0, we separate observed values of x0 into two parts, and set one of them as imputation targets xta0 and the other as conditional observations xco0 . We choose the targets x ta 0 through a target choice strategy, which is discussed in Section 4.3. Then, we sample noisy targets xtat and train θ by solving Eq. (7). We summarize how we set x co 0 and xta0 for training and sampling in Table 1. We also provide the algorithm of training and sampling in Appendix B.1.

heading: 4.3 Choice of imputation targets in self-supervised learning
text: In the proposed self-supervised learning, the choice of imputation targets is important. We provide four target choice strategies depending on what is known about the missing patterns in the test dataset. We describe the algorithm for these strategies in Appendix B.2.
(1) Random strategy : this strategy is used when we do not know about missing patterns, and randomly chooses a certain percentage of observed values as imputation targets. The percentage is sampled from [0%, 100%] to adapt to various missing ratios in the test dataset.
(2) Historical strategy: this strategy exploits missing patterns in the training dataset. Given a training sample x0, we randomly draw another sample x̃0 from the training dataset. Then, we set the intersection of the observed indices of x0 and the missing indices of x̃0 as imputation targets. The motivation of this strategy comes from structured missing patterns in the real world. For example, missing values often appear consecutively in time series data. When missing patterns in the training and test dataset are highly correlated, this strategy helps the model learn a good conditional distribution.
(3) Mix strategy: this strategy is the mix of the above two strategies. The historical strategy may lead to overfitting to missing patterns in the training dataset. The Mix strategy can benefit from generalization by the random strategy and structured missing patterns by the historical strategy.
(4) Test pattern strategy: when we know the missing patterns in the test dataset, we just set the patterns as imputation targets. For example, this strategy is used for time series forecasting, since the missing patterns in the test dataset are fixed to given future time points.

heading: 5 Implementation of CSDI for time series imputation
text: In this section, we implement CSDI for time series imputation. For the implementation, we need the inputs and the architecture of θ.
First, we describe how we process time series data as inputs for CSDI. As defined in Section 3.1, a time series is denoted as {X, M, s}, and the sample space X of X is RK×L. We want to handle X in the sample space RK×L for learning dependencies in a time series using a neural network, but the conditional denoising function θ takes inputs xtat and x co 0 in varying sample spaces that are a part of X as shown in white areas of xtat and xco0 in Figure 2. To address this issue, we adjust the conditional denoising function θ to inputs in the fixed sample space RK×L. Concretely, we fix the shape of the inputs xtat and x co 0 to (K × L) by applying zero padding to xtat and xco0 . In other words, we set zero values to white areas for xtat and x co 0 in Figure 2. To indicate which indices are padded, we introduce the conditional mask mco ∈ {0, 1}K×L as an additional input to θ, which corresponds to xco0 and takes value 1 for indices of conditional observations. For ease of handling, we also fix the output shape in the sample space RK×L by applying zero padding. Then, the conditional denoising function θ(xtat , t | xco0 ,mco) can be written as θ : (RK×L × R | RK×L × {0, 1}K×L)→ RK×L. We discuss the effect of this adjustment on training and sampling in Appendix D.
Under the adjustment, we set conditional observations xco0 and imputation targets x ta 0 for time series imputation by following Table 1. At sampling time, since conditional observations xco0 are all observed values, we set mco = M and xco0 = m
co X where represents element-wise products. For training, we sample xta0 and x co 0 through a target choice strategy, and set the indices of x co 0 as mco. Then, xco0 is written as x co 0 = m
co X and xta0 is obtained as xta0 = (M−mco) X. Next, we describe the architecture of θ. We adopt the architecture in DiffWave [13] as the base, which is composed of multiple residual layers with residual channel C. We refine this architecture for time series imputation. We set the diffusion step T = 50. We discuss the main differences from DiffWave (see Appendix E.1 for the whole architecture and details).
Attention mechanism To capture temporal and feature dependencies of multivariate time series, we utilize a two dimensional attention mechanism in each residual layer instead of a convolution architecture. As shown in Figure 3, we introduce temporal Transformer layer and a feature Transformer layer, which are 1-layer Transformer encoders. The temporal Transformer layer takes tensors for each feature as inputs to learn temporal dependency, whereas the feature Transformer layer takes tensors for each time point as inputs to learn temporal dependency.
Note that while the length L can be different for each time series as mentioned in Section 3.1, the attention mechanism allows the model to handle various lengths. For batch training, we apply zero padding to each sequence so that the lengths of the sequences are the same.
Side information In addition to the arguments of θ, we provide some side information as additional inputs to the model. First, we use time embedding of s = {s1:L} to learn the temporal dependency. Following previous studies [29, 30], we use 128-dimensions temporal embedding. Second, we exploit categorical feature embedding for K features, where the dimension is 16.

heading: 6 Experimental results
text: In this section, we demonstrate the effectiveness of CSDI for time series imputation. Since CSDI can be applied to other related tasks such as interpolation and forecasting, we also evaluate CSDI for these tasks to show the flexibility of CSDI. Due to the page limitation, we provide the detailed setup for experiments including train/validation/test splits and hyperparameters in Appendix E.2.

heading: 6.1 Time series imputation
text: Dataset and experiment settings We run experiments for two datasets. The first one is the healthcare dataset in PhysioNet Challenge 2012 [1], which consists of 4000 clinical time series with 35 variables for 48 hours from intensive care unit (ICU). Following previous studies [7, 8], we process the dataset to hourly time series with 48 time steps. The processed dataset contains around 80% missing values. Since the dataset has no ground-truth, we randomly choose 10/50/90% of observed values as ground-truth on the test data.
The second one is the air quality dataset [2]. Following previous studies [7, 21], we use hourly sampled PM2.5 measurements from 36 stations in Beijing for 12 months and set 36 consecutive time steps as one time series. There are around 13% missing values and the missing patterns are not random. The dataset contains artificial ground-truth, whose missing patterns are also structured.
For both dataset, we run each experiment five times. As the target choice strategy for training, we adopt the random strategy for the healthcare dataset and the mix of the random and historical strategy for the air quality dataset, based on the missing patterns of each dataset.
Results of probabilistic imputation CSDI is compared with three baselines. 1) Multitask GP [31]: the method learns the covariance between timepoints and features simultaneously. 2) GP-VAE [10]: the method showed the state-of-the-art results for probabilistic imputation. 3) V-RIN [32]: a deterministic imputation method that uses the uncertainty quantified by VAE to improve imputation. For V-RIN, we regard the quantified uncertainty as probabilistic imputation. In addition, we compare CSDI with imputation using the unconditional diffusion model in order to show the effectiveness of the conditional one (see Appendix C for training and imputation with the unconditional diffusion model).
We first show quantitative results. We adopt the continuous ranked probability score (CRPS) [33] as the metric, which is freuquently used for evaluating probabilistic time series forecasting and measures the compatibility of an estimated probability distribution with an observation. We generate 100 samples to approximate the probability distribution over missing values and report the normalized average of CRPS for all missing values following previous studies [34] (see Appendix E.3 for details of the computation).
Table 2 represents CRPS for each method. CSDI reduces CRPS by 40-65% compared to the existing baselines for both datasets. This indicates that CSDI generates more realistic distributions than other methods. We also observe that the imputation with CSDI outperforms that with the unconditional model. This suggests CSDI benefits from explicitly modeling the conditional distribution.
We provide imputation examples in Figure 4. For the air quality dataset, CSDI (green solid line) provides accurate imputations with high confidence, while those by GP-VAE (gray dashed line) are far from ground-truth. CSDI also gives reasonable imputations for the healthcare dataset. These results indicate that CSDI exploits temporal and feature dependencies to provide accurate imputations. We give more examples in Appendix G.
Results of deterministic imputation We demonstrate that CSDI also provides accurate deterministic imputations, which are obtained as the median of 100 generated samples. We compare CSDI with four baselines developed for deterministic imputation including GLIMA [21], which combined recurrent imputations with an attention mechanism to capture temporal and feature dependencies and showed the state-of-the-art performance. These methods are based on autoregressive models. We use the original implementations except RDIS.
We evaluate each method by the mean absolute error (MAE). In Table 3, CSDI improves MAE by 5-20% compared to the baselines. This suggests that the conditional diffusion model is effective to learn temporal and feature dependencies for imputation. For the healthcare dataset, the gap between the baselines and CSDI is particularly significant when the missing ratio is small, because more observed values help CSDI capture dependencies.

heading: 6.2 Interpolation of irregularly sampled time series
text: Dataset and experiment settings We use the same healthcare dataset as the previous section, but process the dataset as irregularly sampled time series, following previous studies [22, 35]. Since the dataset has no ground-truth, we randomly choose 10/50/90% of time points and use observed values at these time points as ground-truth on the test data. As the target choice strategy for training, we adopt the random strategy, which is adjusted for interpolation so that some time points are sampled.
Results We compare CSDI with two baselines including mTANs [22], which utilized an attention mechanism and showed state-of-the-art results for the interpolation of irregularly sampled time series. We generate 100 samples to approximate the probability distribution as with the previous section. The result is shown in Table 4. CSDI outperforms the baselines for all cases.

heading: 6.3 Time series Forecasting
text: Dataset and Experiment settings We use five datasets that are commonly used for evaluating probabilistic time series forecasting. Each dataset is composed of around 100 to 2000 features. We predict all features at future time steps using past time series. We use the same prediction steps as previous studies [34, 37]. For the target choice strategy, we adopt the Test pattern strategy.
Results We compare CSDI with four baselines. Specifically, TimeGrad [25] combined the diffusion model with a RNN-based encoder. We evaluate each method for CRPS-sum, which is CRPS for the distribution of the sum of all time series across K features and accounts for joint effect (see Appendix E.3 for details).
In Table 5, CSDI outperforms the baselines for electricity and traffic datasets, and is competitive with the baselines as a whole. The advantage of CSDI over baselines for forecasting is smaller than that for imputation in Section 6.1. We hypothesize it is because the datasets for forecasting seldom contains missing values and are suitable for existing encoders including RNNs. For imputation, it is relatively difficult for RNNs to handle time series due to missing values.

heading: 7 Conclusion
text: In this paper, we have proposed CSDI, a novel approach to impute multivariate time series with conditional diffusion models. We have shown that CSDI outperforms the existing probabilistic and deterministic imputation methods.
There are some interesting directions for future work. One direction is to improve the computation efficiency. While diffusion models generate plausible samples, sampling is generally slower than other generative models. To mitigate the issue, several recent studies leverage an ODE solver to accelerate the sampling procedure [12, 38, 13]. Combining our method with these approaches would likely improve the sampling efficiency.
Another direction is to extend CSDI to downstream tasks such as classifications. Many previous studies have shown that accurate imputation improves the performance on downstream tasks [7, 18, 22]. Since conditional diffusion models can learn temporal and feature dependencies with uncertainty, joint training of imputations and downstream tasks using conditional diffusion models would be helpful to improve the performance of the downstream tasks.
Finally, although our focus was on time series, it would be interesting to explore CSDI as imputation technique on other modalities.
Acknowledgements and Disclosure of Funding
This research was supported by NSF(#1651565, #1522054, #1733686), ONR (N000141912145), AFOSR (FA95501910024), ARO (W911NF-21-1-0125) and Sloan Fellowship.

Review summaries:
This work proposes a probabilistic data imputation model for multivariate time series. The proposed method uses a conditional diffusion model to estimate the (masked) conditional distribution of missing data given observed values. To utilise information from observed values, the authors employ a neural attention mechanism to capture spatiotemporal dependencies of time series. The authors use self-supervised learning for training the model.  In experiments, the author demonstrate competitive performance of their method against a number of comparable methodologies on time series imputation and forecasting tasks. 

This paper focuses on the problem of the probabilistic imputation of missing data in time series. This work proposes a novel approach for probabilistic imputation based on the score-based diffusion models. Experiments on time series imputation tasks show that the proposed method achieves better performance than the current SOTA methods in this space.