Leapfrog Diffusion Model for Stochastic Trajectory Prediction

In plain words

Predicting where a person or a basketball player moves next has to allow for several possible futures. Diffusion models can produce such a spread of futures. They start from random noise and clean it up step by step. Each prediction, though, can take many steps, which is too slow for a system that must update several times a second.

The authors' leapfrog diffusion model (LED) skips most of those steps. A small trained network proposes a set of starting future paths that are already close to a plausible answer. Only a few clean-up steps then run on those starting paths. The starting paths are proposed together, so they can be spread out over different possibilities rather than landing at random.

The authors report that LED predicts faster than the standard diffusion approach, and that its average errors are the best or tied-best on each of the four datasets they tested. The numbers on this card come from the paper and its code.

TL;DR

LED is a trajectory predictor that samples several possible futures with a diffusion model. A trainable leapfrog initializer produces K = 20 correlated samples at denoising step tau = 5 of Gamma = 100, so the first 95 steps are replaced. Author-reported NBA results at 4.0 s are 0.81/1.10 minADE20/minFDE20 (m) against 0.96/1.27 for MID, at about 46 ms per prediction against about 886 ms for 100-step diffusion. The official code's README reports the NBA numbers, but the code omits the denoiser's training stage, and its sampling departs from the paper in several places (see Paper vs code).

Problem

Stochastic trajectory prediction must output a spread of plausible futures, not one path §1 §3.1. Diffusion models represent such spreads well, but they need an adequate number of denoising steps. The authors state that about 100 steps are needed for decent NBA performance, which takes about 886 ms per prediction, while the next frame arrives every 200 ms §1. Independent random starting samples can also miss important futures, so a few samples may cover too little of the distribution §1.

The problem is stated for an agent's past trajectory X over T_p steps, its neighbours' past trajectories, and the ground-truth future Y over T_f steps §3.1. The aim is a model whose K samples include one close to Y eq. 1.

Method

LED keeps the standard forward (noising) process of a diffusion model and changes only where reverse denoising starts §4.1 eq. 3. A standard model draws K i.i.d. Gaussian samples at step Gamma and runs Gamma denoising steps eq. 2. LED instead uses a trainable leapfrog initializer that directly produces K samples of the distribution expected at step tau, then runs tau denoising steps eq. 3 alg. 1.

1. Problem and metric. Predict K trajectories; one of them should be near the ground truth §3.1 eq. 1. Results use minADE_K and minFDE_K §5.3.
2. Standard diffusion (baseline). Forward noising of Y, then K i.i.d. draws from N(0, I) and Gamma conditional denoising steps §3.2 eq. 2.
3. Leapfrog diffusion. Same forward process; the step-Gamma samples are replaced by the leapfrog initializer's K samples at step tau, followed by tau denoising steps §4.1 eq. 3. The authors argue the representation capacity is unchanged because the forward process is the same §4.1.
4. Leapfrog initializer, reparameterization. The initializer predicts a shared mean trajectory mu, a scalar standard deviation sigma, and K normalized sample positions S_k. The K trajectories are Y_k^tau = mu + sigma * S_k §4.2 eq. 4. Splitting the target this way is meant to ease learning §4.2.
5. Mean module. A social embedding by multi-head attention over the neighbours, a temporal embedding from a 1D convolution followed by a GRU, and an MLP that fuses the two §4.2 eq. 5.
6. Scale and sample modules. The scale module outputs sigma. The sample module encodes sigma and outputs all K normalized positions from the same features, so the samples are generated jointly and can be spread apart §4.2.
7. Denoising module. A context encoder over the past trajectories, a noise estimator, and the standard DDPM-style update §4.3 eq. 6.
8. Training stage 1: denoiser. Train the denoising module with the usual diffusion noise-estimation loss over random steps §4.4 eq. 7.
9. Training stage 2: initializer. Freeze the denoiser and train the initializer. The stage-2 loss has a best-of-K distance term weighted by w, plus an uncertainty term that divides the average error by sigma squared and adds log sigma squared as a regularizer §4.4 eq. 8. The authors split training in two because the second stage is more stable when the target distribution is fixed §4.4.
10. Inference. Compute mu, sigma and S_k, form the K samples eq. 4, then run the remaining tau denoising steps for gamma = tau-1 down to 0 §4.5 alg. 1 eq. 6.

Inference, as in the paper's Algorithm 1 alg. 1:

# Leapfrog diffusion inference  [alg. 1]
# Input:  past trajectories (ego and neighbours), leapfrog step tau
# Output: K future trajectories, K = 20 on NBA  [§5.3]

def predict(past, neighbours, tau, num_samples):
    # Shared mean trajectory  [§4.2] [eq. 5]
    mean = mean_module(past, neighbours)               # (T_f, 2)
    # Scalar standard deviation of the leapfrog distribution  [§4.2]
    sigma = scale_module(past, neighbours)             # scalar
    # Normalized sample positions, all K produced together  [§4.2]
    normalized = sample_module(past, neighbours, sigma)
    # (K, T_f, 2); reparameterize into correlated samples  [eq. 4]
    samples = mean + sigma * normalized
    # Refine with tau ordinary denoising steps, gamma = tau-1 ... 0  [§4.5]
    for gamma in reversed(range(tau)):
        samples = denoise_step(samples, past, neighbours, gamma)
    return samples

