💻 computer science

Nucleus-Image: Sparse MoE for Image Generation

Nucleus-Image is a fully open-source, 17B-parameter sparse Mixture-of-Experts diffusion transformer that achieves state-of-the-art image generation quality with only 2B active parameters per forward pass through architectural innovations like Expert-Choice Routing and decoupled timestep modulation, all without relying on post-training optimization.

Original authors: Chandan Akiti, Ajay Modukuri, Murali Nandan Nagarapu, Gunavardhan Akiti, Haozhe Liu

Published 2026-04-15

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Chandan Akiti, Ajay Modukuri, Murali Nandan Nagarapu, Gunavardhan Akiti, Haozhe Liu

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you want to build a super-smart artist who can draw anything you describe. Usually, to make an artist this good, you need a massive studio with thousands of painters working on every single detail of every single image. This is expensive, slow, and requires a huge amount of electricity.

Nucleus-Image is like a revolutionary new art studio that changes the rules. Instead of hiring thousands of painters to work on everything at once, they built a team of 64 specialized experts (like a master of landscapes, a wizard of faces, a guru of text, etc.). But here's the magic trick: For every single image, only about 2 or 3 of these experts actually wake up and do the work. The rest stay asleep.

This allows the studio to have the knowledge of a 17-billion-parameter giant (a massive team), but it only pays the electricity bill of a tiny 2-billion-parameter team.

Here is a breakdown of how they did it, using simple analogies:

1. The "Sleeping Giants" Strategy (Sparse MoE)

Think of the model as a massive library with 17 billion books. In old models, to answer a question, you had to read every single book to find the right page. That's slow.
In Nucleus-Image, they use a Smart Librarian (The Router). When you ask for a picture of a "cat," the Librarian instantly knows, "Oh, we only need the 'Animal Experts' and 'Fur Texture Experts'." It wakes up just those few experts and ignores the rest.

The Result: You get the wisdom of the whole library, but you only pay for the few books you actually needed.

2. The "Decoupled" Brain (Routing vs. Computing)

In previous versions of this technology, the Librarian got confused. The "time of day" (the noise level in the image) was so loud that the Librarian would just say, "Everyone wake up at 9 AM, everyone sleep at 5 PM," regardless of what the picture actually needed.

The Fix: Nucleus-Image separated the Librarian's brain from the painters' hands.
- The Librarian looks at the content (is it a cat?) and the time (is it early or late in the drawing process?) to decide who works.
- The Painters get the full instructions to do their job.
- This prevents the Librarian from getting overwhelmed by the "time" signal and ensures the right experts are always chosen for the right job.

3. The "Text Ghost" Trick

Usually, when an AI draws a picture based on text, it has to carry the text instructions through every single step of the drawing process, like a heavy backpack.

The Fix: Nucleus-Image treats the text instructions like a ghost. It reads the instructions once at the very beginning, memorizes the "Key" and "Value" (the essence of the request), and then lets the text "ghost" float away. The actual drawing process only deals with the image pixels.
The Result: The drawing process becomes incredibly fast and efficient because it's not carrying that heavy backpack anymore.

4. The "Progressive" School Curriculum

You wouldn't teach a baby to run a marathon on day one. You start with walking, then jogging, then running.

The Fix: The model was trained in three stages:
1. 256 pixels: Learning to draw rough shapes (like sketching).
2. 512 pixels: Learning details and textures.
3. 1024 pixels: Learning fine-grained perfection.
- Crucially, they adjusted the "Expert Capacity" (how many experts wake up) as the model got smarter. Early on, they woke up more experts to learn the basics. Later, they woke up fewer, specialized experts to perfect the details.

5. The "No-Shadow" Training Method

Most AI models keep a "shadow copy" of themselves (a backup version) that updates slowly to make the final result smoother. This doubles the memory cost.

