One-sentence Summary

The Microsoft Mage Team's Mage-Flow is an efficient 4B model family that co-designs a low-cost diffusion-based tokenizer (Mage-VAE) with anchor-latent regularization and a native-resolution multimodal transformer trained via rectified flow matching, boosts throughput by $2.5\times$2.5× through resolution packing and kernel fusion, and produces Turbo variants that achieve 0.59s generation and 1.02s editing at $1024^2$10242 on a single A100 GPU.

Key Contributions

• Mage-VAE is a lightweight latent tokenizer that combines one-step diffusion-style encoding/decoding with anchor-latent regularization, matching the reconstruction fidelity of strong public VAEs while reducing tokenization cost by more than an order of magnitude.
• The stack integrates a Native-Resolution Multimodal Diffusion Transformer (NR-MMDiT) with native-resolution packing and fused CUDA kernels, enabling flexible-resolution training and improving end-to-end training throughput by approximately 2.5×.
• A complete generation-and-editing model family is built with Base, RL-aligned, and Turbo variants. Diffusion-NFT is applied as a post-training stage to enhance prompt adherence, text rendering, aesthetic quality, and editing fidelity, and Decoupled DMD distillation with Vision Foundation Model-based adversarial perceptual guidance yields 4-step Turbo models; these models achieve competitive performance on standard benchmarks and deliver generation in0.59s and editing in1.02s at1024² resolution on a single NVIDIA A00 GPU.

Introduction

Visual generative models are now expected to synthesize high-fidelity images, follow compositional text prompts, render multilingual text, and preserve spatial layouts across applications like design, prototyping, and interactive editing. While closed-source systems provide strong performance, they lack transparency, and open-source alternatives rely on backbones scaling from 6B to 80B parameters, making them impractical for fine-tuning, controlled ablations, and deployment on realistic compute budgets. The authors introduce the Mage-Flow stack, a compact 4B-scale foundation that combines a lightweight one-step diffusion tokenizer with a native-resolution multimodal diffusion transformer and fused CUDA kernels, enabling competitive text-to-image generation and instruction-based editing with significantly lower latency and memory than much larger models.

Dataset

The authors construct a large-scale training corpus from two complementary pipelines: a text-to-image generation dataset and an instruction-based image editing dataset. Both pipelines are designed to improve visual quality, safety, semantic alignment, diversity, and coverage of under‑represented capabilities.

Text-to-Image Generation Dataset

• Raw source: about 10 billion image–text pairs collected from open‑source datasets.
• Sample‑level filtering: removes corrupted files, low‑resolution images, extreme aspect ratios, and incorrectly oriented samples; then scores images for brightness, saturation, blur, entropy, watermark, aesthetic quality, OCR content, and NSFW safety. Thresholds are progressively tightened across training stages (256² → 512² → 1024² → SFT) so that early stages keep broad coverage while later stages use cleaner data.
• Cross‑sample deduplication: each image is encoded with a copy‑detection descriptor (SSCD) and indexed with FAISS. Within a dataset, pairs with cosine similarity above 0.9 are grouped, keeping only the best‑quality representative. Across datasets, a persistent index rejects near‑duplicates. A held‑out benchmark index is also matched to reduce evaluation contamination.
• Multi‑granularity captioning: retained images are captioned by a large vision‑language model at four levels: phrase, entity, composition, and photographic descriptions. For text‑rich images the captioner additionally converts visible text into rendering instructions. During training the model samples from these channels to learn to follow prompts of varying length and specificity.
• Concept‑aware synthesis and balancing: supplemental data is generated for long‑tail concepts such as long‑text rendering, rare objects, uncommon attributes, structured layouts, and under‑represented styles (e.g.synthetic text‑rendering samples with diverse fonts, layouts, languages, colors, and backgrounds). All supplemental images pass through the same quality filters and are merged into the main corpus. Phrase‑level captions are used to estimate the concept distribution, and concept‑aware sampling during training reduces the dominance of frequent objects, natural scenes, and common photorealistic styles.
• Final scale: approximately 1.3 billion high‑quality image–text pairs are retained; stage‑wise pre‑training subsets are drawn from this pool with progressively tighter filtering and stronger reweighting.

