Training Text-to-Image Models Without a VAE

Pyramid-JiT: Low-Res Drafts, High-Res Images | Field Notes by Linum

Introduction: The Parameter Allocation Dilemma

When designing a model, a critical question arises: Where should the parameters go?

---

Table of Contents

• The Core Problem
• The Evolution of JiT
• The Encoder-Decoder Struggle
• Introducing Pyramid-JiT (P-JiT)
• Performance Benchmarks
• Open Source Release

---

The Core Problem <a name="the-core-problem"></a>

Current high-end generative models (like Sora or Seedance) are prohibitively expensive because they perform attention calculations on massive token sequences.

For instance, in the Linum v2 model, a mere 5-second video at 720p requires $\approx 110,000$ tokens. In stark contrast, most Large Language Models (LLMs) spend $97\%$ of their training on context windows of $\le 8,000$ tokens.

The Bottleneck: Attention complexity scales quadratically with the number of tokens. To make affordable animation tools a reality, we must reduce the token count by an order of magnitude.

To solve this, Linum uses text-to-image generation as a laboratory for context-window condensation. The logic is simple: a model must master nouns (static objects/scenes) before it can grasp verbs (motion/action).

The Evolution of JiT <a name="the-evolution-of-jit"></a>

Historically, Latent Diffusion Models (LDMs) hit a performance ceiling at $16 \times 16$ token compression unless the auto-encoder is heavily modified. To bypass this, Linum moved toward pixel space models (JiT).

VAE-based compression $\rightarrow$ Pixel Space Compression

By unifying generation and compression into one model, they achieved $32 \times 32$ token compression. This led to the JiT-DDT, a pixel-space encoder-decoder architecture that trained $3.6\times$ faster than previous LDMs.

The Encoder-Decoder Struggle <a name="the-encoder-decoder-struggle"></a>

Scaling encoder-decoder architectures presents a difficult trade-off regarding parameter budgets:

1. Expand the Encoder? This improves structural learning but feels wasteful to spend massive parameters on low-resolution predictions.
2. Expand the Decoder? This improves fine-grained detail, but might cripple the model's ability to understand the overall framing of the scene.

The team looked at the Decoupled Diffusion Transformer (DDT), which suggested a $4:1$ encoder-to-decoder ratio. However, for JiT-DDT, a $1:1$ ratio actually performed better.

This led to a fundamental realization: the industry has largely shifted toward decoder-only models (like GPT) because they are significantly easier to scale without architectural bottlenecks. The goal then became: Can we bring the benefits of encoder-decoder structures into a vanilla, decoder-only JiT?

Introducing Pyramid-JiT (P-JiT) <a name="introducing-pyramid-jit-p-jit"></a>

The solution is Pyramid-JiT (P-JiT). Instead of a separate encoder, P-JiT is a decoder-only architecture that predicts the target image at multiple resolutions throughout the DiT (Diffusion Transformer) trunk.

Essentially, the model creates a "low-res draft" and iteratively refines it into a "high-res image."

Conceptual Logic:

def pyramid_jit_forward(latent, timestep):
    # Predict at multiple scales along the trunk
    low_res_draft = trunk.layer_block_1(latent) 
    mid_res_refinement = trunk.layer_block_2(low_res_draft)
    high_res_final = trunk.layer_block_3(mid_res_refinement)
    return high_res_final

Performance Benchmarks <a name="performance-benchmarks"></a>

P-JiT demonstrates massive efficiency gains over Linum v2.

Efficiency Metrics

Quality Analysis

The team used FD-DINOv2 (where lower is better) to measure quality against 30K real images.

• Real-vs-Real Floor: $25.4$ (The theoretical best score).
• Linum v2: $79.9$
• JiT-DDT: $39.1$
• P-JiT: $27.9$

Key Achievements:

• Match Linum v2's final FD-DINOv2 score.
• Drastically reduce training sample requirements.
• Scale pixels without proportional cost increases.

Caption: P-JiT reaches the quality floor significantly faster than previous iterations.

Open Source Release <a name="open-source-release"></a>

In the spirit of community growth and the pursuit of efficient training, Linum has released the P-JiT artifacts.

• License: Apache 2.0
• Availability: Model weights and code are available.
• Status: This is a research artifact, not a production-ready model release.

These findings serve as a stepping stone toward Linum v3, with the ultimate goal of creating accessible, high-quality animation tools for everyone.