arXiv CSopen access

Transformer-based Diffusion models for Hydrological Time Series Probabilistic Imputation and Forecasting

Transformer-based Diffusion models for Hydrological Time Series Probabilistic Imputation and Forecasting Ferdinand Bhavsar1 , Lionel Benoit1 , Maxime Savatier2 , and Edith Gabriel1 1

Biostatistics and Spatial Processes (BioSP), INRAE, Avignon 84914, France 2 Scientific and Technical Department, ANDRA, Bure 55290, France July 24, 2026

arXiv:2607.21200v1 [stat.ML] 23 Jul 2026

Abstract The modeling of hydrometeorological time series with limited observations is a key challenge in the monitoring of hydro-systems and water resources, as well as for flood or drought risk assessment. Due to the high variability of the underlying processes and the sparsity of available measurements, traditional statistical approaches often struggle to accurately represent their dynamics. In this context, recent advances in deep learning offer a promising direction for improving the representation and generation of complex temporal processes sampled at several observation sites. This study investigates the application of transformer-based diffusion models to the simulation and reconstruction of hydrological time series. The proposed framework is applied to the joint modeling of water quantity and quality at six sites spread across three adjacent headwater catchments located in North-East France on a limestone plateau covered by forests and field crops. The model is calibrated and validated using available observational data, which has been quality controlled and corrected for sensor drift and malfunction through collaborative efforts by LNE metrology expertise and Andra monthly quality control over more than 15 years. Its performance is compared with several established baseline approaches commonly used for time series modeling. Quantitative evaluation metrics are employed to assess the ability of the proposed method to reproduce key temporal characteristics of the observed signals in two settings: the imputation of incomplete time series and the forecasting of upcoming hydrological conditions. Results support the effectiveness of the transformer-based approach and highlight its capacity to capture and simulate the complex patterns present in hydrological data. In particular, the results indicate that diffusion models can efficiently sample realistic time series distributions under observation settings with variable missing data for both forecasting and imputation.

1

Introduction

In hydrological sciences, there is a growing need for the joint modeling of the quantity of water in rivers (i.e., streamflow) and the quality of this water (i.e., chemical, physical, and sometimes biotic signature). This applies in particular in the field of ecohydrology where the quantity and quality of water jointly impact stream ecosystems and stream-dependent organisms (Teurlincx et al., 2019), and in the field of catchment hydrology where physical and chemical tracers (e.g., isotopic signature) are used to identify the origin, storage, and path of water through different compartments of the catchment (Guo et al., 2020). Water quality and quantity are increasingly monitored by automatic and multi-sensor stream gauges, creating a growing database of multivariate time series often recorded at several locations of the river network (von Freyberg et al., 2017; Moiroux-Arvis et al., 2023). Such datasets often include gaps caused by sensor malfunctions, as well as heterogeneous start and end dates due to diverse station setup and discontinuation. This makes hydrological time series heterogeneous, and calls for methods of data augmentation (e.g. gap filling, spatial interpolation), either before integration in process-based models, or for other purposes linked to water management (decision support, early warning, digital twins, etc). Indeed, process-based models usually require structured input information representative of the different water storages and fluxes throughout the catchment of interest to properly depict 1

streamflow generation processes, and available measurements rarely allow to fully quantify them (Rashid et al., 2026; Alornyo et al., 2026). Data-based approaches are particularly well suited for data augmentation, and also allow for data emulation (e.g., stochastic weather (Obakrim et al., 2025) or streamflow generators (Brunner et al., 2019)) as well as now-casting. We will therefore focus on this class of models, and more specifically on ensemble (a.k.a. probabilistic, generative) approaches allowing for the simulation of large ensembles of equally likely outcomes that can be used to assess the variability of the process at hand and to quantify the uncertainty of the simulations. Over the years, powerful parametric probabilistic models based on assumptions about the underlying probabilistic structure of the data have been developed. These models include Markov chains (Shevnina and Silaev, 2019), Autoregressive Integrated Moving Average (ARIMA), (Tunnicliffe Wilson, 2016), and Gaussian random fields (Allcroft and Glasbey, 2003). The main limitation of such models is their dependence on a prior parameterization that restricts their flexibility and, thereby, their ability to capture complex real-world phenomena realistically. To avoid this pitfall, non-parametric resampling methods get rid of explicit assumptions about the underlying distribution, and rely instead on the stochastic resampling of a training dataset. For example, MultiPoints statistics (MPS) simulation is a sequential simulation technique based on pattern matching (Guardiano and Srivastava, 1993; Strebelle and Journel, 2000; Mariethoz et al., 2010). Advances in deep learning have recently drawn considerable attention because these methods effectively capture complex and nonlinear relationships in a wide range of applications (Cybenko, 1989), and because they can be trained efficiently and flexibly when large quantities of data are available (Goodfellow et al., 2016). Over the past few years, long short-term memory (LSTM) models (Hochreiter and Schmidhuber, 1997) have been extensively explored for hydrological time series modeling, including the imputation of missing values (Ren et al., 2022; Weilisi and Kojima, 2022) and forecasting (Feng et al., 2020; Kratzert et al., 2018). While they are theoretically able to model some long-range temporal dependencies (year long dependencies in hydrological time series), in practice LSTM models process time series sequentially and therefore struggle to capture complex global interactions. In contrast, the more recent transformer-based models leverage self-attention to model long-term dependencies more efficiently and in parallel (Vaswani et al., 2023). In hydrology, Hu et al. (2024) used a transformer-based model for predictive forecasting of dissolved oxygen in rivers with greatly improved performance over baselines. The deep-learning methods discussed above do not explicitly account for uncertainty, which is nevertheless essential in hydrology because the complexity of hydrological processes and the limited observations introduce substantial uncertainty in time series modeling. Rather than relying on deterministic models that attempt to recover a single estimate of the real time series, it is therefore preferable to endorse probabilistic methods that model uncertainty explicitly. In the field of deep learning, this is the approach followed by deep generative models, which are a new alternative class of simulation methods that sample from a simple and known latent distribution, and learn through optimization how to transform this latent distribution into a complex and unknown target distribution. Deep generative models have gained significant traction in the last few years for their outstanding performance. They learn to reproduce the complexity of the phenomenon using a training dataset large enough to be representative of its variability and complexity, and exhibit great results in learning to generate synthetic (i.e., never observed) realizations of complex multivariate distributions hardly distinguishable from observations (Arjovsky and Bottou, 2017). Deep generative models have been applied in generating realistic pictures such as face images (Brock et al., 2019; Zhu et al., 2020) and art images (Karras et al., 2019). For a long time Generative Adversarial Networks (GANs) (Goodfellow et al., 2014) have been the leading generative models, and have been successfully applied in hydrology; for example, Chen et al. (2025) designed and trained a custom GAN model for flood forecasting. However, the practical application of GANs often suffers from stability issues, and training them can be a challenge (Goodfellow et al., 2014; Arjovsky et al., 2017; Arjovsky and Bottou, 2017; Bhavsar et al., 2024). Due to these drawbacks, GANs have recently been overshadowed by denoising diffusion generative models (Sohl-Dickstein et al., 2015; Ho et al., 2020). Owing to their greater training stability and improved performance compared to GANs, these models have emerged as the state-of-the-art across a wide range of applications, including image generation and manipulation (Gu et al., 2022; Kawar et al., 2023; Saharia et al., 2022a; Zhang et al., 2023; Ruiz et al., 2023; Dhariwal and Nichol, 2021), time series generation (Yang et al., 2023; Tashiro et al., 2021; Bedin et al., 2025) and time series anomaly detection (Yang et al., 2023). Diffusion models have also attracted increasing attention in geosciences, including hydrology. For instance, C. M. and Thangarasu (2025) compared multiple 2

models for dam inflow forecasting. The authors found that a hybrid architecture transformer-LSTM performed better than diffusion and transformers architectures alone. Meanwhile, Ou et al. (2025) designed a diffusion architecture with a Transformer backbone that outperformed LSTM baselines for forecasting using meteorological forecasts as covariates, the imputation of missing values being performed beforehand using a reanalysis dataset. In Li et al. (2024) the authors used a diffusion model with a Transformer backbone for precipitation nowcasting, outperforming the performances of U-NET architectures. It appears that most existing hydrology-oriented diffusion approaches are designed for specific tasks, typically focusing on either forecasting or data imputation in isolation. This task-specific design limits the flexibility and in turn the broader potential of diffusion models, which are inherently capable of jointly addressing multiple objectives within a unified probabilistic framework, including generation, forecasting, and missing-value imputation. In contrast to hydrology, multi-task diffusionbased models have been successfully applied to the modeling of electrocardiogram time series (Alcaraz and Strodthoff, 2023), stock market and electricity production data (Yuan and Qiao, 2024), as well as air quality and traffic time series (Tashiro et al., 2021; Wang et al., 2023). In the present study we propose a modified diffusion model based on an existing transformer architecture to tackle both imputation and forecasting with a single model for applications in hydrology. We find that the resulting model is competitive, achieving lower error than multiple baselines on two datasets The remainder of this document is organized as follows. First, we introduce some background on diffusion models and, more specifically, on conditional score-based diffusion models. We then show what modifications have been implemented to improve the architecture of the model, explaining our choices. Next, we present results on forecasting and imputation tasks on two datasets, one real and one synthetic, and compare these results to existing baselines. Finally, we provide some conclusions about the application of transformer-based diffusion models to hydrological time series probabilistic imputation and forecasting, and we propose some lines for future research.

2

Background

2.1

Example dataset

Figure 1: OPE hydrological monitoring network (Marne catchment - North-East France): (a) Site description (blue triangles denote stream gauge locations); (b) Multivariate time series profile for some recorded features across 500 timesteps, including data gaps; (c) Percentage of missing data per variable at station 0PE90012. This study addresses the challenge of modeling hydrological time series within the framework of the Perennial Observatory of the Environment (in French Observatoire Pérenne de l’Environnement OPE), a long-term research infrastructure dedicated to monitoring environmental and climate changes around the future Cigéo deep geological repository for radioactive waste. Data are collected from six 3

