Mage-Flow: An Efficient Native-Resolution Foundation Model for Image Generation and Editing

Xinjie Zhang, Peng Zhang, Shicheng Zheng, Jinghao Guo, Zhaoyang Jia, Yifei Shen, Xun Guo, Yuxuan Luo, Jiahao Li, Wenxuan Xie, Fanyi Pu, Xiaoyi Zhang, Kaichen Zhang, Zongyu Guo, Tianci Bi, Dongnan Gui, Zhening Liu, Zimo Wen, Zihan Zheng, Senqiao Yang, Xiao Li, Jinglu Wang, Bin Li, Yan Lu

Mage-Flow is a 4B-scale generative stack that co-designs a lightweight tokenizer and native-resolution transformer for efficient image generation and editing.

How can we build a compact, 4B-parameter generative model that handles native-resolution image generation and instruction-based editing efficiently?

Modern visual generators rely on large backbones and heavy tokenizers, making high-resolution generation and repeated editing computationally expensive and difficult to study under practical budgets. The authors introduce Mage-Flow, a 4B-parameter stack that replaces standard autoencoders with a lightweight, one-step diffusion-based tokenizer (Mage-VAE) and uses a native-resolution transformer (NR-MMDiT) trained with rectified flow matching. This co-design, supported by stack-level CUDA kernel fusion, achieves competitive performance against much larger models while reducing inference latency and memory footprint for interactive use.

Paper Primer

The Mage-VAE tokenizer acts as a learned image codec: it uses a one-step diffusion-style encoder and decoder to map pixels to latents, regularized by an anchor-latent KL objective toward a high-capacity public VAE. This symmetric, fully convolutional design avoids global attention, making tokenization cost nearly linear with image resolution.

The NR-MMDiT backbone removes the traditional "bucket" restriction by packing variable-length image and text sequences into a single batch. This allows the model to process heterogeneous resolutions natively and enables packed classifier-free guidance, where conditional and unconditional branches are evaluated in a single forward pass.

Mage-Flow significantly improves training and inference efficiency compared to larger open-source baselines.

The stack-level kernel fusion increases model FLOP utilization (MFU) from ~14% to ~29% and provides a 2.5× end-to-end training speedup. At 1024² resolution, the 4-step Turbo variants generate images in 0.59s and perform edits in 1.02s on a single A100 GPU.

Why is the tokenizer a bottleneck in modern generative pipelines?

Standard VAEs often use architectures with global attention that scale poorly with resolution; at high resolutions, their encoding and decoding latency can approach that of the diffusion backbone itself.

What is the advantage of native-resolution packing over traditional bucket-based training?

Packing removes the need to discretize image sizes into fixed buckets, allowing the model to observe a richer distribution of aspect ratios and resolutions during every update, which improves generalization to flexible output sizes.

Introduction

We expose the efficiency gap in large visual generators and introduce the compact Mage‑Flow stack.

Large‑scale visual generators are increasingly capable but their training, fine‑tuning, and deployment costs remain prohibitive for most research groups.

We introduce Mage‑Flow, a compact 4 B‑parameter generative stack that couples a lightweight high‑fidelity VAE (Mage‑VAE) with a native‑resolution diffusion transformer (NR‑MMDiT) to close this efficiency gap while preserving image quality.

Redesigning the tokenizer as a one‑step diffusion encoder/decoder with anchor‑latent regularization cuts tokenization cost by more than tenfold, and native‑resolution packing together with fused CUDA kernels yields roughly a 2.5× end‑to‑end training speedup; the resulting Turbo variants generate $1024^2$ images in 0.59 s and edit them in 1.02 s on a single A100 with modest memory.

Our contributions are: (1) the Mage‑Flow stack, (2) the Mage‑VAE tokenizer, (3) the Native‑Resolution MMDiT backbone, (4) a full generation‑and‑editing model family, and (5) a quality‑speed‑memory trade‑off that rivals much larger open‑source systems.

The stack is a tightly coupled pair of a fast latent tokenizer (Mage‑VAE) and a native‑resolution diffusion transformer (NR‑MMDiT) that together enable high‑quality generation and editing at 4 B scale.

