PaperlayerView exact revision on arXiv

Source-linked AI summary

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

arXiv:2607.19064v2cs.CVcs.AIcs.LGcs.MMeess.IV

TL;DR

Strong visual generators are costly to train, fine-tune, and deploy, while competitive open models often rely on large backbones. Mage-Flow offers a compact 4B stack for generation and editing; its models are competitive across standard benchmarks, and its 4-step editing Turbo runs in 1.02s.

  • Problem

    The paper addresses how to make visual generation practical to study, modify, and deploy under realistic compute budgets, given the scale of competitive open-source models.

    Source: §1 · pp. 1, 5

  • Method

    Mage-Flow combines a lightweight high-fidelity tokenizer with a 4B native-resolution diffusion transformer for text-to-image generation and instruction-based editing.

    Source: §1 · p. 5

  • Results

    1.02s is the inference time of 4-step Mage-Flow-Edit-Turbo; across standard generation and editing benchmarks, the models achieve competitive performance against substantially larger open-source systems.

    Source: §1 · p. 6, §7 · p. 32

  • Takeaways & Limitations

    The results support a 4B-scale stack as a practical, research-friendly foundation for image generation, controllable editing, post-training alignment, and vertical-domain applications.

    Source: §7 · p. 32

    The four-step model's gains on general editing benchmarks are benchmark-dependent rather than uniform, with broad editing robustness limited on ImgEdit and no uniform improvement on GEdit.

    Source: §5.3 · pp. 24–25

Abstract

from arXiv · show

Large-scale visual generators are increasingly capable but costly to train, fine-tune, and deploy. We introduce Mage-Flow, a compact 4B-scale generative stack for efficient text-to-image generation and instruction-based image editing. The stack is built from two co-designed components: Mage-VAE, a lightweight high-fidelity latent tokenizer, and a Native-Resolution Multimodal Diffusion Transformer trained with rectified flow matching. Mage-VAE uses one-step diffusion-style encoding and decoding with anchor-latent regularization, preserving the reconstruction quality of strong public VAEs while reducing tokenization cost by more than an order of magnitude. Together with native-resolution packing and stack-level CUDA kernel fusion, the stack supports flexible-resolution training and improves end-to-end training throughput by about $2.5\times$. Built on this foundation, we develop a complete model family with Base, RL-aligned, and Turbo variants for both generation and editing. Diffusion-NFT improves prompt following, text rendering, aesthetic quality, and editing fidelity, while few-step distillation with adversarial perceptual guidance produces 4-step Turbo models for low-latency inference. Despite its compact scale, Mage-Flow and Mage-Flow-Edit achieves competitive performance across standard generation and editing benchmarks. More importantly, the Turbo variants make high-resolution generation and editing practical for interactive use: at $1024^2$ resolution on a single NVIDIA A100 GPU, Mage-Flow-Turbo generates an image in 0.59s, and Mage-Flow-Edit-Turbo edits an image in 1.02s, while maintaining a small memory footprint. These results show that careful tokenizer--backbone--system co-design can deliver strong high-resolution generation and editing within an efficient 4B model family.

1 Introduction

To address practical compute barriers to visual generation, the paper presents Mage-Flow, a compact 4B stack for text-to-image generation and instruction-based editing. It co-designs a lightweight tokenizer, native-resolution backbone, and fused-kernel infrastructure for efficient flexible-resolution use.

