Improved Masked Image Generation with Knowledge-Augmented Token Representations

TL;DR

KA-MIG addresses the core limitation of MIG—weak modeling of semantic dependencies among visual tokens—by injecting explicit, data-derived token priors via three graphs: co-occurrence, semantic similarity, and position-token incompatibility. A graph-aware encoder learns enriched token and positional representations, which are fused into existing MIG backbones with a lightweight additive–subtractive mechanism, enabling parallel decoding with improved fidelity and diversity. Across ImageNet-256 and COCO-30K, KA-MIG yields consistent gains with minimal overhead (≈20M additional parameters) and demonstrates strong, model-agnostic applicability to multiple MIG backbones. The approach provides a practical, scalable way to enhance token-level semantics in MASKED image generation, narrowing the gap to diffusion-based methods while preserving MIG’s speed advantages.

Abstract

Masked image generation (MIG) has demonstrated remarkable efficiency and high-fidelity images by enabling parallel token prediction. Existing methods typically rely solely on the model itself to learn semantic dependencies among visual token sequences. However, directly learning such semantic dependencies from data is challenging because the individual tokens lack clear semantic meanings, and these sequences are usually long. To address this limitation, we propose a novel Knowledge-Augmented Masked Image Generation framework, named KA-MIG, which introduces explicit knowledge of token-level semantic dependencies (\emph{i.e.}, extracted from the training data) as priors to learn richer representations for improving performance. In particular, we explore and identify three types of advantageous token knowledge graphs, including two positive and one negative graphs (\emph{i.e.}, the co-occurrence graph, the semantic similarity graph, and the position-token incompatibility graph). Based on three prior knowledge graphs, we design a graph-aware encoder to learn token and position-aware representations. After that, a lightweight fusion mechanism is introduced to integrate these enriched representations into the existing MIG methods. Resorting to such prior knowledge, our method effectively enhances the model's ability to capture semantic dependencies, leading to improved generation quality. Experimental results demonstrate that our method improves upon existing MIG for class-conditional image generation on ImageNet.

Figures (12)

• Figure 1: Motivation and framework overview. (a) The training pipeline of Masked Image Generation (MIG). (b) The inference process in MIG, where existing works primarily focus on improving sampling quality. (c) Our proposed framework incorporates external prior knowledge graphs to learn richer representations for enhancing MIG.
• Figure 2: The illustration of our proposed KA-MIG framework. We first construct three types of prior knowledge graphs from the training dataset, i.e., $\mathcal{G}_{co}$, $\mathcal{G}_{s}$, and $\mathcal{G}_{p}^{c}$. These graphs are then fed into a graph-aware encoder to learn richer token- and position-aware representations, i.e., $C_{co}$, $C_{s}$, and $P^{c}$. Finally, the enriched prior representations are integrated into the existing MIG framework through a lightweight fusion network.
• Figure 3: Visualization of semantic similarity token via reconstruction. Peak Signal-to-noise Ratio (PSNR) fardo2016formal measures the pixel-level similarity between two images. The red box indicates the location of token (1013). We compare three versions of the same image: (a) the original reconstruction, (b) replacing token (1013) with its most similar token (463) based on our similarity graph $\mathcal{G}_s$, and (c) replacing it with the least similar token (149). The perceptual fidelity of (b) supports the semantic closeness, while (c) introduces noticeable degradation.
• Figure 4: Visualizations for sampling images on ImageNet-256 using selected classes (270: White wolf, and 96: Toucan).
• Figure 5: Visualization of token embeddings from (a) AutoNAT, (b) semantic similarity graph, and (c) co-occurrence graph.