monitoring stations located along the Orge, Saulx, and Ormain rivers, with a temporal resolution of 4 hours, resulting in an extensive dataset for analysis. Figure 1 gives an overview of the dataset, which consists of water-quality time series covering an observation period from 2012 to 2024, i.e., 25980 timesteps. Each station records 13 variables. Among these, 10 were measured for the entire period: fluorescent dissolved organic carbon (fDOC, also referred to as fDOM, or fluorescent dissolved organic matter), dissolved oxygen in mg/L (DO), oxygen saturation in % (SAT), electrical conductivity at 25°C (EC), Polycyclic aromatic hydrocarbons (PAH), nitrate concentration (NO3), water pH (pH), water temperature (TEM), turbidity (TUR), and water level (HTE). Additionally, 3 variables were only measured during the specific years of the measurements: Cyanobacteria (CYA), chlorophyll-a (CHL), dissolved chlorides (CL). The time series exhibit heterogeneous magnitudes: for example, pH remains within a narrow near-neutral range (mean=7.90, max=9.65), while turbidity spans several orders of magnitude (mean=22.52, max=3000.00). Data was corrected for bias by the French National Laboratory of Metrology and Testing (LNE) using monthly quality control by the OPE technicians and contractors (Guigues et al., ress). Some sensor (such as Chl-a, CYA, figure 1) were only deployed during the first years of the station life and for some stations the measurement timesteps were modified across the station life to ensure sufficient power stability on the stations in the early deployment periods between 2012 and 2017. Furthermore, frequent droughts and occasional sensor failure led to added missing data. Finally, in some cases for the period before 2017, sensors timesteps were not always synchronized between station. As a result, the proportion of missing data calculated in this work (which does not account for timesteps changes) can be substantial for some sensors, ranging from 25% to 90%. It should be noted, that the high proportion of missing values reported for the OPE900XX stations is not fully representative of the actual sensor downtime or of the fraction of the year effectively covered by measurements. Rather, it also results from the non-homogeneous acquisition timestep of the raw data (ranging from 4 h to 8 h) which inflate the apparent missing-data rate when computed as a simple ratio of populated rows to expected time steps. This significantly increases the difficulty of modeling the hydrological and biogeochemical dynamics of the site. Another key characteristic of the dataset is its complexity: it exhibits temporal correlations across multiple time scales, cross-correlations between variables, and non-Gaussian distributions of most target variables. In this context, two main tasks, imputation and forecasting, will be considered as goals for the study. Imputation aims to reconstruct missing observations in time series, providing experts with indicative information for more robust hydrological modeling, in particular when models require data with constant time steps and without gaps, which can be most challenging to have for long term measurements in rural or remote rivers and streams, especially when using self-powered stations (solar or battery powered). Forecasting helps to anticipate upcoming hydrological conditions, for instance accurate hydrometeorological forecasts can help predict floods, and better assess the transport and dispersion of pollutants linked to those episodes.

2.2

Diffusion Models

Over the past few years, diffusion models have become one of the dominant generative modeling frameworks, in image synthesis and are attracting increasing attention for probabilistic forecasting and data imputation in scientific time series applications, including hydrology. Previous generative methods learned the parameters of a single transformation linking an easy-to-sample distribution to the target distribution. Variational Auto-Encoders (VAEs) learn this transformation using variational inference, mapping the data into a latent distribution and then back to the data space, thus learning the relation between the latent and data distributions (Kingma and Welling, 2014). Generative Adversarial Networks (GANs) learn the transformation by training a generator model to generate realistic data and a discriminator model to distinguish generated from real data, with both models trained alternately in an adversarial game. In contrast to the above models involving a single transformation, diffusion models, inspired by non-equilibrium thermodynamics, generate data by learning an iterative denoising process that transform a simple, easy to sample noise distribution into the data distribution. Diffusion models have supplanted GANs and VAEs due to their performance in learning complex distributions and because their training is stable, i.e., less prone to issues such as mode collapse, adversarial optimization instabilities, and highly sensitive convergence dynamics (Sohl-Dickstein et al., 2015; Ho et al., 2020; Dhariwal and Nichol, 2021). Denoising diffusion works by approximating the reverse of a noising process. This noising pro4

Figure 2: Illustration of the diffusion process, where t ∈ {0, 1, . . . , T } indexes the diffusion timestep, with t = 0 corresponding to the original data distribution and t = T to the fully noised (Gaussian) distribution. Top row: marginal distributions p(x0 ), p(x1 ), . . . , p(xT ) evolving over time. Middle row: Simplified representation of a data distribution p(x0 ) that starts with two separates modes representing its structure, and which is slowly noised until it becomes a standard Gaussian distribution (p(xT )). Bottom row: example trajectory of a time series at different noise levels, illustrating the iterative noising process. cess is usually Gaussian, and therefore a Gaussian noise is added following a given schedule, until the noisy data approximately follow an isotropic Gaussian distribution. This process is called the forward process, and is detailed hereafter. Let x0 ∈ Rn be a data point sampled from a real, unknown distribution q. A Gaussian noise ϵt with covariance matrix βt In , where In denotes the n × n identity matrix, is iteratively added to x0 at step t during T evenly-spaced steps, with βt ∈ [0, 1], t ∈ {0, 1, . . . , T }, T bounded. These noising steps define the following forward process: p p ϵt ∼ N (0, In ) (1) xt = 1 − βt xt−1 + βt ϵt , where each step in the forward process depends solely on the previous one, thus defining a Gaussian Markov Chain: p(x1:T |x0 ) =

T Y

p(xt |xt−1 ), where p(xt |xt−1 ) = N

p

1 − βt xt−1 , βt In



(2)

t=1

where βt is arbitrarily fixed by defining the noising schedule. Denoising diffusion relies on the choice of a well-behaved noising schedule, and it is therefore necessary to choose an increasing variance schedule such that when t − → T the sample distribution is approximately the one of a white noise vector (i.e., xT ∼ N (0, In )). In Ho et al. (2020), the authors chose to have a linear schedule from β1 = 10−4 to βT = 0.02, but Nichol and Dhariwal (2021) introduced a cosine schedule which puts more focus on the beginning and on the end of the diffusion process, which yielded improved results. From equation 2 one can compute the marginal distribution q(xt |x0 ) in its closed form and this reparameterization allows to find xt without applying the entire forward process: p √ (3) xt = ᾱt x0 + (1 − ᾱt )ϵ, Qt where αt = 1 − βt , ᾱt = s=1 αs , and ϵ is a standard normal variable (i.e., ϵ ∼ N (0, In )). It can be shown that a reverse process exists, and that it transforms the white noise at time T into a data point at time 0 by reversing the noising process through a sequence of denoising steps (Ho et al., 2020). Conditioned on x0 , this reverse process follows a Gaussian distribution:
p(xt−1 | xt , x0 ) = N µt (xt , x0 ), σt2 In . (4) 5

Therefore, if x0 is known, the denoising step at time t can be computed as: xt−1 = µt (xt , x0 ) + σt ϵ

(5)

where ϵ is Gaussian white noise and σt is a standard deviation that depends only on βt . Using equation 5 the reverse process can now be sampled if µt (xt , x0 ) can be computed. However, x0 is not accessible during sampling, and an approximation has to be found instead. Using the reparameterization from Ho et al. (2020), µ can be re-written in the following form: !   p 1  βt 1  ϵ (6) xt − (1 − ᾱt )ϵ =√ xt − p µt (xt , x0 ) = µ̃t xt , √ αt αt (1 − ᾱt ) A neural network is trained to approximate the reverse process, which slowly removes the noise from the noisy data point until its original structure is recovered. This network (with weights θ) takes as input the noisy time series xt and the step t: 1 βt µθ (xt , t) = √ (xt − p ϵθ (xt , t)) αt (1 − ᾱt )

(7)

To approximate the reverse distribution it can be shown that the neural network only needs to be trained to learn how to predict ϵθ , which is an approximation of the real noise ϵ added to the original data-point x0 . The following loss function is used for the training: # "
√ √ 1 − ᾱ 1 2 t (8) βt ϵ − ϵθ ᾱt x0 + 1 − ᾱt ϵ, t 2 + C. θ∗ = Ex0 ,ϵ, t̸=0 2 1 − ᾱt−1 Ho et al. (2020) observed that the weighting terms have little impact on the quality of the generated samples. This leads to the following noise prediction loss: " #
2 √ √ ∗ θ = argmin Ex0 ,ϵ,t ϵ − ϵθ ᾱt x0 + 1 − ᾱt ϵ, t 2 , (9) θ

Given the approximation, a new data point can be generated. Starting at time T , xT is sampled from a normal distribution. Then, working backward through time, the noise z is repeatedly sampled and each previous state xt−1 is calculated using the statistical functions learned by the model: z ∼ N (0, I), xt−1 = µθ (xt , t) +

p

Σθ (xt , t) z.

(10)

The reverse process continues iteratively until it reaches x0 , the newly sampled data-point. Because the example dataset introduced in section 2.1 encompasses a large amount of missing data, we focus in the following on a diffusion model that can be trained while taking gaps into account.

2.3

Conditional Score-based Diffusion Models

Conditional Score-based Diffusion Models (CSDI). Conditional Score-based Diffusion Models (CSDI) are a type of diffusion models specifically designed for time series with missing values (Tashiro et al., 2021). They are inspired by masked language modeling and are tailored to explicitly condition on observed values when predicting the missing ones. CSDI learns to fill in missing values in time series while respecting their temporal structure as well as relationships between variables. In Tashiro et al. (2021) the authors designed both the architecture and the training paradigm of CSDI to handle incomplete time series. The model takes observed values as conditional inputs and predicts missing values accordingly. In the original paper, the authors applied CSDI to both imputation and forecasting tasks, using separate models for each task.

6

d Problem setup. Given conditional observations xco 0 ∈ R , with 0 ≤ d < n, and imputation targets ta x0 ∈ n − d, noisy targets are sampled as: √ √ xta ᾱt xta 1 − ᾱt ϵ, (11) t = 0 +

where ϵ ∼ N (0, In ) and αt follows a predefined noise schedule. The denoising network ϵθ is trained by minimizing the following loss: h
2i co min L(θ) = Ex0 ∼q(x0 ), ϵ∼N (0,I), t ϵ − ϵθ xta . (12) t , t | x0 θ

Transformers. The backbone of the model is based on the Transformer architecture (Vaswani et al., 2023), which has demonstrated strong performance in modeling long-range dependencies through selfattention. Attention is a powerful operation that allows to model dependencies within sequential data, with no limitations regarding to their distance in the sequence (Wang et al., 2014; Bahdanau et al., 2015). It is defined as: ! QK ⊤ V, Attention(Q, K, V ) = softmax √ dk (13) Q = XWQ ,

K = XWK ,

V = XWV ,

d×dk

