Summary (Overview)

• Compact 4B-scale generative stack for text-to-image generation and instruction-based image editing, co-designed for efficiency without scaling to tens of billions of parameters.
• Mage-VAE: a lightweight latent tokenizer using one-step diffusion-style encoding/decoding with anchor-latent regularization, reducing tokenization cost by >10× compared to strong public VAEs while preserving reconstruction quality.
• Native-Resolution Multimodal Diffusion Transformer (NR-MMDiT) trained with rectified flow matching, supporting flexible resolutions (512–2048) and aspect ratios via packing; combined with fused CUDA kernels achieving ~2.5× end-to-end training speedup.
• Complete model family (Base, RL-aligned, Turbo) for both generation and editing; Diffusion-NFT post-training improves prompt following and text rendering; 4-step Turbo models achieve 0.59s generation and 1.02s editing at 1024² on a single A100 GPU with ~18–20 GB peak memory.
• Competitive performance against much larger open-source systems (e.g., Qwen-Image, FLUX.2, FireRed-Image-Edit) across standard benchmarks, establishing a strong quality–speed–memory frontier.

---

Introduction and Theoretical Foundation

Large-scale visual generators (e.g., GPT-Image, FLUX.2, HunyuanImage-3.0) increasingly rely on backbones with 6B–80B parameters, creating barriers to training, fine-tuning, and deployment under realistic compute budgets. The paper argues that strong generation does not require continued backbone scaling; instead, system-level co-design of the tokenizer, backbone, and training infrastructure can deliver high-quality results at a compact 4B scale.

The theoretical foundation builds on two pillars:

1. Rectified flow matching – The backbone is trained to predict a velocity field $v_\theta(z_t, t, \tau)$ that transports noise $\epsilon \sim \mathcal{N}(0,I)$ to data $z$ along a straight-line interpolation $z_t = (1-t)z + t\epsilon$. The objective is:

   $$\mathcal{L}(\theta) = \mathbb{E}_{(x,\tau),t,\epsilon}\left[ \| v_\theta(z_t, t, \tau) - (z - \epsilon) \|_2^2 \right].$$

2. Latent diffusion with efficient tokenization – Public VAEs (e.g., FLUX.2-VAE) are optimized for pixel-level reconstruction but have encoder/decoder costs that scale poorly at high resolution, becoming a pipeline bottleneck. Mage-VAE replaces heavy autoencoder components with lightweight one-step diffusion-style encoding/decoding, regularized toward an anchor latent distribution from FLUX.2-VAE.

---

Methodology

Mage-VAE: Lightweight Image Tokenizer

Architecture. The decoder is a fully convolutional one-step pixel diffusion model (inspired by CoD-Lite), avoiding global attention. The encoder is its architectural dual: a one-step diffusion model generating latents from pixels. Both use stacked convolutional diffusion blocks and a decoupled pixel diffusion head.

Anchor-latent regularization. Instead of a standard Gaussian prior, the posterior $q_\phi(z|x)$ is regularized toward the latent distribution $q_a(z|x)$ of FLUX.2-VAE:

$$\mathcal{L}_{\text{KL}} = \mathbb{E}_x\left[ D_{\text{KL}}\big( q_\phi(z|x) \,\|\, q_a(z|x) \big) \right].$$

Mage-VAE directly produces $16\times$ downsampled latents with 128 channels (internalizing the $2\times$ patchification of FLUX.2 latents).

Training stages.

• Stage I: Multi-step flow-matching pre-training (encoder predicts anchor latents, decoder predicts pixels).
• Stage II: One-step decoder distillation with $\ell_1$, LPIPS, DMD, and DINOv2-projected GAN losses.
• Stage III: Joint one-step encoder–decoder optimization with anchor-latent KL.

Native-Resolution MMDiT (NR-MMDiT)

• 4B-parameter MMDiT with modality-specific normalization and joint self-attention.
• Native-resolution packing: images of different resolutions are flattened into variable-length latent token sequences and packed together with variable-length text tokens under a fixed token budget. FlashAttention with per-sample cumulative offsets enables efficient variable-length attention.
• Packed CFG inference: conditional and unconditional branches are packed into a single forward pass, achieving 1.09×–1.15× speedup.
• Conditioning: frozen Qwen3-VL-4B-Instruct text encoder. For editing, source image latents are concatenated with target latents using a 3D RoPE with a frame dimension to distinguish source/target tokens.

Training Infrastructure

• Stack-level CUDA kernel fusion fuses normalization–activation–residual chains in Mage-VAE, adaptive normalization/RoPE/gating in Qwen3-VL and NR-MMDiT.
• Table 4 shows full system achieves 77.26% MFU (vs. 33.20% without optimizations) and 2.49× training speedup.

Data Pipelines

• Generation: ~10B raw image–text pairs → sample-level filtering (progressive thresholds), cross-sample deduplication (SSCD + FAISS), multi-granularity captioning (Qwen3-VL-32B), concept-aware synthesis → ~1.3B high-quality pairs.
• Editing: ~90M triples (50M open-source + 40M in-house) → VLM-based voting filter (three Qwen3.5-9B experts, majority vote) → 45M retained → edit-type tagging and balancing across 19 categories.

