MIMFlow: Integrating Masked Image Modeling with Normalizing Flows for End-to-End Image Generation

Yang Chen; Xiaowei Xu; Shuai Wang; Xinwen Zhang; Qiushi Guo; Tiezheng Ge; Limin Wang

Source-linked AI summary

MIMFlow: Integrating Masked Image Modeling with Normalizing Flows for End-to-End Image Generation

Yang Chen, Xiaowei Xu, Shuai Wang, Xinwen Zhang, Qiushi Guo, Tiezheng Ge, Limin Wang

arXiv:2606.26016v3Jun 24, 2026cs.CV

TL;DR

Normalizing Flows can overprioritize low-level pixel details, while MIM-based representation learning has remained disconnected from generation. MIMFlow jointly trains masked semantic latents, reconstruction, and flow modeling, reducing FID to 2.50 versus 3.72 for a similar-scale NF baseline on ImageNet 256×256.

Problem
Strict invertibility makes Normalizing Flows prioritize low-level pixel details over high-level semantics, while MIM representation learning remains largely modular and separate from generation.
Method
MIMFlow jointly optimizes masked-image semantic representation, pixel reconstruction, and flow-based density estimation, assigning semantics to the NF and high-frequency synthesis to a decoder.
Results
FID improves from 3.72 to 2.50 versus the similar-scale SimFlow-L baseline, a 32.8% reduction on ImageNet 256×256.
Takeaways & Limitations
MIMFlow demonstrates that masked semantic bottlenecks can improve similar-scale latent Normalizing Flow generation on ImageNet while using 128 tokens.
Takeaways & Limitations
The evidence is limited to class-to-image generation on ImageNet and does not establish that MIM generalizes across diffusion or autoregressive models.

Abstract

from arXiv · show

Normalizing Flows (NFs) are powerful generative models capable of exact density estimation and sampling. However, their strict invertibility often forces the model to exhaust its capacity on low-level pixel details, hindering the capture of high-level semantic structures. While Masked Image Modeling (MIM) has excelled in representation learning, its integration into generative pipelines has remained largely modular and disjointed. In this paper, we propose MIMFlow, a unified end-to-end framework that jointly optimizes latent semantics, pixel reconstruction, and generative flow. By employing a VAE encoder to infer semantic latent from masked images, MIMFlow achieves a principled decoupling of the generative task: the Normalizing Flow focuses on modeling a simplified, low-frequency semantic manifold, while a specialized decoder handles high-frequency synthesis. This design effectively resolves the inherent capacity bottleneck of NFs, allowing the model to prioritize global structural coherence over redundant noise. Empirical results on ImageNet 256$\times$256 show that MIMFlow-L reaches 71.3\% linear probing accuracy and an FID of 2.50. Despite using only 128 tokens (50\% fewer than standard models), it yields a 32.8\% performance gain over similar-scale NF baselines. Our code is available at https://github.com/MCG-NJU/MIMFlow.

1 Introduction

MIMFlow addresses the capacity bottleneck of strictly invertible normalizing flows by integrating masked image modeling into an end-to-end generative framework. It separates semantic manifold modeling from pixel reconstruction so flows prioritize global structure while the decoder handles image details.

Normalizing flows provide exact density estimation and sampling, but strict invertibility prioritizes low-level pixel details over high-level semantics.
MIMFlow jointly optimizes latent semantics, pixel reconstruction, and generative flow in an end-to-end framework.
Masked image modeling has excelled in self-supervised representation learning, yet its role in generative modeling remains under-explored.
Existing MIM-based generative approaches primarily improve visual tokenizers while separating representation learning from the core generative process.
MIMFlow uses a masked-image VAE encoder, an exact-density-estimating normalizing flow, and a reconstruction decoder to bias latents toward global structural coherence.The encoder uses learnable tokens to extract stable representations from masked images.

2 Related Work

Related work spans Normalizing Flows for bidirectional data–latent mappings and Masked Image Modeling for self-supervised representation learning. MIMFlow distinguishes itself by integrating MIM objectives directly into end-to-end NF training.

