PaperlayerView exact revision on arXiv

Source-linked AI summary

Accelerating Masked Image Generation by Learning Latent Controlled Dynamics

Kaiwen Zhu, Quansheng Zeng, Yuandong Pu, Shuo Cao, Xiaohui Li, Yi Xin, Qi Qin, Jiayang Li, Yu Qiao, Jinjin Gu, Yihao Liu

arXiv:2602.23996v1cs.CV

TL;DR

Masked image generation is inefficient because existing feature-reuse methods have limited expressivity and ignore sampling information. MIGM-Shortcut learns lightweight latent controlled dynamics from previous features and sampled tokens, achieving over 4× speedup on Lumina-DiMOO text-to-image generation with negligible performance drop.

  • Problem

    Existing masked image generation acceleration methods mainly reuse previous features with handcrafted rules, while failing to model sampling-controlled dynamics and limiting acceleration.

  • Method

    MIGM-Shortcut learns a lightweight neural shortcut model that uses previous features and sampled tokens to predict the direction of future feature evolution.

  • Results

    Over 4× speedup was achieved for Lumina-DiMOO text-to-image generation with negligible performance drop, while experiments on MaskGIT and Lumina-DiMOO validated the method.

  • Takeaways & Limitations

    Learning latent feature dynamics provides a simple and effective route to practical acceleration and exposes inherent computational redundancy in masked generative models.

  • Takeaways & Limitations

    The approach assumes that time can be implicitly captured by the previous feature and that feature dynamics have approximately uniform local Lipschitz behavior.

Abstract

from arXiv · show

Masked Image Generation Models (MIGMs) have achieved great success, yet their efficiency is hampered by the multiple steps of bi-directional attention. In fact, there exists notable redundancy in their computation: when sampling discrete tokens, the rich semantics contained in the continuous features are lost. Some existing works attempt to cache the features to approximate future features. However, they exhibit considerable approximation error under aggressive acceleration rates. We attribute this to their limited expressivity and the failure to account for sampling information. To fill this gap, we propose to learn a lightweight model that incorporates both previous features and sampled tokens, and regresses the average velocity field of feature evolution. The model has moderate complexity that suffices to capture the subtle dynamics while keeping lightweight compared to the original base model. We apply our method, MIGM-Shortcut, to two representative MIGM architectures and tasks. In particular, on the state-of-the-art Lumina-DiMOO, it achieves over 4x acceleration of text-to-image generation while maintaining quality, significantly pushing the Pareto frontier of masked image generation. The code and model weights are available at https://github.com/Kaiwen-Zhu/MIGM-Shortcut.

1. Introduction

MIGMs achieve strong and multimodal generation capabilities, but their multi-step computation limits efficiency. MIGM-Shortcut addresses this by learning lightweight latent controlled dynamics that incorporate feature history and sampled tokens, achieving over 4× speedup on Lumina-DiMOO with negligible text-to-image quality loss.

  • Motivation: MIGMs generate images by progressively predicting discrete tokens from a masked sequence and achieve performance comparable to cutting-edge continuous diffusion models.Their formulation can also be unified with generation in other modalities.
  • Efficiency challenge: Reducing MIGM generation steps is hindered by the difficulty of modeling multiple tokens jointly, while reducing per-step cost motivates caching-based approaches.These efficiency challenges become more important as MIGMs are integrated into large-scale foundation models.
  • Latent trajectory: Consecutive MIGM features are extremely similar and follow a regular trajectory, suggesting that latent feature dynamics may be structured enough to model.The observation comes from temporal self-similarity analysis and t-SNE visualization of Lumina-DiMOO’s last-layer features.
  • Latent controlled dynamics: Continuous-diffusion trajectory predictors cannot be directly applied because MIGM dynamics depend on exogenous randomness from sampled tokens rather than history alone.MIGMs begin from a fully masked sequence without randomness, so sampling-token randomness provides generation diversity.
  • MIGM-Shortcut: Over 4× speedup is achieved for Lumina-DiMOO with negligible performance drop on text-to-image generation using MIGM-Shortcut.The method is instantiated on MaskGIT and Lumina-DiMOO, using a lightweight neural network to learn the latent controlled dynamics.

2. Related Work