Introduction

  • Mage-VAE matches FLUX.2-VAE reconstruction fidelity while using approximately 12× fewer encoding MACs and 22× fewer decoding MACs per pixel.Its one-step diffusion-style encoder and decoder use anchor-latent KL regularization toward a strong public VAE’s latent distribution, reducing tokenization overhead at high resolution.Source: §1 · pp. 5–6
  • Native-resolution packing combines variable-length image and text sequences across arbitrary resolutions and aspect ratios, removing bucket restrictions and enabling packed classifier-free-guidance inference.The 4B NR-MMDiT backbone is trained with rectified flow matching and uses per-sample 2D rotary embeddings.Source: §1 · pp. 5, 7
  • Fused CUDA kernels increase MFU from approximately 14% to 29% and deliver about a 2.5× end-to-end training speedup.They fuse memory-bound operator chains across Mage-VAE, the Qwen3-VL text encoder, and NR-MMDiT, reducing activation-memory traffic and kernel-launch overhead.Source: §1 · p. 5
  • Mage-Flow’s 4B family achieves competitive or superior benchmark performance against much larger open-source systems, with 4-step Turbo inference at 0.59s for generation and 1.02s for editing.At 1024^2 resolution on one NVIDIA A100, base Mage-Flow generates in 4.37s and Mage-Flow-Edit runs in 10.55s; peak GPU memory is around 18–20 GB, the lowest among compared models.Source: Figure 4 · p. 6, §1 · p. 6
  • The shared stack yields Base, aligned, and Turbo variants for both generation and editing, with editing training mixing generation data to preserve the open-ended generative prior.Diffusion-NFT [14] aligns Mage-Flow for prompt following, aesthetics, bilingual text rendering, and preference alignment; Mage-Flow-Turbo uses 4-step Decoupled-DMD guidance [15] with adversarial perceptual guidance [16].Source: §1 · pp. 6–7

2 Related Work

Related work spans large diffusion-transformer generators and instruction editors, unified multimodal models, tokenization, post-training alignment, and few-step distillation. Mage-Flow takes a complementary compact 4B MMDiT approach, pairing specialized generation and editing with efficient tokenization, Diffusion-NFT, and VFM-guided distillation.

2.1 Image Generation Foundation Models

  • Mage-Flow contrasts with generation systems often sized 6B–80B by co-designing a compact 4B generator with a lightweight VAE for efficient, extensible research.Prior work progressed from latent diffusion, including Stable Diffusion [17] and SDXL [18], to diffusion transformers such as SD3 [20] and SANA [21].Source: §2.1 · p. 7
  • Mage-Flow-Edit provides a compact 4B instruction-based editor, while recent MMDiT editors often rely on substantially larger backbones.Editing progressed from per-image inversion and attention manipulation [24] [25] to instruction-tuned and larger-scale editors such as InstructPix2Pix [28] and OmniGen [32].Source: §2.1 · p. 7
  • Unlike unified models such as Janus-Pro [42], Transfusion [43], BAGEL [44], and Emu3.5 [45], Mage-Flow specializes a 4B MMDiT stack for generation and editing.The paper treats image understanding as an orthogonal capability that can be incorporated through external encoders or downstream systems.Source: §2.1 · pp. 7–8

2.2 Tokenization for Image Generation

  • Mage-VAE addresses costly high-resolution VAE encoding and decoding with a one-step pixel-diffusion encoder–decoder and anchor-latent KL regularization.Its fully convolutional decoder uses compression-oriented pretraining and single-step distillation; the resulting generation-ready latent space has 16× spatial reduction and 128 channels.Source: §2.2 · p. 8

2.3 Post-training for Image Generation

  • Mage-Flow and Mage-Flow-Edit use Diffusion-NFT after supervised fine-tuning on curated capability-targeted data to improve prompt adherence, text rendering, compositional control, and editing fidelity.Diffusion-NFT performs negative-aware fine-tuning on the forward process without likelihood estimation or a specific solver.Source: §2.3 · pp. 8–9

2.4 Distillation for Image Generation

  • Mage-Flow and Mage-Flow-Edit combine Decoupled DMD with VFM-based adversarial guidance to produce 4-step Turbo variants for generation and editing.The cited distillation approaches include progressive distillation [60], consistency models [61] [62], and distribution matching methods [64] [65]; SenseFlow [16] introduces VFM-based adversarial guidance for flow generators.Source: §2.4 · pp. 9–10

3 Mage-Flow Stack · 3.1 Mage-VAE: Lightweight Image Tokenizer

Mage-Flow pairs a 4B native-resolution diffusion transformer with Mage-VAE, a lightweight tokenizer designed to preserve the generation-ready latent structure of a strong public VAE. Mage-VAE uses one-step diffusion-style encoding and decoding with anchor-latent regularization to reduce inference cost while maintaining reconstruction fidelity and downstream compatibility.

