Matryoshka Query Transformer for Large Vision-Language Models

Thank you for your review!

1. Measure of TFLOPs

We measured the TFLOPs by incorporating everything, i.e, including the cost of vision encoder and LLMs. Thank you for pointing out, we would like to incorporate this detail in our revised version.

2. Scenario of just obtain 256 tokens and take the first 8.

During the inference process, the inputs to MQT are queries and visual features from vit, then the cross attention mechanism will output the learned visual feature before feeding into LLM. So there is no softmax involved in this process. We also experimented to find out that indeed no performance changed in this scenario.

3. Sampling vs joint training.

For joint training, we can choose multiple $m$ tokens jointly. We think this design can slightly increase the computation of each training step, but may optimize faster from a global steps view. It’s also interesting to consider the design of how and where to choose the multiple $m$ tokens, for example choose $m$ from a local region (i.e, each $m$ is close with each other) or from a more global region (i.e, joint the the first token and all the 256 tokens). Thanks for your suggestion, we will leave this exploration to our future work.

4. Why log based is worse than fine-grained, not 3.5% as mentioned in the paper but 2.1% from the table.

We apologize for this confusion. We think the 2.1% absolute performance drop which is observed from the table is the same as 3.5% mentioned in the paper, which represents the relative performance change compared to our method. 3.5% is (our score (59.4) - abalated model score (57.3)) / 59.4. We only employed this relative score comparison in our ablation study table. We will incorporate this explanation into our revised version. As we discussed in line 249, this validated our hypothesis that gradually compressing the visual tokens helps the model perform better than log-based choices. Although the tested scenario is 256 tokens, the first few tokens of 256 are still “log-based” trained in the ablated model (for example, still large gaps between 64 and 128 tokens or 128 and 256 tokens), which experimentally found to be not as good as our final model.

5. How is a random token picked in the visualization?

For example, we provided the visualization when our model inference with 8 tokens, please refer to Figure 1 in our general response. We will also incorporate this visualization to our revised version. Here, say we pick the 6th token which is on the 0.75 percentile location. For other settings such as 16 tokens or 256 tokens, we fixed it to randomly pick one visual token for visualization purposes.

Thank you for your review!

1. What scale to choose for the best trade-off?

The trade-off visualization is from Figure 5 and the complete version is from Figure 8 in our paper’s appendix. For the best tradeoff, as we mentioned in line 215, we observed a ''turning point'' in many benchmarks. This point has a relatively lower computation but keeps a high performance. After this point, the performance will decrease more drastically although with a much lower computation. We think this point can be a good answer for the best trade-off.

2. Training Computation

Our main contribution is to achieve a single model for elastic inference of any visual tokens (line 7-9 in our Abstract). As mentioned in our methods (see line 102), we achieved this by choosing $m$ from {2, 4, 6, …, 252, 254, 256}, which means only the first 2 tokens are always supervised for the entire training data, whereas for the rest of visual tokens, less supervision is applied, especially for the visual tokens at the tailing positions. But our gradually reducing token numbers strategy helps stabilize this issue as we demonstrated robust performance in our experiments. Therefore, we apply 2 epochs of training data to help this issue. Gladly, LLaVA instruction tuning is very computational friendly. We only trained for 15 hours (one epoch) more on one A6000 machine, which is a cost with less concern.

3. Study on Q-Former

Concerning the Q-Former layers. We studied the layers for MQT transformers at the early stage of our research. We didn’t find improved performance with multiple layers. This similar design is also employed by Qwen-VL (Bai et al., 2023) which is a strong LVLM in the open-source community. Our extensive experiments in Table 1 and comparison to Q-Former in Figure 3 also demonstrated the sufficient capabilities of MQT design. We also incorporated LLaVA + vallina Q-Former results in the Table 1 of our general response.

Concerning Q-Former at all. Our MQT design is to achieve a single model which is trained only once for elastic inference of any visual tokens (line 7-9 in our Abstract). Query Transformer design has a learned number of visual tokens, therefore, during inference, an elastic number of visual tokens can be choosed for various downstream tasks to achieve our goal.

