Towards Improved Text-Aligned Codebook Learning: Multi-Hierarchical Codebook-Text Alignment with Long Text

TL;DR

The paper tackles the limitation of concise image captions in text-aligned codebook learning by generating long, detailed descriptions with visual-language models. It introduces TA-VQ, a framework that encodes long text at word, phrase, and sentence granularity through a hierarchical encoder and uses a sampling-based optimal transport strategy to align multi-hierarchical image codes with multi-granularity text. The approach yields superior reconstruction and improved performance on downstream tasks (image generation, grounding, and reasoning) across multiple datasets, while reducing computational overhead via sampling and probabilistic modeling. This work advances cross-modal codebook learning by leveraging richer textual semantics and structurally aligned representations, offering practical gains for multi-modal generation and understanding.

Abstract

Image quantization is a crucial technique in image generation, aimed at learning a codebook that encodes an image into a discrete token sequence. Recent advancements have seen researchers exploring learning multi-modal codebook (i.e., text-aligned codebook) by utilizing image caption semantics, aiming to enhance codebook performance in cross-modal tasks. However, existing image-text paired datasets exhibit a notable flaw in that the text descriptions tend to be overly concise, failing to adequately describe the images and provide sufficient semantic knowledge, resulting in limited alignment of text and codebook at a fine-grained level. In this paper, we propose a novel Text-Augmented Codebook Learning framework, named TA-VQ, which generates longer text for each image using the visual-language model for improved text-aligned codebook learning. However, the long text presents two key challenges: how to encode text and how to align codebook and text. To tackle two challenges, we propose to split the long text into multiple granularities for encoding, i.e., word, phrase, and sentence, so that the long text can be fully encoded without losing any key semantic knowledge. Following this, a hierarchical encoder and novel sampling-based alignment strategy are designed to achieve fine-grained codebook-text alignment. Additionally, our method can be seamlessly integrated into existing VQ models. Extensive experiments in reconstruction and various downstream tasks demonstrate its effectiveness compared to previous state-of-the-art approaches.

Figures (18)

• Figure 1: Example of comparing origin caption and long text. The original caption is brief and primarily provides a general description of the image. It lacks detail about the background and certain key elements within the image. In contrast, the long text can offer a more comprehensive image description, including zebra body color, position, mane color position, and background description. This additional text offers richer context, contributing to more robust and effective text-aligned codebook learning.
• Figure 2: The illustration of our proposed TA-VQ framework. The image is first fed to VLM to generate a more detailed text description, and then the text is split into multiple granularities for encoding, i.e., word ($t_w$), phrase ($t_p$), and sentence semantics ($t_s$). Subsequently, the multi-hierarchical encoder encodes and quantizes the image into multi-hierarchical code representation, i.e., $Z_{f_1}$, $Z_{f_2}$, and $Z_{f_3}$. The sampling-base alignment module is employed to achieve $Z_{f_1}$, $Z_{f_2}$, $Z_{f_3}$ and $t_w$, $t_p$, $t_s$ alignment. Finally, the decoder is used to reconstruct the origin image using $Z_{f_3}$.
• Figure 3: Illustration of the Sampling-based Alignment Strategy.
• Figure 4: Examples of unconditional generation on CelebA-HQ. More examples are provided in supplementary materials.
• Figure 5: Examples of semantic synthesis (row 1), text-to-image (row 2), and image completion (row 3). More examples are provided in supplementary materials.

Theorems & Definitions (1)

• Theorem 1