3 Mage-Flow Stack

  • The Mage-Flow stack combines Mage-VAE with a 4B Native-Resolution Multimodal Diffusion Transformer trained using rectified flow matching.The transformer models packed latent sequences with flexible resolutions and aspect ratios, supported by infrastructure for native-resolution 4B-scale learning.Source: §3 · p. 10

3.1 Mage-VAE: Lightweight Image Tokenizer

  • Mage-VAE targets high-resolution VAE latency and memory bottlenecks with lightweight one-step diffusion-style encoding and decoding, regularized toward a strong public VAE’s latent space.For 4-step 1K-resolution generation with FLUX.2-Klein-4B, VAE decoding accounts for 14% of total generation time.Source: §3.1 · p. 10

3.1.1 Architecture

  • Mage-VAE uses a fully convolutional one-step pixel-diffusion decoder and a symmetric one-step encoder, avoiding global attention and keeping decoding cost nearly linear in image resolution.The decoder reconstructs RGB pixels directly from latents; the encoder uses patch embedding followed by stacked convolutional diffusion blocks.Source: §3.1.1 · p. 11

3.1.2 KL regularization with an anchor latent distribution

  • Mage-VAE replaces standard-Gaussian prior matching with KL regularization toward patchified FLUX.2-VAE anchor latents, producing Transformer-ready latents with 128 channels and 16× spatial downsampling.It internalizes the usual 2× patchification of FLUX.2-VAE latents, which have 32 channels and 8× spatial reduction.Source: §3.1.2 · p. 12

3.1.3 Training

  • Training proceeds through separate multi-step encoder and decoder pre-training, one-step decoder distillation, and joint one-step encoder-decoder fine-tuning with anchor-latent KL regularization.Decoder distillation combines reconstruction loss, DMD loss with a compression-oriented diffusion teacher, and DINOv2-projected GAN loss.Source: §3.1.3 · p. 12

3.1.4 Evaluation

  • Mage-VAE cuts FLUX.2-VAE computational complexity by 12.3× for encoding and 22.3× for decoding, with lower latency and memory across tested resolutions, including 4096×4096.At 4096×4096, baseline VAEs can become extremely slow or run out of memory, while Mage-VAE remains efficient for encoding and decoding.Source: Table 1 · p. 11, Figure 7 · p. 11, §3.1.4 · p. 13
  • Swapping Mage-VAE and FLUX.2-VAE while holding generation or editing backbones fixed yields similar or comparable benchmark performance, supporting downstream latent compatibility.This holds both for Mage-Flow Turbo models using FLUX.2-VAE and FLUX.2-Klein-4B using Mage-VAE.Source: Table 2 · p. 12, §3.1.4 · p. 13

3.2 Native-Resolution MMDiT

Mage-Flow uses a 4B Native-Resolution MMDiT with packed variable-length text and image tokens to train across native resolutions and support efficient inference. The same rectified-flow backbone handles generation and editing, with frame-aware positional embeddings distinguishing source and target images for editing.

3.2.1 Model Architecture and Native Packing

  • The 4B backbone follows MMDiT: concatenated text and image tokens use joint self-attention, with modality-specific normalization and projection layers.Cross-modal interaction occurs through self-attention over the combined sequence.Source: §3.2.1 · pp. 13–14
  • Native-resolution packing removes single-bucket training, mixes heterogeneous image sizes and aspect ratios per update, and supports flexible output sizes from 512 to 2048 per dimension.Variable-length text and image sequences are packed under a fixed token budget, avoiding text padding while retaining each image’s native latent grid.Source: §3.2.1 · p. 14
  • Packed classifier-free guidance accelerates inference by 1.09×–1.15× across Mage-Flow and Mage-Edit variants while preserving the original denoising trajectory.Conditional and unconditional branches are evaluated in one forward pass rather than separately.Source: §3.2.1 · p. 14