Normalizing Flows: Normalizing Flows provide mathematically principled bidirectional mappings between data and latent spaces.RealNVP and Glow established exact log-likelihood estimation through coupling layers.
Normalizing Flows: RealNVP and Glow established exact log-likelihood estimation through coupling layers but struggled with high-resolution synthesis.Recent work has revitalized the field through architectural scaling and post-hoc alignment.
Masked Image Modeling: MIMFlow is presented as the first framework to integrate MIM objectives directly into end-to-end NF training.It unifies self-supervised representation learning with generative flow optimization for robust feature extraction and high-quality synthesis.
Masked Image Modeling: MAE and SimMIM demonstrate Masked Image Modeling as a dominant paradigm in self-supervised learning through masked-input reconstruction.The supplied passage identifies masked-input reconstruction as the shared strategy of these methods.

3 Method

MIMFlow combines masked encoding, a fixed-dimensional latent bottleneck, and a latent Normalizing Flow in a VAE-based framework. Joint reconstruction, density-estimation, and semantic objectives shape a robust latent space, followed by decoder fine-tuning for photorealistic detail.

Architecture: MIMFlow uses a Masked Encoder, Latent Normalizing Flow, and decoder to connect masked-image semantics, probabilistic modeling, and pixel reconstruction.The flow maps latent representations to a Gaussian prior for exact likelihood estimation and reverse sampling, while the decoder reconstructs the image.
Latent Token Bottleneck: The Learnable Token Bottleneck extracts K refined query tokens from N masked-image tokens, producing a fixed-dimensional latent z ∈ R^K×D with K < N.The bottleneck stabilizes the flow target across random masks and encourages compression of local pixel redundancies into high-level structural semantics.
Variational Formulation: MIMFlow treats its masked-image encoder and decoder as a VAE conditioned on ˜x = x ⊙(1 − m), replacing the standard Gaussian prior with a Normalizing Flow.The ELBO combines reconstruction likelihood with a KL divergence against the flow-based prior, and the fixed-variance posterior simplifies the KL term to a flow-prior cross-entropy.
Joint Optimization: The invertible flow computes exact latent likelihood using the Gaussian base density and log-determinant Jacobian, learning a semantic manifold with global image dependencies.The reconstruction objective combines ℓ2 and perceptual losses, while the joint loss is Lprob = Lrec + βLNF and jointly optimizes encoder, decoder, and flow parameters.
Semantic Supervision and Fine-Tuning: Auxiliary prediction of DINO or CLIP features from z enriches semantic representations, and adversarial fine-tuning targets missing high-frequency detail in reconstructions.The end-to-end objective adds γLaux, while a patch-based discriminator is introduced during targeted pixel-decoder fine-tuning.

4 Experiment

Experiments on ImageNet 256 × 256 evaluate MIMFlow’s generative quality, semantic representations, efficiency, and component contributions. MIMFlow-L achieves strong performance with compressed 128-token latents, while masking improves both generation and semantic quality and stabilizes flow mappings.

Evaluation Setup: Experiments use ImageNet 256 × 256 and report FID, IS, Precision, Recall, reconstruction FID, and linear probing for semantic quality.Generative metrics use the ADM evaluation suite, while reconstruction FID is computed on the ImageNet validation set.
Main Results: MIMFlow-L produces high-fidelity images with consistent global structures and demonstrates significant performance within the normalizing-flow paradigm.The comparison covers pixel-space, latent autoregressive, latent diffusion, and latent normalizing-flow models.
Efficiency and Token Compression: 128 tokens halve sequence length relative to standard VAE-based models, reducing flow complexity while encouraging compact representation of global image structure.The masked bottleneck limits exposure to high-frequency noise and supports more efficient semantic modeling.
Hardware Efficiency: 37.6GB training memory, 3.11 iterations per second, and 0.011s per-image sampling time improve over 52.3GB, 2.83 iterations per second, and 0.020s under matched hardware.These correspond to a 28% memory reduction, a 10% throughput improvement, and nearly halved sampling time.
Ablations and Stability: gFID 12.82 and Acc. 71.3% are achieved with a 0.4–0.6 masking ratio, while MIMFlow’s Jacobian has a larger minimum singular value and lower, more concentrated log-condition number than STARFlow.The masking ablation identifies a substantial, consistent information bottleneck as important for stabilizing the semantic manifold; Jacobian analysis indicates improved numerical conditioning.