**Figure 3.** Qualitative showcase of *Mage-Flow-Edit* (II). Additional examples cover material and style changes, subject removal, human retouching, rainy weather, meme creation, depth, tone, viewpoint, grayscale conversion, extraction, Canny edges, and multi-type editing. All showcased edit types are supported bidirectionally.

Efficiency and native‑resolution handling are the primary drivers of Mage‑Flow’s value, delivering fast, memory‑light generation and editing at 4 B scale.

Related Work

We situate Mage‑Flow within prior image generation, editing, tokenization, post‑training, and distillation research.

Text‑to‑image generation has moved from pixel‑space diffusion and U‑Net latent diffusion toward large diffusion transformers and rectified‑flow models. Stable Diffusion introduced compressed VAE latents, SDXL scaled the U‑Net design, and Diffusion Transformers shifted the backbone to transformers; SD3 and SANA further improved high‑resolution efficiency. Mage‑Flow counters this trend by fixing the generator at a compact 4 B scale and co‑designing it with a lightweight VAE.

Instruction‑based editing evolved from global stochastic denoising (SDEdit) to localized cross‑attention manipulation (Prompt‑to‑Prompt, Null‑text Inversion) and explicit conditioning (ControlNet, IP‑Adapter). Large‑scale editors such as MagicBrush, OmniGen, and FLUX.1 Kontext achieve strong results but rely on much larger backbones. Mage‑Flow‑Edit offers a 4 B editor that preserves editing quality while dramatically lowering adaptation cost.

Unified multimodal models (Janus‑Pro, Transfusion, BAGEL, Emu3.5) aim to share a single transformer for understanding, generation, and editing, typically at 6 B–80 B scales. These architectures are architecturally distinct from MMDiT‑based generators. Mage‑Flow instead specializes a 4 B MMDiT stack for efficient generation and editing, treating perception as an orthogonal add‑on.

Image tokenization largely determines the efficiency frontier. Continuous‑latent VAEs (e.g., Latent Diffusion, FLUX.1) dominate diffusion pipelines, while discrete VQ‑VAE/GAN tokenizers enable autoregressive generation. Multi‑step diffusion codecs achieve low bitrate but are too slow; one‑step codecs improve latency but need heavy backbones. Building on CoD‑Lite’s observation that lightweight convolutional decoders suffice after compression‑oriented pre‑training, Mage‑VAE trains a fully‑convolutional one‑step diffusion VAE and its dual encoder, achieving 16× spatial reduction with an anchor‑latent KL.

Post‑training aligns generators with human preferences. Preference‑based methods (Diffusion‑DPO, D3PO) fine‑tune on collected comparisons, while RL‑based approaches (DDPO, Flow‑GRPO) optimize non‑differentiable rewards. Diffusion‑NFT sidesteps likelihood estimation by fine‑tuning the forward process, and Mage‑Flow adopts this technique to improve prompt adherence, text rendering, and editing fidelity.

Distillation reduces the costly iterative sampling of diffusion models. Training‑free solvers (DDIM) and progressive distillation cut steps but degrade at extreme budgets. Consistency and latent‑consistency models, InstaFlow, and distribution‑matching distillation (DMD, DMD2) learn few‑step mappings; adversarial variants (Hyper‑SD) preserve detail. Decoupled DMD, combined with VFM‑based adversarial guidance (CLIP, DINOv2) as in SenseFlow, forms the backbone of Mage‑Flow’s 4‑step Turbo generation and editing.

The Mage-Flow Architecture

The Mage‑Flow stack packs native‑resolution latents and a lightweight VAE into a single diffusion pipeline.

High‑resolution diffusion models inherit a hidden latency bottleneck: the VAE that encodes and decodes images can consume as much time as the diffusion transformer itself, especially when scaling to 1K‑2K pixels.

We treat Mage‑VAE as a compact, diffusion‑style encoder/decoder that yields latent tokens, and Native‑Resolution MMDiT as a 4 B‑parameter transformer that attends over those tokens without resampling them to a fixed grid.