3.2.2 Conditioning and rectified-flow training

  • Generation and editing share the rectified-flow objective and NR-MMDiT backbone, using text prompts for generation and instruction-plus-source-image conditioning for editing.Qwen3-VL supplies conditioning embeddings; editing concatenates these with source-image latents and noisy target latents.Source: §3.2.2 · pp. 14–15
  • Frame-aware 3D RoPE assigns visual tokens (h, w, f) positions to distinguish source and target roles while preserving spatial correspondence; loss is computed only on target tokens.This design lets Mage-Flow-Edit initialize directly from Mage-Flow without separate editing modules.Source: §3.2.2 · p. 15

3.3 Training Infrastructure

Mage-Flow’s training efficiency comes from co-designing lightweight tokenization with stack-level kernel fusion. Fusion reduces memory movement and kernel-launch overhead across repeatedly executed modules, improving overall training throughput and memory use.

3.3 Training Infrastructure

  • Fusion targets memory-bound operator chains in Mage-VAE, the frozen Qwen3-VL encoder, and 4B NR-MMDiT blocks, keeping intermediates on-chip and reducing activation movement and kernel launches.The chains include normalization, adaptive modulation, RoPE, gating, activation, and residual addition; Mage-VAE fusion combines normalization, activation, and residual operations, while Transformer-side fusion combines adaptive normalization, rotary embeddings, and gated residual updates.Source: §3.3 · p. 15
  • 2.48× faster per-step training, 29.28% MFU, and 141.44 GB peak per-GPU memory are achieved by the full system, versus 13.88% MFU and 175.45 GB memory.The full system's results show the value of stack-level kernel optimization alongside an efficient tokenizer and architecture.Source: §3.3 · p. 16
  • Replacing FLUX.2-VAE with Mage-VAE reduces per-step time from 1.9285 s to 1.3647 s, a 1.41× speedup.Fusing Mage-VAE and Qwen3-VL gives moderate additional gains; fusing the repeated NR-MMDiT blocks yields the largest improvement because the 4B backbone dominates the training step.Source: §3.3 · p. 16
  • Efficient native-resolution generation requires co-designing lightweight tokenization, which reduces arithmetic cost, with kernel fusion that removes overhead from repeatedly executed VAE, text-encoder, and MMDiT blocks.Because the blocks run many times per forward and backward pass, local operator fusion produces a substantial stack-level throughput gain.Source: §3.3 · p. 16

4 Data Collection and Curation

Mage-Flow curates separate generation and editing datasets to improve data quality, diversity, semantic alignment, safety, and capability coverage. The pipelines combine filtering and balancing with targeted captioning or synthesis.

4.1 Text-to-Image Data Collection and Filtering

  • From roughly 10B raw image–text pairs, four curation stages retain approximately 1.3B high-quality pairs for stage-wise pre-training.The stages are sample-level filtering, cross-sample deduplication, multi-granularity captioning, and concept-aware synthesis.Source: §4.1 · p. 17
  • Generation filtering removes corrupted, unsuitable, low-quality, or unsafe images, with thresholds tightened from 256^2 through 512^2 and 1024^2 to SFT.Filtering checks file properties and image content; early stages preserve broad coverage, while later stages emphasize cleaner, higher-quality data.Source: §4.1 · p. 17, Table 5 · p. 17
  • SSCD descriptors and FAISS nearest-neighbor search remove within- and cross-dataset duplicates above 0.9 cosine similarity and reject matches to a held-out benchmark index.Within-dataset clusters retain the highest-quality representative, with very large clusters capped to limit repeated web templates.Source: §4.1 · p. 17
  • Multi-granularity captions cover concepts, entities, composition, and photographic details, while targeted synthesis adds long-text rendering, rare concepts, and compositional cases.Concept-aware sampling reduces dominance by frequent objects, natural scenes, and common photorealistic styles; filtering and reweighting become stronger across training stages.Source: §4.1 · p. 18, Figure 9 · p. 18