One denoising step, as written in the paper eq. 6:

# One denoising step gamma  [eq. 6]
# Input: noisy samples y_next, shape (K, T_f, 2)

def denoise_step(y_next, past, neighbours, gamma):
    # Context from the past trajectories  [eq. 6a]
    context = context_encoder(past, neighbours)
    # Estimated noise in y_next  [eq. 6b]
    eps = noise_estimator(y_next, context, gamma + 1)
    # Remove the predicted noise  [eq. 6c]
    alpha = alphas[gamma]
    alpha_bar = alpha_bars[gamma]
    mean = (y_next - (1 - alpha) / sqrt(1 - alpha_bar) * eps) / sqrt(alpha)
    # Add fresh Gaussian noise scaled by sqrt(1 - alpha)  [eq. 6c]
    return mean + sqrt(1 - alpha) * standard_normal(y_next.shape)

Stage-2 loss, as written in the paper eq. 8:

# Stage-2 loss for the leapfrog initializer; denoiser frozen  [eq. 8]
# Input: final (K, T_f, 2) after tau denoising steps; future (T_f, 2)

def initializer_loss(final, future, sigma):
    # Euclidean error of each of the K trajectories  [eq. 8]
    errors = norm(final - future)                      # (K,)
    # Best of K, weighted by w = 50  [eq. 8] [§5.2]
    distance_term = w * min(errors)
    # Mean error over K, divided by sigma squared  [eq. 8]
    uncertainty_term = mean(errors) / sigma ** 2
    # Regularizer that stops sigma from growing without bound  [eq. 8]
    regularizer = log(sigma ** 2)
    return distance_term + uncertainty_term + regularizer

The authors considered explicit supervision of the initializer's statistics and did not use it. It needs M >> K denoised samples per example, which the authors estimate at about six days per epoch on NBA §4.4.

Hyperparameters

Results

Author-reported. Setup: four datasets, baselines of ten methods per dataset [tabs. 1-4]; metric minADE20/minFDE20 in metres; NBA inference timings in tab. 6.

Per-subset ETH-UCY results for LED (ADE/FDE): ETH 0.39/0.58, Hotel 0.11/0.17, Univ 0.26/0.43, Zara1 0.18/0.26, Zara2 0.13/0.22 tab. 4. MID* is a re-run of MID's public code under the paper's protocol tab. 3 §5.3. The SDD protocol predicts the destination with LED and fills the rest of the trajectory with a multi-layer perceptron [supp. §2].

Ablations (author-reported, five trials, NFL, K = 20 tab. 5):

Steps on NBA tab. 6: standard diffusion with 100 steps gives 0.94/1.21 at about 886 ms; LED with tau = 3, 5 and 10 gives 0.84/1.10 at about 30 ms, 0.81/1.10 at about 46 ms, and 0.82/1.08 at about 89 ms. Compared with other fast samplers at similar times tab. 7: progressive distillation and DDIM (several step counts) give 0.91 to 1.02 ADE at 54 to 530 ms.

Figures: fig. 1 schematic of the speed-up; fig. 2 inference architecture; fig. 3 best-of-20 NBA predictions against two baselines; fig. 4 learned mean and variance on NBA, where the authors say the variance follows scene complexity; fig. 5 correlated versus i.i.d. samples on NFL with K = 4.

Paper vs code

Code: official repository MediaBrain-SJTU/LED at the commit above. The owner is the MediaBrain group of the authors' institution and the repository is linked from the abstract. The repository has no LICENSE file, so its code is cited by path and line only.

