JLT: Clean-Latent Prediction in Latent Diffusion Transformers

TL;DR

JLT asks whether the JiT story — that clean-data ( $x_0$ ) prediction beats velocity ( $v$ ) prediction in pixel space because it exploits low-dimensional manifold structure — still applies after a learned VAE has already compressed that structure away. The answer, on a 130M DiT over frozen FLUX.2 VAE latents on ImageNet 256², is yes: matched-recipe $x_0$ -prediction beats $v$ -prediction by a large gap, reaching FID-50K 2.50 with CFG (JLT-B/1). The mechanism is a local Gaussian analysis: velocity regression inherits an isotropic target-covariance floor that amplifies the low-variance latent directions a VAE prior wants damped, while clean-latent prediction damps them. The framing: prediction targets are representation-dependent geometric choices, not interchangeable algebraic reparameterizations.

Key claims

On ImageNet 256² with a 130M DiT over frozen FLUX.2 VAE latents, $x_0$ -prediction reaches FID-50K 2.50 with classifier-free guidance, with a “large matched-target gap” over $v$ -prediction under identical representation, backbone, training, and sampling settings [Abstract].
$x_0$ , $\epsilon$ , and $v$ are linearly convertible at any fixed corruption time, so the gap is not algebraic — it comes from how the target’s local geometry interacts with regression error in a learned latent space [Abstract].
A local Gaussian analysis shows velocity regression carries an isotropic target-covariance floor: low-variance latent directions get amplified by the regression process, the opposite of what a compressed latent prior wants [Abstract].
Clean-latent prediction damps the same low-variance directions, matching the geometry of a VAE-compressed latent space — this is offered as the mechanistic reason behind the empirical gap [Abstract].
The principle behind JiT — clean-data prediction wins when input dimensionality exceeds what the manifold actually needs — survives translation into latent space; the “VAE already removed the redundancy” intuition is wrong as a justification for $v$ -prediction [Abstract].

Method

JLT is a 130M latent diffusion Transformer that operates directly on frozen FLUX.2 VAE codes. Two variants of the same backbone are trained under matched settings — representation, architecture, optimizer, schedule, sampler — differing only in the regression target: clean-latent $\hat{x}_0$ vs velocity $\hat{v}$ . The paper grounds the comparison with a local Gaussian analysis of the target covariance: for a fixed corruption time, the conditional distribution of $v$ inherits an isotropic noise-covariance floor that prevents low-variance latent components from being suppressed during regression, whereas the conditional distribution of $x_0$ shrinks those components toward zero. Empirical evaluation is on ImageNet 256² with FID-50K, both unconditional (matched-target gap) and with classifier-free guidance (headline number 2.50). The naming “JLT” mirrors “JiT” — a Just image Transformer over pixels — but in the latent counterpart over FLUX.2 codes.

Results

FID-50K 2.50 on ImageNet 256² with JLT-B/1 (130M, $x_0$ -prediction) under classifier-free guidance [Abstract].
A “large matched-target gap” between clean-latent and velocity prediction at fixed representation, backbone, and training settings [Abstract]. Exact magnitude of the gap, no-CFG numbers, and ablation on intermediate parameterizations are not stated in the abstract.

Why it’s interesting

Two reasons. First, the JiT thesis (Back to Basics: Let Denoising Generative Models Denoise) — that prediction targets are geometric choices that depend on the input space’s dimensionality and structure — was the prescription for pixel-space diffusion, where the manifold is most obviously embedded in a high-dimensional ambient space. The intuition that VAE compression makes the input “small enough” to make the choice algebraic was a reasonable conjecture; JLT falsifies it directly under matched training. Second, this complements Asymmetric Flow Models, which argued that the prediction target is the bottleneck in pixel-space JiT and showed that a rank-asymmetric velocity ( $x_0$ is the rank-0 case, $v$ is the full-rank case) Pareto-dominates both endpoints — JLT now extends that story into the latent regime. The natural next question, untested here, is whether AsymFlow’s intermediate-rank parameterization also wins in latent space, or whether the FLUX.2 VAE’s compression makes pure $x_0$ already optimal. Contrasts with FLUX.2: Analyzing and Enhancing the Latent Space of FLUX — Representation Comparison on the same FLUX.2 VAE, which approached the same “how does latent design interact with the diffusion objective” question from the encoder side.