Post-training and Distillation

Diffusion-NFT (RL alignment). Operates on the forward process with reward-weighted flow matching:

$$\mathcal{L}_{\text{NFT}}^{(s)}(\theta) = \mathbb{E}_{c, x_{0,i}, t}\left[ r_i^{(s)} \| v_\theta^+(x_{i,t}, t, c) - v_{i,t} \|_2^2 + (1 - r_i^{(s)}) \| v_\theta^-(x_{i,t}, t, c) - v_{i,t} \|_2^2 \right].$$

Reward evaluators: PaddleOCR-VL-1.5 (text), Qwen3.5-27B (aesthetic/semantic), RationalRewards (editing).

Few-step distillation (Turbo models). Combines Decoupled DMD with adversarial perceptual guidance:

$$\nabla\mathcal{L}_{\text{Total}} = \underbrace{\Delta_{\text{CA}} + \Delta_{\text{DM}}}_{\nabla\mathcal{L}_{\text{D-DMD}}} + \lambda_{\text{GAN}} \nabla\mathcal{L}_{\text{GAN}},$$

where $\Delta_{\text{CA}} = (w-1)(T_{\text{cond}}^{\text{real}} - T_{\text{uncond}}^{\text{real}})$ and $\Delta_{\text{DM}} = T_{\text{cond}}^{\text{real}} - T_{\text{cond}}^{\text{fake}}$. Adversarial guidance uses a feature discriminator on frozen DINOv2/CLIP features.

---

Empirical Validation / Results

Text-to-Image Generation (Table 8, summary)

• Mage-Flow achieves the best GenEval score among all compared systems (0.90) and strong text rendering (CVTG-2K 0.887, LongText-EN 0.944).
• Turbo variant preserves most quality with only 4 steps.

Instruction-Based Image Editing (Table 13, summary)

• Mage-Flow-Edit achieves best open-source GEdit-CN score and competitive ImgEdit/TextEdit scores.
• Turbo variant maintains performance with 4 steps.

Efficiency (Fig. 4, Tables 1, 4)

• Mage-VAE: 12.3× fewer encoding MACs/pixel and 22.3× fewer decoding MACs/pixel vs. FLUX.2-VAE, with comparable reconstruction (PSNR 36.61 vs. 36.88 on CLIC 2020).
• Training: 2.49× speedup, MFU 77.26%, memory 141.44 GB vs. 175.45 GB.
• Inference: Mage-Flow-Turbo 0.59s/image, Mage-Flow-Edit-Turbo 1.02s/image at 1024² on A100 (80GB), peak memory ~18–20 GB.

Ablation Studies

• Anchor-latent tokenizer compatibility (Table 2): Swapping Mage-VAE and FLUX.2-VAE in downstream backbones yields similar benchmark performance, confirming latent space compatibility.
• Adversarial perceptual guidance (Table 6): Improves generation and text editing under 4-step distillation (e.g., TIIF-Short 83.58 vs. 80.99); gains on general editing are benchmark-dependent.
• Generation data in editing training (Table 7): Helps ImgEdit for Turbo variant (4.38 vs. 4.20); GEdit changes are mixed.
• Native-resolution packing (Table 3): Packed CFG speeds up inference by 1.09×–1.15×.

---

Theoretical and Practical Implications

Theoretical Insights

• Tokenizer–backbone–system co-design is a viable alternative to backbone scaling: a lightweight tokenizer (Mage-VAE) removes the high-resolution bottleneck, while native-resolution packing and kernel fusion make training efficient.
• Anchor-latent regularization provides a principled way to transfer a generation-ready latent space from an expensive VAE to a lightweight one without breaking downstream compatibility.
• Native-resolution packing turns resolution diversity into a training signal, improving flexibility and enabling efficient packed CFG inference.

Practical Implications

• Accessibility: The entire 4B model family fits in ~18–20 GB GPU memory, enabling local desktop deployment and downstream research on a single consumer GPU.
• Speed: 4-step Turbo models (0.59s generation, 1.02s editing) make high-resolution interactive editing practical.
• Extensibility: The stack is designed for easy adaptation (e.g., scientific diagram generation in Appendix D, fine-tuned from Mage-Flow-Base with a single-stage data mixture).
• Reproducibility: Open-source code, models, and project page provide a research-friendly baseline.

---

Conclusion

The Mage-Flow generative stack demonstrates that careful co-design of a lightweight tokenizer, native-resolution diffusion backbone, and optimized training infrastructure can deliver strong high-resolution generation and editing at a compact 4B scale. Key contributions include:

• Mage-VAE: a high-fidelity one-step encoder–decoder with anchor-latent regularization, reducing tokenization cost by >10×.
• NR-MMDiT: a flexible backbone supporting native-resolution packing and packed CFG inference.
• Complete model family (Base, RL-aligned, Turbo) for both generation and editing, achieving competitive or superior performance against much larger open-source systems while being faster and more memory-efficient.

Future directions include more robust multi-image editing, stronger multilingual long-text rendering, and tighter integration with agentic visual creation workflows. The results show that strong visual generation does not necessarily require tens-of-billions-parameter backbones—a well-designed 4B stack can serve as a practical and research-friendly foundation.