LogoAIRosetta

JiT: Back to Basics in Image Generation — Analysis of Kaiming He's New Paper

arXiv

SummaryKaiming He and Tianhong Li propose JiT (Just image Transformers), replacing traditional noise prediction with x-prediction (direct clean-image prediction). Using a pure-pixel Transformer with no tokenizer or pretraining, JiT attains near-SOTA FID on ImageNet with fewer parameters and compute, advocating high-dimensional robustness and minimalist design.

As generative AI models grow ever more complex, a paper titled "Back to Basics: Let Denoising Generative Models Denoise" offers a refreshing reexamination of the core assumptions behind diffusion models. Posted to arXiv on 2025‑11‑16 (arXiv:2511.13720), the work is coauthored by Kaiming He, Associate Professor in MIT's Department of Electrical Engineering and Computer Science, and PhD student Tianhong Li. Kaiming He, known for ResNet and other influential contributions to computer vision, follows a recurring philosophy of simplifying design: from residual connections in ResNet, to masked autoencoding in MAE, and now to JiT (Just image Transformers), he again leverages minimalism to shift paradigms.

JiT is not a brand‑new model family so much as a disruptive rethink of the diffusion‑model paradigm. It challenges the prevailing "noise prediction" strategy and advocates direct clean‑image prediction (x‑prediction), building a pixel‑level, high‑resolution generation framework that the authors describe as "pure". The paper's motto ⁠—⁠ "let denoising models denoise" ⁠—⁠ stresses that indirectly predicting noise often fails catastrophically in high‑dimensional pixel spaces, whereas directly predicting target data enables networks that seem capacity‑limited to operate effectively. On ImageNet, using large patches (patch sizes 16 and 32) at 256×256 and 512×512 resolutions, JiT achieves FID scores comparable to SOTA models such as DiT and PixArt‑α (for example, JiT‑B/16 attains FID 7.48 at 256 resolution) without tokenizer, pretraining, or auxiliary loss terms. This result supports the "simplicity is powerful" philosophy and suggests a new direction for more efficient, broadly applicable generative AI. Below we analyze the paper across background, core innovations, technical details, experiments, and impact.

Background: the "noise trap" in diffusion models and the need to return to basics

Since DDPM (Denoising Diffusion Probabilistic Models) by Ho et al. in 2020, diffusion models have dominated image generation. Their core is a forward noising process that gradually adds Gaussian noise to clean data until it becomes pure noise, and a learned reverse denoising process that restores images. Popular systems like Stable Diffusion and Imagen rely on U‑Net or Transformer backbones and predict noise ε or the noisy image x_t during the reverse process to estimate the clean image x_0.

He and colleagues point out that this "noise prediction" paradigm departs from classical denoising principles. Traditional denoising methods (e.g., Wiener filtering) directly output the clean signal, while modern diffusion models effectively learn to predict noisy quantities. This approach can work in lower‑dimensional domains (such as text) but tends to fail in high‑dimensional pixel spaces. High‑resolution images (e.g., 512×512) contain millions of pixels; predicting noise requires capturing high‑frequency detail and can lead to catastrophic failures: at high noise levels the model collapses, producing blurry or distorted outputs. The paper cites prior efforts (e.g., Chen 2023; Hoogeboom 2023, 2024) that try to mitigate these issues by altering noise schedules or other tricks but do not address the fundamental mismatch of predicting noise rather than target data, which violates the natural‑data‑manifold assumption of generative modeling.

JiT is proposed as a response to this "noise trap." It assumes generation should operate directly on the data manifold ⁠—⁠ the model should output the clean image x_0 rather than indirectly estimate it via noise. This simplifies training and improves stability in high‑dimensional settings. As the paper states: "Predicting clean data allows seemingly capacity‑limited networks to operate efficiently in extremely high‑dimensional spaces." This perspective echoes He's prior work: ResNet bypassed optimization issues with residuals, MAE used masked reconstruction, and JiT uses a pure Transformer to return to pixel‑level generation. Social posts on X (formerly Twitter) echoed the sentiment; a popular post by an account named @jiqizhixin summarized: "JiT is another strong contribution from Kaiming He: large‑patch Transformer operating directly on pixels, no tokenizer required." The post received broad engagement, reflecting community resonance.

Core innovations: the paradigm shift from ε‑prediction to x‑prediction

JiT's innovations are twofold: redefining the prediction target and adopting an extremely simple architecture.

First, the prediction paradigm: conventional ε‑prediction models learn ε_θ(x_t,t) ≈ ε, where x_t = √(α_t) x_0 + √(1‑α_t) ε, and the model's output noise is used to indirectly recover x_0 = (x_t ‑ √(1‑α_t) ε_θ)/√α_t. This indirection magnifies errors in high‑dimensional settings, especially at large noise levels, because the model must effectively invert the noise distribution.

JiT adopts x‑prediction (clean prediction): it directly learns x_θ(x_t,t) ≈ x_0, using the simple relation x_0 = x_θ(x_t,t). This aligns with classical denoising intuition: the model focuses on the target manifold rather than on the noise. The paper argues mathematically that x‑prediction is more robust in high dimensions because it avoids Gaussian noise "explosion" ⁠—⁠ noise dimensionality scales with the number of pixels while the clean data lies on a lower‑rank natural‑image manifold. The x‑prediction loss is L = E[||x_0 ‑ x_θ(x_t,t)||^2], with t ∼ Uniform(0,T) and x_t sampled from the Markov chain. This also simplifies sampling: reverse steps can be written as x_{t‑1} = x_θ(x_t,t) plus a noise adjustment, avoiding more complex inversion formulas.

