Paper: Unified Language-Vision Pretraining in LLM with Dynamic Discrete Visual Tokenization (ICLR-2024)

Reviewer comment (Writing):
The structure of the writing, especially the method section, is very confusing to me. In fact, the key to the proposed method is the learning of the vision token codebook. This codebook and token selector are built on the first stage of learning. However, the authors' description of this part makes it look like an end-to-end training framework, which confuses me a lot. I would suggest the authors carefully revise the structure of this section.

Author response:
* Sorry for the confusion caused by our description of the methodology section. We will make further revisions to the presentation of this section based on your constructive suggestions. We would appreciate it if you have any other detailed revision suggestions.

Reviewer comment (Experiments):
The authors claim their method can achieve in-context learning capability. However, I cannot find any experiments pertaining to this merit.

Author response:
* Thanks for pointing out this. The "in-context capability" mentioned in this work means that the model can take multi-modal prompt input as context and successfully generate context-related contents. We have evaluated LaVIT’s few-shot in-context learning ability on the downstream multi-modal understanding benchmarks. Specifically, we sample several examples from the training data to make task-specific prompts in the form of interleaved [image, text]. For a fair comparison, the few-shot samples are selected by employing the Retrieval In-Context Example Selection (RICES) strategy [3] following the existing methods [1,2]. As observed, LaVIT can learn to complete specific tasks given only a few examples in the form of demonstration. Furthermore, the performance of LaVIT will improve as the number of support samples increases. This phenomenon is consistent across the three benchmarks, demonstrating its remarkable multi-modal in-context learning potential. 

| Method |  | VQAv2 |  | | Vizwiz |  |
| :-----:| :----: | :----: | :----: | :----:|:----: |:----: |
|  | k=2 | k=4 | k=8 | k=2 | k=4 | k=8 |
| KosMos-1 | 51.4 | 51.8 |  51.4 | 31.4 | 35.3 | 39.0 |
| Flamingo-9B | - | 56.3 | 58.0 | - |  34.9 | 39.4 |
| Emu | 56.4 | 58.4 | 59.0 | 37.8 | 41.3 | **43.9** |
| Ours | **66.5** | **67.2** | **67.8** | **38.5** | **41.7** | 42.1 |


* We also provide more visualization examples of the multi-modal in-context image synthesis results in Figure 9 of the appendix. It can be seen that LaVIT is capable of understanding the semantics of multi-modal prompts and generating the context-related images without any other fine-tuning.

Reviewer comment (Reproducibility):
More details about the language-vision pre-training part are required.

Author response:
* The detailed training hyperparameter settings are reported in Table 9 of our appendix. Please let us know if there are any training details that we have not expressed clearly.

Reviewer comment (Experiments):
Is it possible that the impressive performance of the proposed method is due to more training data, especially for the vision tokenization codebook?

Author response:
* For a fair comparison, we employ the same multi-modal pre-training data as BLIP-2 to exclude the performance gains from additional training data. All the used training data (including about 93M samples for understanding and 100M samples for image generation) are from the publicly available datasets introduced in Section 3.3. This data scale is smaller than that of most other methods. For example, the concurrent work EMU use 2.6B data, KOSMOS adopts 3.1B image-caption pairs, and CM3Leon-7B leverage 340M closed-source data. In comparison, our LaVIT has achieved better performance while utilizing relatively less multi-modal data. Besides, the visual codebook is trained using pure images from the same multi-modal data mentioned above without the corresponding captions. This stage does not require additional data.

Reviewer comment (Experiments):
Why do the authors still use the unselected tokens as keys and values? One intuitive approach for these tokens is simply discarding them. Moreover, these tokens should contribute minorly to the final performance and are redundant as authors claimed.

Author response:
* Admittedly, merely using the selected tokens can approximately encode the global semantics of the image content. However, in our experiments, we have found that directly discarding the unselected patches may result in some information loss, such as the semantics of some background objects, which may impact the performance of multimodal understanding. For instance, in visual question answering tasks, if the question refers to these background objects, the model may be unable to provide the correct answer. Therefore, we design the token merger to compress the semantics of unselected patches onto the retained ones instead of directly discarding them. This process is fulfilled according to the semantic similarity between these patches, so the unselected tokens serve as keys and values. We also conducted the ablation study to explore the role of token merger. As shown by the below Table, we can observe a distinct performance drop in multi-modal understanding tasks without the token merger. In summary, the combination of token selector and merger contributes to reducing the redundancy among tokens while maximally preserving the image details.

