Discretizing Continuous Time Series for Imputation with Masked Diffusion Training

AI summary

Overview

• Research area: Time series analysis, specifically missing-value imputation, drawing on discrete masked diffusion models from natural language processing and on Transformer architectures.
• Technical level: Advanced. The paper assumes familiarity with diffusion models, absorbing-state discrete diffusion (D3PM, MDLM), Transformer backbones, vector quantization, and spectral (Fourier) regularization.
• Scope: The paper introduces and evaluates MDTIM (Masked Diffusion Time-series Imputation Model), a masked-diffusion framework for imputing continuous multivariate time series, with a stochastic tokenization pipeline that preserves ordinal structure and continuous reconstruction via expectation decoding.

What This Paper Is About

Time series collected in the real world are frequently incomplete due to sensor failures or transmission errors, and recovering the missing values accurately matters for everything downstream. The authors argue that existing methods have two weaknesses: missing and observed values are embedded in the same representation space (with zero placeholders mixed in with valid data), and continuous diffusion models such as CSDI are trained to predict the Gaussian noise that was added rather than the original signal. The paper's goal is to adapt masked diffusion, where a special [MASK] token is structurally separate from valid values and the model directly predicts the original content, to the continuous and ordinal nature of time series.

Key Contributions

1. A masked diffusion framework for time series imputation. The authors apply the MDM training paradigm to imputation so that the [MASK] token is structurally orthogonal to valid observations, separating missing and observed states in the representation and having the model directly predict original values rather than added noise.
2. Stochastic Discretization. Continuous values are normalized per instance using only observed values (to avoid leakage), projected onto a grid of K bins, and assigned tokens after injecting uniform noise ε ∼ U(−0.5, 0.5) before rounding, so that the token assignment is unbiased in expectation.
3. Ordinal-Aware Soft Labeling. Instead of one-hot targets, targets come from a truncated Gaussian kernel over token indices (window size w = 2, with zero probability assigned to the [MASK] token and distant bins), preserving the ordinal relationships between neighboring bins.
4. Expectation-based unmasking and spectral consistency regularization. Continuous values are reconstructed as the probability-weighted expectation over the predicted token distribution, averaged over M = 10 noise samples, and training adds an L1 loss between the Fourier transforms of the reconstruction and the ground truth (adapted from DiffusionTS).

Main Findings

