🤖 AI

Discretizing Continuous Time Series for Imputation with Masked Diffusion Training

The paper proposes the Masked Diffusion Time-series Imputation Model (MDTIM), which improves time series imputation by structurally separating masked and observed values and introducing Stochastic Discretization to adapt masked diffusion training for continuous, ordinal data, thereby achieving superior robustness and scalability over existing methods.

Original authors: Dongbin Kim, Seungyun Lee, Geonwoo Shin, Jaewook Lee

Published 2026-08-20

📖 4 min read☕ Coffee break read

CC BY 4.0

Original authors: Dongbin Kim, Seungyun Lee, Geonwoo Shin, Jaewook Lee

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Time series data is the rhythm of the modern world, a continuous stream of numbers recording everything from the temperature in a city to the heartbeat of a patient. These records are rarely perfect; sensors fail, signals drop, and gaps appear in the data. To make sense of the whole picture, scientists must fill in these missing pieces, a task known as imputation. The challenge lies in the nature of the data itself: it is continuous, flowing smoothly from one moment to the next, yet it is also ordered, where a value at one second is intimately connected to the value at the next. Traditional methods often struggle here, either treating the missing spots as simple zeros that confuse the pattern or using complex mathematical models that try to guess the missing value by predicting the noise that was added to it, rather than the value itself.

A team of researchers at Seoul National University has developed a new approach to this problem, one that changes how computers view the missing data. They created a system called the Masked Diffusion Time-series Imputation Model, or MDTIM. Instead of trying to predict the noise, this model treats the missing data like a blank space in a sentence that needs to be filled with the correct word. In language models, a special symbol like [MASK] is used to hide a word, and the model learns to guess the original word based on the context. The researchers realized this same logic could work for time series, but only if they could solve a fundamental mismatch: time series are smooth and continuous, while language words are distinct and separate.

To bridge this gap, the researchers introduced a method called Stochastic Discretization. Imagine trying to describe a smooth, flowing river using only a limited set of stepping stones. If you simply round the water level to the nearest stone, you lose the subtle variations of the flow. The new method adds a tiny, controlled amount of randomness when converting the smooth data into these stepping stones. This randomness ensures that, on average, the information is preserved even though the data has been simplified into categories. Furthermore, the model understands that these categories are not just random labels; they have an order. A value that is slightly higher than the current one is more likely to be the correct guess than a value that is wildly different. By teaching the model to respect this order, the system can reconstruct the smooth flow of the original data with high precision.

The results of this approach are striking. The researchers tested their model on a variety of real-world datasets, including electricity consumption, weather patterns, and hospital patient records. In every case, the new model outperformed existing state-of-the-art methods. It was particularly effective when large chunks of data were missing, such as when a sensor failed for a long period. In these difficult scenarios, where other models struggled to guess the missing values, the new system maintained its accuracy. It also proved to be much faster, completing its calculations in seconds where older, similar methods took minutes or even hours.

The study confirms that treating missing time series data as a structured puzzle to be solved, rather than a noisy signal to be cleaned, leads to better results. By combining the logic of masked language models with a clever way of handling continuous numbers, the researchers have created a tool that is both more accurate and more efficient. This work suggests that the future of data analysis may lie in adapting techniques from one field, like language processing, to solve deep problems in another, provided the underlying differences are carefully bridged. The model does not just fill in the blanks; it reconstructs the story of the data with a clarity that was previously out of reach.

Technical Summary: Masked Diffusion Time-series Imputation Model (MDTIM)

Problem Statement

Time series imputation is essential for reliable analysis, yet real-world data often suffer from missing values due to sensor failures or transmission errors. The paper identifies two primary limitations in existing approaches:

Representation Misalignment: Current methods, including Transformer-based discriminative models and continuous diffusion models, embed missing and observed values within the same data manifold. Even with attention masks, placeholder values (e.g., zeros) occupy the same embedding space as valid observations, potentially introducing spurious correlations.
Objective Misalignment: Continuous diffusion models (e.g., CSDI) are typically trained to predict the Gaussian noise added to the data rather than the original signal itself. This indirect objective is suboptimal for a task explicitly aimed at recovering the original values.

Furthermore, while Masked Diffusion Models (MDM) have shown promise in Natural Language Processing (NLP) by structurally separating missing tokens (via a [MASK] symbol) and directly predicting original values, applying them to time series is challenging. Time series data are inherently continuous and ordinal, whereas language tokens are discrete and categorical. Simple quantization severs ordinal relationships and loses fine-grained information.

Methodology

The authors propose the Masked Diffusion Time-series Imputation Model (MDTIM), a framework that bridges the gap between discrete masked diffusion and continuous time series modeling. The core components are:

1. Stochastic Discretization