| Setting  | Flickr | VQAv2 | OKVQA | GQA |
| :-----:| :----:|:----: |:----: |:----: |
| w/o Token Merger | 70.8 |  51.9  |  42.1 |  39.7  | 
| w/ Token Merger | **72.7**   |  **56.4**   | **46.2** | **42.3** |

Reviewer comment (Theory):
For the causal attention, which tokens are seen as previous ones? As there is actually no sequence for an image, using causal attention for image encoding is thus not convincing.

Author response:
* In the causal attention, the remaining merged visual tokens are flattened into the 1D sequence in the raster order, where the previous tokens of each one are originally on its top left side. Though the original input image is 2D shaped, the usage of causal attention will contribute to converting 2D visual features into a 1D sequence with causal dependency in the output space of token merger. To validate the effectiveness, we replace causal attention with bi-directional attention in the token merger. The results are shown in Table 6 of the appendix. Compared to bi-directional attention, adopting causal attention achieves superior performance across all multi-modal understanding tasks. This phenomenon is reasonable as the next-token prediction objective of LLM requires the causal dependency between input tokens, using causal attention will ensure the consistency with the LLM.

| Setting  | Flickr | VQAv2 | OKVQA |
| :-----:| :----:|:----: |:----: |
| Bi-directional Attention | 69.9 |  55.6  |  44.1 | 
| Causal Attention | **73.2** | **57.3** | **46.4** |

Reviewer comment (Theory):
What is the function of the second half of Eqn. 4?

Author response:
* The second half of Eq-4 serves as a regularization term to impose constraints to minimize the number of retained tokens. We expect to preserve the whole semantic information of an image with as few tokens as possible.

Reviewer comment (Novelty):
The novelty of this paper is limited and the ideas of the paper have already been explored in previous work. For example, end-to-end training LLM has been explored in Emu[MR1], and the transformation of visual information into discrete tokens have been explored in [MR2, MR3, MR4, MR5].

Author response:
We would clarify the unique technical contributions of LaVIT as follows: 
* (1) This work mainly focuses on exploring what makes a good visual token design when extending text-oriented LLM to deal with multimodal inputs. Although the vector quantization for visual features was explored in some previous works [2,3,4,5], directly applying these codebook techniques to tokenize images into discrete tokens trained with LLM are sub-optimal. As a key difference from textual tokens, visual tokens (image patches) exhibit a notable interdependence. In other words, the mutual information between different visual tokens is bigger compared to text tokens, making the semantics of a visual token can be deduced from its adjacent patches. This is also why the majority of mask image modeling methods [4,5] employ an exceptionally high mask rate for image patches. Analogously, such visual interdependence will render the next-token paradigm less effective through self-supervision. Hence, the proposed dynamic visual tokenizer introduces the token selector and merger to pick the most informative tokens and merge the trivial ones. This mechanism contributes to reducing the redundancy and interdependence among visual patches, rendering the unified next-token prediction objective more effective. We have conducted experiments about different visual tokenization techniques on the same 30M pre-training data to validate this. As observed, these methods are not specially tailored for LLM and thus achieve inferior performance compared to LaVIT. 

| Method | Codebook Size |  Flickr | GQA | VQAv2 | OKVQA |
| :-----:| :----: |  :----: | :----:|:----: |:----: |
| VQ-VAE | 8192 | 60.5 |  35.0  | 48.8 | 40.2 |
| BEIT-2 | 8192 | 65.8  | 38.3 | 51.1 | 41.4 |
| Ours | 8192 | **70.0** |  **40.8**  | **53.0** | **44.5** |

* (2) The developed visual tokenizer can produce tokens with a dynamic sequence length varying from the complexity of the image content itself. However, the existing methods tokenize the image into a long and fixed sequence length (e.g., 512 or 1024). Given that the attention computation in LLM exhibits a quadratic relationship with respect to the token length, a long sequence will invariably reduce the computation efficiency. The tokens sparsification of LaVIT's tokenizer thus enables efficient multi-modal training and inference (as shown in Table 3 (b)).

* (3) The same discrete format for two modalities' tokens contributes to the unification of both training and inference. For training, both visual and textual tokens can be optimized under a unified next-token classification objective indiscriminately. In inference, the LaVIT can take any modality combinations as input and produce both visual and textual content by generative sampling. Although the concurrent work Emu [1] also proposes to unlock the LLM during pre-training, it optimizes the visual modality by regressing the next visual features. As validated by Table 3 (a) in Section 4.3, the inconsistent optimization objectives for image and text are not conducive to unified multi-modal modeling within LLM. Therefore, our LaVIT can achieve superior performance than Emu using less training data and smaller model size.