Image Editing Dataset

• Raw triples: about 90 million (source image, edit instruction, target image) triples. Roughly 50M are aggregated from open‑source instruction‑based editing datasets; the remaining ≈40M are synthesized in‑house — 10M low‑level image‑processing triples and 30M general semantic‑editing triples.
• VLM‑based filtering: each triple is evaluated by three independent visual‑language model experts with different judging criteria. An expert checks whether the edit is correctly applied, whether unrelated regions are preserved,, and whether the result is visually plausible.. A triple is kept only when at least two of the three experts vote to pass.. After filtering, roughly 20M open‑source triples and 25M synthesized triples survive,, yielding a 45M‑triple retained pool..
• Edit‑type tagging and balancing: a manual taxonomy of 19 edit categories is defined.. Data units (whol.e datasets, sub‑datasets,s or samples sharing an explicit edit‑type field)are mapped to these categories.. Sampling rates are then adjusted so that the training mixture provides broad, well‑proportioned coverage across the taxonomy,, preventing frequent operations from dominating while keeping rare edit types sufficiently represented..
• Usage: the final balanced editing set is used directly for editing model training..

Method

The authors propose the Mage-Flow stack, which is built around two core components: Mage-VAE, a lightweight high-fidelity latent tokenizer, and a 4B Native-Resolution Multimodal Diffusion Transformer (NR-MMDiT) trained with rectified flow matching. Mage-VAE provides a compact generation-ready latent space, while NR-MMDiT models packed latent sequences with flexible resolutions and aspect ratios.

To address the latency bottleneck of traditional VAEs at high resolutions, the authors design Mage-VAE as a lightweight pixel-diffusion-based tokenizer. It replaces heavy high-resolution autoencoder components with lightweight one-step diffusion-style encoding and decoding. The architecture features a symmetric design: the decoder uses stacked convolutional diffusion blocks and a decoupled pixel diffusion head to reconstruct RGB pixels directly from latents, while the encoder acts as a latent generator conditioned on pixels. To ensure the latent space is generation-ready, the authors regularize the posterior toward an anchor latent distribution induced by the FLUX.2-VAE.

The training pipeline consists of three stages. In Stage I, the encoder and decoder are pre-trained separately as multi-step diffusion models. In Stage II, the decoder is distilled into a one-step model using reconstruction loss, a DMD loss, and a DINOv2-projected GAN loss. In Stage III, the one-step encoder and decoder are jointly fine-tuned with the anchor-latent KL regularizer.

The generative backbone is a 4B-parameter NR-MMDiT. Unlike standard bucket-based training, NR-MMDiT trains directly on native-resolution sequences. Images with different resolutions and aspect ratios are encoded by Mage-VAE, flattened into variable-length latent token sequences, and packed into a single contiguous batch under a fixed token budget. The same packing principle applies to text conditions. This native-resolution formulation removes the single-bucket restriction, exposing the model to a richer resolution distribution and allowing a single checkpoint to generalize to flexible output sizes, ranging from 512 to 2048 in both height and width.

Text conditioning is provided by a frozen Qwen3-VL text encoder. The model is trained in the Mage-VAE latent space with the rectified flow-matching objective:

$$\mathcal{L}(\theta) = \mathbb{E}_{(x, \tau), t, \epsilon} \left[ \| v_{\theta}(z_t, t, \tau) - (z - \epsilon) \|_2^2 \right]$$L(θ)=E(x,τ),t,ϵ​[∥vθ​(zt​,t,τ)−(z−ϵ)∥22​]

where $z_t = (1 - t)z + t\epsilon$zt​=(1−t)z+tϵ and $\epsilon \sim \mathcal{N}(0, I)$ϵ∼N(0,I). For editing tasks, the model extends the 2D rotary positional embedding with an additional frame dimension to distinguish source and target visual tokens.

