Back to Basics: Let Denoising Generative Models Denoise

TL;DR

JiT (“Just image Transformers”) argues that the long-running pixel-vs-latent and noise-vs-image debates in diffusion both bottom out in one prescription: when input dimensionality exceeds network hidden width, the network must predict clean data ( $x_0$ ) rather than noise ( $\epsilon$ ) or velocity ( $v$ ). Under the manifold assumption, clean images lie on a low-dimensional manifold while noised quantities span the full high-dimensional space; a width-768 transformer can compress to the former but not represent the latter. A plain Vision Transformer with patch size 16 or 32 on raw pixels, no VAE, no tokenizer, no perceptual or representation-alignment loss, reaches competitive ImageNet 256/512 FID purely by switching the prediction target — and a low-rank bottleneck at the input projection helps further. Conceptually nothing more than DiT-on-pixels with $x_0$ -prediction.

Key claims

Among the nine combinations of (loss space, prediction space) ∈ { $x_0, \epsilon, v$ }², only $x_0$ -prediction works on ImageNet-256 with JiT-B/16 (768-d patches into a 768-d transformer); $\epsilon$ - and $v$ -prediction fail catastrophically across all three loss-space reweightings (FID ≈ 8.6 vs 372 / 96) [Table 2(a), §4.2].
The failure mode is dimensionality-gated, not loss-weighting-gated: at JiT-B/4 on ImageNet-64 (48-d patches into 768-d transformer) all nine combinations work with marginal FID gaps [Table 2(b)], showing the failure appears precisely when patch dimension approaches or exceeds hidden width.
Noise-schedule shift alone cannot rescue $\epsilon$ / $v$ -prediction at high input dimension; $\epsilon$ -prediction stays > 350 FID across all tested logit-normal shifts $\mu \in \{0, -0.4, -0.8, -1.2\}$ while $x_0$ -prediction merely benefits modestly from higher noise [Table 3, §4.2].
Network width does not have to match or exceed patch dimension under $x_0$ -prediction: JiT/32 at 512² (3072-d patches) and JiT/64 at 1024² (12288-d patches) train successfully with standard B/L/H widths just by scaling noise proportionally [§4.2, “Increasing hidden units is not necessary”].
A low-rank linear bottleneck at the input projection — two linear layers mapping patch → $d_b$ → hidden, with $d_b$ as small as 16 to 32 of the 768 patch dimensions — does not collapse training and in fact improves FID by up to 1.3 points across $d_b \in [32, 512]$ on JiT-B/16 [Fig. 4, §4.2 “Bottleneck can be beneficial”].
A toy experiment with $d$ -dim data embedded in $D$ -dim space ( $D \gg d$ ) via a fixed orthogonal projection reproduces the same pattern: a 256-d MLP succeeds at $x_0$ -prediction even when $D$ is large but the network is under-complete, while $\epsilon$ - and $v$ -prediction fail as $D$ grows [§3.3, Fig. 2].
Adopting general-purpose Transformer improvements (SwiGLU, RMSNorm, RoPE, qk-norm, multi-token in-context class conditioning) on top of the baseline JiT compounds to FID 5.49 → 3.39 for JiT-B/16 → JiT-L/16 at 256² without CFG [Table 4, §4.4].

Method

The architecture is a textbook ViT/DiT applied to raw pixel patches: divide an image into non-overlapping patches of size $p \times p$ , linearly embed each to the transformer’s hidden width, add positional embeddings, run through a stack of standard transformer blocks with adaLN-Zero conditioning on timestep and class label, and linearly project each output token back to a $3p^2$ pixel patch. No VAE, no pre-trained tokenizer, no auxiliary loss, no perceptual loss, no representation alignment. Conceptually identical to DiT operating directly on pixels rather than VAE latents.