4. TextQA and ChartQA

Please refer to the general response’s Table 2 for added TextVQA results. We observe a very steady performance trend even when reducing the number of visual tokens to 16 tokens. In our paper’s Table 1, we have tested on 4 comprehensive evaluation benchmarks for Large Multimodal Models including MME, MMMU, MMBench, and MM-Vet, where our model demonstrated robust performance on these tasks. All 4 benchmarks involving evaluation on OCR and text related tasks, we report the specific scores for these sub categories in our general response’s Table 3, please refer to it. In summary, our model performs the best in these text related tasks with 3.2 points higher than LLaVA-1.5 on average.

As for ChartQA, since LLaVA-1.5 didn’t incorporate any training data for chart understanding, their paper didn’t evaluate on this dataset. When we downloaded their weight and conducted evaluation, it reveals LLaVA-1.5 performs very poorly on ChartQA. Therefore, this task may not be suitable for use as a reference for comparison. But we also provided these results here for your reference. LLaVA-1.5 with 576 tokens achieved 18.2 on ChartQA, and our model with only 256 tokens, reducing to [144, 64, 36, 16, 8, 4] achieved 14.3, 14.1, 14.0 , 13.6 , 13.7 , 12.7 , and 12.0 separately. This also indicates the strong robustness of our model in reducing numbers of visual tokens.

5. Attention map visualization and quantitative results.

We have conducted extensive quantitative evaluation in our paper’s Table 1 and Table 2 and demonstrated strong capabilities of our model. We conducted attention map visualization in order to understand what visual information is the model focusing on when using a lower number of visual tokens. The visualization on 1 visual token is not cherry-picked but randomly picked. For example, we provided the visualization when our model inference with 8 tokens, please refer to Figure 1 in our general response. We will also incorporate this visualization to our revised version. For other settings such as 16 tokens or 256 tokens, we fixed it to randomly pick one visual token for visualization purpose.

Thank you for your review!

1. Comparison with dynamic baselines.

For comparison with LLaVA-PruMerge, please refer to our Table 1 in general response. In summary, Our model significantly outperformed LLaVA-PruMerge in 5 out of 6 tasks, and by a 2.3 absolute score on average. Moreover, LLaVA-PrugeMerge employs a interquartile range (IQR) method to identify import visual tokens, the identified number is “fixed” in IQR. Whereas our MQT can enable a dynamic number of any visual tokens during inference.

As for dynamic in ViT encoder, the following discussion is directly cited from LLaVA-PruMerge paper Section 3.5 which corresponds with our thoughts as well. ''Block by block, tokens are gradually reduced in number, correspondingly reducing the computation cost in the internal ViT'', ''LMMs usually require a large stack of visual tokens… Thus, using previous token merging methods to obtain one refined class token as representation of visual input is not consistent with the literature of large multimodal models''.

To the best of our knowledge, we didn’t find a dynamic visual tokens baseline in ViT for us to compare. Besides, training a dynamic ViT for this purpose is very computationally intensive. This highlights the efficiency and lightweight nature of our MQT design.

2. Improvement in Low FLOPs regime.

Our main contribution is not to train a model which performs better than independently trained low FLOPs models. But rather to have a model that can train only once that can enable elastic inference of any number of visual tokens. Please refer to line 12-15 in our abstract and line 43-44 in our introduction section. The illustration of Figure 1 demonstrated that our model is at least as good as training (exhaustively 256) possible models by picking several choices of visual token numbers in this range for evaluation and comparison. And we are glad to find out that our model even achieved better performance in high FLOPs regime.

3. How is $m$ chosen at inference? from the method section?

Yes, your assumption is correct. Please also refer to line 50 in our introduction section: ''During inference, we have the flexibility to selectively utilize solely the initial $m$ visual tokens''. In the method section, we mentioned how $m$ is chosen for training time such as at line 99, therefore how $m$ is chosen at inference time is intuitive. We will incorporate this detail to the method section in our revised paper. Thank you.

