Composable Visual Tokenizers with Generator-Free Diagnostics of Learnability
CompTok employs diffusion-based visual tokenization with InfoGAN objectives, enhancing compositionality and predictability of the token space, achieving SOTA in image generation.
Key Findings
Methodology
This work introduces CompTok, a diffusion decoder coupled with a recognition model trained via an InfoGAN-style objective to ensure token utility. It incorporates swap-based regularization and adversarial flow constraints to promote compositional control and realism. The framework involves a two-stage pipeline: an encoder E mapping images to tokens, and a diffusion decoder D reconstructing images conditioned on these tokens. The recognition head Qϕ predicts tokens from decoded images, enforcing causal effectiveness. Swap regularization expands the token space's compositionality, while AFM ensures generated samples stay on the natural image manifold. The evaluation introduces AvgIG and MC metrics to quantify the local sensitivity and global connectivity of the token space, predicting downstream generator performance.
Key Results
- CompTok achieves a rFID of 1.27 on ImageNet, outperforming existing 2D grid and 1D sequence tokenizers. Its gFID surpasses baseline models, indicating higher generation quality. The AvgIG and MC metrics strongly correlate with downstream task performance, outperforming traditional reconstruction metrics. Ablation studies confirm that mutual information supervision and swap regularization are critical for enhancing compositionality. The model demonstrates high-level semantic editing via token swapping, showing robust control.
- Across remote sensing, medical imaging, and OCR tasks, CompTok’s AvgIG and MC metrics correlate more with task performance than rFID, indicating better preservation of semantic information. Experiments show improved efficiency, requiring fewer steps for high-quality image synthesis. The learned token space exhibits high connectivity, facilitating interpolation and editing, which are vital for practical applications. The results validate the framework’s applicability in diverse domains.
- Ablation results highlight the importance of mutual information supervision and adversarial flow regularization in maintaining a well-structured token space. The combination of these techniques results in a more learnable, compositional latent space that supports complex editing and generation tasks. The experiments demonstrate the method’s superiority over existing approaches in both quantitative metrics and qualitative control, establishing a new benchmark for visual tokenizers.
Significance
This research advances the state-of-the-art in visual tokenization by explicitly optimizing for compositionality and learnability, addressing key limitations of prior methods. The introduction of predictive metrics (AvgIG and MC) provides practical tools for evaluating token spaces beyond reconstruction fidelity. The framework significantly enhances the controllability and diversity of generative models, facilitating high-level semantic editing and multi-task adaptability. Its broad applicability across domains like remote sensing, medical imaging, and OCR underscores its potential to transform intelligent visual systems, enabling more interpretable and flexible image synthesis. This work lays a foundation for future exploration of structured, controllable, and semantically rich token representations.
Technical Contribution
The core technical innovation lies in integrating diffusion models with InfoGAN-style objectives to enforce token utility and compositionality. The swap regularization mechanism systematically expands the valid token combinations, while AFM adversarial flow regularizer ensures the generated samples remain on the natural image manifold. The introduction of AvgIG and MC as metrics offers a novel quantitative assessment of the token space’s structure, predictive of downstream generator success. These contributions collectively enable the creation of a structured, highly controllable latent space that surpasses existing tokenization approaches in both quality and interpretability.
Novelty
This work is the first to embed InfoGAN-style recognition objectives within a diffusion-based visual tokenizer, explicitly targeting the enhancement of compositionality and learnability. The combined use of swap regularization and adversarial flow constraints to shape the token space’s geometry is novel, addressing the fragmentation issues seen in prior models. Additionally, the proposed AvgIG and MC metrics provide new tools for evaluating the latent space’s structure, bridging the gap between reconstruction fidelity and downstream controllability. These innovations collectively push the frontier of structured, controllable image generation.
Limitations
- Despite its advances, the framework requires complex multi-stage training involving regularizers and metrics, which increases computational costs and implementation complexity. The reliance on adversarial training can lead to instability, especially in high-dimensional spaces. The current model’s performance may degrade on highly diverse or out-of-distribution datasets due to assumptions about natural image manifolds. Further research is needed to improve training efficiency and robustness across broader domains.
Future Work
Future directions include integrating multi-modal signals (text, audio) to enrich token semantics, developing more efficient training algorithms, and extending the metrics to evaluate multi-scale and hierarchical structures. Exploring unsupervised or weakly supervised approaches could reduce reliance on labeled data. Additionally, applying the framework to video and 3D data would broaden its impact. Enhancing the interpretability of token representations and integrating user-guided editing tools are promising avenues to make the technology more accessible and controllable for practical applications.
AI Executive Summary
The rapid evolution of visual generative models has underscored the importance of effective tokenization strategies. Traditional approaches, such as VQ-VAE and its variants, primarily focus on reconstruction fidelity, often neglecting the structure and utility of the latent token space. This limitation hampers the ability of downstream generators to perform fine-grained control and semantic editing, especially in complex, high-dimensional visual domains. Recognizing this gap, the authors introduce CompTok, a novel framework that leverages diffusion models combined with InfoGAN objectives to produce structured, compositional visual tokens.
CompTok’s architecture involves an encoder that maps images into a sequence of tokens and a diffusion decoder trained to reconstruct images conditioned on these tokens. A recognition head ensures each token’s causal effectiveness, while swap-based regularization and adversarial flow constraints promote a well-connected, realistic token space. The key innovation lies in the introduction of two metrics—AvgIG and MC—that quantify the local sensitivity and global connectivity of the token space, respectively. These metrics serve as predictors of the downstream generator’s performance, addressing the limitations of traditional reconstruction metrics.
Experimental results on datasets like ImageNet demonstrate that CompTok surpasses existing tokenizers in both reconstruction quality and generative control. It achieves a rFID of 1.27, outperforming models like VQGAN and ViT-VQGAN, while also enabling high-level semantic editing through token swapping. The metrics correlate strongly with downstream task performance, validating their utility. Ablation studies confirm the importance of mutual information supervision and swap regularization in shaping a learnable, compositional latent space.
This work significantly advances the field by providing a structured, controllable tokenization framework that bridges the gap between reconstruction fidelity and generative controllability. Its broad applicability across remote sensing, medical imaging, and OCR tasks highlights its potential to revolutionize visual AI. Future research will focus on multi-modal integration, efficiency improvements, and expanding to video and 3D data, paving the way for more flexible, interpretable, and high-quality visual synthesis.
Deep Dive
Abstract
We introduce CompTok, a training framework for learning visual tokenizers whose tokens are enhanced for compositionality. CompTok uses a token-conditioned diffusion decoder. By employing an InfoGAN-style objective, where we train a recognition model to predict the tokens used to condition the diffusion decoder using the decoded images, we enforce the decoder to not ignore any of the tokens. To promote compositional control, besides the original images, CompTok also trains on tokens formed by swapping token subsets between images, enabling more compositional control of the token over the decoder. As the swapped tokens between images do not have ground truth image targets, we apply a manifold constraint via an adversarial flow regularizer to keep unpaired swap generations on the natural-image distribution. The resulting tokenizer not only achieves state-of-the-art performance on image class-conditioned generation, but also demonstrates properties such as swapping tokens between images to achieve high level semantic editing of an image. Additionally, we propose two metrics that measures the landscape of the token space that can be useful to describe not only the compositionality of the tokens, but also how easy to learn the landscape is for a generator to be trained on this space. We show in experiments that CompTok can improve on both of the metrics as well as supporting state-of-the-art generators for class conditioned generation.