Just Image Transformers (JiT) Overview

• Just Image Transformers (JiT) are pixel-space diffusion models that denoise raw images directly, avoiding the reconstruction bottleneck of latent compression.
• JiT employs an end-to-end training regime with x-prediction and velocity-matching objectives, demonstrating scalable improvements in high-resolution synthesis and convergence speed.
• PixelREPA introduces a Masked Transformer Adapter that refines semantic alignment by preventing feature hacking, ensuring effective denoising while maintaining full image details.

Just Image Transformers (JiT) are pixel-space diffusion transformer models that perform denoising directly in raw image space rather than in a compressed latent space. Proposed by Li and He, JiT is characterized by end-to-end training on raw images without any pretrained tokenizer or autoencoder and without auxiliary perceptual or adversarial losses. Within this formulation, the removal of the latent tokenizer is presented as a way to avoid the reconstruction bottleneck of latent diffusion models, whose compression stage can suppress fine textures and small structures, bound achievable output quality by tokenizer fidelity, add an extra component to train and maintain, and incur decoding overhead at sampling time. A 2026 analysis further situates JiT as a test case for whether training accelerants developed for latent-space diffusion transfer to pixel space, concluding that naive Representation Alignment (REPA) does not do so straightforwardly and motivating the PixelREPA variant instead (Shin et al., 15 Mar 2026).

1. Pixel-space diffusion and the rationale for JiT

JiT belongs to the class of pixel-space diffusion transformers: the denoiser operates in raw image space rather than over a compressed latent produced by a pretrained tokenizer. The immediate architectural consequence is that the model is, in principle, a fully self-contained generation pipeline. The reported motivation is not merely aesthetic simplification; it is tied to a specific critique of latent diffusion. In latent diffusion models, the tokenizer reduces computational cost, but the compression stage introduces a reconstruction bottleneck that can suppress fine textures and small structures, cap output quality by tokenizer fidelity, and add both maintenance cost and sampling-time decoding overhead (Shin et al., 15 Mar 2026).

The significance of JiT is therefore tied to high-resolution synthesis. Because denoising remains in image space, the model is presented as able, in principle, to preserve arbitrary high-frequency details. This makes pixel-space diffusion especially compelling at high resolution, where those details matter most. A plausible implication is that JiT should be evaluated not only as a generative architecture but also as a position in the broader design space between latent efficiency and pixel fidelity.

2. Objective structure and training regime

The paper describing REPA failure in JiT states that JiT uses $x$-prediction, or clean image prediction, with a velocity-matching objective. In the same study, auxiliary alignment is introduced only as an added regularizer, so the full modified objective is written as

$L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$

with $\lambda > 0$, and the reported experiments use $\lambda = 0.1$ (Shin et al., 15 Mar 2026).

The experimental setting is centered on ImageNet $256 \times 256$, with additional analysis on ImageNet $32 \times 32$ to show a resolution dependence in alignment behavior. The implementation is reported to follow JiT exactly except for the Masked Transformer Adapter used by PixelREPA. Training uses 8 NVIDIA H200 GPUs, Adam, learning rate $2 \times 10^{-4}$, batch size 1024, $(\beta_1,\beta_2)=(0.9,0.95)$, and EMA decays 0.9996, 0.9998, and 0.9999. Evaluation uses FID and Inception Score computed with 50K samples; sampling uses a 50-step Heun ODE solver and classifier-free guidance interval $[0.1,1.0]$. The evaluated JiT scales are B/16, L/16, H/16, with comparison to JiT-G/16.

These details matter because the article’s central claims about JiT are not purely architectural. They concern the interaction between JiT’s pixel-space training dynamics and external semantic supervision under a fixed high-resolution ImageNet protocol.

3. Representation Alignment in JiT and the failure of naive transfer from latent space

Representation Alignment (REPA) is presented as a training acceleration method for diffusion transformers that aligns intermediate diffusion features with features from a frozen pretrained semantic encoder, specifically DINOv2. In the reported formulation, a projection head $h_\phi(\cdot)$ maps diffusion hidden states into the semantic feature space, and those projected features are matched to encoder outputs $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$0 using cosine similarity. In latent-space diffusion, this is described as effective because both the diffusion state and the semantic target are already compressed, so alignment is performed between two information-bottlenecked spaces of roughly matched complexity (Shin et al., 15 Mar 2026).

The same paper argues that this logic breaks in JiT. On ImageNet $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$1, JiT + REPA may improve early in training, but later degrades and ends up worse than vanilla JiT. The proposed diagnosis is an information asymmetry: denoising occurs in a high-dimensional image space with $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$2 degrees of freedom, whereas the semantic target is strongly compressed. Many pixel-distinct images can therefore map to very similar semantic features. Under direct regression to that compressed target, the model is said to engage in feature hacking: it overfits to the narrow semantic space, collapses diverse pixel configurations that share similar embeddings, and thereby harms diversity and final image quality.

The paper’s subset analysis sharpens this argument. On subsets that are tightly clustered in semantic feature space but visually diverse in pixels, JiT + REPA performs worse than vanilla JiT. On subsets with more spread-out semantic features, REPA is less harmful or even beneficial. The failure is therefore not uniform; it is concentrated where the semantic target is most ambiguous relative to the richness of pixel space. This is also why ImageNet $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$3 is informative: the study reports that REPA can work at low resolution but fails at high resolution, suggesting that the pathology depends on the gap between semantic compression and pixel-level variability.

4. PixelREPA: transformed alignment targets and constrained prediction