4.2 Image-Editing Data Collection and Filtering

  • 45M editing triples survive majority-vote filtering by three Qwen3.5-9B experts, from a raw pool of roughly 90M triples.The retained pool comprises roughly 20M open-source and 25M synthesized triples; each expert checks edit execution, preservation of unrelated regions, and visual plausibility.Source: §4.2 · pp. 18–19, Figure 10 · p. 19
  • In-house editing synthesis supplies about 10M low-level processing triples and 30M semantic-editing triples, spanning operations such as replacement, removal, style transfer, text editing, and restoration.Type-specific pipelines combine generation, inpainting, segmentation, image processing, and template-based instruction generation to supplement sparse or unreliable open-source categories.Source: §4.2 · p. 19
  • A manually defined 19-category taxonomy maps retained editing data to edit types, and sampling rates are adjusted to balance frequent and rare operations.Annotation units follow each source’s organization, and the resulting mixture aims to preserve broad coverage across the taxonomy.Source: Figure 9 · p. 18, §4.2 · pp. 19–20

5 Training

Mage-Flow trains generation and editing in a shared Mage-VAE latent space with the same 4B Native-Resolution MMDiT and rectified-flow objective, then uses progressive pre-training and fine-tuning, Diffusion-NFT alignment, and four-step distillation. The recipe adapts conditioning, data, and rewards to each task.

5.1 Pre-training and Supervised Fine-tuning

  • Generation training progresses from fixed-resolution pre-training through native-aspect-ratio stages to 1024-pixel supervised fine-tuning, producing Mage-Flow-Base.Pre-training uses 1.2B image–text pairs at 256 × 256, then 600M at 512 pixels and 300M at 1024 pixels; fine-tuning uses a curated 150M high-quality subset.Source: §5.1 · p. 20
  • Two-stage editing adaptation of Mage-Flow-Base mixes editing triples with generation pairs and yields Mage-Flow-Edit-Base while preserving open-ended synthesis ability.The stages use 35M editing triples and 35M generation pairs, followed by 20M triples and 10M pairs; multi-image edits comprise at most 0.5% of editing data. The resulting model supports semantic edits, appearance transformations, restoration, and structure-aware outputs through a unified one-to-many interface.Source: §5.1 · p. 21

5.2 Diffusion-NFT Post-training

  • Diffusion-NFT aligns Mage-Flow using capability-tagged prompts and rubric-specific rewards, with a two-stage schedule that shifts emphasis toward harder text-rendering prompts.The approximately 20K-prompt pool covers text rendering, aesthetics, and semantic understanding; training runs 140 steps at a 1:1:1 mixture, then 60 steps at 2:4:1, retaining all three groups.Source: §5.2 · pp. 21–22

5.3 Few-step Distillation

  • Adversarial perceptual guidance benefits generation and text editing under four-step compression, but its effects on general editing benchmarks are mixed rather than uniform.The reported comparison says it slightly improves ImgEdit but does not uniformly improve GEdit.Source: Table 6 · p. 23, §5.3 · p. 24

5.2 Diffusion-NFT Post-training

  • Editing Diffusion-NFT interleaves editing and generation updates at 4:1, using task-specific rewards to produce Mage-Flow-Edit while preserving generation capabilities.The editing stream draws from approximately 30K prompts across edit tasks and uses RationalRewards to assess instruction adherence, source-content preservation, plausibility, and rendered-text quality; post-training runs for 300 optimizer steps.Source: §5.2 · p. 23

5.3 Few-step Distillation

  • Adding generation data to four-step editing distillation improves ImgEdit from 4.20 to 4.38, while the two GEdit splits remain comparable and do not improve uniformly.The full-step editor changes only marginally; the reported benefit is broad category-wise editing robustness under four-step compression, not a uniform gain across editing metrics.Source: Table 7 · p. 24, §5.3 · p. 25
  • Four-step generation distillation combines Decoupled DMD with adversarial perceptual guidance to retain the teacher’s visual quality and prompt following with a shorter trajectory.The student is distilled from the frozen Mage-Flow teacher on roughly 200K curated prompt–image pairs spanning six content categories.Source: §5.3 · pp. 24–25
  • Editing distillation trains Mage-Flow-Edit-Turbo on editing and generation data, using edited targets for the discriminator’s real branch to preserve edit following and generation ability.The training mixture is 3:1 editing to generation data, with roughly 250K editing samples across six categories; both Turbo variants inherit teacher quality and alignment with four inference steps.Source: §5.3 · p. 25