Second, architectural design: JiT is a "Just image Transformer," discarding tokenizers (e.g., VAE encoders) and auxiliary losses (e.g., perceptual losses). Inputs are raw RGB pixel patches (patch sizes 16×16 or 32×32) directly embedded into a Transformer. The embedding uses two linear layers with a bottleneck (bottleneck dimension d' < 768, e.g., d'=256) to reduce parameters while preserving information (see paper Fig. 4). The Transformer body follows a ViT variant: SwiGLU activations, RMSNorm normalization, RoPE positional encoding, and QK‑Norm for query‑key normalization. These components, though originating in language models, adapt well to images: the paper's Table 4 shows RoPE improving FID by ~0.5 and QK‑Norm enhancing stability.

Conditioning uses an in‑context class‑conditioning scheme: rather than a single class token, 32 class tokens are appended (inspired by Li 2024's autoregressive approaches), improving class control. During training, logit‑normal sampling is used at distribution tails for numerical stability. Model sizes are compact: JiT‑B/16 (Base, Patch 16) has ~300M parameters, JiT‑H/16 (High) ~600M, both significantly smaller than DiT's 1B+ models.

Technical details: training, sampling, and extensions

Training is straightforward and efficient. The dataset is ImageNet‑1K (1.28M images) at 256 or 512 resolution. The noise schedule is linear β_t from 1e‑4 to 0.05 with T=1000 steps. The loss is plain MSE, without a VLB lower bound or extra regularizers. Optimizer is AdamW with learning rate 5e‑4, batch size 128, trained for 600 epochs (≈500K iterations). The paper emphasizes "no pretraining": models are trained from random initialization, and the use of large patches accelerates convergence via locality.

Sampling uses a DDIM‑like variant: starting from pure noise x_T, iterate x_{t‑1} = √α_{t‑1} x_θ(x_t,t) + √(1‑α_{t‑1}) η, where η controls residual noise. Classifier‑Free Guidance (CFG) scales of 2.2⁠–⁠2.9 ensure strong conditional fidelity. Cross‑resolution generation is a notable capability: a model trained at 256 can be upscaled to 512 with only a slight FID increase (≈1 point, Table 3) and no retraining.

Extensions explore the bottleneck embedding: original patch vectors of 768 dims (16×16×3) compressed to d'=32 raised FID by only ~0.2, demonstrating x‑prediction's low‑rank friendliness. A classifier head attached at layer 4 or 8 (global average pooling + linear layer) with auxiliary CE loss improved generation quality by ≈10% (Table 4). The paper also examines noise scaling: a negative‑mean noise setting (P_mean=‑0.8, P_std=0.8) improved high‑frequency detail.

A GitHub implementation (LTH14/JiT) provides PyTorch code supporting multi‑GPU training. Example command: torchrun main_jit.py ‑‑model JiT‑B/16 ‑‑img_size 256 ‑‑epochs 600, facilitating reproducibility.

Experimental validation: benchmarks and ablations

The paper evaluates comprehensively on ImageNet using FID, IS (Inception Score), and precision/recall metrics. JiT‑B/16 (256) achieves FID 7.48 and IS 350.2, outperforming PixArt‑α (FID 8.1) and EDM (FID 9.2), and approaching DiT‑XL/2 (FID 6.5). At 512 resolution, JiT‑B/32 attains FID 12.3, beating noise‑prediction baselines by ~20%, which supports the claim of improved high‑dimensional robustness.

Ablation studies are systematic: Table 1 compares x‑prediction versus ε‑prediction, showing a ~15% lower FID for x‑prediction. Table 2 analyzes noise schedule: ε‑prediction collapses at high noise (FID >50) while x‑prediction remains stable. Table 4 details "Just Advanced" components: RoPE adds +0.5 FID improvement, SwiGLU +0.3. Cross‑resolution results (Table 3) show 256→512 FID only rising to ~8.2. Qualitative figures (Fig. 3) display JiT‑generated samples with cinematic detail in lighting and texture.

Compared with SOTA, JiT avoids tokenizer‑induced VAE artifacts, requires no pretraining (saving ~80% compute), and uses about 50% fewer parameters while matching FID. Industry commentary (e.g., Bohrium AI X posts) observed: "JiT shows chasing noise misses the real signal⁠—⁠returning to basics matters."

Impact and outlook: minimalist design reshaping generative AI

JiT's impact is broad. Theoretically, it validates that direct prediction can be superior in high‑dimensional spaces and challenges the centrality of noise prediction in diffusion models, potentially inspiring new paradigms such as x‑prediction for video or 3D generation. Practically, its minimalist architecture lowers barriers: the open‑source code aids reproduction and the model is more suitable for edge deployment (e.g., on‑device image generation). JiT is also computationally efficient: training a 256 model reportedly needs 8×A100 GPUs for ~1 week at 600 epochs.

Challenges remain: x‑prediction requires stronger generalization in low‑noise regimes and future work may integrate multimodal signals. In a broader sense, He's approach may herald a "post‑diffusion era" where simpler mathematical principles replace engineering complexity. As one X post by @jiqizhixin put it: "JiT demonstrates the power of large‑patch Transformers⁠—⁠no frills, direct pixel generation." In the generative AI arms race, JiT is a reminder that innovation can come from returning to fundamentals rather than merely scaling complexity.

In summary, JiT marks a new chapter in He's simplification‑driven research. By combining x‑prediction with a pure Transformer, it shows that less can be more: direct denoising outperforms indirect noise chasing in high‑dimensional image generation. This insight improves generation quality and promotes AI democratization. JiT may become a benchmark that helps shift the field from a "noise era" toward "clean generation."

CopyrightStatement

Comment

Please login before commenting

NoComments

RelatedNews

LoadingMore

JiT: Back to Basics in Image Generation — Analysis of Kaiming He's New Paper - Airosetta News