Related work has improved discrete image generation through masked prediction and acceleration techniques, but existing approaches face inefficiency, multimodality, differentiability, or feature-approximation limitations. MIGM-Shortcut instead learns a moderate-complexity feature dynamics model applicable to existing trained base models.

  • Masked Image Generation: Masked image generation predicts multiple tokens per network evaluation, avoids rigid generation order, supports tuning-free image editing, and benefits representation learning.MIGMs use randomly masked token sequences and flexible conditioning masked positions.
  • Generation Acceleration: Few-step methods reduce generation steps in continuous and discrete diffusion models, but discrete approaches still struggle with the multi-modality problem.Other methods reduce per-step cost through pruning or quantization, while some require time-differentiable features unavailable in discrete diffusion.
  • Feature Caching: TeaCache and Block Caching refine cached-feature approximations with learned polynomial coefficients, whereas MIGM-Shortcut learns feasible feature dynamics with moderate complexity for existing trained base models.The proposed contrast emphasizes broader applicability across well-trained base models.

3. Method

MIGM-Shortcut models masked image generation as latent controlled dynamics, learning a lightweight transition from previous features and newly sampled tokens to bypass repeated base-model computation. Its design uses a compact attention bottleneck, while training regresses the next feature with frozen-base-model targets and inference alternates full and shortcut steps.

  • Shortcut model formulation: The method treats the unmasking model’s hidden features as a latent controlled trajectory driven by sampled token observations.At step i, f_ti is the unmasking model’s last hidden state, while x_ti is the generated token sequence.
  • Shortcut model formulation: The shortcut transition predicts f_ti+1 from f_ti, x_ti, and time, replacing the cumbersome base-model computation with a lightweight neural network.The transition is formulated as f_ti+1 = f_ti + Sθ(f_ti, x_ti, ti) + ϵ, with normally distributed error.
  • Motivation and empirical evidence: Feature pairs typically have cosine similarity above 0.95, while local difference ratios concentrate around a constant, supporting a low-complexity approximation.The observed approximately uniform local Lipschitz behavior suggests the shortcut model need not have high-frequency representational capacity.
  • Network implementation: The backbone uses cross-attention over newly decoded tokens followed by self-attention, with dimensionality bottlenecks and time conditioning to reduce cost and aid convergence.Inputs are projected down and back up, while sinusoidal time embeddings modulate features through adaptive layer normalization.
  • Training: Training minimizes MSE on collected tuples (f_ti, x_ti, t_i, f_ti+1), with only θ trainable and the base model M frozen.More elaborate distribution-matching and self-prediction strategies were tested, but MSE alone performed equally well.
  • Inference: Inference computes the first feature with the base model and can replace subsequent heavy feature computations with the shortcut model.This creates an inference workflow that alternates full base-model steps with shortcut steps.

4. Experiments

Experiments on MaskGIT and Lumina-DiMOO show that MIGM-Shortcut accelerates generation while preserving or improving quality. Ablations further support incorporating sampling information and balancing shortcut-model complexity.

  • MaskGIT: MaskGIT-Shortcut consistently generates better images faster than vanilla MaskGIT, even surpassing vanilla’s optimal performance on ImageNet-512.Quality is evaluated by FID across different configurations.
  • Lumina-DiMOO: At budget B = 14, DiMOO-Shortcut achieves the highest ImageReward and UniPercept-IQA, the second highest CLIPScore, and a 4 acceleration rate.The comparison uses Pareto-optimal configurations of multiple acceleration methods.
  • Lumina-DiMOO: 4.0–5.8 acceleration rate is achieved by all three DiMOO-Shortcut configurations while approaching vanilla Lumina-DiMOO performance.Acceleration is measured relative to Lumina-DiMOO using 64 steps.
  • Ablation studies: Removing cross-attention drastically degrades performance, producing over-smoothed images because the model predicts the expectation over possible sampling results.Cross-attention incorporates sampled-token information into feature prediction.
  • Ablation studies: Shortcut-model complexity must balance accuracy and computation: insufficient complexity cannot model feature dynamics, whereas excessive complexity wastes computation like the base model.The ablation varies bottleneck ratio R and attention-layer count D; the default setting is R = 2, D = 1.

5. Conclusion

MIGM-Shortcut accelerates masked image generation with minimal performance drop by exploiting information in continuous features that sampling otherwise discards. It learns latent feature dynamics from previous features and sampled tokens, enabling shortcuts that skip the computationally heavy base model.

  • Conclusion: MIGM-Shortcut achieves remarkable acceleration with minimal performance drop by leveraging information wasted in continuous features during MIGM sampling.The method targets reduced computational complexity by recovering useful information from continuous features.
  • Conclusion: Existing methods that approximate future features from previous features suffer from limited expressivity and neglect sampling information.These limitations motivate incorporating sampled tokens into the learned dynamics model.
  • Conclusion: The proposed lightweight model combines previous features and sampled tokens to learn latent feature evolution dynamics, then skips the heavy base model through a learned shortcut.Experiments on MaskGIT and Lumina-DiMOO demonstrate the method’s efficiency-oriented application.

