Skip to content

Training Text-to-Image Models 3.6× Faster (JiT-DDT)

Linum’s field-notes post introduces JiT-DDT, a pixel-space text-to-image architecture that trains a 512² model with 3.6× fewer GPU-hours than their Linum v2 latent-diffusion baseline at 256² while generating 4× more pixels. The recipe layers three things onto the JiT (Li & He, 2025) pixel-space backbone: (a) an encoder-decoder split borrowed from DDT, where the encoder predicts an 8×-downsampled version of the image to explicitly learn structure and the decoder handles high-frequency detail; (b) x-prediction + v-loss to sidestep the curse-of-dimensionality that killed v-prediction under aggressive 32×32 patchification; (c) Z-Image-style architecture tweaks (single-stream MM refiners, SwiGLU, Muon on 2D weights, sigmoid attention gating, PixelREPA into DINOv3, Qwen3.5-4B text embeddings). Weights and code released under Apache 2.0 as a research checkpoint on the path to Linum v3.

  • JiT-DDT trains a 512² text-to-image model in 3.6× fewer GPU-hours than the Linum v2 latent-diffusion baseline at 256², while generating 4× more pixels per sample [§Introduction, §From Linum v2 to JiT-DDT].
  • 32×32 token reduction is achievable by pushing compression into the DiT via linear patchification instead of a VAE; the standard CNN-VAE stack (FLUX, Ideogram, Z-Image) empirically caps out at 16×16 [§Hitting the VAE compression wall].
  • Aggressive patchification only works under x-prediction, v-loss, not v-prediction, v-loss: velocity u = ε − x₀ implicitly forces the network to fit high-dimensional Gaussian noise, which fails as channel dimension grows; x-prediction lets the network exploit the low-dimensional image manifold [§V-prediction and the curse of dimensionality].
  • Vanilla JiT struggles with fine-grained detail. JiT-DDT recovers detail by splitting the DiT into an encoder that predicts a 64×64 downsampled x₀ (structure-only) and a decoder that predicts the full 512×512 x₀ (detail), each with its own patchification and output head [§JiT-DDT: our encoder-decoder pixel-space architecture].
  • Unlike the original DDT, JiT-DDT keeps encoder and decoder equally sized (encoder’s task is easier at 8× downsample) and retains the text condition on the decoder (open-world T2I has detail annotated in the prompt) [§JiT-DDT].
  • Tuning the noise schedule was the load-bearing change that made JiT-DDT beat vanilla JiT — none of the other DDT-port ablations moved the needle until the schedule was widened; the post explicitly notes no prior work has published a schedule dialed for pixel-space [§Adjusting the noise schedule].
  • Plain REPA (align an early DiT layer to DINOv3) actively hurt the pixel-space JiT; PixelREPA (mask x% of the visual hidden state, push through a shallow transformer, then align to DINOv3 on a matched-token-count downsampled image) is the working variant. Ablating PixelREPA on/off midway through training is a wash — the masking may be functioning as a regularizer against DINO-space overfitting [§JiT setup, §Why does the JiT-DDT work?].
  • The single-stream MM-DiT architecture increases attention-sink risk; sigmoid attention gating (23 gates per block, one per head) neutralizes it [§JiT setup].
  • The post frames JiT-DDT alongside Self-Flow (BFL) and RAE v2 as three variants of the same underlying diagnosis: early DiT layers have a gradient-propagation problem and need an extra loss injection to learn structure; JiT-DDT injects it via the encoder’s low-res x₀ prediction, Self-Flow via later-layer self-distillation, RAE v2 via REPA-as-x-prediction in DINO space [§Why does the JiT-DDT work?].

Backbone changes vs Linum v2. The block structure moves from alternating self-attention (image) + cross-attention (text) to a single stream where image and text tokens are concatenated and go through joint self-attention. The transformer becomes wider and shallower — 23 layers × 2944 hidden — following Z-Image. Activations are SwiGLU. All 2D matrices (Q/K/V, FFN weights) are trained with Muon; input/output layers, scales, and AdaLN biases stay on AdamW.