The Fix: Nucleus-Image didn't keep a shadow. Instead, they saved many "snapshots" of the model during training. At the very end, they mathematically merged these snapshots together to create the final, super-smooth model.
The Result: They saved massive amounts of computer memory during training without losing any quality.

6. The "Perfect Recipe" (Data & Optimization)

They didn't just use random internet images. They built a 700-million-image library that was:

Cleaned: Removed blurry, bad, or duplicate pictures.
Graded: Sorted into "A1" (Masterpiece) to "A5" (Okay) tiers.
Captioned: Given multiple descriptions (short, medium, and detailed) so the model learns different ways to describe things.
Synthetic Text: They even generated images with perfect text on them to teach the model how to write words inside pictures (a common struggle for AI).

The Bottom Line

Nucleus-Image proves that you don't need to brute-force your way to the best AI. By being smart about who works (Sparse MoE), efficient about how they work (Decoupled routing, Text Ghosts), and disciplined about what they learn (Curriculum and Data), they created a model that rivals the biggest, most expensive AI in the world, but runs on a fraction of the power.

And the best part? They gave the recipe, the ingredients, and the final dish to the public for free. No secret sauce, no hidden tricks—just open science.

1. Problem Statement

The field of text-to-image generation faces a fundamental trade-off between generation quality and computational efficiency. State-of-the-art models (e.g., SD3.5, Flux) often require massive parameter counts and active computation budgets, leading to high inference costs and memory footprints. While Sparse Mixture-of-Experts (MoE) architectures have successfully scaled Large Language Models (LLMs) by activating only a fraction of parameters per token, applying MoE to Diffusion Transformers (DiT) presents unique challenges:

Routing Instability: In diffusion models, adaptive modulation (scaling based on timesteps) varies significantly across the denoising process. If the router sees modulated inputs, it may collapse into "timestep-choice routing" rather than "content-choice routing," destroying expert specialization.
Gradient Sparsity: At low resolutions, the token sequence is short. Standard MoE routing can lead to insufficient tokens per expert, causing unstable gradient estimates.
Inference Overhead: Text tokens in standard DiTs often pass through MLP/MoE layers, preventing efficient caching and increasing latency.
Training Complexity: Maintaining Exponential Moving Average (EMA) weights doubles memory usage, and standard learning rate schedules lack flexibility for curriculum-based training.

2. Methodology

Nucleus-Image addresses these challenges through a holistic approach spanning data, architecture, optimization, and systems engineering.

A. Data Pipeline & Curriculum

Corpus Construction: Built a 1.5B image-caption pair dataset from 700M unique images. The pipeline includes multi-stage filtering, deduplication (SHA-256 + pHash), and aesthetic tiering (A1–A5).
Multi-Granularity Captions: Images are captioned at three levels (short, medium, detailed) to support diverse training objectives.
Synthetic Text Rendering: A dedicated pipeline generates images with controlled text to improve legibility and typography.
Curriculum Learning: Training follows a progressive resolution curriculum (256 → 512 → 1024) with multi-aspect-ratio bucketing at every stage. A "progressive sparsification" strategy reduces the expert capacity factor as resolution increases.

B. Model Architecture

Sparse MoE DiT: A 32-layer Transformer with a hidden dimension of 2048.
- Total Parameters: ~17B.
- Active Parameters: ~2B per forward pass.
- Experts: 64 routed experts + 1 shared expert per MoE layer (29 layers use MoE; the first 3 are dense for stability).
Expert-Choice Routing: Unlike token-choice routing, experts select their top-k tokens. This guarantees uniform expert utilization without auxiliary load-balancing losses.
Decoupled Routing & Computation: A critical innovation.
- Router Input: Uses the unmodulated token representation concatenated with the timestep embedding. This prevents the timestep-dependent modulation scale from dominating routing decisions.
- Expert Input: Uses the fully modulated representation for computation.
- Result: Stable routing based on content and timestep, while computation remains fully conditioned.
Text Handling: Text tokens are excluded from the transformer backbone's MLP and MoE layers. They participate only as Key-Value (KV) pairs in Joint Attention. This allows for full KV caching of text across denoising steps.
Grouped Query Attention (GQA): Uses a 4:1 ratio (16 Q heads, 4 KV heads) to reduce KV cache memory by 4x.