A. Selecting which Layer’s Feature to Investigate

The study selects the feature layer for MIGM-Shortcut by seeking the last layer while favoring the highest cosine similarity between consecutive steps; an exemplary Lumina-DiMOO sample satisfies both conditions simultaneously.

  • A. Selecting which Layer’s Feature to Investigate: MIGM-Shortcut investigates which layer’s feature is most suitable for modeling.The analysis uses an exemplary sample generated by Lumina-DiMOO.
  • A. Selecting which Layer’s Feature to Investigate: The desired feature is from the last layer to enable higher acceleration rates.Layer selection is motivated by the goal of accelerating MIGM-Shortcut.
  • A. Selecting which Layer’s Feature to Investigate: The desired feature dynamics also have the highest cosine similarity between consecutive steps, making them more amenable to learning.Figure 12 plots similarity against layer depth for the exemplary Lumina-DiMOO sample.
  • A. Selecting which Layer’s Feature to Investigate: The last-layer choice and highest consecutive-step similarity are achieved simultaneously in the examined sample.This joint outcome is reported from the layer-depth similarity analysis.

B. More Exploration of Training Strategies

The authors explore alternatives to simple MSE training for MIGM-Shortcut, including distribution matching and exposure to its own predictions. Neither strategy yields notable benefits, supporting the assumption that smooth dynamics make the target easy to learn with MSE.

  • B. More Exploration of Training Strategies: Adding a KL-divergence term matches logits from shortcut-model features with those from the base model’s features.The base model’s classification head H is frozen.
  • B. More Exploration of Training Strategies: Training the shortcut model through several rollout steps exposes it to its own predicted features, addressing potential exposure bias at inference.Inference may provide predicted rather than true base-model features as input.
  • B. More Exploration of Training Strategies: Neither distribution matching nor rollout-based self-exposure brings notable benefits, as shown in Fig. 13.The authors suggest smooth dynamics make the target easy enough for MSE loss, consistent with their core assumption.

C. Details of Human Study

The human study uses Rapidata4 for pairwise image-quality comparisons between acceleration methods and vanilla outputs under a unified instruction. Each paired datapoint receives three independent annotations, and method win rates are aggregated from pairwise wins against vanilla.

  • Study protocol: Rapidata4 evaluates image quality through pairwise comparisons of each method’s output against vanilla output.Trials use A/B presentation under a fixed unified instruction asking which image has better overall visual quality.
  • Study protocol: Each acceleration method is evaluated on all paired datapoints, with three independent raters annotating every datapoint.
  • Win-rate calculation: For each method M, the reported win rate against vanilla aggregates pairwise wins for the method and vanilla across all comparisons.The aggregation uses the total numbers of pairwise wins for M and vanilla.

D. More Visualization

Additional visual examples show that DiMOO-Shortcut consistently generates high-quality images while operating at high speed.

  • D. More Visualization: DiMOO-Shortcut consistently generates high-quality images with high speed across the examples in Fig. 14-22.The figure compares DiMOO-Shortcut with other methods.

E. Complete Results of Ablation Study

The ablation study identifies N = 1, R = 2 with cross attention as the default DiMOO-Shortcut configuration offering the optimal quality-speed tradeoff. It also defines the architectural variants and compares feature similarity and training strategies.

  • Variant performance: N = 1, R = 2 with cross attention achieves the optimal quality-speed tradeoff among DiMOO-Shortcut variants.Table 3 reports metric values for the different variants.
  • Variant definitions: N denotes the number of cross-attention and self-attention layers, while R denotes the bottleneck ratio.The “No cross attention” variant replaces the cross-attention layer with a self-attention layer.
  • Feature dynamics: Feature cosine similarity between consecutive steps is evaluated across layer depth, averaging over all tokens with a 25%–75% token-range shade.The evaluation visualizes similarity trends and token-level variation across layers.
  • Training strategies: Training strategies are compared using dist matching and Rollout2/3, where the former adds a KL divergence loss term and the latter rolls out 2/3 steps with the shortcut model.These definitions specify the alternative training procedures shown in Figure 13.
Loading 2602.23996v1…