stzhao/rice-range/tmp/video-velocity-model - Storage Bucket

Autoregressive Video Generation

This repo learns video dynamics in latent space rather than pixel space. Each frame is encoded into a latent space, and an autoregressive model is trained to predict how that latent evolves over time. Decoding the predicted latent trajectory back through the decoder yields video.

Two modeling families share the same data pipeline, transport, and training loop:

1. Flow / velocity matching with a two-branch DiT (DiTwDDTHead) — predicts the instantaneous velocity of the latent trajectory and rolls it out with an Euler ODE.
2. Block-causal autoregressive GPT (NanoGPTVideoVelocity) — a spacetime transformer that predicts the next frame's latent (or its velocity) one frame at a time, with KV-cached rollout.

This document is a technique reference for the data-processing and training machinery. Operational, copy-pasteable data-processing commands (preprocessing, UMAP, interpolation, PPL) live in data_processing/README.md. The autoregressive Gaussian/GMM theory is derived in RAE/transformer_autoregressive_gaussian.md; a Chinese-language code walkthrough is in overview.md.

Environment

cd video-velocity-model

conda create -n vvm python=3.11
conda activate vvm
pip install -r requirements.txt

python -m pip install --upgrade uv
uv sync

Data representation

The stage-2 model never sees pixels at train time — it consumes precomputed latents.

• Backbones. dinov2, mae, siglip2 (ViT encoders, latent 768×16×16) and flux2_ae (FLUX.2 AE, latent 128×16×16). RAE decoder checkpoints and latent normalization stats come from RAE-collections; flux2_ae encodes/decodes with the FLUX.2 AE directly (data_processing/flux2_ae/wrapper.py maps [0,1]↔[-1,1]).
• Latent file (.npz). features of shape (T, C, H, W) and timesteps of shape (T,). Files are stored class-organized (<class>/<video>_patch_tokens.npz) so the dataloader can read labels from the directory name.
• Stage-1 role at train time. The RAE is instantiated with load_encoder=False — the encoder is never run during training; the decoder exists only to turn predicted latents (and ground-truth windows) back into mp4 for inspection.

Producing these latents, plus the analysis tools (UMAP trajectory plots, latent interpolation, PPL smoothness metric) are documented in data_processing/README.md.

Trajectory construction (TrajectoryPlan)

When transport.use_trajectory=true, the transport replaces noise→data flow matching with a trajectory plan over a window of seq_len latent frames. The window is placed on a uniform time grid linspace(0, 1, seq_len); the stored timesteps are ignored (del timesteps in the loop) so supervision depends only on frame order + uniform spacing. A time t is sampled (uniform or truncated logit-normal), the enclosing frame interval is found, and (x_t, u_t) are constructed by one of three modes (path_video.py):

• discrete — x_t snaps to the left frame; the target is still the finite-difference interval velocity u_t = Δx / dt.
• continuous — x_t is the linear interpolation of the two frames; u_t is the piecewise-constant interval velocity.
• smooth_continuous — per-frame velocities are estimated on the whole grid (central difference interior, one-sided 2nd-order at the ends), then a local cubic Hermite segment gives both x_t and u_t. The result is C¹, so the velocity target no longer jumps at frame boundaries — added because sharp piecewise-constant targets made continuous overfit/rollout brittle.

time_axis_mode sets the units of the velocity target while keeping the model time t∈[0,1]:

• window (legacy): u_t ∝ Δx · (N−1), so its magnitude scales with window length.
• frame: u_t = Δx, magnitude decoupled from N. The sampler mirrors this with a matching dt_step so pred · dt_step ≈ Δx in both modes.
• physical: reserved, not implemented.

Training objective

Two trajectory training modes (transport.trajectory_training_mode):

• point — sample one random t per window, build (x_t, u_t) from the plan, and match. With a point head the loss is plain velocity-matching MSE mean_flat((model(x_t,t,y) − u_t)²).
• autoregressive_discrete — teacher-forcing over the whole window at once: feed frames 0..N−2, target frames 1..N−1. With prediction: 'x' the target is the next frame itself; with 'velocity' it is (x_{i+1}−x_i)/dt_target. This is the mode the NanoGPT configs use.

Probability heads (distribution_heads.py) decouple the loss from the target type:

head	output channels	loss
point	C	MSE (μ−x)²
gaussian_diag	2C	per-channel ½[(x−μ)²/σ² + log σ²]
gaussian_spherical	C+1	shared scalar σ²: ‖x−μ‖²/2σ² + ½C·log σ²
gmm	K + 2KC	−logsumexp_k[log π_k + log 𝒩_k]