Pixel-space objective. Patchification is 32×32 (so a 512×512 image is 256 tokens). Training is x-prediction, v-loss. Perceptual losses (LPIPS + P-DINO) are applied when timestep t → 0. Sigmoid attention gating (one gate per head per block) mitigates attention sinks in the single-stream design.

Text conditioning. Qwen3.5-4B replaces T5-XXL. The DiT takes hidden states from three of Qwen’s full-attention blocks (post-sparse/linear layers) and learns a projection to combine them into a single text condition.

PixelREPA. For a 512² image with 32×32 patches (256 visual tokens), the auxiliary loss downsamples the image to 256² so DINOv3’s 16×16 patches also yield 256 tokens. The DiT’s hidden state at block 8 has x% of tokens masked, passes through a shallow transformer, and all output tokens are aligned to the DINOv3 tokens by cosine distance.

JiT-DDT encoder-decoder split. The DiT is cut in half. The encoder patchifies the noised image at 32×32, gets the text condition, and predicts a 64×64 (8×-downsampled) version of x₀ through its own output head. The decoder receives (a) the encoder’s final hidden states, (b) its own view of the noised image via a separate patchification, and (c) the same text condition; it predicts the full 512×512 x₀. Encoder and decoder are equally sized (~11 layers each of the 23-layer trunk).

Refiners. Two Z-Image-style MM-DiT refiner blocks massage image and text tokens independently before the shared trunk. The other Z-Image tweaks (AdaLN truncation, post-norm gate + tanh, RMSNorm placement) did not move the needle and were dropped.

  • Headline: 3.6× fewer GPU-hours than Linum v2 to reach comparable quality, at 512² vs Linum v2’s 256² (i.e., 4× more pixels per sample).
  • Qualitative comparison on 26 prompts shows JiT-DDT is (a) far more faithful to art styles (charcoal, Roman mosaic, oil painting), (b) generates more realistic lighting (Linum v2 was prone to over-saturation), (c) still struggles with human faces when they are not the main subject (small faces in scenes with a chef, a florist, a street musician).
  • Ablations reported qualitatively (no FID/GenEval numbers in the post): the noise-schedule change was decisive; PixelREPA on/off midway was a wash; dropping the decoder’s text conditioning was a wash; refiners helped; the other Z-Image knobs did not help.
  • Released: Apache-2.0 weights + code on HuggingFace + GitHub, explicitly framed as a research checkpoint, not a product release. Model card and repo scaffolding written by Claude Fable 5.1 from the authors’ internal experiment repo.

This is the first filed text-to-image instance of the “pure JiT-line + encoder-decoder split” recipe at production scale — it plants a flag between the small end (MiniT2I at 258M, MiniT2I: Pixel-Space Text-to-Image with MM-JiT at 258M Parameters) and the well-resourced end (PixelDiT-T2I 1.3B from NVIDIA, PixelDiT: Pixel Diffusion Transformers for Image Generation; AsymFLUX.2 klein 9B fine-tune, Asymmetric Flow Models). The three-way framing of JiT-DDT / Self-Flow (Self-Flow: Self-Supervised Flow Matching for Scalable Multi-Modal Synthesis) / RAE v2 as “gradient-propagation fixes for early DiT layers” is the wiki’s cleanest published articulation yet of why so many disparate 2026 recipes converge on the same delta over vanilla flow-matching DiT. It also directly stress-tests SRA’s (No Other Representation Component Is Needed: Diffusion Transformers Can Provide Representation Guidance by Themselves (SRA)) “DiT can supply its own representation guidance” thesis — the Linum team tried JiT + Self-Flow first and it lost to JiT + PixelREPA, suggesting external DINO alignment is still helpful in the low-sample (100–150M) regime where SRA and Self-Flow do most of their work. Finally, the noise-schedule-as-load-bearing-knob observation is a concrete gap in the pixel-space literature that other groups can now go dial in.