6 Experiments

Mage-Flow’s 4B-scale models deliver competitive text-to-image generation and instruction-based editing, with Turbo variants retaining much of that performance at four denoising steps. The experiments also assess the models across benchmark suites and illustrate generation and editing at varied resolutions and task types.

6.1 Experimental Setup

  • Generation benchmarks use native 1024^2 resolution; Base models are evaluated with 30 denoising steps, Mage-Flow with 20, and Turbo with four unless otherwise specified.Editing evaluation covers ImgEdit-Bench, GEdit-Bench, and TextEdit-Bench; Base and Mage-Flow-Edit use 30 steps, while Edit-Turbo uses four.Source: §6.1 · p. 26, Instruction-based image editing · p. 26

6.2.1 Text-to-image generation

  • Mage-Flow achieves the strongest overall GenEval score and remains competitive on DPG-Bench and TIIF-Bench; Mage-Flow-Turbo retains competitive performance with four sampling steps.The overall comparison also reports strong performance against larger open-source specialist models across fine-grained generation and text-rendering benchmarks.Source: Table 9 · p. 27, §6.2.1 · pp. 27–28, Table 10 · p. 27
  • Mage-Flow attains near-best open-source multi-region text-rendering accuracy on CVTG-2K and strong English long-text rendering, while Chinese LongText remains an area for improvement.Mage-Flow variants also maintain strong alignment and text-rendering ability on both OneIG language splits, supporting fine-grained generation.Source: Table 11 · p. 28, Table 12 · p. 28, §6.2.1 · p. 29

6.2.2 Instruction-based image editing

  • Mage-Flow-Edit delivers competitive editing quality at 4B parameters, including the best open-source GEdit-Bench Chinese score and a near-best English score; Edit-Turbo preserves most performance in four steps.The models remain competitive on ImgEdit-Bench and TextEdit-Bench against substantially larger open-source systems.Source: Table 13 · p. 29, §6.2.2 · p. 29
  • On GEdit-Bench, Mage-Flow-Edit scores 8.127 on English and 8.123 on Chinese, while Edit-Turbo scores 8.271 and 8.264, respectively, using four denoising steps.Both variants outperform Qwen-Image-Edit-2511 [5], FireRed-Image-Edit-1.0 [8], and LongCat-Image-Edit [22] under the same protocol.Source: Table 15 · p. 30, §6.2.2 · p. 31
  • Mage-Flow-Edit reaches 14.142 on TextEdit-Bench Synthetic and 16.255 on Real-World, while precise layout preservation and complex text replacement remain improvement areas.TextEdit-Bench isolates text-editing ability; the real-world subset uses non-reference evaluation, while the synthetic subset uses reference-based evaluation.Source: Table 16 · p. 31, §6.2.2 · p. 31

6.3.1 Text-to-image visualizations

  • Generation visualizations show photorealistic and imaginative scenes, portraits, and legible English and Chinese text on packaging, posters, magazine covers, and signage.The galleries span native resolutions and aspect ratios.Source: §6.3.1 · p. 31

6.3.2 Instruction-based editing visualizations

  • Editing visualizations cover content, scene, appearance, human-centered, creative, and low-level-control edits, including paired degradation–restoration tasks and bidirectional edit types.The galleries pair source images with edited outputs and include composite multi-task examples.Source: §6.3.2 · pp. 31–32

7 Conclusion

Mage-Flow is a compact, efficient 4B-scale stack for text-to-image generation and instruction-based editing, built through joint tokenizer, native-resolution model, and systems design. Its model family combines alignment and distillation methods to deliver competitive generation and editing performance.