To make native-resolution 4B-scale learning practical, the authors implement stack-level kernel fusion. They fuse dominant operator chains inside the repeated blocks of Mage-VAE, the Qwen3-VL text encoder, and NR-MMDiT, keeping intermediate values in on-chip memory to reduce activation memory movement and kernel launches.

The Mage-Flow generation corpus is curated from roughly 10B raw image-text pairs. The pipeline includes sample-level filtering, cross-sample deduplication, multi-granularity captioning, and concept-aware synthesis. Sample-level filters remove corrupted, low-quality, or unsafe images. Cross-sample deduplication uses SSCD descriptors to suppress near-duplicate visual modes. Multi-granularity captioning assigns phrase-level, entity-level, composition-level, and photographic captions to standardize textual supervision.

Concept-aware synthesis supplements long-tail images. The resulting data composition ensures broad coverage across object, scene, design, and synthetic domains, while concept-aware sampling reduces the dominance of frequent objects during training.

The authors train the Mage-Flow model family with a unified recipe. Pre-training follows a progressive curriculum: starting at 256x256 resolution for broad visual-language alignment, moving to 512-pixel native-aspect-ratio regimes, and finally 1024-pixel regimes for fine details. Supervised fine-tuning is then performed on a high-quality subset. For instruction-based editing, the model is initialized from the text-to-image base checkpoint and adapted using a balanced mixture of editing triples and generation pairs.

Post-training applies Difusion-NFT, a negative-aware fine-tuning method operating on the forward process of flow-matching generators. It uses online rollout groups and capability-routed reward evaluators to optimize a reward-weighted flow-matching objective with implicit positive and negative policies. Finally, to reduce inference cost, the authors distill the aligned checkpoints into 4-step Turbo models using Decoupled DMD combined with adversarial perceptual guidance.

Experiment

Mage-VAE is evaluated as a lightweight tokenizer distilled from FLUX.2-VAE, showing that anchor-latent supervision preserves reconstruction quality and latent compatibility while reducing encoding cost by over an order of magnitude. The Mage-Flow and Mage-Flow-Edit models are further distilled into 4-step Turbo variants using Decoupled DMD with adversarial perceptual guidance, which retains generation and editing quality, with adversarial gains especially evident for generation and text editing. Overall, the 4B Mage-Flow achieves competitive text-to-image and instruction-based editing performance against much larger models, and the 4-step versions largely maintain this capability.

Mage-VAE achieves reconstruction fidelity on CLIC 2020 that is comparable to the strongest FLUX.2-VAE baseline while attaining the best or near-best quality on FFHQ. It reduces computational complexity by roughly an order of magnitude for encoding and over twenty times for decoding, enabling efficient handling of high resolutions where other VAEs become impractically slow or exceed memory limits. The lightweight design retains generation-compatible latent structure distilled from the expensive teacher, offering a superior quality–efficiency trade-off. Mage-VAE matches FLUX.2-VAE's PSNR, SSIM, and LPIPS on CLIC 2020 with a much lower parameter count and kMACs/pixel. On FFHQ at 1024×1024, Mage-VAE reaches the highest PSNR and SSIM among all compared VAEs, outperforming both SD variants and the teacher FLUX.2-VAE. Encoding complexity is reduced by about 12× and decoding complexity by about 22× relative to FLUX.2-VAE, while still maintaining high reconstruction quality. At 4096×4096 resolution, baseline VAEs either run out of memory or become extremely slow, whereas Mage-VAE remains efficient for both encoding and decoding. Despite the computational savings, the distilled tokenizer preserves the latent geometry required by FLUX.2-style diffusion generators, as indicated by cross-tokenizer compatibility in downstream tasks.