where WQ , WK , WV ∈ R are matrices with trainable coefficients, and X ∈ Rn×d is an input multivariate time series. Unlike recurrent architectures, Transformers process the full sequence in parallel, which captures global temporal interactions efficiently. Given an input sequence representation x ∈ RT ×d , multi-head attention allows the model to attend to multiple representation subspaces simultaneously: MultiHead(X) = Concat(head1 , . . . , headH )WO ,

(14)

where headi , 0 ≤ i ≤ H is an independent attention operation with its own trainable set of weights: headh = Attention(Qh , K h , V h )

(15)

h , WVh ∈ Rd×dk and WO ∈ RHdk ×d are learned parameters. where WQh , WK

After multi-head attention a two-layer fully connected network called a Multi-Layer Perceptron (MLP) is applied independently to each timestep. The MLP is a per-timestep nonlinear layer that converts the multi-head attention outputs, i.e. temporal and features context, into richer latent features and stores learned patterns in its weights.: MLP(x) = σ(xW1 + b1 )W2 + b2 ,

(16)

where σ(·) denotes a non-linear activation function. Model architecture. Unlike most diffusion models, CSDI does not rely on a U-Net architecture (Ronneberger et al., 2015). Instead, it consists of a stack of layers grouped in blocks operating all at a fixed resolution throughout the network. Each block contains a skip connection, which allows the input of the block to bypass the intermediate layers and be added directly to the block’s output. In CSDI’s architecture, blocks are therefore called residual blocks (He et al., 2015). The model takes two input channels. The first is the main input, defined as the concatenation of co the noisy target values (xta t ) and the conditional observations (x0 ):  ta co  x = xt , x 0 . (17) The second input channel consists of side information, which includes supplementary data that can improve prediction quality, such as the observation mask and optional temporal or feature embeddings. To capture temporal dependencies and inter-feature relationships, each residual block contains two Transformer layers: one operating along the temporal axis and another one operating along the feature axis. These two layers enable the model to learn complex dependencies across time and between variables. At each generation step, a residual block produces two types of outputs: (1) a residual output which is the input of the next residual block, (2) a secondary output, which is stored (skip connection). The final output of the network uses all the skip connections that are aggregated and are then passed through final convolutional layers. 7

Training strategy. The training procedure of CSDI is inspired by masked language modeling and differs from that of standard diffusion models. The reverse denoising process cannot be directly approximated during the inference because the true target values xta 0 are unknown. Instead, during training, artificial missingness is introduced by masking observed values. Hence, given a sample x0 from the dataset, the observed entries are randomly split into two disjoint subsets: one subset is co treated as imputation targets xta 0 , while the other subset is used as conditional observations x0 . In the original CSDI paper, multiple masking strategies for training were proposed. The random strategy selects imputation targets by randomly masking a percentage of observed values, while the historical strategy leverages patterns from training data by aligning observed and missing indices between samples to capture structured missingness. The mixed strategy combines both approaches, improving generalization while learning realistic missing patterns. In this work, a combination of random masking and forecasting masking strategies is employed. In the latter, future time steps are masked to train the model for prediction tasks.

3

Methods

3.1

Augmented CSDI for hydrometeorological time series

Due to a limited training sample size, a high dimensionality and an heterogeneous scaling of the hydrological signals involved, the original CSDI architecture presented in section 2.3 performed poorly when applied to the test dataset introduced in section 2.1. The main challenges faced by the model were (see Appendix D for details): • the training failed to converge to a proper solution due to training instability and poor convergence, • the simulated time series ignored the observed values, failing to capture local relationships and resulting in jagged artifacts, • the simulations collapsed toward a single mean, with individual time series showing little variation around it, • the diurnal (and multi-day) cycles were somewhat captured, but the timing of sporadic peaks often did not align with observations. To overcome these limitations we designed the updated CSDI architecture detailed below. Custom CSDI architecture The modification of the original CSDI architecture to better comply with hydrological data followed a step by step approach, each step being guided by the analysis of the limitations listed above. See Appendix B for a comparison of the two architectures. To improve the simulation of local patterns during imputation and enhance the conditioning to observed values, multiple convolutional layers have been added and combined with a multi-scale layer in each residual block in the form of a Residual U-block (RSU), based on the nested U-shaped design introduced in U²-Net (Qin et al., 2020). Each RSU consists of a lightweight U-Net–style convolutional sub-network embedded within the residual mapping. The time series are down-sampled twice along the temporal axis and processed at progressively lower temporal resolutions before being up-sampled and fused back. Internal skip connections within the RSU preserve fine-grained temporal information while keeping broader contextual information across scales. This design allows each residual block to jointly model local and global temporal patterns in a parameter-efficient manner, thanks to the pooling layers. To improve the inclusion of the side information (i.e., the timesteps encoding and the mask) to the simulations, we changed the manner in which this information is integrated into the model. The side information was revamped, with three objectives: (1) side information should integrate the observed values and (2) it should have increased control over the hidden time series states, while (3) the timesteps should encode periodic signals explicitly. For the first objective, the observations xco 0 were simply appended to the side information. Although these observations are already present in the residual stream, this results in the information now being input twice, with each pathway serving a distinct functional purpose. The residual branch integrates observed values through convolution and attention, capturing both local and long-range patterns. In contrast, the side channels leverage the observed values alongside the rest of the side 8

information to modulate the residual stream in a complementary manner (see Figure 12 for more details). For the second objective, instead of inputting the data from the side information through a simple additive biasing, a modulation mechanism inspired by StyleGAN (Karras et al., 2019) was adopted. Specifically, the model learns two conditioning-dependent components: a multiplicative scaling term and an additive offset, which together modulate the activations and enable more expressive and reliable use of the conditioning information. Given an input sequence representation x ∈ RT ×c and the side information representation y ∈ RT ×d the following conditioning is applied: Conditioning(x, y) = x ⊙ α(y) + β(y) T ×c

(18)

T ×c