How does the Mage‑Flow stack differ from conventional diffusion pipelines that use fixed‑resolution token buckets?

Conventional pipelines first rasterize every image to a common resolution, then split it into uniform patches, which forces padding and wastes compute on empty regions. Mage‑Flow skips this step: each image retains its original pixel grid, Mage‑VAE emits a token per pixel (or patch), and the transformer attends over the exact native layout via 2D‑RoPE. The result is less padding, lower memory, and unchanged spatial fidelity.

Image A tokens: $t_{A,1}, t_{A,2}, t_{A,3}, t_{A,4}$; Image B tokens: $t_{B,1}, t_{B,2}, t_{B,3}$.

Pack sequence for Sample 1: [text 1, text 2, text 3, text 4, $t_{A,1}$…$t_{A,4}$] → length 8.

Pack sequence for Sample 2: [text 1, text 2, text 3, text 4, $t_{B,1}$…$t_{B,3}$] → length 7.

Apply per‑sample 2D‑RoPE: each image token receives a positional pair (row, col) matching its native coordinates (e.g., $t_{A,3}$ gets (2, 1)). Text tokens receive a separate 1‑D position stream.

FlashAttention processes the batch of packed sequences, padding only the final token of the longer sample (1 token) to align lengths.

Because padding is limited to the difference in total token counts, the memory overhead grows with the *maximum* sequence length, not with the sum of all pixel counts, yielding a substantial saving for heterogeneous image sizes.

**Figure 5.** *Mage-Flow* architecture. *Mage-Flow* uses Qwen3-VL to encode text prompts and *Mage-VAE* to map images into compact transformer-ready latents. Images of different resolutions and aspect ratios are flattened into variable-length latent token sequences and packed together with variable-length text tokens in a single batch. The 4B Native-Resolution MMDiT processes these packed sequences with per-sample 2D rotary positional embeddings and FlashAttention variable-length attention, avoiding fixed resolution buckets while preserving each sample’s native spatial layout. The right panel illustrates the MMDiT block, where text and image streams use modality-specific normalization and projection layers but interact through joint self-attention.

By eliminating resolution‑bucket padding and using a lightweight VAE, the Mage‑Flow stack delivers high‑fidelity 1K–2K generation with a fraction of the compute that traditional diffusion stacks require.

Mage-VAE: Efficient Latent Tokenization

Designing a lightweight VAE that inherits FLUX.2’s latent geometry while cutting compute.

High‑capacity VAEs such as FLUX.2‑VAE deliver excellent latent representations but are too heavy for real‑time deployment. Mage‑VAE tackles this bottleneck by swapping the expensive bitrate control of CoD‑Lite codecs for a KL regularizer that pulls the posterior toward an FLUX.2‑derived anchor distribution.

Mage‑VAE is a one‑step encoder‑decoder that mimics the FLUX.2 latent space by regularizing its posterior toward an anchor distribution and by emitting latents already down‑sampled for the diffusion transformer.

Patchify $4\times4$ image → four $2\times2$ patches.

Each patch is projected to a $32$‑dim vector (FLUX.2 style).

Spatial reduction by $8\times$ collapses the four patches into a single spatial location.

Mage‑VAE’s encoder directly maps the original $4\times4$ image to a $1\times1$ latent of $128$ channels, matching the final shape of the FLUX.2 pipeline.

KL regularizer computes $\text{KL}(q_{\phi}(z\mid x)\,\|\,p_{\text{anchor}}(z))$ on this $128$‑dim vector, pulling it toward the anchor Gaussian.

This toy example shows that Mage‑VAE eliminates the intermediate patch‑wise encoding while still producing a latent that lives in the same high‑dimensional anchor space, enabling a single‑step diffusion downstream.

How does Mage‑VAE differ from a conventional VAE that uses a standard normal prior?

Conventional VAEs push $q_{\phi}(z\mid x)$ toward a fixed isotropic Gaussian, which discards the spatial‑channel structure needed by diffusion transformers. Mage‑VAE instead regularizes toward the FLUX.2‑VAE anchor distribution, preserving the 16× down‑sampled, 128‑channel layout that the downstream MMDiT expects.