Reviewer comment (Evaluation):
The generation process relies on a stable diffusion model, which makes the verification of the visual generation capabilities of the proposed model limited. Do the image generation results rely on strong stable diffusion, or are they more attributable to the discrete labels generated by the model? It is unclear to what extent the good results achieved are due to stable diffusion.

Author response:
* The denoising U-Net in LaVIT serves as a pixel decoder to recover image pixels (rendering) from the corresponding discrete visual token. In other words, this U-Net is responsible for mapping from discrete space to pixel space rather than generating new image contents based on input prompts. As shown in Figure 7 of the appendix, the reconstructed images from token IDs by the U-Net exhibit a high degree of semantics and general structure to the original images. Therefore, the quality of the generated image (i.e., prompt-following, aesthetic, diversity) is predominantly determined by the discrete visual token generated by the capability of LaVIT rather than this decoder. This claim is further supported by evidence from the following two aspects:


* (1) Despite the weight of conditional U-Net in LaVIT is initialized from the Stable Diffusion v1.5, the LaVIT achieves the superior text-to-image generation performance than the SD v1.5 (7.4 (Ours) v.s. 9.9 (SD 1.5) in FID metric as in Table 2). Hence, this improvement stems from the LaVIT's capability to better correlate the visual and textual tokens. 

* (2) Both GILL[6] and Emu[1] have leveraged the same pre-trained weight of Stable Diffusion v1.5 like ours as the weight initialziation for the decoder. After training on the same data (i.e., the open-source LAION-Aesthetics), they still achieved inferior performances than LaVIT (e.g., 7.4 (Ours) v.s. 12.2 (GILL) and 11.7 (Emu)) as shown in Table 2. 

* Therefore, all the above evidence can fully demonstrate that the good zero-shot image generation results are attributed to the strong capability of LaVIT.

Reviewer comment (Experiments):
Visual token merger part is not well validated. How much additional complexity does this stage bring to the model? Are different codebooks needed for data in different domains (animation, painting, natural landscape, etc.)? Why is a complex merger mechanism needed? Can a simple attention map from a pre-trained model indicate the importance of individual image patches?

Author response:
* **Why is a complex merger mechanism needed?** The token merger aims to compress the semantics of unselected visual patches onto the retained ones according to their feature similarity. Compared to directly discarding these unselected patches, this merge mechanism not only reduces the redundancy among patches but also maximally preserves the visual details. The ablation study to explore the role of token merger is shown as follows. As observed, the introduction of this module can significantly improve the performance of multimodal understanding.

| Setting  | Flickr | VQAv2 | OKVQA | GQA |
| :-----:| :----:|:----: |:----: |:----: |
| w/o Token Merger | 70.8 |  51.9  |  42.1 |  39.7  | 
| w/ Token Merger | **72.7**   |  **56.4**   | **46.2** | **42.3** | 

* **attention map from a pre-trained model.** Although the attention map from the pre-trained ViT can reflect the importance of individual image patches to some extent, it is difficult to manually determine a threshold based on attention weights to select which tokens should remain as input to LLM. In contrast, the proposed token select and merge mechanism can dynamically pick the most informative visual tokens based on the complexity of the image content itself.

* **Additional complexity.** The token merger consists of 12 transformer blocks. It has approximately 280 million parameters, comprising just 3% of the total model. This allocation of parameters does not significantly increase the complexity of the model.

* **Different codebooks needed.** Our LaVIT does not require different codebooks for data from different domains. The learned codebooks can be adapted to diverse domains. As shown in Figures 10 and 11 of the appendix, LaVIT can generate images with different styles, such as oil painting, DSLR photos, animation, and landscape.

Reviewer comment (Novelty):
The processing of the decoder (using the stable diffusion model) is the same as Emu's [MR1] strategy. A reference to the corresponding location and some justifications are required.

Author response:
* Thank you for pointing out the lack of discussion with the Emu decoder. There are the following main differences between these two decoders. (1) The Emu's decoder takes the continuous visual embeddings yielded by LLM as the conditions to generate images. In contrast, our pixel decoder is conditioned on the discrete visual token IDs. It contains an additional module to recover the dynamic discrete token sequence back to the corresponding visual semantic features $x_{\text{rec}}$. The pixel decoding procedure of LaVIT is based on the reconstructed $x_{\text{rec}}$. (2) The training strategies for two decoders are different. the training of Emu's decoder requires the image-text pairs and LLM, during which text is first fed into LLM to generate the visual embeddings. However, our decoder only takes pure image data as input and does not need the combination with LLM during training (only at first stage of training), which reduces the computational overhead and the requirement of the rare multi-modal training data. We have added these discussions in the appendix following your valuable suggestions.

Reviewer comment (Evaluation):
In the experimental results in Table 3, using a larger number of tokens (fixed) achieves worse results, and more explanations need to be added.

Author response:
* The fixed tokenizer variant transforms all the image patch features into the visual tokens as input to LLM during pre-training.  As mentioned above, the redundancy and high interdependence between the image patches may affect the effectiveness of next-token prediction through self-supervision. In comparison, the proposed dynamic tokenizer can reduce the trivial visual redundancy and thus achieve superior zero-shot performance.

Reviewer comment (Experiments):
Compare the performance of text generation using merged continuous embeddings and quantized visual tokens. The loss of detailed information due to quantization is likely to affect the performance of text generation. However, it would be interesting to see how much of a difference there is.

Author response:
* Yes, the loss of detailed information due to quantization will affect the multi-modal understanding performance. The comparisons between using merged continuous embeddings and quantized visual tokens have been presented in Table 7 (a) of the appendix. For your convenience, we have listed the detailed results in the following Table:

| Setting |   Flickr | VQAv2 | OKVQA |
| :-----:| :----:|:----: |:----: |
| Quantized Visual Tokens | 70.3 |  53.0  | 44.8 |    
| Merged Continuous Embeddings | **72.7**   |  **56.4**   | **46.2** |  

* As observed, the continuous visual features achieve better multi-modal understanding performance than quantized visual codes. This phenomenon is reasonable as quantized embeddings represent the centroid of continuous visual features, potentially leading to the loss of visual details. Therefore, when visual input serves as the condition (i.e., sequence like [image, text]), we use the merged continuous visual features as the input to LLM.

Reviewer comment (Experiments):
Show the performance difference with and without token merger. It would help to understand the importance of the token merger module.

Author response:
* Thanks for your constructive suggestions. We have added the ablation study about the effect of the proposed token merger. From the presented results, it can be observed a distinct performance drop in multi-modal understanding tasks when eliminating the token merger module. Especially for visual question answering tasks that need fine-grained visual information as the reference, directly dropping the unselected visual tokens will lose some visual details. The proposed token merger can preserve the image details to some extent by compressing the unselected image patches onto the retained ones.

| Setting  | Flickr | VQAv2 | OKVQA | GQA |
| :-----:| :----:|:----: |:----: |:----: |
| w/o Token Merger | 70.8 |  51.9  |  42.1 |  39.7  | 
| w/ Token Merger | **72.7**   |  **56.4**   | **46.2** | **42.3** |

Reviewer comment (Writing):
Figure 1 (a): `Adater-Style` -> `Adapter-Style`

Author response:
* Thanks for pointing out these typos in the draft. We have corrected them in the revision following your valuable suggestions.

Reviewer comment (Writing):
The first line in Section 3: `the Large language model` -> `the large language model`

Author response:
* Thanks for pointing out these typos in the draft. We have corrected them in the revision following your valuable suggestions.

Reviewer comment (Writing):
The last line in `Token Selector` of Section 3.1 : `the binary decision M` -> `the binary decision mask M`

Author response:
* Thanks for pointing out these typos in the draft. We have corrected them in the revision following your valuable suggestions.

Reviewer comment (Theory):
Equation 3: $||l_{2}(\hat{x_i})-l_{2}(c_{j})||$ -> $||\hat{x_i}-c_{j}||_2$

Author response:
* Thanks for pointing out these typos in the draft. We have corrected them in the revision following your valuable suggestions.

Reviewer comment (Experiments):
The authors train a visual codebook to assign discrete labels to visual tokens, which can be computationally costly. In the context of efficiently adapting LLM for vision-language tasks, it would be interesting to explore whether existing codebooks [1, 2, 3] can be utilized to provide useful supervision for training LLM.

Author response:
* We fully agree with the reviewer that exploring the influence of existing codebooks is meaningful. Actually, in our earlier experiments, we have tested the impact of different codebooks (e.g., VQ-VAE [1,2] and BEIT-2 [3]), but found that they were not effective in serving as a good visual tokenizer for multimodal LLM. The detailed comparisons between these codebooks and our proposed dynamic one are shown in the following Table. All these ablations are conducted on the same 30M pre-training data with clip ViTL/14 as the visual encoder.