4. Addressing Low FLOPs regime and MQT model starting with a lower $M$

Please first see the response to (2) above. For your question about the base MQT starting with a smaller $M$ and lower sparsity ratio, we believe you are probably referring to the baseline model with ''Fixed #Visual Tokens'' as illustrated in our Figure 1. Its results are indicated by the orange curve. If we misunderstand or you have further questions, please respond to our rebuttal and we are very happy to address your questions.

Thank you for your review!

1. How many layers does MQT transformer need as compared to Q-Former?

We studied the layers for MQT transformers at the early stage of our research. We didn’t find improved performance with multiple layers. This similar one layer design is also employed by Qwen-VL (Bai et al., 2023) which is a strong LVLM in the open-source community. Our extensive experiments in Table 1 and comparison to Q-Former in Figure 3 (''With only 2 visual tokens, MQT-LLAVA outperforms InstructBLIP on all 8 tasks'') have demonstrated the sufficient capabilities of MQT design.

2. Should text prompts be involved in the visual feature extraction process?

Thanks for the detailed observation of our model. Most current LVLMs (except BLIP2 has an extra text encoder) do not use text prompts as input for extracting visual information. All the visual features are provided to LLM, together with the input question. LLM can then employ the textual information and identify useful visual features in answering the text question. We acknowledge this is a good suggestion in the lower number of visual tokens scenarios, whereas the model can pre-select the useful visual feature with textual information before extracting them. But an extra text understanding model needs to be properly designed for this purpose. We will leave this suggestion to future study of our work, but we believe this is orthogonal to the main claims of our paper.

3. TextVQA Performance

Please refer to the general response’s Table 2 for added TextVQA results. We observe a very steady performance trend even when reducing the number of visual tokens to 16 tokens. In our paper’s Table 1, we have tested on 4 comprehensive evaluation benchmarks for Large Multimodal Models including MME, MMMU, MMBench, and MM-Vet, where our model demonstrated robust performance on these tasks. All 4 benchmarks involving evaluation on OCR and text related tasks, we report the specific scores for these sub categories in our general response’s Table 3, please refer to it as well. In summary, our model performs the best in these text related tasks with 3.2 points higher than LLaVA-1.5 on average.

4. How does the model perform when the number of visual tokens exceeds the original 576 tokens?

Technically, the number of visual tokens is bounded by the original number since that’s the maximum visual information that a vision encoder can represent. But we did experiment with this interesting idea at an early stage. When we initialize the query tokens as sinusoidal embeddings similar to positional embeddings, then we can extrapolate the number of query tokens beyond the original limit. However, we didn’t find significant improvement in performance and the extrapolated visual tokens are also non intuitive in their meanings.

5. Adaptive Pooling or Vanilla Q-Former baseline.

Our main idea is to train a model once and enable elastic inference of any number of visual tokens (under the maximum) for various computation resources and downstream tasks (line 7-9 and 12-14 in our Abstract). Adaptive pooling based approach will limit the choice of visual tokens to only several options. Whereas for Vanilla Q-Former, we conducted preliminary experiments at the starting point of our project, please refer to the detailed scores in our general response’s Table 2. In summary, vanilla Q-Former doesn’t demonstrate strong capabilities for our tasks.

6. Evaluation on more extensive heavy computation cases like LLaVA-Next.

Our MQT design can be served as a plug-and-play module to other models (as mentioned by reviewer jmsV, too), we adopted LLaVA-1.5 as the base model to validate our methods and have demonstrated strong capabilities. We plan to extend our MQT method to videos and Interleaved image text tasks like LLaVA-Next as well. Recent work such as SlowFast-LLaVA (Xu et al., 2024) from Apple have demonstrated promising results in this direction. However, these are out of the scope of this paper and can be considered as a future work direction. We have released source code and models for the community to try out the ideas for other models as well.