**Figure 6.** *Mage-VAE* is a lightweight pixel-diffusion-based VAE distilled from the FLUX.2-VAE latent space. FLUX.2-VAE provides the anchor latent distribution used for diffusion pre-training and anchor-latent regularization. The *Mage-VAE* decoder is distilled to reconstruct images from FLUX.2-aligned latents, while the *Mage-VAE* encoder is trained to produce latents compatible with the same anchor space. This design yields an efficient one-step encoder-decoder that preserves a generation-ready latent structure and can be interchanged with FLUX.2-VAE in downstream generators.

**Figure 7.** Inference cost of *Mage-VAE*. Encoding and decoding latency and memory is tested across resolutions on 80GB A100 GPU. SD-VAE, SD-3.5-VAE, and FLUX-VAE are not shown, as their model structures are almost identical to FLUX.2-VAE, causing their inference costs to completely overlap.

Native-Resolution MMDiT Backbone

Native‑resolution packing lets a single update learn from arbitrarily sized images and prompts.

Standard diffusion training forces every step to draw samples from a single resolution bucket, which discretises the native size distribution and makes extreme aspect ratios cumbersome.

Instead of forcing all images in a batch to share a fixed grid, the model packs each image’s latent token stream together with its prompt under a global token budget, letting every update see a mix of native sizes.

Compute cumulative offsets: Sample A starts at index 0, Sample B starts at index 40 (30 + 10).

Concatenate tokens → sequence length = 30 + 10 + 12 + 5 = 57, but we truncate to the 50‑token budget, discarding the last 7 tokens of Sample B’s prompt.

FlashAttention receives the 50‑token tensor together with offset array [0, 40] and processes each sample independently.

During CFG, the unconditional copy of Sample A (30 + 10 tokens) and Sample B (12 + 5 tokens) are appended after the conditional copies, yielding a single 100‑token forward pass that produces both outputs.

Native packing turns resolution diversity into a usable signal: larger images automatically contribute more visual tokens, while shorter prompts free budget for higher‑resolution samples.

How does native‑resolution packing differ from simply padding all sequences to the same length?

Padding adds dummy tokens that still participate in attention, wasting compute and memory. Packing reorders the real tokens and supplies per‑sample offsets so attention never sees padding, keeping the operation strictly proportional to the true token count.

When source and target images share a batch, each visual token gets a third coordinate (frame) so the attention mechanism can tell “this token belongs to the source” versus “this token belongs to the target”.

Why can’t we reuse the ordinary 2‑D RoPE for source‑target editing?

Without a frame dimension the model cannot distinguish which tokens belong to the source image and which belong to the target, so the attention updates would mix them indiscriminately, breaking the spatial correspondence needed for precise editing.

Beyond the model itself, training speed hinges on fusing the many memory‑bound operator chains inside VAE, the frozen text encoder, and the 4 B NR‑MMDiT blocks.

**Figure 1** Qualitative showcase of *Mage-Flow*. Native-resolution samples spanning editorial design, photorealism, food photography, bilingual text rendering, portraits, and stylized art. *Mage-Flow* supports flexible generation with height and width from 512 to 2048, including extreme 4:1 aspect ratios such as 512 × 2048 and 2048 × 512. The examples demonstrate coherent layouts, fine visual details, and legible English and Chinese text across diverse resolutions and aspect ratios.

Finding 2: Removing the single‑bucket restriction lets each update mix heterogeneous sizes, turning resolution diversity into a useful training signal while enabling packed CFG inference that speeds up guidance by up to 1.15×.

Finding 3: Co‑designing the model (lightweight Mage‑VAE, native‑resolution NR‑MMDiT) with system‑level kernel fusion cuts per‑step time by 2.48× and raises MFU to 29 %—the dominant efficiency gains stem from fusing the 4 B diffusion backbone.

Data Collection and Curation

How the Mage‑Flow stack gathers and curates data for generation and editing.

