Improved Masked Image Generation with Knowledge-Augmented Token Representations

TL;DR

This paper introduces KA-MIG, a framework that incorporates explicit knowledge graphs of token dependencies to enhance masked image generation, resulting in higher quality image synthesis.

Contribution

It proposes a novel knowledge-augmented approach using three types of token knowledge graphs to improve semantic dependency learning in masked image generation.

Findings

Improved image generation quality on ImageNet dataset.

Effective integration of knowledge graphs enhances semantic understanding.

Outperforms existing MIG methods in class-conditional image synthesis.

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.