5 Limitation

The study validates masked semantic bottlenecks for latent normalizing-flow models but does not establish that MIM generalizes across generative families. Its scope is limited to ImageNet class-to-image generation, while text-to-image extension remains future work.

Model-family scope: MIMFlow’s findings do not support a general claim about MIM across diffusion or autoregressive generative models.The experiments specifically validate MIM-style masked semantic bottlenecks for latent normalizing-flow models.
Generation setting: The study focuses on class-to-image generation on ImageNet rather than text-to-image generation.Text-to-image extension remains future work because prompt alignment and compositional semantics are central challenges.
Evaluation scope: Evaluation uses standard ImageNet generation metrics, including FID, IS, precision, and recall.The supplied passage introduces these metrics while discussing the study’s evaluation scope.

6 Conclusion

MIMFlow unifies Masked Image Modeling and Normalizing Flows in an end-to-end generative framework, decoupling semantic manifold modeling from high-frequency texture synthesis through a masked bottleneck with learnable tokens.

MIMFlow integrates Masked Image Modeling with Normalizing Flows in a unified end-to-end generative framework.
A masked bottleneck with learnable tokens decouples generative tasks between semantic modeling and texture synthesis.
The Normalizing Flow models the low-frequency semantic manifold, while a specialized decoder synthesizes high-frequency texture.

A Implementation Details

MIMFlow’s implementation specifies its architecture and training hyperparameters in dedicated tables, while using distinct loss regions to balance global reconstruction with masked-context inference. The adversarial refinement stage follows a GAN-based formulation and discriminator design.

Optimization objectives: Pixel-level MSE is computed over the entire image to promote global reconstruction fidelity.The loss covers all image pixels rather than only masked regions.
Optimization objectives: The auxiliary semantic loss is restricted to masked patches to enforce inference of missing structural context.This focuses semantic learning on information that must be inferred from the visible image.
Configuration: Detailed architectural specifications and training hyperparameters are summarized in Table 6 and Table 7, respectively.Table 6 describes MIMFlow-L architecture configurations, where K denotes latent tokens and D denotes latent dimensionality.
Adversarial refinement: The adversarial refinement stage adopts a GAN loss formulation and discriminator architecture.The supplied passage identifies the adversarial objective and discriminator design but does not provide their detailed specifications.

B Spectral Analysis of the Jacobian … C Extended Experimental Results

The paper evaluates NF Jacobian geometry with matrix-free automatic-differentiation methods, estimating extreme singular values and local conditioning without explicitly forming the full Jacobian. Extended experiments report training hyperparameters and progressive ImageNet ablations under a specified evaluation protocol.

B Spectral Analysis of the Jacobian: Matrix-free automatic-differentiation methods estimate the Jacobian’s extreme singular values without explicitly computing or storing its D × D matrix.The approach uses iterative procedures based on AD primitives because the data-space Jacobian is computationally intractable to construct directly.
C Extended Experimental Results: Training uses two phases: joint optimization of the VAE and NF, followed by adversarial refinement of the decoder.These phase definitions are provided in the training-hyperparameters table.
B.1 Estimation of the Maximum Singular Value (σmax): The maximum singular value σmax(J) is estimated by power iteration on the operator J^T J using Jacobian-vector and vector-Jacobian products.The iteration normalizes J^T(Jv_k), and convergence yields σmax(J) = ∥Jv∥.
B.1 Estimation of the Maximum Singular Value (σmax): Power iteration avoids explicit matrix construction by computing Jv_k through a JVP and applying J^T through a VJP.These automatic-differentiation primitives implement the required operator products directly.
B.2 Estimation of the Minimum Singular Value (σmin): The minimum singular value is estimated with inexact shifted inverse iteration using A = J^T J + αI, where α > 0 provides Tikhonov regularization.The shift ensures strict positive definiteness and improves numerical stability.
B.2 Estimation of the Minimum Singular Value (σmin): The shifted linear system is solved by conjugate gradients in a matrix-free procedure, with operator-vector products evaluated through a JVP-VJP sequence.The converged minimum eigenvalue λmin of A recovers the smallest singular value as σmin(J) = sqrt(max(λmin − α, 0)).
B.2 Estimation of the Minimum Singular Value (σmin): The NF’s local conditioning is evaluated with the condition number κ(J) = σmax/σmin.This combines the estimated maximum and minimum singular values into a local conditioning measure.
C Extended Experimental Results: Progressive ablations are conducted on ImageNet 256 × 256 with 10K samples and evaluated without CFG.The extended-results table specifies this evaluation setting.