PixelREPA is introduced to make representation alignment compatible with pixel-space diffusion. Its central move is to transform the alignment target rather than directly force JiT features to match the frozen encoder output. Instead of the REPA-style MLP projection head, PixelREPA uses a Masked Transformer Adapter (MTA) composed of a shallow Transformer adapter and partial token masking. The adapter is specified as a 2-block Transformer with self-attention, and it is trained to transform an intermediate JiT feature map toward the semantic representation space of the frozen encoder (Shin et al., 15 Mar 2026).

Two stated benefits motivate this replacement. First, the adapter learns a structured transformation from JiT features to semantic features instead of forcing JiT itself to collapse into semantic space. Second, self-attention enables contextual aggregation, so each token prediction can use neighboring tokens rather than relying on a point-wise mapping. This is presented as a direct response to the claim that the original REPA-style MLP is too easy to exploit via a near-trivial shortcut.

The adapter alone is reported to be insufficient. PixelREPA therefore adds random partial masking of adapter input tokens during training. Masking serves two roles in the paper’s interpretation: shortcut prevention and information bottlenecking. By breaking direct token-wise correspondence, it forces genuine contextual reasoning; by reducing the effective information available to the alignment branch, it narrows the gap between rich pixel features and compressed semantic targets. Crucially, masking is applied only in the alignment branch, not in the main denoising pathway. JiT therefore retains full-information denoising while semantic alignment is regularized.

The reported PixelREPA loss replaces direct projection with adapter-based prediction under a patch-wise mask $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$4, and the default mask ratio is $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$5. The paper’s interpretation of this value is explicit: too little masking leaves shortcut learning available, while too much masking blocks too much supervision and harms learning.

5. Quantitative characterization across model scales

The principal ImageNet $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$6 results compare vanilla JiT and PixelREPA-enhanced JiT across multiple scales. The reported metrics are FID and Inception Score.

Model	FID	IS
JiT-B/16	3.66	275.1
PixelREPA-B/16	3.17	284.6
JiT-L/16	2.36	298.5
PixelREPA-L/16	2.11	309.5
JiT-H/16	1.86	303.4
PixelREPA-H/16	1.81	317.2
JiT-G/16	1.82	292.6

For JiT-B/16, PixelREPA reduces FID from 3.66 to 3.17 and improves IS from 275.1 to 284.6; the paper states this as a 13.4% improvement in FID and the abstract reports more than $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$7 faster convergence. For JiT-H/16, PixelREPA-H/16 achieves FID $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$8 and IS $L = L_{\text{JiT}} + \lambda L_{\text{PixelREPA}},$9, slightly outperforming vanilla JiT-H/16 and also exceeding JiT-G/16 in both reported metrics (Shin et al., 15 Mar 2026).

The study further states that PixelREPA is scalable and inference-free: it improves JiT across model sizes and affects training only, so it adds no inference overhead. This suggests that the reported gains are attributed to training dynamics rather than a change in the deployed sampler.

6. Ablations, misconceptions, and terminological ambiguity

The ablation results are used to isolate why PixelREPA differs from both vanilla JiT and naive REPA. On B/16 at 200 epochs, the reported FID values are 4.37 for JiT, 5.14 for JiT + REPA, 4.68 for PixelREPA$\lambda > 0$0 (adapter without masking), and 4.00 for full PixelREPA. The conclusion drawn in the paper is threefold: standard REPA is harmful, the Transformer adapter alone helps but is still insufficient, and masking is essential. A second ablation over mask ratio gives FID 4.26 at $\lambda > 0$1, 4.00 at $\lambda > 0$2, 4.38 at $\lambda > 0$3, 4.32 at $\lambda > 0$4, and 4.58 at $\lambda > 0$5, again supporting the interpretation that $\lambda > 0$6 best balances bottlenecking and learnability (Shin et al., 15 Mar 2026).

These experiments address a common misconception: the study does not argue that semantic alignment is intrinsically unsuitable for JiT. Its stated conclusion is narrower and more technical. Standard REPA is not directly suitable for high-resolution pixel-space diffusion transformers because the alignment target must respect the information structure of the denoising space. Another misconception is that masking is merely generic regularization. In the paper’s framing, masking is a compatibility mechanism that narrows the information gap between pixel-space denoising and compressed semantic targets.

A separate source of confusion is terminological. The acronym “JiT” is also used by “Just-in-Time: Training-Free Spatial Acceleration for Diffusion Transformers,” which denotes a training-free spatial acceleration framework for diffusion transformers built around anchor tokens, a spatially approximated generative ODE, and a deterministic micro-flow. That work is unrelated to Just Image Transformers despite the shared abbreviation (Sun et al., 11 Mar 2026). In current usage, disambiguation therefore depends on expansion: “Just Image Transformers” refers to pixel-space diffusion transformers, whereas “Just-in-Time” refers to inference-time spatial acceleration.

7. Broader significance for pixel-space generative modeling

Within the evidence presented, JiT functions as a case study in the non-transferability of methods between latent-space and pixel-space diffusion. The paper’s broader implications are stated explicitly. Alignment must be space-aware; direct regression to a compressed semantic space can cause feature hacking in high-resolution pixel generation; a transformed alignment target is better than a direct projection; and masking should be understood not only as regularization but as a mechanism for making semantic alignment feasible without sacrificing the full denoising pathway (Shin et al., 15 Mar 2026).

For the study of pixel-space diffusion transformers more generally, this positions JiT at the center of a methodological distinction. Latent diffusion benefits from compression during both denoising and alignment; JiT removes that compression from the main generative pathway in order to avoid the reconstruction bottleneck. The resulting gain in representational richness also changes which auxiliary objectives remain well-posed. PixelREPA is presented as one practical answer to that change: it preserves JiT’s pixel-space denoising while constraining only the auxiliary alignment branch. This suggests that future work on JiT-like systems will likely depend less on whether external semantic supervision is used at all than on how carefully the supervision is made commensurate with image-space information.