7 Conclusion

  • Mage-Flow combines Mage-VAE, native-resolution diffusion Transformer modeling, and fused CUDA kernels to reduce tokenization cost and training overhead while supporting flexible resolutions and aspect ratios.Together, these designs make large-scale native-resolution training practical under a fixed compute budget.Source: §7 · p. 32
  • Across standard generation and editing benchmarks, Mage-Flow achieves competitive performance against substantially larger open-source systems while maintaining low memory usage and fast inference.The reported scope includes both generation and editing settings.Source: §7 · p. 32
  • The family includes Base, RL-aligned, and Turbo variants for generation and editing; Diffusion-NFT improves alignment, while Decoupled-DMD with adversarial perceptual guidance produces 4-step Turbo models.The alignment targets prompt following, text rendering, aesthetic quality, and editing preferences; the Turbo models target low-latency inference.Source: §7 · p. 32
  • The conclusion argues that strong visual generation need not require tens-of-billions-parameter backbones and identifies multi-image editing, multilingual long-text rendering, and agentic creation workflows as future directions.The authors position the 4B-scale stack as a practical, research-friendly foundation for generation, controllable editing, post-training alignment, and vertical-domain applications.Source: §7 · p. 32

A Detailed Mage-VAE Training Process

Mage-VAE is trained in three stages: multi-step flow-matching pre-training, one-step decoder distillation, and joint one-step VAE optimization. The process yields one-step encoding and decoding with high-fidelity reconstruction and a regularized latent space compatible with the pretrained anchor space.

Stage I: Multi-step flow-matching pre-training

  • Encoder and decoder are pre-trained as multi-step diffusion models with flow matching in the X-prediction parameterization.The decoder learns image reconstruction from anchor latents, while the encoder learns to produce patchified FLUX.2-VAE anchor latents from images.Source: A Detailed Mage-VAE Training Process · p. 44

Stage II: One-step decoder distillation

  • The one-step decoder is distilled from the multi-step decoder using pixel and perceptual reconstruction objectives, a DINOv2-projected GAN loss [76], and DMD [64].DMD transfers the compression-oriented diffusion teacher’s distribution prior [75]; the teacher and DINOv2 feature extractor remain frozen during distillation.Source: A Detailed Mage-VAE Training Process · p. 44

Stage III: Joint one-step VAE optimization

  • After encoder warm-up, the decoder is unfrozen and the complete one-step encoder–decoder pipeline is fine-tuned end to end.During warm-up, gradients pass through the fixed decoder to train the encoder without perturbing the decoder’s reconstruction ability.Source: A Detailed Mage-VAE Training Process · p. 45
  • Joint one-step optimization uses anchor-latent KL regularization to limit latent-space drift and yields a compact, regularized Mage-VAE for Mage-Flow training.The encoder is first warm-started while the distilled decoder is fixed, then both are jointly fine-tuned so they can co-adapt while retaining anchor-space compatibility.Source: A Detailed Mage-VAE Training Process · p. 45

B VAE Latency Across Resolutions

Mage-VAE has the lowest reported latency at every tested resolution, with its advantage widening sharply at high resolution. Its single-step, memory-efficient design supports practical native-resolution training and high-resolution generation.

B VAE Latency Across Resolutions

  • Latency comparisons span 512^2–4096^2 on an 80GB NVIDIA A100 GPU under bf16 precision; Mage-VAE’s single-step, memory-efficient design drives near-flat scaling.The passage connects this scaling to the practicality of native-resolution packing and high-resolution generation under a fixed compute budget.Source: B VAE Latency Across Resolutions · pp. 45–46
  • At 1024^2, Mage-VAE encodes / decodes in 19.1/33.8 ms, versus 83.8/153.3 ms for FLUX.2-VAE and 88.3/137.9 ms for Qwen-Image-VAE.At 512^2, the methods are broadly comparable, while the gap at 2048^2 reaches roughly an order of magnitude.Source: B VAE Latency Across Resolutions · p. 45
  • At 4096^2, Mage-VAE encodes/decodes in 149.6/375.0 ms, versus 47.3/183.9 s for FLUX.2-VAE and 106.7/192.0 s for Qwen-Image-VAE; HunyuanVideo-VAE and HunyuanImage-3.0-VAE run out of memory.Table 17 reports encoding / decoding latency in milliseconds, measured on an 80GB NVIDIA A100 GPU under bf16 precision; the prose gives the cited 4096^2 values in seconds for the two comparison VAEs.Source: B VAE Latency Across Resolutions · p. 45, Table 17 · p. 46