Training high‑fidelity diffusion models demands massive, clean image‑text data, yet raw web collections are noisy, duplicated, and unsafe. Mage‑Flow mitigates this by building two complementary pipelines that filter, deduplicate, and balance data for generation and instruction‑based editing.

The model receives a source image, a natural‑language edit instruction, and a target image, learning to apply precise edits conditioned on language rather than learning a generic image‑to‑image mapping.

Encode the source image into a latent representation.

Encode the instruction “replace the apple with a green pear” into token embeddings.

Cross‑attention aligns the word “apple” with the red region and “pear” with the intended replacement region.

At each diffusion step, the model updates only the aligned region, gradually turning the red apple into a green pear while leaving the background unchanged.

After the final timestep, decode the latent to obtain the edited image matching the target.

Because the model learns to localize edits via language, it can perform precise modifications without an explicit mask, unlike conventional inpainting that requires a binary mask.

The generation pipeline starts from ~10 B raw image–text pairs, applies staged sample‑level filters, deduplicates near‑duplicate visuals, captions images at multiple granularities, and finally augments long‑tail concepts with targeted synthesis.

The editing pipeline assembles ~90 M source‑instruction‑target triples, filters them with three independent VLM experts, tags each triple with one of 19 edit categories, and balances the taxonomy to ensure diverse edit coverage.

**Figure 9. Data composition for generation and editing pre-training.** (a) Concept distribution of the curated generation corpus after merging filtered web data with targeted supplemental data. (b) Final composition of the editing pre-training mixture after adjusting sampling rates across constituent editing datasets.

**Figure 10** Editing data filtering pipeline. Raw editing triples are collected from open-source editing datasets and in-house synthesis. Each triple is evaluated by three Qwen3.5-9B experts with different system prompts and partially overlapping rubrics. Each expert analyzes the source image, target image, and edit instruction, and its reasoning output is parsed into a pass/fail decision. A triple is retained only if it receives a majority vote from the three experts. The surviving 45M triples are then assigned to a manually defined edit-type taxonomy and reweighted across categories to form the final editing training set.

Training Recipe

Unified training pipeline that scales from low‑resolution pre‑training to fast 4‑step Turbo models.

The training pipeline proceeds in three logical phases: progressive pre‑training, RL‑guided post‑training, and few‑step distillation. Figure 11 visualizes this flow across the generation and editing branches.

**Figure 11** Overview of the *Mage-Flow* training pipeline. The shared backbone is first trained with progressive text-to-image pre-training, summarized as low-resolution pre-training followed by native-resolution pre-training, and then supervised fine-tuning produces the base generation checkpoint *Mage-Flow-Base*. From this checkpoint, the text-to-image branch applies Diffusion-NFT post-training to obtain *Mage-Flow* and 4-step distillation to obtain *Mage-Flow-Turbo*. The editing branch forks from *Mage-Flow-Base*: continued source-conditioned editing training produces *Mage-Flow-Edit-Base*, mixed editing-and-generation post-training produces *Mage-Flow-Edit*, and 4-step distillation produces *Mage-Flow-Edit-Turbo*.

Start cheap at low resolution to learn coarse semantics, then gradually increase resolution so the model can capture fine details without exploding compute.

Stage 1 processes Image A: 1 M pixels total, gradient update cheap.

Stage 2 processes Image B: 2 ×  more pixels, re‑uses Stage 1 weights, cost roughly doubles.

Stage 3 processes Image C: 4 ×  more pixels than Stage 2, but the model already knows coarse shapes, so only fine‑detail gradients are large.

After Stage 3, the model can render 1024‑pixel text without having spent 1.2 B × 1024² pixel‑updates.

Gradually increasing resolution amortizes the expensive high‑resolution learning over earlier cheap stages, yielding high‑fidelity output with far less total compute.

Why not train the full 1024‑pixel model from the start?

Training at full resolution would require orders of magnitude more FLOPs per step, making the 1.2 B‑sample curriculum infeasible; the progressive schedule concentrates compute where it matters most—first on semantics, then on detail.