• Best MAE on real-world benchmarks. In Table 1 (sequence length L = 48, MAE averaged over 3 random seeds), MDTIM achieves the lowest MAE on Energy, ETTh, and Weather under both Uniform and Geometric masking. On Energy at 30% Uniform it reaches 0.044 (70% Uniform: 0.085); on ETTh 0.127 at 30% Uniform (70%: 0.211); on Weather 0.032 at 30% Uniform.
• Frequency-domain priors win on synthetic data but degrade on real data. FGTI is best on the synthetic Sine dataset, where MDTIM trails marginally (0.002 vs. 0.001 at 30% Uniform). But on ETTh, FGTI degrades substantially (0.218 vs. MDTIM's 0.127 at 30% Uniform).
• Robustness to contiguous missingness. Under Geometric masking, most baselines suffer significant degradation while MDTIM on Energy at 30% increases only from 0.044 to 0.053.
• Advantage grows with sequence length. At L = 192 with 30% Uniform missing, MDTIM records MAE 0.123 versus BRITS 0.192 and SAITS 0.145.
• Better probabilistic calibration. In CRPS comparisons against CSDI (Table 3), MDTIM is lower in every setting, e.g. ETTh at 30% Uniform: 0.0863 vs. 0.1120; Energy at 30% Uniform: 0.0429 vs. 0.0629. On the periodic Sine dataset MDTIM's CRPS is nearly constant across missing rates (0.0020).
• Effective under natural real-world missingness. On PhysioNet 2012, which has roughly 80% natural missingness from irregular ICU sampling, MDTIM achieves 0.236 MAE at 30% Uniform versus SAITS at 0.283. FGTI degrades the most here, even underperforming BRITS (0.338 vs. 0.333).
• Large efficiency gains over continuous diffusion. On Energy (uniform missing, inference on the full test set with M = 10 noise samples for MDTIM and T = 50 diffusion steps for CSDI), the Small MDTIM runs in 4.05s with 0.65M parameters, against CSDI at 401.69s with 1.19M parameters and up to 894.76s at the Large scale. The 0.65M Small MDTIM also surpasses the 88.24M Large SAITS in MAE (0.046 vs. 0.074 at 30%).
• Imputation quality transfers to forecasting. With a fixed pre-trained PatchTST forecaster, MDTIM reduces forecasting MAE by 19–26% over SAITS and 16–19% over CSDI across horizons H = 4, 6, 8, 12 on Energy.
• Both proposed components help, but soft labeling is data-dependent. Expectation-based unmasking consistently reduces error across all datasets. Soft labeling adds a large gain on the periodic Sine dataset (0.0071 → 0.0022 at 30% Uniform), a milder gain on ETTh, and no further improvement beyond expectation unmasking on Energy.
• Vocabulary size shows a trade-off. On ETTh, K = 20 underperforms from high quantization error (0.138 at 30% Uniform), K = 60 slightly degrades from harder classification (0.129), and K = 40 is best (0.127).

Methodology in Plain English

The authors take an approach designed for text—masked diffusion, where tokens are progressively replaced with a [MASK] symbol and the model learns to restore them—and adapt it to numbers.

The challenge is that text tokens are discrete categories, while time series values are continuous and ordered. To bridge this, they first normalize each input window using statistics computed only from observed values. Each value is then placed on a grid of K bins. Rather than rounding deterministically, they add random noise drawn uniformly between −0.5 and 0.5 before rounding. A value at grid coordinate 3.6, for example, becomes token 4 with 60% probability and token 3 with 40% probability, so the expected token equals the original coordinate and no information is systematically lost to the grid.

Because bins have an order (bin 5 is closer to bin 6 than to bin 15), the training target is not a single correct class but a soft distribution shaped like a truncated Gaussian centered on the true bin. This tells the model that near-misses are much better than far-misses, which plain cross-entropy would not.

The backbone is a Transformer that alternates between processing the time axis and the channel (variate) axis, built on the Diffusion Transformer design with rotary position embeddings, and its output vocabulary spans 1.5K entries covering the range [−1.5, 1.5] to absorb distribution shift. Training combines the masked-diffusion cross-entropy with the soft labels and a spectral loss that compares the Fourier transform of the current reconstruction against that of the clean signal, encouraging global temporal structure to be preserved rather than only local accuracy.

At inference, because the tokenizer injected randomness, the model runs multiple times with different noise samples, averages the predicted probability distributions, and returns the expected value over bin centers, which is then denormalized back to the original scale.

Why This Matters

Impact on research. The paper offers a concrete argument that the dominant continuous-diffusion recipe for imputation is misaligned with the task: predicting added noise is an indirect route to recovering original values, and the paper's masked-diffusion alternative is both more direct and, in these experiments, faster and more accurate. It also shows that masked diffusion—previously studied almost exclusively in discrete modalities such as NLP—can be extended to a continuous, ordinal domain without abandoning the discrete training paradigm.

Real-world applications:

• Clinical monitoring and ICU analytics, where vital-sign streams are irregularly sampled and often heavily incomplete (the PhysioNet 2012 benchmark reflects this setting).
• Energy grid and smart-meter analytics, where sensor dropouts and transmission errors corrupt consumption or load readings.
• Environmental and weather sensing networks, where station failures create gaps in multivariate measurements.
• Industrial predictive maintenance, where long gaps from sensor malfunction must be filled before downstream forecasting or anomaly detection.

Industry relevance. The efficiency results matter for deployment: MDTIM's Small variant completes inference in 4.05s on Energy where CSDI requires 401.69s, and a 0.65M-parameter MDTIM outperforms an 88.24M-parameter SAITS. That combination of accuracy, low latency, and small model size is relevant to production pipelines where imputation is a preprocessing step rather than the final product.

Future Directions

• Data-adaptive discretization. The optimal vocabulary size K appears to depend on the range and granularity of the underlying signal; automatically selecting the resolution per dataset is proposed as a way to improve generalization.
• Adaptive frequency objectives. The spectral consistency loss assumes stable frequency components and was disabled on PhysioNet; designing a frequency objective that modulates its own influence based on local signal regularity is an open problem.
• Extremely long and non-stationary sequences. The current evaluation covers standard benchmarks of moderate length; scaling MDTIM to very long sequences and to multi-scale or non-stationary domains such as financial tick data remains future work.
• Understanding when ordinal regularization pays off. The ablation shows soft labeling helps greatly on periodic Sine, mildly on ETTh, and not at all on Energy—identifying in advance which signals benefit from ordinal-aware guidance is left unresolved.

Target Audience

Researchers and graduate students working on time series analysis, generative modeling, or missing-data problems, particularly those already familiar with diffusion models and Transformer architectures. It is also relevant to applied machine learning engineers and practitioners in healthcare, energy, and environmental monitoring who need fast, accurate imputation as part of a larger pipeline, and to anyone interested in transferring discrete masked-diffusion methods from language modeling into continuous domains.

Authors’ abstract

Time series imputation is a crucial area for reliable time series analysis, yet it remains challenging due to the complex temporal dynamics and noise of real-world data. Existing approaches, however, exhibit two limitations: missing and observed values are embedded within the same representation space without explicit structural separation, and continuous diffusion-based methods are trained to predict added noise rather than the original signal. To address these, we propose the Masked Diffusion Time-series Imputation Model (MDTIM), which leverages the training paradigm of masked diffusion model for imputation tasks. The MASK token is structurally orthogonal to valid observations, and the model directly predicts the original values, naturally aligning both the representation and the learning objective with the imputation task. To bridge the gap between discrete masked diffusion and the continuous, ordinal nature of time series, we further introduce Stochastic Discretization, which maps continuous values to ordinal-aware tokens while preserving continuous dynamics. Our experiments on diverse benchmarks confirm that MDTIM achieves superior robustness and scalability, consistently outperforming state-of-the-art deterministic and generative baselines across various missing scenarios.