To handle the continuous nature of time series within a discrete token space without losing information, the authors introduce a stochastic tokenization pipeline:

Instance-Adaptive Normalization: Input windows are normalized using statistics computed solely from observed values to prevent data leakage.
Stochastic Token Generation: Normalized values are projected onto a discrete grid of size $K$ . Before rounding to the nearest integer, uniform noise $\epsilon \sim U(-0.5, 0.5)$ is injected. This ensures that the expected value of the discrete token matches the continuous input ( $E[z] \approx x$ ), preserving information in expectation and allowing the model to learn dynamics beyond fixed grid resolution.

2. Ordinal-Aware Soft Labeling

Standard classification treats tokens as independent classes, ignoring the ordinal nature of time series (e.g., token $k$ is closer to $k+1$ than $k+10$ ). MDTIM employs Ordinal-Aware Soft Labeling:

Instead of one-hot targets, the model is trained with soft target distributions derived from a truncated Gaussian kernel centered on the ground-truth token.
This assigns higher probability to neighboring tokens, enforcing continuity. Crucially, the [MASK] token is assigned zero probability in the target distribution, maintaining structural separation between missing and observed states.

3. Model Architecture and Training

Backbone: A Factorized Temporal-Variate Transformer based on the Diffusion Transformer (DiT). It processes time ( $T$ ) and feature ( $C$ ) axes via interleaved blocks, capturing both temporal dependencies and cross-variable correlations.
Training Objective: The model minimizes a dual-domain loss:
- Ordinal-Aware Masked Diffusion Loss: A weighted cross-entropy loss using the soft labels described above, adapted from the continuous-time training paradigm of MDLM.
- Spectral Consistency Regularization: An auxiliary loss minimizing the $L_1$ distance between the Fast Fourier Transform (FFT) of the reconstructed sequence and the ground truth, ensuring global temporal coherence.
Continuous Reconstruction: During inference, the model generates a probability distribution over discrete tokens. The final continuous value is recovered by computing the probability-weighted expectation of the predicted distribution (Expectation-based Unmasking), followed by denormalization.

Key Contributions

Masked Diffusion Framework for Time Series: The first adaptation of the masked diffusion paradigm to time series imputation, providing explicit structural separation between missing and observed representations.
Novel Quantization Pipeline: The introduction of Stochastic Discretization and Ordinal-Aware Soft Labeling, which preserves ordinal relationships and continuous dynamics within a categorical token space without information loss.
Precise Reconstruction Strategy: The use of an Expectation-based Unmasking strategy and Spectral Consistency Regularization to enable precise reconstruction of continuous time series from discrete outputs.

Experimental Results

The authors evaluated MDTIM on four benchmarks: Energy, ETTh, Weather, and a synthetic Sine dataset, under various missing patterns (Uniform and Geometric) and ratios (30%, 50%, 70%).

Imputation Accuracy: MDTIM consistently achieved the lowest Mean Absolute Error (MAE) across real-world datasets (Energy, ETTh, Weather) compared to 19 baselines, including RNNs (BRITS, GRU-D), Transformers (SAITS, PatchTST), and generative models (CSDI, GP-VAE).
Robustness: The model demonstrated superior robustness under geometric masking (consecutive missing values), where many baselines degraded significantly.
Scalability: Performance advantages grew with sequence length (up to $L=192$ ), indicating effective capture of long-range dependencies without attention dilution.
Probabilistic Quality: In terms of Continuous Ranked Probability Score (CRPS), MDTIM outperformed the diffusion-based baseline CSDI, providing better-calibrated uncertainty estimates.
Efficiency: MDTIM significantly reduced inference time compared to continuous diffusion models (e.g., CSDI), completing inference in seconds versus minutes for large-scale models, while maintaining higher accuracy.
Downstream Impact: When used as a pre-processing step for a fixed PatchTST forecaster, MDTIM-imputed data led to substantially lower forecasting errors compared to other imputation methods.

Significance and Claims

The paper claims that MDTIM addresses the fundamental misalignment in existing time series imputation methods by aligning both the representation (via orthogonal [MASK] tokens) and the learning objective (direct prediction of original values) with the task of imputation.

The authors assert that their method successfully bridges the gap between discrete masked diffusion and continuous time series data. By leveraging stochastic discretization and ordinal-aware objectives, MDTIM achieves state-of-the-art robustness and scalability across diverse missing scenarios. The work suggests that the masked diffusion paradigm, previously confined to discrete modalities like NLP, is highly effective for continuous data when paired with appropriate discretization strategies that preserve ordinal and spectral properties.

The paper remains modest regarding limitations, noting that the optimal vocabulary size may vary by dataset, the spectral consistency loss is less effective on highly irregular signals (like PhysioNet), and future work is needed to extend the model to extremely long sequences and non-stationary domains.