where α ∈ R and β ∈ R are the outputs of a convolution layer applied to the side information, and ⊙ is the element-wise multiplication. Finally, for the third objective, inspired by feature engineering from time series literature (Bansal et al., 2025), we include explicit encoding of some cycles. Previous deep learning for hydrology literature recommends similar Fourier encoding to help capture seasonal pattern dependencies within the time series properties and enhance the model’s performance (He et al., 2025). Given the timesteps t ∈ N of the time series, Fourier encoding is used to represent periodic structure at multiple scales. The Fourier timesteps mapping is defined as:   2π 2π t), cos( α t), . . . , sin( α2πN t), cos( α2πN t) ∈ R2N Φ(t) = sin( α 1 1 where {αi }N i=1 denote a set of predefined cycle lengths. In our use case, we use cycles of 1 day, 7 days, 30 days, 90 days and 365 days. The resulting representation is fixed and augments the original temporal input with explicit phaseaware periodic features, enabling the model to capture both short- and long-term temporal dependencies associated with the specified cycles. While the above architecture modifications led to improvements in simulated time series, it also led to instabilities during the training, with the gradients diverging to extremely high values. To alleviate this issue, Root Mean Squared Normalization (RMSNorm, Zhang and Sennrich (2019)) was applied in each residual block. RMSNorm focuses on re-scaling invariance and regularizes the summed inputs according to the root mean square (RMS) statistic: RMSNorm(xi ) = q

xi PC 1

ϵ+ C

⊙ γi ,

for i ∈ {1, . . . , C}

(19)

2 i=1 xi

where ϵ << 1, C ∈ N is the number of channels in the residual block, and ⊙ is the element wise product. This addition of RMSNorm has already been identified in the literature for stabilization and performance improvement (Jones et al., 2026), and is related to AdaLN in Diffusion Transformers (DiT, (Peebles and Xie, 2023)). Covariate Augmented Custom CSDI model The custom architecture offers satisfactory results for imputation but its forecasting performance remains limited when unforeseen external events (e.g., a rainfall event generating a streamflow response, and in turn a change in the chemical composition of stream water) induce distributional shifts in the hydrological time series. However, such events and in particular droughts and rains can be forecasted using numerical weather models, and the associated information will be leveraged in the Covariate Augmented Custom CSDI model by simulating the time series x knowing the covariate data y derived from numerical weather forecasts. This conditional distribution can be modeled directly in the generator network. This idea was first exploited in conditional GANs by Mirza and Osindero (2014) to generate according to some conditioning parameters. This has since been adapted to generate conditional images and text with denoising diffusion models, see e.g., Saharia et al. (2022b); Dieleman et al. (2022). In practice, the denoising network is extended to incorporate covariates as additional inputs: x̂0 = Dθ (xt , t, xco 0 , y),

(20)

where xt denotes the noisy sample at diffusion step t, xco 0 represents the observations, and y is the covariate time series. The model implicitly learns the conditional distribution through the standard 9

diffusion objective, without modifying the loss function. In other words, conditioning is achieved purely through the architectural augmentation.

3.2

Baselines

To assess the performances of the custom architecture, we compare the new model against three types of baselines. First, we consider classic statistical methods, including Gaussian Processes and Multi-point Statistics. These methods are well established for capturing correlations in space, time and between variables, and also for quantifying uncertainty. Refer to Appendix A for more details on the implementation of the Gaussian process in the present setting. Second, we compare the custom CSDI model to other neural networks dedicated to time series processing. These include the base CSDI architecture described in section 2.3 as well as a U-NET architecture, which is commonly used in image diffusion models (Ronneberger et al., 2015; Zhang et al., 2024). Third, forecasting tests on a synthetic dataset include two additional forecasting baselines: TSMixer (Chen et al., 2023) and NHits (Challu et al., 2022). These models are not evaluated on imputation tasks because they do not handle missing values. While NHits is deterministic, TSMixer has a probabilistic and deterministic version. We use both versions as baselines, naming them respectively TSMixer-P and TSMixer-D. This allows us to compare both point-wise prediction models and probabilistic baseline neural-networks. We use the Darts implementation for these last two baselines (Herzen et al., 2022).

4

Results

4.1

Evaluation settings

All experiments were conducted on an NVIDIA RTX 2000 Ada Generation GPU equipped with 16 GB of RAM. To evaluate our model, we use two datasets: (1) the OPE dataset introduced in section 2.1, and (2) a synthetic time series dataset that seeks to mimic hydrological dynamics, which has been designed to test how the models perform on a complete dataset and how they adapt to new data. In all cases all variables were normalized using a standard normalization procedure. For the covariates-augmented model, we use weather covariates derived from the SAFRAN reanalysis (QuintanaSegui et al., 2008). For each OPE station the covariates time series were extracted at the closest 8km grid point in the SAFRAN dataset. The resulting series are normalized and passed to the model as a secondary input as described in Section 3.1.

Synthetic time series dataset: Inspired by Rattanatheerawon and Fooprateepsiri (2025), we generated a synthetic dataset that seeks to mimic the hydrological time series of the example dataset. The synthetic dataset consists of 8 concatenated multivariate time series, with a timestep of 1h, and with three components for each multivariate time series: (1) a dynamic component, (2) a periodic component correlated with the dynamic one that comprises two periods of approximately 24h and 30 days, and (3) a periodic component independent of the others with periods of approximately 48h and 25 days. Finally, a univariate time series with regime switches is used to model a station experiencing droughts. A more detailed description is available in Appendix E. The performance on the synthetic dataset was evaluated using all models. In contrast, the performance on the real-world dataset was assessed only for the models capable of both imputation and forecasting because this dataset encompasses missing values.

Evaluation metrics: We adopt both deterministic (Root Mean Squared Error, Mean Absolute Error, SMAPE) and probabilistic (CRPS) metrics for model evaluation. The deterministic metrics measure the point-wise accuracy of the models, while the CRPS evaluates the fit of an estimated probability distribution to a ground-truth observation.

10

Figure 3: Synthetic data generated to test models. In blue we have a dynamic component, in orange a periodic component linked to the dynamic component, in green an independent periodic component, and in red a single time series with drought like periods (see Annex E for more details). Test Sets Configuration: For imputation ground truth, we randomly remove 10%, 25%, 50% and 90% values on 24 timesteps test time series. The points to mask were all selected with the same seed for every model, so that the comparison was based on the same set of observed values. A different seed was used for each hydrological variables, meaning that the missing steps varied along the last axis. For forecasting, we remove the 6, 9, 12 last timesteps of 24 timesteps test time series. Evaluations are computed between the masked ground truth and the median of the simulations for probabilistic models, or the predicted timeserie for predictive models. See Appendix 4 for an overview of the evaluation framework.

Figure 4: The evaluation setting of our diffusion model for imputation (top) and forecasting (bottom).

4.2

Performance on the synthetic dataset

The models are trained and tested on the synthetic dataset described in Section 4.1. Figure 5 displays the performance of the different models for the two tasks of interest, namely imputation and forecasting. 11

(a) Imputation Metrics Visualization

(b) Forecasting Metrics Visualization

Figure 5: Comparison of model performance on the synthetic dataset across missing ratios and prediction horizons. For clarity the metrics of the U-NET model are omitted or left out of bounds when it substantially underperforms compared to other models. Results in Figure 5.a show that for imputation all the metrics support that the custom CSDI model outperforms the three baselines. The GP benchmark leads to competitive results when the ratio of missing values is limited (typically < 50%), but the custom CSDI stands out for the stability of its accuracy when the ratio of missing values increases. It is worth noticing that the original CSDI model leads to disappointing results when applied to this synthetic dataset mimicking high-resolution hydrological observations, most likely because of the limited training sample size, combined with the high dimensionality (75 variables) and the heterogeneous scaling of hydrological signals, as described in section 3.1. Results in Figure 5.b focus on the relative forecasting performance of the different models across multiple forecast horizons, and show that for this task too the custom CSDI model almost consistently outperforms all the baselines (see Appendix Tables 15 for more detailed results). For short-term prediction the GP model performs well (second behind custom CSDI) but as the horizon grows the probabilistic TS-Mixer model (TS-Mixer P) gets better and equals the custom CSDI model for the 18 time steps horizon, even marginally surpassing it in terms of RMSE. The good performance of these two models for this prediction horizon could be explained by the fact that the multi-layer perceptron architecture which makes up TS-Mixer is more stable during training than the base CSDI. To refine the evaluation of the custom CSDI model, Figure 6 investigates how it captures periodic signals through the comparison of the power spectral density (PSD) in simulations and in the synthetic data. For a time series x ofPlength T the PSD is defined as the squared magnitude of the t T Fourier coefficient, i.e., P SD(x) = | t=0 x(t)e−j2π T |2 . Spectra are normalized and a one dimensional Gaussian filter kernel is applied along the frequency axis to reduce noise. For the generated time series, the PSD is computed for each individual simulated time series, then the mean PSD over 100 simulations is computed and reported in the figure.

12

Figure 6: Power spectral density comparison between synthetic data (blue) and the mean of 100 simulations performed by the custom CSDI model (orange). The dotted vertical blue lines show the expected frequencies of the periodic signal embedded in the synthetic data, i.e., 24h for variable 15 and at 48h for variable 18. The results in Figure 6 show that the custom CSDI model successfully captures and reproduces the daily (dependent periodic variable, top in Figure 3) and two-days (independent periodic variable, bottom in Figure 3) periodic signals present in the data. Both the amplitude and the frequency are properly modeled. Since the targeted tasks are imputation within short windows (48 timesteps, with 90% of values missing at random) and short-term forecasting (horizons of at most 18 steps), every evaluation window is shorter than the dominant seasonal period. We therefore do not assess the ability of the custom CSDI model to reproduce low-frequency components such as seasonal or pluriannual cycles. Finally, Table 1 evaluates how imputation accuracy scales when the length of the missing segments (i.e., the size of the data gaps) increases from 50 to 180 consecutive time steps (i.e. from 8.3 days to 30 days). The results reveal a clear performance inflection for the custom CSDI model based on gap size. For short-to-medium horizons (50 and 100 steps), the custom CSDI model outperforms the GP benchmark, achieving a SMAPE of 9.52% for missing segments of size 100. However, when the gaps exceed 140 steps, the performance of the custom CSDI model sharply decreases with SMAPE rising from 9.52% (gap size = 100) to 17.78% (gap size = 140). For such long gaps the GP benchmark achieves better accuracy. These results could be explained by the fact that the custom CSDI model efficiently leverages the local context (i.e., the data available before and after a gap) for the imputation of missing values, but is less effective in capturing the large-scale context. In case of data comparable to sub-daily resolution hydrological time series (i.e., the synthetic dataset presented in this section) the above features make the custom CSDI model well suited for the imputation of gaps of moderate size and for short- to medium-term forecasting (less than two weeks). When local information becomes rare or absent a decrease in performance is expected.

4.3

Performance on the OPE dataset

The OPE dataset is partitioned chronologically into training, validation, and test sets of respectively 4152, 20, and 28 elementary time series, each comprising 24 timesteps (4 days). Figure 7 illustrates imputation results for four different models: the original CSDI model, the custom CSDI model, the custom CSDI model augmented with covariates, and the GP benchmark model.

13

Table 1: Imputation performance comparison between the proposed Custom Model and Gaussian Processes (GP) across varying missing gap sizes on synthetic time series. Model / Metric

50 Timesteps

100 Timesteps

140 Timesteps

180 Timesteps

Custom Model RMSE MAE SMAPE (%) CRPS

0.0405 ± 0.0006 0.0250 ± 0.0002 8.3586 ± 0.0726 0.0555 ± 0.0009

0.0355 ± 0.0003 0.0241 ± 0.0001 9.5221 ± 0.0387 0.0570 ± 0.0005

0.0718 ± 0.0006 0.0471 ± 0.0004 17.7802 ± 0.1686 0.1173 ± 0.0009

0.0737 ± 0.0011 0.0494 ± 0.0004 19.8093 ± 0.2360 0.1199 ± 0.0008

Gaussian Processes RMSE MAE SMAPE (%) CRPS

0.0494 ± 0.0020 0.0347 ± 0.0008 11.1288 ± 0.2047 0.0828 ± 0.0014

0.0612 ± 0.0010 0.0447 ± 0.0007 13.7765 ± 0.3650 0.1021 ± 0.0013

0.0637 ± 0.0008 0.0423 ± 0.0007 13.9469 ± 0.2436 0.1096 ± 0.0010

0.0660 ± 0.0010 0.0477 ± 0.0009 16.7788 ± 0.4187 0.1216 ± 0.0016

Figure 7: Illustration of imputation results over 24 timesteps (4 days), on one sample of the test dataset, across four models for dissolved oxygen at station OPE90013 showed as normalized values. Red crosses denote observed timesteps, blue dots indicate masked ground truth, and the simulated distribution is represented by its median (green line) and 5-95th percentile band (light green band). The visual inspection of Figure 7 indicates that both versions of the custom CSDI model, and to a lesser extent the GP model, successfully capture the main temporal patterns of the time series. In contrast the original CSDI model tends to overestimate the temporal variability of the underlying process. Figure 8 synthesizes the evaluation metrics for both imputation and forecasting, and detailed results are reported in Appendix Tables 16.

14

(a) Imputation Metrics Visualization

(b) Forecasting Metrics Visualization

Figure 8: Performance comparison across missing ratios and prediction horizons on the OPE dataset. U-Net is omitted, as its poor performance would distort the scale of the plots. Results in Figure 8 show that across all metrics, the custom CSDI architecture (with and without covariates) consistently outperforms all baselines. As expected, the covariate-augmented version performs slightly better than the variant without covariates, highlighting the additional predictive value of weather auxiliary data. Interestingly, as was the case for the synthetic dataset, the GP model outperforms the original CSDI architecture. This suggests that the original CSDI model may not be suited for hydrological data. In contrast, a custom architecture specifically designed and tailored to the characteristics of the dataset demonstrates superior performance. The covariate-augmented model performs similarly to the model without covariates in terms of imputation; however, it performs significantly better for forecasting tasks. This improvement is particularly visible for the simulations displayed in Figure 9. In this configuration a sudden meteorological event (here a rainfall event) breaks the stationarity of the hydrometeorological system, and the baseline model fails to capture the sudden peak of stream water level observed on 05/01/2024 whereas the model incorporating covariates successfully reproduces this feature thanks to the auxiliary information brought by the meteorological covariates.

15

Figure 9: Compared forecasting results for four tested models, Custom with covariates (CovAug), Custom without covariates (Custom), Gaussian Processes (GP) and Original CSDI (Original). Red crosses show observed timesteps, blue dots indicate masked ground truth, and the simulated distribution is represented by its median and 5-95th percentile band. Results are shown for the normalized water level variable (HTE) at station OPE90012. To evaluate the consistency of multivariate time series simulation the Figure 10 displays the covariance structure of the 75 variables, grouped by station. Parameters are grouped into blocks to facilitate the visualization but this display does not reflect any block-wise modeling assumptions.

16

Figure 10: Comparison between the Pearson covariance matrix of the training dataset (top-left) and generated simulations by: the original CSDI (top-right), the custom covariate augmented model (bottom-left), and the custom model (bottom-right). The correlations are grouped by station. Results in Figure 10 show that the original CSDI (top-right) struggles to capture the correlation structure. Some emerging structure is visible, with parallel positive diagonals slightly appearing, but overall the simulated correlation seems weaker than the one from the real data. In contrast, the custom model (bottom-right) better reproduces the overall correlation structure across stations. The presence of parallel positive diagonals inside the non-diagonal blocks shows correlations between the same parameters measured at different locations, which suggests that the custom CSDI model performs satisfactorily for multisite modeling. However, the contrast between station is reduced, with weaker correlation magnitudes in custom CSDI simulations compared to observations, suggesting that the generative model partially underestimates the strong positive and negative correlations. In particular, we note that the negative correlations are not well captured. The possible cause for this has already been identified in the literature, as Qin et al. (2022) showed that since softmax attention weights are strictly non-negative, the cross-station conditioning mechanism struggles to encode negative relationships between variables. Finally, the covariate-augmented CSDI model (bottom-left), which has the same backbone as the custom CSDI model, seems to capture the correlations more strongly. It also captures the parallel positive diagonals inside the blocks, showing correlations between the same parameters measured at different locations. However, in this case, the simulated correlation seems slightly stronger than in the real data. To further explore the contribution of adding covariates to the custom CSDI model, Figure 11 displays the distribution of simulated values with respect to the observed ones on a forecasting scenario. The red dashed line represents a perfect predictor and is shown for reference.

17

Figure 11: Visualization of the distribution of predicted values with respect to their real value, without covariates (left) and with covariates (right). All values are normalized. The red dashed line represents a perfect predictor. EC stands for electric conductivity, and HTE for water level. To evaluate the ability of the model to generate realistic predictive distributions, this experiment considers the same forecasting task as Figure 9, simulating 15 days given 16 days. Results in Figure 11 highlight the impact of incorporating covariates on the forecasting performance. Without covariates (left), the model has access only to the historical distribution of the target variables and therefore has limited ability to predict the true values. As a consequence, the generated predictions show larger variability and weaker alignment with the ground truth. However, the larger spread also provides broader coverage of possible outcomes, which is desirable in applications where uncertainty quantification is important. In contrast, the model with covariates demonstrates the opposite tradeoff. Predictions are more concentrated around the true values, leading to improved predictive accuracy, but with reduced variability.

5

Discussion

In this study we proposed a Transformer-based diffusion model for the joint imputation and forecasting of multivariate hydrological time series. The resulting model architecture, showcased in 3.1, is based on the addition of layers from recent literature to the original CSDI framework in order to tailor the model to the targeted application. We discuss below the strengths and limitations of the custom CSDI model. Uncertainty quantification. Using a diffusion generative model presents the advantage of probabilistic formulation. The model can generate a large ensemble of equally likely simulations, allowing users to quantify uncertainty. This feature is important in the context of hydrology where decision-making is better informed when information about model uncertainty is available. This also allows the model to be used for anomaly detection when values are measured outside the range of simulated values. Domain specific architecture. Through the comparison with the original CSDI and the UNET architecture, we show that the proposed ad-hoc architecture performs better on both intended tasks, namely imputation and forecasting. Good performance has been reached by drawing on the deep learning literature applied to hydrological time series to guide design choices, such as periodic encoding. At the same time it is important not to lose sight of the conventional generative modeling in order to benefit from research in this community. Modeling inter-variable dependencies. One limitation identified during the evaluation of the custom CSDI model is its limited ability to capture negative cross-variable correlations. While the 18

Transformer architecture has demonstrated strong performance in modeling long-range dependencies through self-attention, the results suggest that the implementation we used has limited capability to capture negative relationships. An interesting direction to address this would be to replace the standard attention mechanism with cosFormer, which has been shown to better handle negative correlations (Qin et al., 2022). Due to time constraints, this direction was not explored in the current study and is left for future work Choice of covariates. The covariate-augmented model relies on meteorological information from the SAFRAN reanalysis. While this provides a useful way to assess the potential contribution of exogenous meteorological drivers, it does not fully correspond to an operational forecasting setting where future covariates would have to be provided by numerical weather forecasts and would therefore be uncertain. The operational performance of the model may thus depend on the quality of the available meteorological forecasts. This should be evaluated explicitly before using the covariateaugmented custom CSDI model for operational hydrometeorological forecasting. Data Quality. We note that the OPE dataset used for training was corrected for sensor drift and malfunction through a joint LNE/Andra metrology effort, involving monthly quality control and regular maintenance over 15 years (Guigues et al., ress). This correction is non-trivial: in natural streams, long-term sensor deployment is susceptible to drift from biofouling, sediment, and vegetation debris accumulating on the sensor surface, which can substantially degrade data quality without rigorous, expert-driven QC. While real-time forecasting with the trained model remains possible using raw, uncorrected sensor data, we expect simulation quality to degrade to the extent that inputs deviate from the drift-corrected training distribution.

6

Conclusion

In this work, we introduced a Transformer-based architecture able to perform simultaneously the imputation and forecasting of multivariate hydrological time series. The model is based on the CSDI architecture, which has been improved with convolutions, multi-scale layers, and Root Mean Squared Normalization. The resulting architecture allows for a stable training and performs well in modeling synthetic and real datasets. In particular, the custom CSDI model better captures uncertainty than a set of baselines, and is competitive on predictive metrics. Despite an imperfect modeling of the correlation between variables, the proposed approach offers a promising direction for data-driven modeling of hydrological data with sub-daily resolution across more than 15 years. By jointly addressing imputation and forecasting within a single Transformerbased model using a long-term water quality and water level dataset, this work contributes to the growing literature on deep learning hydrology. The results support the relevance of the proposed model for short-term forecasting and for the imputation of gaps of size between hours to days. The results support the relevance of the proposed model for short-term forecasting and for the imputation of short to medium-size gaps, but we did not tested the ability of the model to reproduce the hydrological dependencies, seasonal variability, or low-frequency dynamics for periods longer than 15 days. Future work should therefore focus on the evaluation of the model on longer sequences, including long consecutive gaps and temporally structured sensor failures.

7

Acknowledgment

This method has been developed within the scope of the Geolearning chair. The authors are grateful to French National Agency for Radioactive Waste Management (ANDRA) partners for support and fruitful discussions. We would especially like to thank their technicians, engineers and researchers responsible for the collection, data curation and quality control of the OPE data, following during more than 15 years a protocol developed in partnership with LNE and ANDRA, with monthly quality control, sensor drift correction, and uncertainty estimation.

8

Code availability

The source code is available on this git repository: link. 19

References Alcaraz, J. M. L. and N. Strodthoff (2023). Diffusion-based time series imputation and forecasting with structured state space models. Allcroft, D. J. and C. A. Glasbey (2003, 09). A latent gaussian markov random-field model for spatiotemporal rainfall disaggregation. Journal of the Royal Statistical Society Series C: Applied Statistics 52 (4), 487–498. Alornyo, S. A., W. R. McGillis, and P. J. Culligan (2026). From sensors to simulations: How ai is transforming water monitoring, management, and policy. Water 18 (14). Arjovsky, M. and L. Bottou (2017). Towards principled methods for training generative adversarial networks. arXiv preprint arXiv:1701.04862 . Arjovsky, M., S. Chintala, and L. Bottou (2017). Wasserstein generative adversarial networks. In International conference on machine learning, pp. 214–223. PMLR. Bahdanau, D., K. Cho, and Y. Bengio (2015). Neural machine translation by jointly learning to align and translate. In Proceedings of the 3rd International Conference on Learning Representations (ICLR 2015). Bansal, A., K. Balaji, and Z. Lalani (2025). Temporal encoding strategies for energy time series prediction. Bedin, L., Y. Janati, G. Victorino Cardoso, J. Duchateau, R. Dubois, and E. Moulines (2025, 06). Reconstructing ecg from indirect signals: a denoising diffusion approach. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences 383 (2299), 20240330. Bhavsar, F., N. Desassis, F. Ors, and T. Romary (2024). A stable deep adversarial learning approach for geological facies generation. Computers & Geosciences 190, 105638. Brock, A., J. Donahue, and K. Simonyan (2019). Large scale GAN training for high fidelity natural image synthesis. Brunner, M. I., A. Bárdossy, and R. Furrer (2019). Technical note: Stochastic simulation of streamflow time series using phase randomization. Hydrology and Earth System Sciences 23 (8), 3175–3187. C. M., N. and N. Thangarasu (2025, 11). Enhancing dam inflow forecasting using hybrid deep learning and diffusion models: a case study on the malampuzha dam. Water Practice and Technology 20 (11), 2276–2293. Challu, C., K. G. Olivares, B. N. Oreshkin, F. Garza, M. Mergenthaler-Canseco, and A. Dubrawski (2022). N-hits: Neural hierarchical interpolation for time series forecasting. Chen, C., F. Wang, Z. Wang, D. Zhang, and L. Xiang (2025, June). A novel flood forecasting model based on TimeGAN for data-sparse basins. Stochastic Environmental Research and Risk Assessment 39 (6), 2267–2280. Chen, S.-A., C.-L. Li, N. Yoder, S. O. Arik, and T. Pfister (2023). Tsmixer: An all-mlp architecture for time series forecasting. Cybenko, G. (1989). Approximation by superpositions of a sigmoidal function. Mathematics of control, signals and systems 2 (4), 303–314. Dhariwal, P. and A. Nichol (2021). Diffusion models beat gans on image synthesis. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan (Eds.), Advances in Neural Information Processing Systems, Volume 34, pp. 8780–8794. Curran Associates, Inc. Dieleman, S., L. Sartran, A. Roshannai, N. Savinov, Y. Ganin, P. H. Richemond, A. Doucet, R. Strudel, C. Dyer, C. Durkan, C. Hawthorne, R. Leblond, W. Grathwohl, and J. Adler (2022). Continuous diffusion for categorical data. ArXiv abs/2211.15089.

20

Feng, D., K. Fang, and C. Shen (2020, September). Enhancing streamflow forecast and extracting insights using long-short term memory networks with data integration at continental scales. Water Resources Research 56 (9). Goodfellow, I., J. Pouget-Abadie, M. Mirza, B. Xu, D. Warde-Farley, S. Ozair, A. Courville, and Y. Bengio (2014). Generative adversarial nets. In Advances in neural information processing systems, pp. 2672–2680. Goodfellow, I. J., Y. Bengio, and A. Courville (2016). Deep Learning. Cambridge, MA, USA: MIT Press. http://www.deeplearningbook.org. Gu, S., D. Chen, J. Bao, F. Wen, B. Zhang, D. Chen, L. Yuan, and B. Guo (2022, June). Vector quantized diffusion model for text-to-image synthesis. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 10696–10706. Guardiano, F. and M. Srivastava (1993). Multivariate geostatistics: Beyond bivariate moments. In A. Soares (Ed.), Geostatistics Troia, pp. 133–144. Dordrecht: Kluwer Academic Publications. Guigues, N., M. Savatier, and B. Lalère (in press). Quantifying measurement uncertainty in real time continuous water quality monitoring: methodological approaches and implications for data interpretation. Environmental Science: Processes & Impacts. In press. Guo, D., A. Lintern, J. A. Webb, D. Ryu, U. Bende-Michl, S. Liu, and A. W. Western (2020). A data-based predictive model for spatiotemporal variability in stream water quality. Hydrology and Earth System Sciences 24 (2), 827–847. He, J., Y.-J. Chen, A. Jafari, A. Idamekorala, and G. Fox (2025). Deep learning foundation and pattern models: Challenges in hydrological time series. He, K., X. Zhang, S. Ren, and J. Sun (2015). Deep residual learning for image recognition. Herzen, J., F. Lässig, S. G. Piazzetta, T. Neuer, L. Tafti, G. Raille, T. V. Pottelbergh, M. Pasieka, A. Skrodzki, N. Huguenin, M. Dumonal, J. Kościsz, D. Bader, F. Gusset, M. Benheddi, C. Williamson, M. Kosinski, M. Petrik, and G. Grosch (2022). Darts: User-friendly modern machine learning for time series. Journal of Machine Learning Research 23 (124), 1–6. Ho, J., A. Jain, and P. Abbeel (2020). CoRR abs/2006.11239.

Denoising diffusion probabilistic models.

Hochreiter, S. and J. Schmidhuber (1997, 12). Long short-term memory. Neural computation 9, 1735–80. Hu, J., P. Wang, D. Li, and S. Liu (2024). A long-term multivariate time series prediction model for dissolved oxygen. Ecological Informatics 82, 102695. Jones, C., E. Noutahi, J. Hartford, and C. Eastwood (2026). Elucidating the design space of flow matching for cellular microscopy. Karras, T., S. Laine, and T. Aila (2019). A style-based generator architecture for generative adversarial networks. In Proceedings of the IEEE/CVF conference on computer vision and pattern recognition, pp. 4401–4410. Kawar, B., S. Zada, O. Lang, O. Tov, H. Chang, T. Dekel, I. Mosseri, and M. Irani (2023, June). Imagic: Text-based real image editing with diffusion models. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 6007–6017. Kingma, D. P. and M. Welling (2014). Auto-encoding variational bayes. Kratzert, F., D. Klotz, C. Brenner, K. Schulz, and M. Herrnegger (2018). Rainfall–runoff modelling using long short-term memory (lstm) networks. Hydrology and Earth System Sciences 22 (11), 6005–6022. Li, C., X. Ling, Y. Xue, W. Luo, L. Zhu, F. Qin, Y. Zhou, and Y. Huang (2024). Precipitation nowcasting using diffusion transformer with causal attention. IEEE Transactions on Geoscience and Remote Sensing 62, 1–16. 21

Mariethoz, G., P. Renard, and J. Straubhaar (2010). The direct sampling method to perform multiplepoint geostatistical simulations. Water Resources Research 46 (11). Mirza, M. and S. Osindero (2014). arXiv:1411.1784 .

Conditional generative adversarial nets.

arXiv preprint

Moiroux-Arvis, L., L. Royer, D. Sarramia, G. De Sousa, A. Claude, D. Latour, E. Roussel, O. Voldoire, P. Chardon, R. Vandaële, et al. (2023). Connecsens, a versatile iot platform for environment monitoring: Bring water to cloud. Sensors 23 (6), 2896. Nichol, A. and P. Dhariwal (2021). CoRR abs/2102.09672.

Improved denoising diffusion probabilistic models.

Obakrim, S., L. Benoit, and D. Allard (2025). A multivariate and space-time stochastic weather generator using a latent gaussian framework. Stochastic Environmental Research and Risk Assessment 39 (9), 3677–3701. Ou, Z., C. Nai, B. Pan, Y. Zheng, C. Shen, P. Jiang, X. Liu, Q. Tang, W. Li, and M. Pan (2025, July). Probabilistic diffusion models advance extreme flood forecasting. Geophysical Research Letters 52 (15). Peebles, W. and S. Xie (2023). Scalable diffusion models with transformers. Qin, X., Z. Zhang, C. Huang, M. Dehghan, O. R. Zaiane, and M. Jagersand (2020). U2-net: Going deeper with nested u-structure for salient object detection. Pattern recognition 106, 107404. Qin, Z., W. Sun, H. Deng, D. Li, Y. Wei, B. Lv, J. Yan, L. Kong, and Y. Zhong (2022). cosformer: Rethinking softmax in attention. Quintana-Segui, P., P. Le Moigne, Y. Durand, E. Martin, F. Habets, M. Baillon, C. Canellas, L. Franchisteguy, and S. Morel (2008). Analysis of near-surface atmospheric variables: Validation of the SAFRAN analysis over France. Journal of Applied Meteorology and Climatology 47 (1), 92–107. Rashid, M., A. Saeed, M. Khalid, A. Murtaza, and M. Waqar Saleem (2026). The transformative role of artificial intelligence in water resources engineering: A comprehensive review. Environmental Modelling & Software 197, 106857. Rattanatheerawon, W. and R. Fooprateepsiri (2025). Diffusion-driven time-series forecasting to support sustainable river ecosystems and sdg-aligned water-resource governance in thailand. Sustainability 17 (22). Ren, H., E. Cromwell, B. Kravitz, and X. Chen (2022). Technical note: Using long short-term memory models to fill data gaps in hydrological monitoring networks. Hydrology and Earth System Sciences 26 (7), 1727–1743. Ronneberger, O., P. Fischer, and T. Brox (2015). U-net: Convolutional networks for biomedical image segmentation. In Medical Image Computing and Computer-Assisted Intervention–MICCAI 2015: 18th International Conference, Munich, Germany, October 5-9, 2015, Proceedings, Part III 18, pp. 234–241. Springer. Ruiz, N., Y. Li, V. Jampani, Y. Pritch, M. Rubinstein, and K. Aberman (2023, June). Dreambooth: Fine tuning text-to-image diffusion models for subject-driven generation. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), pp. 22500–22510. Saharia, C., W. Chan, H. Chang, C. Lee, J. Ho, T. Salimans, D. Fleet, and M. Norouzi (2022a). Palette: Image-to-image diffusion models. In ACM SIGGRAPH 2022 Conference Proceedings, SIGGRAPH ’22, New York, NY, USA. Association for Computing Machinery. Saharia, C., W. Chan, H. Chang, C. A. Lee, J. Ho, T. Salimans, D. J. Fleet, and M. Norouzi (2022b). Palette: Image-to-image diffusion models. Shevnina, E. and A. Silaev (2019). The probabilistic hydrological marcsHYDRO (the markov chain system) model: its structure and core version 0.2. Geoscientific Model Development 12 (7), 2767– 2780. 22

Sohl-Dickstein, J., E. A. Weiss, N. Maheswaranathan, and S. Ganguli (2015). Deep unsupervised learning using nonequilibrium thermodynamics. Strebelle, S. and A. G. Journel (2000). Sequential Simulation Drawing Structures from Training Images. Dissertation. Tashiro, Y., J. Song, Y. Song, and S. Ermon (2021). Csdi: Conditional score-based diffusion models for probabilistic time series imputation. In M. Ranzato, A. Beygelzimer, Y. Dauphin, P. Liang, and J. W. Vaughan (Eds.), Advances in Neural Information Processing Systems, Volume 34, pp. 24804–24816. Curran Associates, Inc. Teurlincx, S., D. van Wijk, W. M. Mooij, J. J. Kuiper, I. Huttunen, R. J. Brederveld, M. Chang, J. H. Janse, B. Woodward, F. Hu, and A. B. Janssen (2019). A perspective on water quality in connected systems: modelling feedback between upstream and downstream transport and local ecological processes. Current Opinion in Environmental Sustainability 40, 21–29. System dynamics and sustainability. Tunnicliffe Wilson, G. (2016, 03). Time series analysis: Forecasting and control,5th edition, by george e. p. box, gwilym m. jenkins, gregory c. reinsel and greta m. ljung, 2015. published by john wiley and sons inc., hoboken, new jersey, pp. 712. isbn: 978-1-118-67502-1. Journal of Time Series Analysis 37, n/a–n/a. Vaswani, A., N. Shazeer, N. Parmar, J. Uszkoreit, L. Jones, A. N. Gomez, L. Kaiser, and I. Polosukhin (2023). Attention is all you need. von Freyberg, J., B. Studer, and J. W. Kirchner (2017). A lab in the field: high-frequency analysis of water quality and stable isotopes in stream water and precipitation. Hydrology and Earth System Sciences 21 (3), 1721–1739. Wang, Q., J. Zhang, S. Song, and Z. Zhang (2014). Attentional neural network: Feature selection using cognitive feedback. In Advances in Neural Information Processing Systems (NIPS 2014). Wang, X., H. Zhang, P. Wang, Y. Zhang, B. Wang, Z. Zhou, and Y. Wang (2023). 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, KDD ’23, New York, NY, USA, pp. 2409–2418. Association for Computing Machinery. Weilisi and T. Kojima (2022). Investigation of hyperparameter setting of a long short-term memory model applied for imputation of missing discharge data of the daihachiga river. Water 14 (2). Yang, C., T. Wang, and X. Yan (2023). Ddmt: Denoising diffusion mask transformer models for multivariate time series anomaly detection. Yang, D., J. Yu, H. Wang, W. Wang, C. Weng, Y. Zou, and D. Yu (2023). Diffsound: Discrete diffusion model for text-to-sound generation. IEEE/ACM Transactions on Audio, Speech, and Language Processing 31, 1720–1733. Yuan, X. and Y. Qiao (2024). Diffusion-ts: Interpretable diffusion for general time series generation. Zhang, B. and R. Sennrich (2019). Root mean square layer normalization. Red Hook, NY, USA: Curran Associates Inc. Zhang, H., J. Zhou, Y. Lu, M. Guo, P. Wang, L. Shen, and Q. Qu (2024). The emergence of reproducibility and generalizability in diffusion models. Zhang, L., A. Rao, and M. Agrawala (2023, October). Adding conditional control to text-to-image diffusion models. In Proceedings of the IEEE/CVF International Conference on Computer Vision (ICCV), pp. 3836–3847. Zhu, J.-Y., T. Park, P. Isola, and A. A. Efros (2020). Unpaired image-to-image translation using cycle-consistent adversarial networks.

23

A

Gaussian Processes details

For our baseline, we use a Gaussian Process with a composite kernel combining a smooth periodic component, a non-periodic short-timescale component, and observational noise. Specifically, the kernel consists of a Matérn-ν = 3/2 covariance modulated by a periodic exponential-sine-squared term to capture quasi-periodic structure, plus an additional Matérn component for medium-scale variability, and a white-noise term for uncorrelated measurement errors.

B

Architecture detail

(a) Original architecture

(b) Modified architecture

Figure 12: Comparison between the original CSDI architecture and the modified architecture. The modified components are highlighted in red. The original architecture is adapted from Tashiro et al. (2021).

24

C

Hyperparameters Model Iterations Recurrent Blocs Batch Size Learning Rate Nb diffusion steps

Custom 100 4 16 1e−3 100

Cov. Augmented 100 4 16 1e−3 100

Original 100 4 8 1e−3 50

Table 2: Hyperparameters used to train the different models.

D

Ablations Tests

Model Original +Cond. before Transformers +StyleGAN Cond+Norm +Convolutions

+Seasonality

H

RMSE

MAE

sMAPE

CRPS

CRPS sum

6 0.3887 ± 0.0024 0.1260 ± 0.0003 22.1140 ± 0.0881 0.1171 ± 0.0005 0.0656 ± 0.0001 12 0.4351 ± 0.0023 0.1489 ± 0.0005 24.8425 ± 0.0694 0.1420 ± 0.0004 0.1020 ± 0.0001 18 0.5019 ± 0.0007 0.1779 ± 0.0004 27.0204 ± 0.0746 0.1747 ± 0.0003 0.1528 ± 0.0002 6 0.3257 ± 0.0008 0.0855 ± 0.0003 17.5146 ± 0.1022 0.0802 ± 0.0004 0.0555 ± 0.0001 12 0.4009 ± 0.0005 0.1234 ± 0.0006 22.8549 ± 0.1559 0.1159 ± 0.0008 0.0887 ± 0.0005 18 0.4733 ± 0.0005 0.1661 ± 0.0007 27.7593 ± 0.1747 0.1558 ± 0.0005 0.1405 ± 0.0005 6 0.3240 ± 0.0007 0.0840 ± 0.0004 16.6410 ± 0.1247 0.0785 ± 0.0005 0.0548 ± 0.0001 12 0.3929 ± 0.0008 0.1163 ± 0.0004 21.5238 ± 0.0816 0.1087 ± 0.0004 0.0844 ± 0.0003 18 0.4610 ± 0.0011 0.1540 ± 0.0010 25.6988 ± 0.0860 0.1429 ± 0.0009 0.1288 ± 0.0009 6 0.3146 ± 0.0008 0.0839 ± 0.0003 17.7016 ± 0.0865 0.0794 ± 0.0002 0.0555 ± 0.0001 12 0.4018 ± 0.0009 0.1208 ± 0.0005 22.0666 ± 0.1432 0.1137 ± 0.0005 0.0865 ± 0.0003 18 0.4830 ± 0.0004 0.1629 ± 0.0004 27.2257 ± 0.0340 0.1544 ± 0.0004 0.1354 ± 0.0003 6 0.3211 ± 0.0021 0.0878 ± 0.0010 17.3480 ± 0.1479 0.0822 ± 0.0009 0.0568 ± 0.0002 12 0.3895 ± 0.0007 0.1140 ± 0.0004 21.1204 ± 0.1045 0.1066 ± 0.0006 0.0846 ± 0.0004 18 0.4526 ± 0.0009 0.1473 ± 0.0004 25.0509 ± 0.0569 0.1374 ± 0.0007 0.1272 ± 0.0006

Table 3: Ablation study results (mean ± std) across forecasting horizons. Bold green indicates the best (lowest) value per horizon/metric pair. In this section, we present some ablation tests to showcase how different layers influence simulation results. Table 3 presents forecasting metrics as we add our custom layers, step by step. Figure 13 illustrates forecasting results for these different results. As discussed in Section 2.3, the vanilla CSDI baseline fails to converge to a meaningful solution: the predicted mean is flat, no temporal structure is captured, and all reported metrics show poor performances. Conditioning placement. Moving the conditioning layer before the Transformer blocks, rather than after like in the original CSDI, is the first modification necessary for convergence. This single change enables the model to produce temporally coherent outputs, as seen in Figure 13, suggesting that early injection of the conditioning signal is critical for the diffusion process to learn meaningful dynamics. RMSNorm and StyleGAN-based conditioning. To improve training stability, we introduce RMSNormalization alongside StyleGAN-style conditioning. Without normalization, gradients explode during training, preventing convergence. We noticed a posteriori that this design is closely related to Adaptive Layer Normalization (AdaLN), a technique widely adopted in Diffusion Transformer (DiT) architectures (Peebles and Xie, 2023) precisely because it stabilizes transformer training in the diffusion setting. Our variant differs in using RMSNorm rather than standard Layer Normalization and omits the learned scaling factor. We are not the only ones to choose RMSNorm, as Jones et al. (2026) independently explored it for diffusion transformers. At this stage, the generated time series exhibit realistic hydrological structure, however predictions are substantially biased and the predictive uncertainty does not cover the ground truth. 25

Convolutional layers. Adding convolutional layers on top of the above yields to a trade-off: the predicted mean becomes slightly flatter, while the confidence band widens considerably, driving the model toward a more uncertain but better-calibrated regime. Qualitatively, the simulated time series still behave like real hydrological time series, yet seasonality remains poorly captured and temporal dependencies appear chaotic (Figure 13). Point-wise metrics (MAE, RMSE) improve at short horizons, but all other losses worsen. Critically, however, removing the convolutional layers degrades the inter-variable correlation structure (Figure 14), indicating that without them the model overfits and fails to reproduce important multivariate dependencies. Encoding periodic signals. Finally, incorporating the encoding of periodic signals yields the best performance on long-horizon forecasting across all metrics. As shown in Figure 13, it is only with this component that the predicted peaks align with those of the ground truth, while maintaining a well-calibrated uncertainty range. This constitutes our final custom architecture.

Figure 13: Ablation study: predicted water temperature forecasts (15-day horizon, conditioned on 16 observed days) at station OPE0013, for each successive architectural modification. Green line and shaded envelope: posterior mean and uncertainty band; blue line: ground truth; red crosses: observed conditioning values. Each panel isolates the effect of one added component, from vanilla CSDI (top) to our full proposed architecture (bottom).

26

Figure 14: Inter-variable correlation matrices estimated, real-data (left) and our custom architecture without convolutions (right).

E

Synthetic dataset details

The synthetic dataset is designed to mimic multivariate hydrological time series at K different stations. It contains three groups of variables per station: a dynamic component, a correlated periodic component, and an independent seasonal component. In addition, a regime-switching process mimics a station with drought behavior. For each station 1 ≤ k ≤ K, we define a baseline mean µk = 0.25 + 0.07k, modeling differences between measurement locations. The latent driver of each dynamic variable follows an autoregressive AR(1) process,
zk,t = max 0.2, µk + ϕk (zk,t−1 − µk ) + ϵk,t + gt ,

ϵk,t ∼ N (0, 0.012 ),

(21)

with AR coefficient ϕk = 0.95 − 0.02k, and gt denotes the coefficient for global extremes shared between all stations ( Exp(0.2) with probability pspike (t), (22) gt = 0 otherwise, where the probability of spike pspike (t) is higher during periods of bursts bt and otherwise low. Bursts evolve according to ( 1 with probability 0.01, bt = (23) bt−1 · β otherwise. where β = 0.93 represents the burst decay. Then we define pspike (t) as: ( 0.08 if bt > 0.2, pspike (t) = 0.01 otherwise.

(24)

First, the observed dynamic variable is obtained by adding small observation noise to the latent, Dk,t = zk,t + ηt ,

ηt ∼ N (0, 0.012 ),

(25)

To mimic periodic cycles of hydrological variable, a correlated periodic counterpart to the observed dynamic variable is defined as Sk,t = 0.6 zk,t + sk,t + ξt ,

ξt ∼ N (0, 0.012 ).

(26)

with sk,t a periodic component composed of two sinusoidal terms with slightly varying periods and random phase shifts ϕ1,k , ϕ2,k ∼ Uniform(0, 2π), defined as     2πt 2πt + ϕ1,k + 0.03 sin + ϕ2,k . (27) sk,t = 0.07 sin p1,k p2,k

27

Each station also includes a simple independent periodic variable, defined as     2πt 2πt Zk,t = 0.15 + 0.08 sin + ϕ̃1,k + 0.04 sin + ϕ̃2,k + ωt , p̃1,k p̃2,k

(28)

with ωt ∼ N (0, 0.012 ). The frequencies and phases of this component are independent of the latent process and global extremes. In addition to station-level variables, a single regime process R[t] models a station with drought dynamics. In the normal regime, the process follows Rt = µN + ϕN (R[t − 1] − µN ) + ϵt + 0.2 gt ,

(29)

with µN = 0.35, ϕN = 0.94, and ϵt ∼ N (0, 0.012 ). During drought periods, the process gradually reverts toward a near-zero mean, Rt = µD + ϕD (R[t − 1] − µD ),

(30)

with µD = 0.02 and ϕD = 0.85. Drought episodes persist for a minimum duration and evolve smoothly to avoid instantaneous drops. Spikes and drought are mutually exclusive: if gt > 0, the regime is forced to normal dynamics. Finally, the data set stacks the three dynamic, three correlated periodic, three independent periodic variables, and the regime process as one multivariate time series.  X = D1 ,

...,

DK ,

S1 ,

...,

SK ,

28

Z1 ,

...,

ZK ,

 R ∈ Rn×(3K+1) ,

(31)

F

Detailed metrics Missing Model Custom GP U-NET 25% NHITS TSMixer-D/P CSDI Custom GP U-NET 50% NHITS TSMixer-D/P CSDI Custom GP U-NET 75% NHITS TSMixer-D/P CSDI Custom GP U-NET 90% NHITS TSMixer-D/P CSDI

RMSE 0.0119 ± 0.0000 0.0272 ± 0.0001 0.1629 ± 0.0004 – – 0.0925 ± 0.0002 0.0125 ± 0.0000 0.0323 ± 0.0001 0.2082 ± 0.0003 – – 0.1269 ± 0.0004 0.0154 ± 0.0002 0.0438 ± 0.0002 0.2342 ± 0.0001 – – 0.1482 ± 0.0006 0.0237 ± 0.0002 0.0687 ± 0.0002 0.2461 ± 0.0003 – – 0.2250 ± 0.0005

MAE 0.0094 ± 0.0000 0.0151 ± 0.0000 0.1372 ± 0.0004 – – 0.0652 ± 0.0001 0.0098 ± 0.0000 0.0175 ± 0.0000 0.1743 ± 0.0003 – – 0.0862 ± 0.0002 0.0110 ± 0.0000 0.0239 ± 0.0000 0.1950 ± 0.0002 – – 0.0998 ± 0.0004 0.0157 ± 0.0001 0.0400 ± 0.0001 0.2048 ± 0.0002 – – 0.1734 ± 0.0005

SMAPE 4.2593 ± 0.0023 5.4961 ± 0.0113 52.4132 ± 0.2071 – – 21.0319 ± 0.0289 4.3723 ± 0.0031 6.1521 ± 0.0150 70.2633 ± 0.3969 – – 24.8444 ± 0.0406 4.7121 ± 0.0033 8.2759 ± 0.0184 81.2690 ± 0.0595 – – 28.4976 ± 0.0813 6.7905 ± 0.0257 14.3499 ± 0.0307 86.5986 ± 0.1939 – – 48.3544 ± 0.1071

CRPS 0.0208 ± 0.0000 0.0354 ± 0.0001 0.2956 ± 0.0004 – – 0.1471 ± 0.0002 0.0218 ± 0.0000 0.0411 ± 0.0001 0.3704 ± 0.0004 – – 0.1987 ± 0.0006 0.0241 ± 0.0000 0.0567 ± 0.0001 0.4147 ± 0.0004 – – 0.2357 ± 0.0008 0.0339 ± 0.0001 0.0959 ± 0.0003 0.4352 ± 0.0003 – – 0.3952 ± 0.0012

(a) Synthetic Imputation Metrics Horizon Model Custom GP U-NET NHITS 6 TSMixer-D TSMixer-P CSDI Custom GP U-NET NHITS 12 TSMixer-D TSMixer-P CSDI Custom GP U-NET NHITS 18 TSMixer-D TSMixer-P CSDI

RMSE 0.0489 ± 0.0000 0.0529 ± 0.0007 0.2404 ± 0.0011 0.0578 ± 0.0000 0.0524 ± 0.0000 0.0588 ± 0.0010 0.1889 ± 0.0005 0.0565 ± 0.0001 0.0770 ± 0.0004 0.2484 ± 0.0008 0.0615 ± 0.0000 0.0614 ± 0.0000 0.0582 ± 0.0001 0.1864 ± 0.0006 0.0714 ± 0.0000 0.0823 ± 0.0002 0.2509 ± 0.0007 0.0764 ± 0.0000 0.0743 ± 0.0000 0.0682 ± 0.0002 0.1831 ± 0.0002

MAE 0.0207 ± 0.0000 0.0284 ± 0.0002 0.2013 ± 0.0013 0.0367 ± 0.0000 0.0274 ± 0.0000 0.0295 ± 0.0004 0.1345 ± 0.0005 0.0263 ± 0.0001 0.0442 ± 0.0002 0.2071 ± 0.0006 0.0385 ± 0.0000 0.0346 ± 0.0000 0.0314 ± 0.0001 0.1332 ± 0.0003 0.0330 ± 0.0001 0.0484 ± 0.0001 0.2080 ± 0.0007 0.0456 ± 0.0000 0.0430 ± 0.0000 0.0361 ± 0.0001 0.1309 ± 0.0001

SMAPE 6.6687 ± 0.0121 9.8296 ± 0.0487 85.5644 ± 0.7897 15.2595 ± 0.0000 10.3289 ± 0.0000 9.7174 ± 0.0434 34.8726 ± 0.1201 8.1881 ± 0.0186 14.5196 ± 0.0939 88.3506 ± 0.4370 15.8211 ± 0.0000 13.3507 ± 0.0000 10.8374 ± 0.0271 35.3875 ± 0.0216 9.5539 ± 0.0330 16.8145 ± 0.0592 88.8613 ± 0.7188 17.2002 ± 0.0000 16.0825 ± 0.0000 11.7899 ± 0.0247 34.7984 ± 0.0380

CRPS 0.0512 ± 0.0002 0.0699 ± 0.0005 0.4319 ± 0.0020 0.1062 ± 0.0000 0.0793 ± 0.0000 0.0893 ± 0.0004 0.3154 ± 0.0011 0.0635 ± 0.0007 0.1047 ± 0.0004 0.4413 ± 0.0010 0.1113 ± 0.0000 0.1000 ± 0.0000 0.0822 ± 0.0001 0.3196 ± 0.0009 0.0794 ± 0.0006 0.1159 ± 0.0003 0.4442 ± 0.0009 0.1318 ± 0.0000 0.1243 ± 0.0000 0.0865 ± 0.0001 0.3155 ± 0.0005

(b) Synthetic Forecasting Metrics Figure 15: Synthetic dataset performance comparison. Best values are highlighted in green; secondbest in blue. Top: Imputation metrics across missing ratios. Bottom: Forecasting metrics across horizons.

29

Missing Model Custom Augmented 25% GP CSDI U-NET Custom Augmented 50% GP CSDI U-NET Custom Augmented 75% GP CSDI U-NET Custom Augmented 90% GP CSDI U-NET

RMSE 0.1676 (±0.0049) 0.1597 (±0.0029) 0.1609 (±0.0008) 0.3962 (±0.0009) 0.7479 (±0.0006) 0.1865 (±0.0012) 0.1850 (±0.0010) 0.1915 (±0.0004) 0.3913 (±0.0007) 0.8144 (±0.0007) 0.2933 (±0.0009) 0.2744 (±0.0018) 0.3839 (±0.0059) 0.4263 (±0.0014) 0.9145 (±0.0006) 0.3563 (±0.0021) 0.3212 (±0.0020) 0.5257 (±0.0012) 0.5608 (±0.0024) 0.9850 (±0.0004)

Horizon Model Custom Augmented 6 CSDI GP U-NET Custom Augmented 12 CSDI GP U-NET Custom Augmented 18 CSDI GP U-NET

RMSE 0.3211 (±0.0021) 0.2679 (±0.0013) 0.3887 (±0.0024) 0.4059 (±0.0072) 0.9606 (±0.0012) 0.3895 (±0.0007) 0.3013 (±0.0024) 0.4351 (±0.0023) 0.4870 (±0.0090) 1.0037 (±0.0006) 0.4526 (±0.0009) 0.3412 (±0.0017) 0.5019 (±0.0007) 0.5640 (±0.0046) 1.0199 (±0.0003)

MAE SMAPE 0.0356 (±0.0003) 7.5701 (±0.0195) 0.0356 (±0.0002) 7.5243 (±0.0260) 0.0457 (±0.0002) 9.9667 (±0.0560) 0.1192 (±0.0004) 19.0637 (±0.0491) 0.5361 (±0.0005) 92.2608 (±0.2457) 0.0432 (±0.0001) 9.4832 (±0.0269) 0.0440 (±0.0001) 9.3706 (±0.0292) 0.0638 (±0.0003) 13.7339 (±0.0531) 0.1239 (±0.0002) 20.3463 (±0.0494) 0.5857 (±0.0010) 101.1553 (±0.3862) 0.0656 (±0.0001) 13.2244 (±0.0311) 0.0633 (±0.0002) 12.0324 (±0.0540) 0.1044 (±0.0009) 18.9232 (±0.0734) 0.1441 (±0.0004) 23.6783 (±0.0289) 0.6431 (±0.0004) 111.1768 (±0.1814) 0.1127 (±0.0006) 22.3026 (±0.1002) 0.0957 (±0.0004) 18.5258 (±0.0629) 0.1596 (±0.0004) 26.2071 (±0.0546) 0.2360 (±0.0005) 39.5478 (±0.0619) 0.6862 (±0.0004) 118.6330 (±0.1639)

CRPS 0.0322 (±0.0003) 0.0320 (±0.0002) 0.0438 (±0.0001) 0.1069 (±0.0004) 0.4638 (±0.0005) 0.0391 (±0.0001) 0.0394 (±0.0001) 0.0602 (±0.0002) 0.1115 (±0.0001) 0.5098 (±0.0008) 0.0588 (±0.0002) 0.0558 (±0.0003) 0.1010 (±0.0006) 0.1324 (±0.0003) 0.5686 (±0.0003) 0.1030 (±0.0005) 0.0852 (±0.0004) 0.1563 (±0.0003) 0.2194 (±0.0003) 0.6157 (±0.0003)

(a) Imputation Metrics MAE SMAPE 0.0878 (±0.0010) 17.3480 (±0.1479) 0.0773 (±0.0002) 15.3664 (±0.0752) 0.1260 (±0.0003) 22.1140 (±0.0881) 0.1126 (±0.0015) 20.2815 (±0.1756) 0.6716 (±0.0010) 116.9822 (±0.4429) 0.1140 (±0.0004) 21.1204 (±0.1045) 0.1038 (±0.0003) 19.5852 (±0.0642) 0.1489 (±0.0005) 24.8425 (±0.0694) 0.1618 (±0.0007) 26.4087 (±0.0674) 0.6962 (±0.0005) 120.2760 (±0.1908) 0.1473 (±0.0004) 25.0509 (±0.0569) 0.1317 (±0.0004) 23.3131 (±0.0377) 0.1779 (±0.0004) 27.0204 (±0.0746) 0.2010 (±0.0009) 30.3275 (±0.0780) 0.7050 (±0.0004) 121.3809 (±0.1661)

CRPS 0.0822 (±0.0009) 0.0702 (±0.0003) 0.1171 (±0.0005) 0.1059 (±0.0007) 0.6058 (±0.0009) 0.1066 (±0.0006) 0.0936 (±0.0003) 0.1420 (±0.0004) 0.1540 (±0.0009) 0.6297 (±0.0006) 0.1374 (±0.0007) 0.1186 (±0.0004) 0.1747 (±0.0003) 0.1968 (±0.0008) 0.6379 (±0.0003)

(b) Forecasting Metrics Figure 16: OPE dataset performance comparison. Best values are highlighted in green; second-best in blue. Top: Imputation metrics across missing ratios. Bottom: Forecasting metrics across horizons.

30