C Recaptioning System Prompt

The recaptioning system uses Qwen3-VL-32B-Instruct [13] to generate four mutually consistent caption layers in one pass, guided by an objective, structured prompt and returned as JSON.

C Recaptioning System Prompt

  • The single prompt produces entity, phrase, composition, and photographic descriptions in one pass as a JSON object.It fixes objectivity and consistency principles, a step-by-step procedure, the four-layer output schema, and a style guide.Source: C Recaptioning System Prompt · p. 46
  • Descriptions must remain visually objective, physically and spatially consistent, organized from overall scene to foreground, and written in present tense with specific language.The prompt directs captions to use concrete visual qualities rather than subjective judgments and precise terms for quantity, size, shape, color, and texture.Source: C Recaptioning System Prompt · p. 46
  • The procedure identifies the main subject and its attributes, describes the background, details composition and style, extracts subject phrases, and synthesizes a summary.Source: C Recaptioning System Prompt · p. 46
  • The JSON schema separates a short main-entity label, concise attribute phrases, a single-sentence contextual composition summary, and a detailed photographic description.The photographic layer includes image category, style, recognized text, entities, lighting and color, camera composition, and photography techniques; style is inferred consistently across composition and photographic layers, defaulting to photorealistic photography when unclear.Source: C Recaptioning System Prompt · p. 46

D Reward Models

The appendix describes four capability-routed reward evaluators for Diffusion-NFT post-training: three score generation and one scores instruction-based editing. Each returns a scalar reward in [0, 1], and rewards are not mixed within a prompt.

Reward evaluators

  • Four evaluators cover text rendering, aesthetic quality, and semantic understanding for generation, plus RationalRewards for instruction-based editing.Each prompt’s capability tag routes it to exactly one evaluator.Source: D Reward Models · p. 47

Text rendering

  • PaddleOCR-VL-1.5 scores normalized target strings by exact substring match or capped character-edit similarity; absent or unrecognizable text scores 0.The OCR reward averages scores across the prompt’s targets; a perfect rendering scores 1.Source: D Reward Models · p. 47

Aesthetic quality

  • The quality criteria span general photographic quality and detailed human-anatomy checks, including both broad photographic-quality considerations and detailed assessments of human anatomy.Source: D Reward Models · p. 48

Semantic understanding

  • The Qwen3.5-27B semantic judge averages yes/no prompt-alignment checks covering objects, attributes, counts, spatial relations, actions, scene, and faithfulness.The fine-grained checks provide an interpretable signal for multi-object scenes.Source: D Reward Models · pp. 48–49

Instruction-based editing

  • RationalRewards rates editing on text faithfulness, image preservation, visual quality, and applicable text rendering, then linearly maps the mean 1–4 score to [0, 1].Text rendering is not applicable when the edit involves no text; all-4 and all-1 ratings map to 1 and 0, respectively.Source: D Reward Models · p. 49

E Application: Scientific Diagram Generation

The section applies Mage-Flow-Base to scientific-diagram generation through a single-stage mixed-data fine-tune, reporting a 20.81-point Overall gain and qualitative improvements in layout, text, and arrows.

E Application: Scientific Diagram Generation

Quantitative Results.

  • 61.61 Overall for Mage-Flow-SciForma exceeds zero-shot Mage-Flow-Base’s 40.80 by 20.81 points, with component and arrow accuracies rising by 13.70 and 24.70 points.Results are on SciFormaBench-2K; the baselines are evaluated zero-shot.Source: Quantitative Results · p. 50, Table 19 · p. 50
  • Mage-Flow-SciForma produces precise layouts, legible mathematical indices and descriptors, and correctly routed arrows, fixing Mage-Flow-Base errors while matching SciForma-9B-Base.Examples include neural-network architectures with skip-connections, unrolled transformer blocks, and annotated processing pipelines.Source: Quantitative Results · p. 50, Figure 27 · p. 51, Figure 26 · p. 51

Related papers

Loading 2607.19064v2…