Losses are normalized to per-element scale so heads are directly comparable: at σ=1 every Gaussian/GMM head equals MSE/2, and point reproduces the legacy deterministic MSE bit-for-bit. log_var_clamp=(-10,10) bounds log σ² before exp. The transport injects head_type / num_mixtures / log_var_clamp into the model so the output projection is sized consistently. The theory (chain-rule factorization, why MSE is the fixed-variance single-Gaussian special case, temperature sampling) is in RAE/transformer_autoregressive_gaussian.md.

Model architectures

DiTwDDTHead (flow / velocity matching)

A two-branch DiT operating on a single latent map [B, C, 16, 16]:

• s-branch (encoder blocks) patchifies the latent, adds 2D positional encoding, runs encoder blocks, and fuses time/label/motion embeddings into a conditioning feature s.
• x-branch (decoder blocks) patchifies the latent and is modulated by s through AdaLN, then final_layer + unpatchify produce the velocity field.

Uses AdaLN-style modulation, RMSNorm, RoPE, SwiGLU FFN, and class dropout (CFG-ready).

NanoGPTVideoVelocity (autoregressive)

A block-causal GPT over spacetime patches. Each frame flattens to H·W tokens; a window of T frames is T·H·W tokens.

• Block-causal attention — intra-frame full attention + inter-frame strictly causal (k_time ≤ q_time), via F.scaled_dot_product_attention. A KVCache enables AR rollout: at step i only Z_i's H·W tokens are fed and concatenated with cached Z_{<i}.
• 3D RoPE with rope_dims=(24,20,20) for (time, height, width) bands; the rope_t_scale knob rescales the time coordinate (see Position Interpolation).
• Configurable internals — norm_type (layer/rms), qk_norm (per-head Q/K norm applied pre-RoPE), use_sandwich_norm, mlp_type (gelu 4× / swiglu 8/3×), and independent bias / out_proj_bias toggles.
• Weight init — Linear std=0.02; residual c_proj scaled by 0.02/√(2·n_layer); probability-head log_var rows zero-init + log_var_init_bias.
• Probability head — same HeadSpec machinery sizes the output projection for point/gaussian_*/gmm.
• Diagnostics — per-layer Q/K/output L2 and attention max/min/entropy, drained by layer_stats(); attention-distribution stats are recomputed (no_grad) only on the 3 key_layers to bound cost.

Representative NanoGPT config (NanoGPT-medium-FLUX2-AE_ar_train_vpred_gmm.yaml): n_layer=24, n_head=16, n_embd=1024, rms norm, qk_norm, swiglu, transport velocity + autoregressive_discrete + time_axis_mode=frame, head=gmm (K=4).

Conditioning

The conditioning stream is t + y + motion (terms summed, then SiLU → AdaLN):

• Time t — Gaussian-Fourier embedding (always on).
• Label y — LabelEmbedder. --disable-label-condition drops y from the forward pass entirely (not the same as class_dropout_prob=1, which keeps the branch but feeds a null label). The y_embedder module is kept for checkpoint compatibility but excluded from optimization, and DDP switches to find_unused_parameters=True so long runs don't error on the unused branch.
• Motion — --use-motion-condition / --motion-dim. A synthetic per-video motion code: all videos share a zero start vector; each video's end vector is a deterministic hash of its sample id; motion(t) is their linear interpolation, recomputed every rollout step. A learned linear layer maps it into the conditioning space. Trajectory-mode only.

CFG and autoguidance forward paths (forward_with_cfg, forward_with_autoguidance) are wired for guidance.scale > 1.

Optimization & training loop

train_video.py provides DDP, EMA, gradient accumulation, gradient clipping, LR scheduling, checkpointing, and wandb logging.