Swapping the tokenizer between Mage-VAE and FLUX.2-VAE while keeping the downstream backbone fixed leads to near-identical performance across generation and editing benchmarks, demonstrating strong cross-tokenizer compatibility. The results confirm that the lightweight Mage-VAE preserves the latent geometry required by FLUX.2-style diffusion generators, making it an efficient plug-in replacement without sacrificing output quality. Under the Mage-Flow-Turbo backbone, both tokenizers achieve identical GenEval scores (0.88) and differ by only 0.01 on DPG-Bench, with other metrics like OneIG-EN and LongText-CN matching or nearly matching. For the FLUX.2-Klein-4B backbone, GenEval is unchanged at 0.83, DPG-Bench varies by at most 0.01, and scores on LongText-EN and LongText-CN remain identical regardless of the tokenizer used. The consistent benchmarks across backbones show that Mage-VAE can replace the computationally heavier FLUX.2-VAE without meaningfully affecting generation or instruction-based editing quality.

Packing the conditional and unconditional branches of classifier-free guidance into a single batch eliminates redundant computation and reduces overall inference latency without changing the denoising trajectory. Across text-to-image and instruction-based editing variants, this optimization consistently improves throughput by 9–15% on a single NVIDIA A100 GPU, with base generation models benefiting the most. Packed CFG achieves speedups ranging from 1.09× to 1.15× across all tested model and step configurations. The base text-to-image model with 20 steps sees the largest relative gain, while the editing variants with longer per-step times see a more modest 1.09–1.10× improvement. All measurements are performed on a single A100 GPU, ensuring the reported speedups reflect practical inference efficiency gains.

Switching from FLUX.2-VAE to the lightweight Mage-VAE tokenizer immediately cuts per-step training time by about 29% and raises model FLOP utilization from 33.20% to 45.63% at essentially the same peak GPU memory. Adding fused CUDA kernels for the VAE and text encoder yields only modest additional speedups, while fusing the NR-MMDiT backbone produces the dominant improvement, lowering memory to 141.44 GB and boosting overall speedup to 2.49×. The full system demonstrates that both a compact tokenizer and stack-level kernel fusion are necessary to remove memory-bound overhead from repeated blocks. Replacing FLUX.2-VAE with Mage-VAE alone provides a 1.41× speedup and raises MFU by over 12 percentage points without changing peak memory. Fusing the VAE and Qwen3-VL text encoder each give incremental speedups (up to 1.46×), while fusing the NR-MMDiT blocks delivers a much larger leap to 2.49× overall speedup. DiT kernel fusion also reduces peak per-GPU memory from approximately 175 GB to 141 GB, showing that memory-bound overhead is concentrated in the diffusion backbone. The full configuration more than doubles MFU (33.20% to 77.26%), confirming that stack-level operator fusion effectively converts hardware capability into usable throughput.

Filtering thresholds become progressively more stringent from the initial pre-training stage to the supervised fine-tuning stage. Minimum resolution and aesthetic quality requirements increase, while the allowed watermark probability drops dramatically, shifting emphasis from broad visual coverage to higher-resolution, cleaner, and more visually appealing images. Aspect ratio, file size, and NSFW thresholds remain constant across all stages. Watermark tolerance tightens sharply from below 0.5 in the 256² stage to below 0.05 in the SFT stage, virtually eliminating watermarked images in later training phases. Aesthetic score floor rises from 4.5 to 6.5 across the stages, ensuring that only visually appealing samples survive the final filtering step.

the paper evaluate Mage-VAE as a lightweight tokenizer, confirming it matches the reconstruction quality of FLUX.2-VAE while reducing encoding and decoding complexity by over an order of magnitude, and it seamlessly replaces the teacher in downstream generation with no impact on quality. Packed classifier-free guidance removes redundant computation to accelerate inference by up to 15%, and fusing the NR-MMDiT backbone more than doubles training throughput while halving peak memory. Data filtering pipelines apply increasingly stringent aesthetic and watermark thresholds from initial pre-training to supervised fine-tuning, ensuring only clean, high-quality images are used in later stages.