Latent Forcing: Reordering the Diffusion Trajectory for Pixel-Space Image Generation

TL;DR

Latent Forcing trains a single diffusion transformer to jointly denoise raw pixels and deterministic self-supervised latents (DINOv2) of the same image, but with separate noise schedules so the latent denoises first and then conditions the pixel denoising. The generated latent is discarded at the end, so the pipeline is lossless at encoding and end-to-end at inference — no VAE, no decoder. On ImageNet-256, the cascaded-schedule Single-Schedule LF-DiT reaches SOTA pixel-diffusion FID at matched compute, and the authors argue that the order of information revealed during denoising — not distillation or compression — is the primary driver of the gains attributed to REPA-style latent supervision.

Key claims

Multi-time, multi-tokenizer diffusion is a valid generative model for joint distributions of pixels and deterministic latents; when the latent is a deterministic function of pixels, optimization is invariant to the choice of schedule ordering, so ordering is purely a training/inference convenience — but empirically the order chosen dominates convergence [§3.2].
Across schedules and across multiple latent encoders (DINOv2-B+Reg, Data2Vec2, downsampled pixels), denoising latents before pixels strictly outperforms denoising pixels first: FID-10K (no guidance) goes from 24.39 (pixel-earlier) to 18.65 (latent-earlier) for DINOv2 and 27.69 → 23.61 for Data2Vec2 [Table 1].
Scaling a tokenizer’s variance is informationally equivalent to time-shifting its noise schedule, so “which tokenizer to combine” and “what schedule to use per tokenizer” are the same design choice [§3.3, Eq. 4].
The cascaded schedule (denoise latent fully, then pixels) gives the best Single-Schedule FID — 12.42 unguided / 6.60 guided at ImageNet-256, beating Linear-Offset (20.98 / 10.73) and Variance-Shift (13.48 / 8.16) at matched training [Table 8].
Cascaded generation suffers cascaded error; injecting up to 25% noise into the latent during pixel-step training (but 0% at inference) improves FID from 14.11 → 10.93 at 200 epochs and prevents late-training degradation [Table 7].
AutoGuidance outperforms CFG for the pixel timesteps because DINOv2 features already probe to the class label, so class conditioning becomes redundant during pixel generation; using CFG-Interval on latents + AutoGuidance on pixels is reported as the best mix [§5.4].
The framing isolates order (not REPA distillation alone) as the primary factor behind pretrained-representation gains in diffusion, reframing the reconstruction–generation tradeoff: only early timesteps need to be compressed, so the global compression rate of a tokenizer is the wrong axis to optimize [§1, §5.5].
A 4-layer output-expert split at the end of the DiT (one expert for latent output, one for pixel) leaves parameter and FLOP counts unchanged but improves FID 13.97 → 12.93 [Table 5, §4.2].

Method

The model is a standard DiT/JiT pixel transformer with two minimal changes. First, the time embedding is doubled: a second 2-layer MLP encodes the latent-modality time variable and is added into the adaLN-Zero conditioning vector (a ~0.5% parameter overhead). Second, pixel and latent inputs are patchified into the same number of tokens (DINOv2 inputs are pre-resized so its 14×14 patches align with the pixel 16×16 grid), each linearly projected to the hidden size, and the two embeddings are added per patch — so the token count and FLOPs match a single-modality DiT. The prediction target is $\hat{x}_0$ (x-prediction) with $1/(t^2+\epsilon)$ loss weighting, following JiT, and the pixel + latent losses are reweighted to equal magnitude.

Training samples a single global time $t \in [0,1]$ uniformly (the “Multi-Schedule” model uses independent per-modality times; the “Single-Schedule” model maps $t$ through fixed schedules $f_i(t)$ ). Variance time-shifts are applied so that uniform sampling on rescaled latents is informationally equivalent to the shifted-logit-normal schedules used in standard DiTs. The cascaded schedule denoises the latent from $t=1$ to $0$ over the first 25 of 50 Heun steps with the pixel held at full noise, then denoises pixels from $t=1$ to $0$ over the next 25 steps with the latent fixed. At inference the generated latent is discarded after pixel denoising completes, leaving a raw RGB image.

The Multi-Schedule variant — which trains on independently sampled $(t_{\text{latent}}, t_{\text{pixel}})$ pairs — is used only to survey the schedule space; the Single-Schedule variant trains on one specific trajectory and is the apples-to-apples comparator for baselines. Two further additions ship in the final model: a 4-layer Mixture-of-Output-Experts at the end of the DiT (one expert routed to latent output tokens, one to pixel tokens), and 25% maximum latent noise sampled during pixel-denoising training steps to combat cascaded error.

Results

On ImageNet-256, LF-DiT-L (cascaded schedule, 80 epochs, ViT/L backbone) reaches FID-10K unguided 12.42 and FID guided 6.60 [Table 8]. At 200 epochs with the 25% noise augmentation, FID-10K drops to 10.93 [Table 7].
System-level: the paper reports a new state-of-the-art for diffusion-transformer-based pixel generation at the authors’ compute scale on ImageNet conditional and unconditional generation [§1, Abstract; Table 11 in the paper PDF].
The schedule-space sweep [Table 1, Fig. 3] shows the ordering effect is robust to (a) latent choice — holds for DINOv2-B+Reg, Data2Vec2 (an SSL model trained only 150 epochs on ImageNet), and even spatially-downsampled pixels acting as the “latent” — and (b) curve shape — both variance-shift and linear-offset trajectories show latent-first dominates, though cascaded wins overall.
AutoGuidance applied to pixel steps and CFG-Interval applied to latent steps is the best-performing guidance mix [§5.4].

Why it’s interesting

Latent Forcing fits cleanly into Luma’s tracking of pixel-space diffusion (Pixel-space diffusion) as a third paradigm alongside PixelGen (perceptual losses) and L2P (LDM-to-pixel transfer): it keeps the pixel-space output but recovers latent-diffusion’s convergence by making the latent a generated scratchpad inside the same model rather than an encoder input or a frozen feature target. The reframing of REPA — “what REPA buys you is mostly ordering, not distillation” — is a sharper claim than the prior literature has made, and if it holds, it shifts the design search from “find a better tokenizer” to “find a better schedule.” The cascaded inference budget (50% of steps spent denoising a latent that’s then thrown away) is also a concrete cost the team should weigh against Diffusion training efficiency-style recipes that don’t allocate inference compute to discarded modalities. Worth comparing head-to-head with PixelGen: Improving Pixel Diffusion with Perceptual Supervision at parity epochs and parity inference budget.