C.1 Progressive ablation.

The progressive ablation isolates component contributions under a 10K-sample evaluation without CFG. Replacing the SD-VAE latent interface with learnable tokens while removing GAN training worsens gFID, whereas adding DINO+CLIP supervision improves semantic confidence as measured by IS.

Progressive ablation: The ablation progressively isolates each component using 10K samples without classifier-free guidance.It starts from an end-to-end SD-VAE baseline with 256 tokens and GAN training.
Progressive ablation: 19.52 gFID results after replacing the latent interface with learnable tokens and removing GAN loss, up from 11.24.This indicates that learnable tokens alone do not explain the final gain.
Progressive ablation: Adding DINO+CLIP auxiliary supervision improves semantic confidence as reflected by IS.The supplied passage does not report the magnitude of this improvement.

C.2 Ablation on Masked Reconstruction Weight

The ablation studies how the masked-region reconstruction loss should be weighted relative to unmasked regions for generative modeling. Equal weighting is optimal, while over-weighting masked patches harms both reconstruction and generation by weakening global structural coherence.

Weighting objective: The MIM weight scales masked-region reconstruction loss relative to unmasked regions, whose weight remains fixed at 1.The study balances standard MIM’s masked-patch objective against the need for global pixel-level fidelity in high-quality synthesis.
Optimal setting: 1 achieves the best performance by assigning equal importance to masked and unmasked regions during joint optimization.This setting is identified as optimal in the ablation shown in Tab. 9.
Over-weighting masked regions: Increasing the MIM weight to 5 or 10 degrades both reconstruction (rFID) and generation (gFID) metrics.Over-weighting masked regions may bias the model toward local patch-filling at the expense of global structural coherence.

C.3 Efficiency Analysis · C.4 Semantic Evolution across Flow Depth

MIMFlow improves efficiency by operating with 128 tokens, substantially fewer than existing 256- or 1024-token approaches. Linear probing shows that NF depth does not improve semantic accuracy, motivating MIMFlow’s division of semantic modeling and density estimation between the encoder and flow.

C.3 Efficiency Analysis: 128 tokens let MIMFlow operate with a substantially reduced token budget compared with existing methods using 256 or 1024 tokens.The passage identifies reduced token count as a key efficiency advantage.
C.3 Efficiency Analysis: 183ms per iteration is the training latency for the 128-token setting on an H20 GPU with batch size 32.This latency isolates the effect of token count using the Improved STARFlow backbone.
C.3 Efficiency Analysis: The latency comparison is conducted with the Improved STARFlow backbone to isolate the effect of token count.The experiment compares token-budget settings rather than changing the backbone.
C.4 Semantic Evolution across Flow Depth: Encoder latent representations achieve high linear-probing accuracy under the MIM objective.The encoder’s latent space provides the strongest semantic classification performance described in the passage.
C.4 Semantic Evolution across Flow Depth: NF depth does not increase linear-probing accuracy, which occasionally plateaus or slightly declines as flow depth increases.The comparison covers intermediate representations at various NF depths.
C.4 Semantic Evolution across Flow Depth: Normalizing Flows perform complex distribution warping through bijective mappings but lack inherent hierarchical feature abstraction or semantic distillation.This characterizes the semantic limitation observed across flow depth.
C.4 Semantic Evolution across Flow Depth: MIMFlow delegates semantic modeling to the MIM-trained encoder while subsequent flow blocks focus on probabilistic density estimation.This division addresses the observed limitation of semantic evolution within NF depth.