The single critical design decision is what the network output represents. The paper formalizes that the three reparameterization choices ( $x_0$ -pred, $\epsilon$ -pred, $v$ -pred) and the three loss-space choices yield nine inequivalent training objectives, none of which are mathematically equivalent under network capacity constraints. The recommended recipe is $x_0$ -prediction with $x_0$ -loss under a linear flow-matching schedule ( $\alpha_t = 1-t$ , $\sigma_t = t$ ), trained with logit-normal $t$ sampling and a noise shift toward higher noise. Sampling uses a 50-step Heun solver. The optional bottleneck is a rank- $d_b$ factorization of the input patch-embedding matrix — purely linear, parameter-saving, and improves quality across a wide $d_b$ range. The motivation invoked is the IRMAE observation (2010.00679) that low-rank linear bottlenecks induce smoother manifold structure in autoencoder latents.

Results

On ImageNet-256, JiT-B/16 (131M params, 25 GFlops) reaches FID 4.37 with the “Just Advanced Transformers” improvements [Table 5]. The paper reports competitive numbers up through JiT-L/16 (FID 3.39 / 2.79 with CFG, no CFG numbers in parens). The headline negative result is more striking than the headline positive: at JiT-B/16 with $\epsilon$ -prediction the model produces FID 372 (vs 8.6 for $x_0$ ), and at JiT/32 on 512² the FID gap between $x_0$ -prediction and the other two parameterizations widens further. The bottleneck ablation drops JiT-B/16 FID from ~6.7 (no bottleneck) to ~5.4 at $d_b = 32$ [Fig. 4], for free — bottleneck layers add only $d_b(3p^2 + h)$ parameters where $h$ is hidden width. The recipe scales analogously at 1024² with JiT/64 (12288-d patches) without architectural changes beyond noise scaling.

Why it’s interesting

This paper is the foundational JiT backbone that several downstream wiki entries build on top of without it being filed yet. PixelGen: Improving Pixel Diffusion with Perceptual Supervision explicitly uses “JiT (Joint image Transformer — a DiT operating on raw pixel patches with no VAE)” as its baseline and adds LPIPS + Perceptual-DINO losses on top; Continuous Adversarial Flow Models post-trains JiT-H/16 with a learned discriminator; Representation Fréchet Loss for Visual Generation post-trains the same JiT-H/16 with Fréchet-distance loss; Asymmetric Flow Models replaces JiT’s $x_0$ -prediction with a rank-asymmetric velocity prediction $u_A = P\epsilon - x_0$ and reports the rank-8 case dominates both pure $x_0$ and pure $u$ . The natural read is that JiT is the canonical pixel-space DiT baseline of 2026 and the field is now layering criterion- and parameterization-side improvements on top of it. The key claim against the parameterization fork in Asymmetric Flow Models — that “small but non-zero noise-prediction rank Pareto-dominates pure $x_0$ ” — is a direct counterfactual to JiT’s ” $x_0$ is the right answer” headline; AsymFlow concedes JiT’s failure-mode story but argues the right fix is a rank-asymmetric velocity, not pure $x_0$ .

For Luma, the second-order claim — that VAEs matter for LDMs not just as sequence-length reducers but as manifold smoothers that enforce low-rank structure on the diffusion target — is the load-bearing one. JiT’s input-bottleneck ablation (FID improves with $d_b = 32$ on 768-d patches) is empirical evidence that the low-rank constraint itself is the useful piece of what a VAE provides, separable from spatial compression. If true, the implication for video diffusion is that the VAE’s role can in principle be replaced by an architectural bottleneck without giving up the latent prior — relevant for Pixel-space diffusion as a whole and especially for video, where temporal VAE artifacts are the dominant failure mode.

One concern Terrance flagged that the paper does not address: does the input-bottleneck trick also rescue $\epsilon$ - and $v$ -prediction? If yes, the ” $x_0$ is special” framing weakens to “you need a bottleneck somewhere — either at the input or implicitly in the prediction target.” Without that ablation, the paper’s two contributions (prediction target + bottleneck) are entangled.