• Stage 1 is absent. The paper trains the denoiser first §4.4 §5.2. The denoising loss exists train_led_trajectory_augment_input.py but fit never calls it; fit trains only the initializer train_led_trajectory_augment_input.py. The denoiser is loaded from a checkpoint that is not in the repository led_augment.yml train_led_trajectory_augment_input.py. README: "add training/evaluation for diffusion models" is listed as a TODO README.md.
• Stage 2 schedule differs. The paper gives 200 epochs, lr 1e-4 and decay 0.9 every 32 epochs §5.2. The config sets 100 epochs, decay step 8 and gamma 0.5 led_augment.yml. The learning rate comes from the command line, whose default is 0.002 main_led_nba.py; the config's lr (1e-3) is never read train_led_trajectory_augment_input.py.
• Optimizer. Paper: Adam §5.2. Code: AdamW train_led_trajectory_augment_input.py.
• Denoiser "frozen". The paper freezes the denoiser in stage 2 §4.4. The code only excludes its weights from the optimizer; the denoiser is left in training mode train_led_trajectory_augment_input.py and its layers include dropout 0.1 model_diffusion.py layers.py.
• Sample scaling. Paper: Y_k = mu + sigma * S_k eq. 4. Code also divides the normalized samples by the mean standard deviation across the K samples train_led_trajectory_augment_input.py; the same rescaling is used at test time train_led_trajectory_augment_input.py. The code's network outputs log sigma squared; the scale is exp(v/2) train_led_trajectory_augment_input.py.
• Noise in denoising. Paper eq. 6c adds sqrt(1 - alpha) z at each step eq. 6. The accelerated sampler multiplies this noise by 1e-5 train_led_trajectory_augment_input.py; it is used in training and testing train_led_trajectory_augment_input.py train_led_trajectory_augment_input.py.
• Distance term. Paper: w * min_k ||Y - Y_k||_2, a whole-trajectory norm eq. 8. Code: per-step L2 errors multiplied by linear weights from 2.0 down to 0.1, averaged over time, then minimum over K train_led_trajectory_augment_input.py train_led_trajectory_augment_input.py. The code's comment calls the weighting "not necessary" train_led_trajectory_augment_input.py.
• Uncertainty term. Paper: mean error over K divided by sigma squared, plus log sigma squared eq. 8. Code: per-step errors averaged over K and time, with no temporal weights train_led_trajectory_augment_input.py.
• Trajectory fed to the loss. The paper's loss uses Y_k without saying whether these are the tau-step outputs or the initializer samples eq. 8. The code uses the outputs after tau denoising steps train_led_trajectory_augment_input.py.
• Sample module. Paper: each of the three modules has its own social and temporal encoder §4.2. Code: one social transformer is shared by all three; temporal encoders are separate model_led_initializer.py.
• Denoiser architecture. Paper: noise estimator implemented as MLPs §4.3. Code: transformer layers with conditional linear layers, width 512 model_diffusion.py.
• Input features. The paper does not state inputs. Code uses 6 features per step: absolute position normalized by mean [14, 7.5] and scale 5, position relative to the last observed point, and velocity train_led_trajectory_augment_input.py led_augment.yml. Targets are displacements from the last observed point, divided by 5 train_led_trajectory_augment_input.py.
• Gradient clipping. Not in the paper; code clips at 1.0 train_led_trajectory_augment_input.py.
• Reproduction. README reports NBA values of 0.1766/0.2694 (1 s), 0.3693/0.5642 (2 s), 0.5817/0.8366 (3 s), 0.8095/1.0960 (4 s) from the released checkpoint README.md. These match tab. 1 within rounding. The checkpoint is not in the repository, so we did not run it.

Pitfalls

• You cannot retrain end to end from this repository. The denoiser checkpoint, the NBA dataset and its loader are absent from the clone; the README points to a Google Drive folder README.md. Only the stage-2 initializer training and NBA evaluation code are present.
• Sampling noise is switched off in practice. Each denoising step adds noise scaled by 1e-5 train_led_trajectory_augment_input.py. The paper's update includes the noise term eq. 6. Treat the code's denoising as nearly deterministic when you reimplement.
• Dropout stays on at evaluation. The model is never switched to eval mode on the NBA test path train_led_trajectory_augment_input.py; the only eval call is in the unused p_sample_loop train_led_trajectory_augment_input.py. Results depend on a fixed seed train_led_trajectory_augment_input.py.
• Hard-coded shapes. The accelerated sampler splits the K = 20 samples into two halves of 10 and hard-codes those sizes train_led_trajectory_augment_input.py model_diffusion.py. Changing K or the batch shape needs code changes. tau is a module constant, not a config value train_led_trajectory_augment_input.py.
• CUDA is required. Several tensors are moved to GPU unconditionally train_led_trajectory_augment_input.py train_led_trajectory_augment_input.py.
• Unclear loss. The paper's eq. 8 does not say which trajectories enter the loss, nor whether the norm is over the whole trajectory. The code's choices are listed under Paper vs code.
• Speed-up figures beyond NBA. The abstract gives speed-ups of 19.3, 30.8, 24.3 and 25.1 times on NBA, NFL, SDD and ETH-UCY. The body reports timings only for NBA tab. 6 tab. 7. The timing hardware is not stated.
• Not like-for-like across datasets. The SDD protocol predicts a destination first [supp. §2]. MID is re-run under the paper's protocol for SDD tab. 3 §5.3. Compare SDD and ETH-UCY numbers with care.
• Averaging over agents. The test metric averages over every agent in the batch; the code builds 11 agents per scene train_led_trajectory_augment_input.py train_led_trajectory_augment_input.py. The paper says these are 10 players and the ball §5.1.
• Explicit initializer supervision is too costly. The authors estimate about six days per epoch on NBA for it, so they dropped it §4.4.
• Variance is a learned log. The network output is log sigma squared, not sigma train_led_trajectory_augment_input.py.

Lineage

• Builds on: the diffusion trajectory predictor MID, with the denoiser trained by the usual diffusion noise-estimation schedule §4.4 §2. The code's README acknowledges borrowing code from MID, NPSN and GroupNet README.md.
• Denoising step follows the DDPM update §4.3.
• Compared against: fast samplers progressive distillation and DDIM tab. 7; ten trajectory baselines per dataset, including GroupNet and MID [tabs. 1-4].
• Superseded by: not stated in this wiki.