Autoregressive Denoising Diffusion Models for Multivariate Probabilistic Time Series Forecasting

In plain words

Forecasting several related series at once is hard, because they move together. A common shortcut is to predict each series alone, which ignores those links. TimeGrad instead predicts all series for the next time step jointly, and it draws a random sample from a learned distribution rather than giving one best guess. Repeating that draw many times gives a spread of plausible futures.

The sampler is a diffusion model. It starts from pure random noise and removes the noise in a series of small steps, with a neural network trained to predict the noise at each step. A recurrent network reads the history and tells the sampler what has happened so far. The paper reports that this beats the other probabilistic forecasters it compares against on five of six public datasets. The cost is speed: every forecast step runs the denoising network once per diffusion step, so sampling is slow.

TL;DR

• Multivariate probabilistic forecaster: an LSTM encodes the past, and a conditional denoising diffusion model samples each future step §3.
• Training uses the standard noise-prediction loss, applied to random diffusion steps §3.1 alg. 1.
• Inference draws 100 sample paths, each with 100 reverse diffusion steps per forecast step §4.2 alg. 2.
• Author-reported CRPS_sum wins on Solar, Electricity, Traffic, Taxi and Wikipedia; on Exchange (the smallest set) it does not tab. 2.
• The official code is the pts package in pytorch-ts, which includes the model, but the paper's experiment configuration is not in the repository.

Problem

Classical forecasters give univariate point forecasts, need hand-built seasonal features and are fitted one series at a time §1. Multivariate models need a tractable output distribution. A full Gaussian over D dimensions needs O(D^2) network outputs and an O(D^3) loss, which is impractical for large D, and it captures only second-order dependencies §3. Low-rank Gaussians, copulas, normalising flows and VAEs each impose structure that can misfit disconnected modes of the data §4.3. Energy-based models are less constrained but are hard to train §1. The method assumes covariates are known for every time point, including the forecast horizon §3.

Method

TimeGrad writes the future of a multivariate series as a product of per-step conditionals. Each step depends on the whole past and the known covariates §3 eq. 8. An RNN carries a state h_t through the series, taking the value and covariates at each step as input eq. 9. Each per-step distribution is a diffusion model whose noise predictor is conditioned on h_{t-1} eq. 10.

The per-step diffusion model follows Ho et al. (2020) §2:

1. Forward process. Gaussian noise is added through a fixed schedule beta_1 ... beta_N. The noisy value at any level has a closed form §2 eq. 3.
2. Reverse process. Each reverse step is Gaussian, with a learned mean and a variance Sigma §2 eq. 1. The mean is parameterised through a noise predictor epsilon_theta, which the paper derives from the closed-form posterior mean §2.
3. Variance. Sigma is fixed to the posterior variance beta_tilde §3.1 eq. 5.
4. Loss. With that variance the weighted objective reduces to an unweighted squared error on the noise, now conditioned on h_{t-1} §3.1 eq. 7.
5. Training. For each time step of the prediction window, draw a diffusion step n and a noise sample, take a gradient step on the noise error, and repeat until convergence §3.1 alg. 1.
6. Inference. Run the RNN over the context window to get h_T. For each future step, draw x^N from white noise and apply the reverse steps down to n = 1, adding noise at every step except the last §3.2 alg. 2. Feed the sample back into the RNN and repeat for the horizon. Repeat the whole procedure S = 100 times to get empirical quantiles §3.2 §4.2.
7. Scaling. Divide each dimension by its context-window mean before the model, and multiply samples back afterwards §3.3.
8. Covariates. Time features such as day-of-week and hour-of-day, and static embeddings where present, plus lagged values chosen from the data frequency §3.4.
9. Denoiser architecture. Eight residual blocks of 1-D dilated convolutions with kernel 3, gated activation sigma(.) * tanh(.), and 8 channels. Dilation alternates between 1 and 2. The noise index goes through a Fourier embedding with N_max = 500 into a 32-dimensional vector. Skip connections are summed for the output §4.2 fig. 2.

The simplified conditional objective that is optimised per time step is:

\mathbb{E}_{\mathbf{x}_t^0, \epsilon, n} \left[ \| \epsilon - \epsilon_\theta (\sqrt{\bar{\alpha}_n} \mathbf{x}^0_t + \sqrt{1 - \bar{\alpha}_n} \epsilon, \mathbf{h}_{t-1}, n ) \|^ 2 \right] §3.1 (simplifies eq. 7)

# Training step on a batch of windows  [§3.1] [alg. 1]
# Input:  windows (B, L, D) with L = context + prediction, already scaled
# Output: scalar loss