Fine‑tune the base model using reinforcement‑learning rewards that are specific to the type of capability being improved (text rendering, aesthetics, or semantics).

Why not simply fine‑tune with a single scalar reward?

Different capabilities produce rewards on incompatible scales; a single scalar would bias the model toward the evaluator with the largest numeric range, harming the other abilities.

Separate the classifier‑free‑guidance signal from the distribution‑matching signal, give each its own noise schedule, and add a lightweight adversarial term to preserve perceptual quality in only four denoising steps.

Compute $\Delta$₍CA₎ = (w − 1) × (0.8 − 0.5) = 0.2 × 0.3 = 0.06.

Compute $\Delta$₍DM₎ = 0.8 − 0.6 = 0.2.

Assume the adversarial gradient ∇L₍GAN₎ = 0.05 and $\lambda$₍GAN₎ = 1.0.

Combine: ∇L₍Total₎ = 0.06 + 0.2 + 0.05 = 0.31, which is the update applied to the student.

Even with tiny numeric values, the CA term nudges the student toward the conditional teacher, while the DM term enforces overall distribution alignment; the adversarial term adds a perceptual safety net.

How does Decoupled DMD differ from standard distribution‑matching distillation?

Standard distillation mixes the guidance and distribution terms into a single loss and shares one noise schedule, which becomes unstable when the diffusion trajectory is aggressively shortened. Decoupled DMD gives each term its own schedule, keeping both signals reliable in the 4‑step regime.

**Figure 13** Few-step distillation framework. The student maps a prepared noised latent $z_t$ and condition $c$ to a predicted clean sample. Decoupled DMD re-noises the student output at two independent noise levels: the classifier-free augmentation branch queries the frozen teacher with and without the condition to form $\Delta_{CA}$, while the distribution-matching branch contrasts the teacher with a trainable fake-score model to form $\Delta_{DM}$. In parallel, generated and real images are encoded by frozen vision foundation models such as DINOv2 and CLIP, and a lightweight feature discriminator provides the adversarial gradient $\nabla\mathcal{L}_{GAN}$. The Decoupled DMD and adversarial gradients are combined to train the 4-step Turbo student.

Experimental Setup

Experimental setup and key findings on distillation and adversarial guidance.

We evaluate Mage‑Flow and its Turbo variants on eight text‑to‑image benchmarks and three instruction‑based editing benchmarks, all at native 1024×1024 resolution. Unless noted otherwise, Mage‑Flow‑Base runs 30 denoising steps, Mage‑Flow uses 20 steps, and the Turbo models run only 4 steps. Adversarial perceptual guidance is applied with guidance scale $w\!=\!7.5$ and adversarial weight $\lambda_{\text{GAN}}\!=\!0.13$ at noise levels $\tau_{\text{ca}}$ and $\tau_{\text{dm}}$.

Adversarial perceptual guidance most benefits generation and text‑editing quality when the diffusion trajectory is compressed to four steps.

Finding 4 reports consistent gains on highly compressed generation and editing pipelines, while gains on broader editing benchmarks vary.

**Table 7. Effect of mixing generation data during editing training.** We compare editing models trained with and without generation data. The effect is modest for the full-step editor; for the 4-step Turbo editor, generation data improves ImgEdit, while GEdit changes are mixed.

**Figure 14.** Composition of the distillation sets. The inner rings show the main categories, and the outer rings show the corresponding sub-categories. (a) The generation distillation set covers diverse content categories such as people, scenes, products, design, food, and synthetic data. (b) The editing distillation set covers various edit-type categories such as object editing, scene and viewpoint changes, text editing, appearance modification, control-map editing, and complex edits.

**Table 8.** Overall comparison. Table 8 summarizes the main text-to-image results across prompt following, fine-grained generation, and text rendering benchmarks.

Quantitative Performance

Mage‑Flow trims compute with a 4B native‑resolution diffusion transformer and a compact VAE, preserving quality.

The experimental setup (Section 7) details training data, evaluation protocols, and hardware; here we focus on the outcomes.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers