Improved Masked Image Generation with Knowledge-Augmented Token Representations

AI summary

Overview

Research area: Computer vision, specifically class-conditional and text-to-image generation with masked (non-autoregressive) visual token models.

Technical level: Advanced. The paper assumes familiarity with VQ-VAE tokenization, masked image generation (MaskGIT-style parallel decoding), graph convolutional networks, and standard generative metrics (FID, IS, Precision, Recall).

One-sentence scope: The paper proposes KA-MIG, a model-agnostic framework that constructs three token-level knowledge graphs from training data and fuses them into existing masked image generation backbones (MaskGIT, AutoNAT, TiTok, MaskGen) to improve generation quality.

What This Paper Is About

Masked image generation (MIG) works by predicting many visual tokens in parallel, which makes it fast but usually less accurate than diffusion models. The paper argues this is partly because individual visual tokens lack clear semantic meaning and token sequences are long (e.g., 256 tokens per image), so transformers struggle to learn the dependencies between tokens directly from data. KA-MIG's goal is to supply that missing structure as explicit prior knowledge graphs extracted from the training set, rather than only refining the sampling/decoding strategy as prior work does.

Key Contributions

1. A knowledge-augmented masked image generation framework (KA-MIG) that injects prior semantic knowledge graphs into the internal token representations of existing MIG models instead of modifying only the decoding strategy.
2. The construction of three token knowledge graphs from training data: two positive priors (co-occurrence graph, semantic similarity graph) and one negative prior (position-token incompatibility graph).
3. A lightweight additive–subtractive fusion mechanism that adds positive prior features to unmasked token representations and subtracts negative prior features at each transformer layer, with a hybrid precompute/compute-online strategy to control cost.
4. Experiments on ImageNet-256 with three backbones plus text-to-image generation on MS-COCO, showing consistent FID improvements, alongside ablations, embedding visualizations, and an efficiency analysis (code released at https://github.com/GuotaoLiang/KA-MIG).

Main Findings

• Improvements on all three ImageNet-256 backbones: MaskGIT improved from FID 6.18 to 5.69 (IS 182.1 to 170.2, Precision 0.80 to 0.81, Recall 0.52 to 0.50); AutoNAT from FID 2.68 to 2.45 (IS 278.8 to 274.1, Precision 0.82, Recall 0.56); TiTok-b64 from FID 2.48 to 2.40 (IS 214.7 to 217.0); TiTok-s128 from FID 1.97 to 1.90 (IS 281.8 to 271.9). Note that IS decreased in the MaskGIT, AutoNAT, and TiTok-s128 cases while FID improved.
• Stronger gains on longer token sequences: The authors attribute the larger improvements on MaskGIT and AutoNAT (256 tokens) versus TiTok (64 or 128 tokens) to the more complex token interactions in longer sequences, which they argue validates the core motivation.
• Competitive against other paradigms: The paper reports that TiTok-s128-KA achieves the best FID among the compared GAN, diffusion, and autoregressive models, and that AutoNAT-KA's FID of 2.45 outperforms LlamaGen-XL (2.62), VAR-d20 (2.57), and TiTok-b64 (2.48).
• Lightweight addition: The method adds approximately 20 million additional parameters overall, with the backbone frozen and only the classification layer plus new parameters fine-tuned.
• Ablation — every graph helps, and they complement each other: On AutoNAT, baseline FID 2.68 / IS 278.8; adding only the semantic similarity graph gives 2.49 / 279.6; only the position-token graph 2.51 / 285.6; only the co-occurrence graph 2.51 / 282.1; any pair gives 2.46–2.48; all three give the best 2.45 / 274.1. The semantic similarity graph produced the largest single-graph FID improvement.
• Efficiency trade-off: Precomputing the co-occurrence and semantic similarity graphs adds only +0.79M parameters each and virtually no extra TFLOPs, while precomputing the class-conditional position-token graph would add +196M parameters. The chosen configuration precomputes the two light graphs and computes the position-token graph online (+15M params, +0.06 TFLOPs), compared with online computation of the other two at +16M params and +0.09 TFLOPs each.
• Negative prior improves token ranking: Under random masking of 10,000 images, NDCG@100 rose from 0.0282 to 0.0284 at mask ratio 0.3 and from 0.0288 to 0.0289 at mask ratio 0.5 when the position-token incompatibility graph was added.
• Semantic similarity validated by reconstruction: Replacing token (1013) with its most similar token (463) yielded no perceptible difference, while the least similar token (149) significantly degraded quality; in the appendix, top-2 (463, 658) was visually indistinguishable, while top-3 (721) and (149) degraded quality across a 1024-image PSNR analysis.
• t-SNE structure: The semantic similarity graph embedding forms small clusters of two or three tokens (consistent with the top-2 construction), and token (1013) sits close to (463) but far from (149). The co-occurrence graph embedding forms multiple coherent clusters, and low-frequency tokens (shown in red) are not isolated but embedded among semantically similar neighbors.
• Text-to-image also benefits: On COCO-30K with MaskGen as backbone, FID improved from 22.27 to 21.01 and CLIP-Score from 25.58 to 26.10 at 16 steps. Only the co-occurrence and semantic similarity graphs were used here, because the position-token graph is class-specific.

Methodology in Plain English

The approach has three steps.

1. Build knowledge graphs from the training data. Because individual tokens have no human-readable meaning, the authors mine structure statistically. The co-occurrence graph is a weighted undirected token-to-token graph where each edge records how often two tokens appear in a first-order local neighborhood (vertical, horizontal, diagonal); low-frequency edges are pruned as noise. The semantic similarity graph treats tokens as similar if they appear at similar spatial positions across many images — each token gets a length-N position-distribution vector, Jensen-Shannon divergence measures distances, and each token keeps its top two most similar tokens as a directed graph. The position-token incompatibility graph is a per-class directed position-to-token graph marking tokens that never appear at a given location for that class (for example, ground or grass tokens in the upper-middle region of "airplane" images).

2. Encode the graphs. Two 3-layer GCNs with ReLU produce global token representations from the co-occurrence and similarity graphs. For the negative graph, position features are formed by averaging the embeddings of the tokens deemed incompatible with each position, passed through a learnable weight matrix, producing a per-position embedding.

3. Fuse into the existing model. A fusion network is inserted before each transformer layer. Positive prior features are added to unmasked token representations via zero convolutions; negative prior features are subtracted, scaled by a hyperparameter alpha. Because the priors are learned independently of the MIG model, they can be precomputed and cached, so inference adds only light additive and subtractive operations. Training keeps the standard masked-token negative log-likelihood objective; the backbone is frozen and only the classification layer and new parameters are tuned.

Experimental setup: Backbones are MaskGIT, AutoNAT, and TiTok, evaluated on ImageNet at 256×256 (AutoNAT and TiTok were not evaluated at 512×512 because public checkpoints are unavailable there). Metrics are FID (FID-50K, following Dhariwal and Nichol 2021), IS, Precision, and Recall. Text-to-image uses MaskGen on MS-COCO. Training configurations differ per backbone (MaskGIT: 1500 iterations, batch 100, learning rate 0.0001; AutoNAT: 1300 iterations, batch 300, learning rate 0.0001; TiTok-b64 and TiTok-s128: 1200 iterations, batch 500, learning rate 2.0e-06, with weight decay 0.03 for TiTok-s128). AutoNAT and MaskGIT training used 4 A100 GPUs; TiTok training used 4 A100 GPUs with generation on 2 4090 GPUs, and TFLOPs were computed on an NVIDIA RTX 4090.

Why This Matters

Impact on research. The paper shifts attention from improving MIG sampling schedules to improving the model's internal representations, and shows that externally supplied token-level priors are a practical, model-agnostic lever. It also offers a data-driven way to give meaning to otherwise uninterpretable codebook tokens, which the t-SNE and reconstruction experiments support.

Real-world applications:

• Faster high-fidelity image synthesis where parallel decoding is attractive, such as interactive design and content creation tools.
• Text-to-image generation with improved fidelity, relevant to creative and advertising workflows (demonstrated on MS-COCO with MaskGen).
• Class-conditional data generation, e.g., synthetic training data for categories such as animals, vehicles, or scenes.
• Deployment-constrained settings where the storage-versus-compute trade-off described in the efficiency analysis matters, since graphs can be cached or computed on demand.

Industry relevance. The approach requires no external annotations or hand-designed rules, is compatible with several existing MIG codebases, adds roughly 20 million parameters, and reuses precomputed artifacts — properties that favor integration into existing generation pipelines. Quantitative latency, memory, or throughput measurements beyond parameter counts and TFLOPs are not reported.

Future Directions

• Extending the negative position-token prior to text-to-image: it was excluded from the MS-COCO experiments because it is defined per class, but the paper does not describe how a text-conditioned version would be built.
• Evaluating at 512×512 for AutoNAT and TiTok once public checkpoints exist; the paper only provides qualitative ImageNet-512 samples.
• Reducing the storage cost of the class-conditional incompatibility graph, which the analysis shows would add +196M parameters if precomputed per class.
• Applying the prior-graph idea to other token-based generation domains (for example video or audio) and exploring additional graph types beyond the three identified here.

Target Audience

Researchers and engineers working on visual token generation, masked/parallel decoding, and efficient generative modeling, particularly those already familiar with MaskGIT-style pipelines, VQ-VAE tokenization, and graph neural networks. Practitioners wanting to retrofit an existing MIG backbone with modest overhead also benefit, as do readers interested in how latent codebook tokens can be given structure without external annotation.

Authors’ 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.