def training_step(windows, observed, covariates):
    # 1. Encode the window with the RNN, with the true values as input
    #    [§3] [eq. 9]
    states = rnn(concat(windows, covariates))  # (B, L, H)

    # 2. Map each state to the conditioning input of the denoiser
    #    (the width is not stated in the paper; see Paper vs code)
    cond = project(states)  # (B, L, C)

    # 3. Draw one diffusion step and one noise sample per time point
    #    Steps are 0-based: index k stands for the paper's n = k + 1
    steps = uniform_int(0, N, shape=(B, L))  # [alg. 1]
    noise = normal(shape=windows.shape)  # epsilon ~ N(0, I)  [§2]

    # 4. Noise the window in closed form  [§2] [eq. 3]
    noisy = (
        sqrt(alpha_bar[steps]) * windows
        + sqrt(1 - alpha_bar[steps]) * noise
    )

    # 5. Predict the noise from the noisy value, the step and the condition
    #    [§3.1]
    predicted = denoiser(noisy, steps, cond)  # (B, L, D)

    # 6. Unweighted squared error on the noise, masked where unobserved
    #    [§3.1] [eq. 7]
    squared = (noise - predicted) ** 2
    loss = masked_mean(squared, observed)
    return loss

# Sample future paths  [§3.2] [alg. 2]
# Input:  context (B, C, D), future covariates (B, P, F), S paths
# Output: samples (B, S, P, D) in original units

def forecast(context, past_covariates, future_covariates, num_samples):
    # 1. Scale each dimension by the mean of its context window  [§3.3]
    scale = mean_abs(context)  # per series and dimension
    # 2. Run the RNN over the context to get the warm-up state h_T  [§3.2]
    state = rnn_over(context / scale, past_covariates)
    paths = []
    for step in range(P):
        # 3. Draw x^N ~ N(0, I) for every path  [alg. 2]
        x = normal(shape=(B, num_samples, D))
        for k in reversed(range(N)):
            # 4. Predict the noise and form the reverse-step mean
            #    [§2] [eq. 1]; beta[k] is the paper's beta_{k+1}
            eps = denoiser(x, k, cond_from(state))
            mean = (x - beta[k] / sqrt(1 - alpha_bar[k]) * eps) / sqrt(alpha[k])
            # 5. Add noise with variance beta_tilde, none at the last step
            #    [§2] [eq. 5] [alg. 2]
            z = normal(shape=x.shape) if k > 0 else zeros_like(x)
            x = mean + sqrt(beta_tilde[k]) * z
        # 6. Keep the sample and feed it back as the next input  [§3.2]
        paths.append(x * scale)
        state = rnn_step(x, future_covariates[:, step], state)
    return stack(paths, axis=2)

Hyperparameters

Results

Author-reported. Setup: six public datasets with the splits and preprocessing of the cited Vec-LSTM work §4.1 tab. 1. Metric: CRPS_sum, computed on the sum over dimensions, averaged over the horizon §4.1. TimeGrad is averaged over 10 re-trainings with its standard error tab. 2. Baselines are classical multivariate models (VAR, VAR-Lasso, GARCH, VES), deep models (KVAE, Vec-LSTM variants, GP variants, and a Transformer with masked autoregressive flows) §4.3.

Source: tab. 2. Bold marks the best value per column in the paper.

• The paper states that TimeGrad is the new state of the art on all but the smallest dataset §4.3. The table agrees: Exchange (8 dimensions) is won by VES, VAR and Transformer MAF at 0.005 tab. 2.
• Ablation on Electricity, varying the number of diffusion steps N over 2, 4, ..., 256 with everything else fixed, reports that N can drop to about 10 with little loss and that the best value is near N = 100 §4.4 fig. 3. Mean and standard error are over 5 runs fig. 3.
• Figure 4 shows predicted medians and 50% and 90% intervals for the first 6 of 963 Traffic dimensions in the first rolling window fig. 4. It has no numbers.

Paper vs code