C. Training & Optimization

Loss Functions:
- Rectified Flow Loss: Primary objective for velocity prediction.
- Z-Loss & Orthogonal Loss: Regularizes router logits and weights to prevent explosion and encourage expert diversity (without standard load-balancing penalties).
- Wavelet Loss: Applied in the final high-resolution stage to emphasize high-frequency details.
Optimizer: Uses Muon (orthogonalized momentum) for attention and FFN weights, combined with AdamW for modulation/patch embeddings.
Learning Rate Schedule: Adopts Warmup-Stable-Merge (WSM).
- No online EMA shadow copy is maintained.
- Training runs with a constant learning rate after warmup.
- The final model is produced by merging the last $N$ checkpoints using inverse-square-root weights, mathematically equivalent to EMA but with zero training memory overhead.
Parallelism: Uses Expert Parallelism (EP) across 64 GPUs (1 expert per GPU) combined with FSDP2. Custom Triton kernels handle grouped matrix multiplication and token permutation to minimize communication overhead.

D. Inference Optimizations

Text KV Caching: Since text tokens are not timestep-conditioned and do not pass through MoE, their KV states are computed once and cached.
Joint Attention: Single Flash Attention kernel call per block for both image and text, scaling as $O(S_i \cdot (S_i + S_t))$ .

3. Key Contributions

First Open-Source MoE Diffusion Model at High Quality: Nucleus-Image is the first fully open-source MoE diffusion model to reach SOTA quality tiers without post-training RL or DPO.
Decoupled Routing Design: Solves the specific instability of MoE in diffusion transformers caused by timestep modulation, enabling stable expert specialization.
Efficient Text Integration: By excluding text tokens from MoE/MLP layers, the model achieves significant inference speedups and memory savings via KV caching.
WSM Training Strategy: Eliminates the memory cost of EMA shadow weights while preserving the benefits of weight averaging through post-training checkpoint merging.
Progressive Sparsification: A novel curriculum that adjusts the MoE capacity factor based on resolution and token density to ensure gradient stability at low resolutions.

4. Results

Nucleus-Image achieves state-of-the-art or near-state-of-the-art performance on major benchmarks while activating only ~2B parameters (vs. 17B total), matching models with significantly larger active budgets.

GenEval: Score 0.87 (matches Qwen-Image, beats SD3.5 Large and Flux.1 Dev). Notably leads in Spatial Position (0.85).
DPG-Bench: Score 88.79 (highest overall, beating Qwen-Image at 88.32 and Seedream 3.0 at 88.27). Strongest in Entity, Attribute, and Relation subcategories.
OneIG-Bench: Score 0.522 (surpasses Imagen4 and Recraft V3). Strong performance in Style and Text Rendering.
Efficiency: Achieves these results with a fraction of the inference cost of dense models due to sparse activation and text KV caching.

5. Significance

The Nucleus-Image report demonstrates that Sparse MoE scaling is a highly effective path for image generation, breaking the traditional trade-off between quality and efficiency.

Scalability: It proves that models can scale to 17B+ parameters while maintaining low inference costs (~2B active), making high-quality generation accessible on consumer-grade or smaller cloud hardware.
Open Science: By releasing the full model weights, training code, and the curated 700M-image dataset, it democratizes access to high-tier diffusion technology.
Architectural Blueprint: The "Decoupled Routing" and "Text-Exclusion" strategies provide a new blueprint for designing efficient, stable diffusion transformers, potentially influencing future multimodal models.
No Post-Training Tuning: The model achieves its performance purely through pre-training, suggesting that high-quality data curation and architectural design are more critical than RLHF for image generation.