• Optimizer — fused AdamW. optimizer.weight_decay_split: gpt applies the GPT-style 2-group split (optim_utils.py): decay group = params with dim≥2 that are not embeddings; no-decay group = biases, norm weights, 1-D scalars, and embeddings.
• Scheduler — linear or cosine with warmup (warmup_from_zero ramps from 0; otherwise constant-lr warmup) decaying to final_lr.
• EMA — on by default (ema_decay); --disable-ema saves ~1× model memory and the per-step update, using the live model for sampling/eval/checkpoints. Old checkpoints without an ema key fall back to live weights.
• Time-distribution shift — time_dist_shift = √(shift_dim/shift_base); with the UCF101 default 4096/4096 it is 1.0 (no reweighting). Time is drawn uniform or truncated-logit-normal.
• Precision — fp32 / fp16 (GradScaler) / bf16; optional torch.compile.
• Single-video overfit — --overfit-video-path + --overfit-start-frame + --overfit-repeat-factor pin one fixed window and virtually repeat it, for end-to-end sanity checks. --save-groundtruth-video decodes the first window to <exp>/generated_videos/groundtruth/groundtruth_decoded.mp4.

Logged diagnostics (every log_interval): train/loss, lr, ut_norm, pred_norm, param_norm, grad_norm + grad_clip_scale; per-key-layer qk/q, qk/k, blk/out L2 and (with --collect-attn-stats) attn/max|min|entropy; and head stats log_var/*, sigma/*, loss/mahalanobis, loss/log_det, mixture/entropy, mixture/top1 for Gaussian/GMM heads.

Sampling (sample_ode_video)

Video generation is a naive Euler rollout from a real first-frame latent (not from noise):

velocity head:  x_next = x_current + v(x_current, t, y, motion) · dt_step · step_scale
x-pred head:    x_next = pred(x_current, t, y, motion)

with dt = 1/(num_frames−1), dt_step matching the training time_axis_mode. Per step it recomputes motion(t) and draws the next latent through the head's sample() with sample_temperature (τ=0 → deterministic μ) and, for GMM, mixture_temperature (0 → argmax component). autoregressive_history=True + use_kv_cache=True feed the growing prefix through the NanoGPT KV cache. Generated latents are stacked [B,T,C,H,W] and decoded frame-by-frame.

Evaluation

• FVD — data_processing/sample_for_fvd.py loads a NanoGPT video-flow checkpoint, shards the test split across DDP ranks, rolls out videos from first-frame latents, and writes mp4s + a manifest; data_processing/cal_fvd.py extracts TorchScript I3D (Kinetics-400) features and computes the Fréchet distance with streaming mean/covariance (float64, sqrtm).
• Extended context / Position Interpolation — --pi-train-frames sets rope_t_scale = pi_train_frames / num_frames, compressing the time-axis RoPE so a model trained on short windows can roll out longer ones.
• PPL smoothness — data_processing/cal_metric.py measures StyleGAN-style perceptual path length (LPIPS) of latent interpolation paths across backbones; see data_processing/README.md.

Running training

Canonical single-GPU invocation:

export PYTHONPATH="$PWD/RAE/src:$PYTHONPATH"
torchrun --standalone --nnodes=1 --nproc_per_node=1 \
  RAE/src/train_video.py \
  --config RAE/configs/stage2/training/UCF101/<config>.yaml \
  --data-path video_features/<feature_dir> \
  --results-dir ckpts/ucf101 \
  --precision fp32 --wandb

Convenience launchers (set EXPERIMENT_NAME, WANDB_MODE, ENTITY, PROJECT as needed; note the launchers hardcode an absolute repo root you may need to adjust):

• RAE/scripts/run_video_flow_train.sh — DINOv2 velocity DiT (discrete).
• RAE/scripts/run_video_flow_train_smooth.sh <dinov2|mae|siglip2|flux2_ae> — smooth_continuous trajectory supervision with the DiTDH-S_*_smooth_exp.yaml configs. Honors DISABLE_LABEL_CONDITION, USE_MOTION_CONDITION, MOTION_DIM, ENABLE_WANDB, COMPILE, CKPT.

Config families under RAE/configs/stage2/training/UCF101/:

• DiTDH-S_*_smooth_exp.yaml — DiT velocity matching, smooth trajectories.
• DiTDH-*_overfit.yaml / DiT-*_overfit.yaml — single-window overfit sanity checks.
• NanoGPT-*-FLUX2-AE_ar_train*.yaml — autoregressive GPT; the vpred/xpred × gauss_diag/gauss_sphere/gmm matrix selects prediction target × probability head.

Notes

• Training requires a GPU. The stage-1 decoder/stat paths come from the config's stage_1 section; keep RAE-collections and backbone weights valid.
• Reconstructions and sampled videos are for qualitative inspection, not benchmark scoring (use FVD for that).
• time_axis_mode='frame' is the convention for the NanoGPT autoregressive configs; the sampler's dt_step is matched to it so rollout displacement is consistent with training regardless of window length.