Code at commit 7860c9693d55b5c086867477cc33c89485ed0167 of zalandoresearch/pytorch-ts (the authors' organisation), MIT licence. The paper says code will be released after review §4. The README lists this paper as an implemented model README.md.

• Loss window. The paper sums the loss over the prediction window only, t in [t_0, T], with the context window providing the state §3.1. The code computes the RNN, the target and the loss over context plus prediction time_grad_network.py time_grad_network.py.
• Scaling fallback. The paper divides by the context mean, or by 1 if it is zero §3.3. The code's MeanScaler uses the mean absolute value over observed points. When a series has no observed points in the context, it falls back to the batch-wide mean for that dimension, and it has a floor of 1e-10 scaler.py scaler.py scaler.py.
• Denoiser input. The paper conditions the noise predictor on the RNN state h_{t-1} §3 eq. 10. The code first projects that state to a vector of width conditioning_length (default 100) through the diffusion output head time_grad_network.py distribution_output.py time_grad_estimator.py.
• Series embedding. The paper describes categorical embeddings for time and static features §3.4. The code's network ignores the embedding_dimension argument (default 5) and sets its series embedding to dimension 1 time_grad_network.py time_grad_estimator.py.
• Time features. The paper names day-of-week and hour-of-day features §3.4. The code builds Fourier time features from the data frequency time_grad_estimator.py.
• Noise-index embedding frequencies. The paper states Fourier positional embeddings with N_max = 500 §4.2. The code multiplies the step index by 10^(4k/d), so its frequencies rise with the channel index k. The usual Transformer encoding uses falling frequencies epsilon_theta.py.
• Variance and loss. The code uses beta_tilde as the sampling variance gaussian_diffusion.py and an unweighted squared error on the noise gaussian_diffusion.py. This agrees with the paper §3.1 eq. 5 eq. 7.
• Diffusion schedule. The code uses a linear schedule from 1e-4 to the end value over N points gaussian_diffusion.py, with 100 steps and end value 0.1 by default time_grad_estimator.py time_grad_estimator.py. This agrees with §4.2.
• Architecture details. The paper gives 8 residual layers and 8 channels §4.2 fig. 2. The code agrees time_grad_estimator.py, and adds details the paper does not state: a 0.4 leaky-ReLU epsilon_theta.py, a skip-sum divided by the square root of the layer count epsilon_theta.py, and a residual sum divided by the square root of 2 epsilon_theta.py.
• RNN cell. The paper states two LSTM layers of width 40 §4.2. The code agrees as defaults time_grad_estimator.py. The example notebook uses a GRU Time-Grad-Electricity.ipynb.
• Training configuration. The paper gives Adam, learning rate 1e-3 and batch size 64 §4.2. The repository has no script that reproduces the table. The trainer defaults are batch 32, 100 epochs and 50 batches per epoch trainer.py. The example notebook uses batch 64, 20 epochs and 100 batches per epoch Time-Grad-Electricity.ipynb.
• Data. The paper says the datasets are preprocessed exactly as in the Vec-LSTM work §4.1. The notebook loads GluonTS's electricity_nips recipe Time-Grad-Electricity.ipynb. Whether that matches the paper's preprocessing is not verified here.

No other discrepancies found for the sampling loop, the noise schedule or the simplified objective.

Pitfalls

• Loss window differs. Implement the paper's prediction-window loss §3.1 or the code's full-window loss, not both. The code's version is the one that reproduces the released model (see Paper vs code).
• Scale fallback. A series that is all zeros in the context is scaled by the batch mean in the code, not by 1. Data with all-zero context windows will differ between the two. The paper also writes "mean", while the code uses the mean absolute value; the two agree for non-negative data such as the Exchange and Solar sets tab. 1.
• Step indexing. The code counts diffusion steps from 0 to N-1, and betas[0] is the paper's beta_1 gaussian_diffusion.py gaussian_diffusion.py. Off-by-one errors are easy when porting the paper's equations.
• Sampling cost. Each forecast step loops N times over the denoiser, so a horizon of P steps costs about P * N denoiser calls per sample batch §6 time_grad_network.py. Training needs one denoiser call per time point §6. The paper suggests a faster-sampling variant from WaveGrad §6, which is not in this code.
• Scaling is in place. log_prob divides its input tensor by the scale in place gaussian_diffusion.py. Pass a copy when reusing the window tensor.
• Circular padding. The input and conditioning convolutions use padding of 2 with circular mode epsilon_theta.py epsilon_theta.py. Internally the series is D + 4 long, and the two unpadded kernel-3 convolutions bring it back to D epsilon_theta.py. The figure caption says the spatial size stays D fig. 2; the output length is D, but the intermediate length is not.
• Unstated values in the code. Dropout 0.1 time_grad_estimator.py, weight decay 1e-6 trainer.py and the conditioning width 100 time_grad_estimator.py are code defaults. The paper does not state them. Whether the paper's experiments used them is not verifiable from the repository.
• Clipping. The sampler does not clip the denoised estimate by default gaussian_diffusion.py. The paper does not mention clipping.
• Traffic scaling. The paper says Traffic does not need scale normalisation §4.1. The code scales whenever scaling=True, which is the default time_grad_estimator.py. Whether scaling was switched off for Traffic in the paper's runs is not stated.
• Stochastic results. Each TimeGrad number is a mean over 10 re-trainings, but the sampler draws 100 random paths per forecast. Compare CRPS_sum with the reported standard error, not single-run values.
• Electricity ablation numbers come from a single dataset §4.4. The paper says the trend holds on others but shows no numbers fig. 3.

Lineage

• Builds on: Ho et al. (2020), Denoising Diffusion Probabilistic Models, for the per-step diffusion model and its training loss §2. Sohl-Dickstein et al. (2015) introduced the diffusion-model framework §2. DeepAR (Salinas et al.) provides the autoregressive recurrent conditioning §3.
• Closest related: WaveGrad and DiffWave, which use conditional denoising diffusion for waveforms without a temporal recurrence §5.1. The paper contrasts its autoregressive treatment with those.
• Compared against: VAR, VAR-Lasso, GARCH, VES, KVAE, Vec-LSTM (Salinas et al. 2019, low-rank Gaussian copula), GP scaling and copula variants, and the Transformer with masked autoregressive flows (Rasul et al., conditioned normalising flows) §4.3 tab. 2.