| Method | Codebook Size |  Flickr | GQA | VQAv2 | OKVQA |
| :-----:| :----: |  :----: | :----:|:----: |:----: |
| VQ-VAE | 8192 | 60.5 |  35.0  | 48.8 | 40.2 |
| BEIT-2 | 8192 | 65.8  | 38.3 | 51.1 | 41.4 |
| Ours | 8192 | **70.0** |  **40.8**  | **53.0** | **44.5** |

* As observed, directly applying these existing codebook techniques to train LaVIT leads to sub-optimal downstream performance. We speculate that this may be due to the fact that the current codebooks were not explicitly tailored to accommodate text-oriented language models. For example, VQ-VAE's codebook is learned by reconstructing image pixels, which may cause the learned code to focus on low-level visual patterns and be incompatible with high-level word tokens. Although BEIT-2's codebook is derived by distilling high-level features from the CLIP teacher model, their visual codes do not have the 1D causal dependency and thus are not well compatible with LLM. Moreover, most of these existing codebooks tokenize an image into a long, fixed sequence (e.g., 1024 for DALLE [2]). Such a long sequence will invariably result in low efficiency, especially for generating images. Therefore, we introduce a well-designed dynamic visual tokenizer for better compatibility with LLM in the unified training.

Reviewer comment (Experiments):
Hard to be completely fair when compared against other baselines. For instance, compared to BLIP-v2, this work adds additional English text corpus and further fine-tunes the LLM. The significant performance improvement over previous work may be attributed to the additional training data and trainable parameters.

Author response:
* In terms of data, for a fair comparison, we employ the same multi-modal pre-training data (about 93M samples) as BLIP-2 [4] to exclude performance gains from additional data usage This data scale is smaller than that of most other methods. For example, the concurrent work Emu [5] uses 2.6B data, and KOSMOS [6] adopts about 3.1B image-caption pairs. Besides, the purpose of including English text corpus during pre-training is to preserve the already learned language understanding ability of LLM (e.g., the performance on linguistic benchmarks like MMLU) when acquiring good multimodal capabilities. The leveraged text corpus is Open-Source data (Redpajama [7]) for training the LLaMA model from scratch, which therefore, does not introduce new textual data and is just for alleviating the forgetting of previously learned knowledge.

* While LaVIT fine-tunes the LLM compared to BLIP-2, our objective is to enable the visual modality to integrate better with LLM like a foreign language. In this way, text-oriented LLM can not only understand visual modality input but also generate visual content like language. In contrast, BLIP-2 can only comprehend the visual inputs but lacks the capability to generate them. Moreover, tuning the LLM instead of the adapter with limited trainable parameters can leverage the remarkable reasoning capabilities of LLM to learn the interaction across different modalities.

Reviewer comment (Reproducibility):
The implementation details of the ablation studies could be clearer. It is confusing why the proposed method achieves different results in Tables 3, 6, and 7 when doing ablation studies.

Author response:
* All the ablation studies were conducted on part of pre-training data (about 30M image-text pairs) by using the ViTL/14 of CLIP as the visual encoder. The language model is scheduled to train for 10K steps in these experiments. All other hyper-parameters are set to the same values as those listed in Table 9 of the appendix. Due to the costly training resources, these ablation studies did not progress to the final 10K steps and were stopped when the model performance did not change much. But it's important to note that we make sure that the same step was taken in each set of experiments for a fair comparison. For example, the experiments in Table 3 (a) are stopped at 9K steps, Table 6 for 9.5K steps and Table 7 (a) for 9.2K steps. As a result, the reviewer will find slight performance differences (within 1% change) in Tables 3, 6, and 7.

Reviewer comment (Evaluation):
The authors claim that the learned codebook encodes high-level semantic concepts. However, the visualization results of codebooks in Figure 5 may suggest that the same code groups image patches based on low-level visual patterns, such as color and textures. It would be interesting to see additional image-level visualization results, similar to those in [4], to demonstrate that the codebook encodes high-level semantic concepts.

Author response:
* Thanks for your constructive suggestions. We have updated the visualizations for the learned codebook in the draft and added more examples of image patches belonging to the same visual code. As observed in Figure 8 of the appendix, the learned discrete visual codes can convey explicit high-level visual semantics. For example, code 13014 represents the human arm, code 7260 shows the part of the car, code 15235 represents the bus, and code 15090 indicates the donuts. These visualizations demonstrate the interpretability of the learned codebook.