Matryoshka Query Transformer for Large Vision-Language Models

Authors: wenbo  hu
zi-yi  dou
liunian harold li
amita  kamath
nanyun  peng
kai-wei  chang
Venue: NeurIPS_cc_2024_Conference
Published: 2024-05-14 19:09:01

Abstract: Large Vision-Language Models (LVLMs) typically encode an image into a fixed number of visual tokens (e.g., 576) and process these tokens with a language model. Despite their strong performance, LVLMs face challenges in adapting to varying computational constraints. This raises the question: can we achieve flexibility in the number of visual tokens to suit different tasks and computational resources? We answer this with an emphatic yes. Inspired by Matryoshka Representation Learning, we introduce the Matryoshka Query Transformer (MQT), capable of encoding an image into $m$ visual tokens during inference, where $m$ can be any number up to a predefined maximum. This is achieved by employing a query transformer with $M$ latent query tokens to compress the visual embeddings. During each training step, we randomly select $m \leq M$ latent query tokens and train the model using only these first $m$ tokens, discarding the rest.
Combining MQT with LLaVA, we train a single model once, and flexibly and drastically reduce the number of inference-time visual tokens while maintaining similar or better performance compared to training independent models for each number of tokens. 
Our model, MQT-LLaVA, matches LLaVA-1.5 performance across 11 benchmarks using a maximum of 256 tokens instead of LLaVA’s fixed 576. Reducing to 16 tokens (8x less TFLOPs) only sacrifices the performance by 2.4 points on MMBench. On certain tasks such as ScienceQA and MMMU, we can even go down to only 2 visual tokens with performance drops of just 3\% and 6\% each.
Our exploration of the trade-off between the accuracy and computational cost brought about by the number of visual tokens facilitates future research to achieve the best of both worlds.

Main paper: 
 introduction: 
 Recent work in Large Vision-Language Models (LVLMs) (OpenAI, 2023; Liu et al., 2023b; Bai et al., 2023) has shown remarkable performance across a broad range of vision-language tasks (Huang et al., 2023; Chen et al., 2023; Cai et al., 2024a; Li et al., 2023b) . These LVLMs typically consist of a vision encoder to embed images into grid features, which are fed into a Large Language Model (LLM) (Touvron et al., 2023; Chiang et al., 2023) for processing and reasoning alongside a text input.A key research question is how to transform these raw visual embeddings into the visual tokens that are fed into the LLM. Prior work either directly projects the grid features with a multi-layer perceptron (MLP) (Liu et al., 2023b) or compresses the grid features into fewer tokens with a query transformer or resampler (Li et al., 2023a; Dai et al., 2023; Bai et al., 2023; Ye et al., 2023; Alayrac et al., 2022) . However, these models all need to pre-determine how many tokens an image is worth, and set a fixed number for all images. Finding a flexible number that adaptively strikes a balance Figure 1 : Our model, MQT-LLAVA, matches LLaVA-1.5 performance on 11 benchmarks using only 256 visual tokens instead of 576. We achieve a 2x speed-up with 256 tokens and 8X speed-up in TFLOPs using 16 tokens with only a 2.4 performance drop compared to LLaVA-1.5 on MMBench.between efficiency and performance is difficult. More visual tokens encode more information, but come at a higher inference cost, as the complexity of the transformers used in these LVLMs scales quadratically with the number of input tokens. Additionally, not all applications require or allow the same token budget: some applications have limited computational resources, necessitating a lower token budget to ensure real-time processing. In practice, most best-performing LVLMs choose a fixed, large number of visual tokens per image (e.g., 576 for LLaVA-1.5) without the ability to adaptively adjust the visual token allocation at deployment time.In this work, inspired by Matryoshka Representation Learning (MRL) (Kusupati et al., 2022; Kudugunta et al., 2023) , we introduce Matryoshka Query Transformer (MQT), a simple way to train a single LVLM that supports adaptively changing the number of visual tokens at inference time. We use a query transformer (Li et al., 2022; Alayrac et al., 2022) with M latent query tokens to transform grid features into visual tokens. Crucially, during each training step, we train the model using only the first m latent query tokens while dropping the rest, where m is randomly selected within the range of M . With such a tail-token dropping strategy, the query tokens form a Matryoshka structure. Intuitively, the significance of each token correlates with its placement within this nested structure. During inference, we have the flexibility to selectively utilize solely the initial m visual tokens.We combine MQT with LLaVA-1.5: the resulting model, MQT-LLAVA, is able to match LLaVA-1.5 performance across 11 benchmarks using only a maximum of 256 tokens, instead of LLaVA's fixed 576. When the maximum number of tokens is dropped drastically to only 2 tokens, MQT-LLAVA performance drops by only 3% on ScienceQA and 6% on MMMU. Finally, we study the performance of 2, 4, 8, 16, 36, 64, 144, and 256 visual tokens during inference across 11 benchmarks, and offer a trade-off in the selection of visual tokens that balances achieving the highest accuracy with minimizing computational costs on different tasks. Interestingly, we find that changing the number of visual tokens impacts different tasks very differently. For instance, tasks involving language-based reasoning and subject-level scientific knowledge can achieve excellent performance with only a few tokens, whereas complex open-ended visual question tasks that involve rich local information details require a larger number of tokens.In summary, we make the following key contributions:• We introduce Matryoshka Query Transformer (MQT), which allows for a flexible choice of the number of visual tokens and accommodates varying computational constraints in different tasks. • Leveraging MQT, we build MQT-LLAVA, a vision-language model that matches the performance of LLaVA-1.5 using less than half the number of visual tokens, and outperforms it in 6 out of 11 benchmarks. • We further explore the performance and computation trade-offs across 11 tasks and demonstrate that a significant speed-up can be achieved with minimal performance drop by reducing the number of visual tokens (e.g., 8X fewer TFLOPs with 2.4 points drop on MMBench).Pre-trained LLMFigure 2 : Our model employs a query transformer to encode images as visual tokens. We randomly select the first m tokens during training, and enable flexible choice of any m number under M during inference, where M is the maximum number of initialized tokens.Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?Answer: [Yes] Justification: We clearly state our main claims in abstract and introduction.Guidelines:• The answer NA means that the abstract and introduction do not include the claims made in the paper. • The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers. • The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings. • It is fine to include aspirational goals as motivation as long as it is clear that these goals are not attained by the paper. 
matryoshka query transformer: 
 Preliminary: Matryoshka Representation Learning (MRL). MRL (Kusupati et al., 2022; Kudugunta et al., 2023) involves training models with nested dimensions to learn representations at multiple granularities, enabling adaptive deployment per computational constraints. MRL defines a series of models f 1 , f 2 , . . . , f M with the same input and output space but growing hidden dimensions.The name "Matryoshka" comes from the fact that the parameters of f m are contained by f m+1 . For example, in Kudugunta et al. (2023) , {f m } are a series of Transformers with the same depth but different widths. Consider a specific Feed Forward Network (FFN) block in f M that has d M neurons in the hidden layer. Then, the FFN block in f m will contain the first d m neurons, andd 1 ≤ d 2 ≤ • • • ≤ d M .MRL then trains these models jointly with the following loss:EQUATIONwhere L is the loss function and y is the ground truth label. Note that for each training step, MRL performs forward and backward passes for all M models, inducing significant training overhead compared to training one model. After training, MRL can perform inference with any hidden dimension d i≤M , enabling flexible deployment based on specific needs. MRL is our motivation to train LVLMs that can perform inference with a flexibly selected number of visual tokens.We first explain how we encode images with a query transformer, then discuss our training paradigm.Encoding images with a Query Transformer. We employ a query transformer-based architecture to extract visual tokens from images following previous work (Li et al., 2022; Bai et al., 2023) . Specifically, an input image x is first processed by an image encoder and are then flattened intoH × W grid features G = [g 11 , • • • , g 1W , • • • , g H1 , • • • , g HW ].Then, a query transformer Q is applied to compress the grid features to M visual tokens. Specifically, Q assumes a set of latent query tokens Z = [z 1 , . . . , z M ] as input, where M is usually smaller than H × W . The query tokens cross-attend to the grid features and compress the information into the query tokens. The final-layer query tokens become the visual tokens V that are fed to a large language model together with the input text tokens. I.e., V = Q(Z, G). A linear projection layer is added in the end to match the hidden size of the language model.1 Matryoshka Query Transformer. To enable elastic inference, given the M latent query tokens Z = [z 1 , . . . , z M ], at each training step, we feed only the first m query tokens to the query transformer Q. Subsequently, we obtain only m visual tokens from the query transformer. m can be any number equal to or smaller than the maximal token number M . In practice, we choose m from a linear set of maximum dimensions, in increments of 2, e.g. m can be any number in {2, 4, 6, . . . , 252, 254, 256} when M = 256. From a training efficiency perspective, our approach uses, on average, half of the visual tokens compared to the original query transformer-based models.Formally, given an input image with its corresponding text question q and answer y, at each training step, we randomly select a m and feed the first m latent tokens Z 1:m and the text question q to the model. We compare the model output and y and minimizeEQUATIONLM is the language model, L is the language modeling loss function, and c m is a constant coefficient to control the weight of different numbers of visual tokens, which is always set to 1 in our setting.Discussion. Here we discuss several interesting properties of MQT. (1) Unlike the original matryoshka representation learning that maintains a nested structure in the parameter space, we specifically target LVLMs and make the visual tokens Matryoshka-like. ( 2 Question: Does the paper discuss the limitations of the work performed by the authors?Answer: [Yes]Justification: See Limitation in Appendix B Guidelines:• The answer NA means that the paper has no limitation while the answer No means that the paper has limitations, but those are not discussed in the paper. • The authors are encouraged to create a separate "Limitations" section in their paper.• The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be. • The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated. • The authors should reflect on the factors that influence the performance of the approach.For example, a facial recognition algorithm may perform poorly when image resolution is low or images are taken in low lighting. Or a speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon. • The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size. • If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and fairness. • While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an important role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations. 
experiments: 
 We first introduce the implementation details of our query transformer architecture ( §3.1). We then show the empirical performance of our approach compared to state-of-the-art models across 11 benchmarks ( §3.2. Finally, we further study the performance-efficiency trade-off ( §3.3).MQT-LLAVA Implementation Details. We implement our models based on LLaVA-1.5 (Liu et al., 2023a) , except that we use our Matryoshka Query Transformer instead of an MLP to obtain the visual tokens. The MQT is a single-layer Transformer with cross-attention. Following Liu et al. (2023a) , we select CLIP ViT-L/14 (Radford et al., 2021) as our vision encoder, supporting 336x336 image resolution, and Vicuna-v1.5 (Chiang et al., 2023) as our LLM. As studied in Hu et al. (2023) ; Zhu et al. (2023) ; Liu et al. (2023b) , we adopt a two-stage training approach. We train only the query transformer in the first-stage alignment, using LLaVA-558K for 1 epoch with a batch size of 256 and a learning rate of 1e-3. We then fine-tune both the query transformer and LLM using LLaVA-665K for 2 epochs with a batch size of 128 and a learning rate of 2e-5. All training is on 8xA6000s, with 4 and 30 hours per stage, respectively. We apply MQT during the second stage (c.f. §4.3).Baselines. As shown in Table 1 , we compare our model with LLaVA-1.5 (Liu et al., 2023a) and our model's baseline LLaVA query transformer (QT-LLaVA), which is trained with a fixed number of 256 visual tokens across all training stages. We also list other models' results for comparison, including BLIP-2 (Li et al., 2023a) , InstructBLIP (Dai et al., 2023) , Shikra (Chen et al., 2023) , IDEFICS (IDEFICS, 2023), and Qwen-VL (Bai et al., 2023) .Evaluation Benchmarks. We evaluate our model across 11 mainstream benchmarks, including VizWiz (Gurari et al., 2018) , ScienceQA-IMG (Lu et al., 2022) , VQA-v2 (Goyal et al., 2017) , GQA (Hudson and Manning, 2019), POPE (Li et al., 2023c) , MME Perception (Fu et al., 2023) , MME Cognition (Fu et al., 2023) , MMBench (Liu et al., 2023c) , LLaVA-Bench (In-the-Wild) (Liu et al., 2023b) , and MM-Vet (Yu et al., 2024) .Table 1 presents the results of MQT-LLAVA with inference visual token budgets of 2, 4, 8, 16, 36, 64, 144, and 256 . We refer to the baseline approach, where the model is trained with a fixed number of visual tokens across all training stages, as LLaVA Query Transformer (QT-LLaVA outperforms the baseline QT-LLaVA with 256 tokens in 9 out of 11 benchmarks. One possible explanation is that by enforcing our model to only see fewer tokens during training, the stricter constraint helps the model generalize better to unseen tasks. This is especially evident in the higher performance on VizWiz. When compared to open-source state-of-the-art models, our model with 256 tokens achieves on par or better than LLaVA-1.5 performance with 576 tokens across 11 benchmarks, outperforming it in 6 out of 11 benchmarks. Even with 64 tokens, our model falls short of LLaVA-1.5 by only 0.9 points on average. When drastically drop to only 2 tokens, our score falls by only 3% on ScienceQA and 6% on MMMU. While directly adding a query transformer to LLaVA degrades the performance, our strategy can achieve comparable or better performance than LLaVA-1.5. We explore performing inference using a variety of numbers of visual tokens, including 1) an extremely low number of tokens; 2) a number of visual tokens unseen during training. As shown in Figure 3 , MQT-LLAVA with only 2 visual tokens outperforms InstructBLIP (Vicuna-7B), which is based on Q-Former (Li et al., 2023a) using 32 visual tokens. This demonstrates the effectiveness of our model in compressing visual information, pointing to its potential use for applications in computation-heavy tasks. For an unseen number of visual tokens, we pick a random number of visual tokens: 77, and include its results in Appendix C. Despite never being explicitly trained for this number of tokens, our model can generalize to any number within 256 during inference, demonstrating a further benefit of our elastic approach.To demonstrate our computational efficiency, we compute TFLOPs when running MQT-LLAVA on MMBench with 8, 16, 36, 64, 144, and 256 visual tokens, compared to LLaVA with 576 tokens. As shown in Figure 1 , we are able to achieve significant speed-ups with little-to-no performance loss: our model with 256 and 144 tokens respectively achieve a 2x and 3x speed-up compared to Figure 4 : Grad-CAM visualization of 1 randomly picked token from using 8, 16, 64, 256 visual tokens, respectively, to encode an image. The model effectively concentrates on high-level concepts using fewer tokens and delves into low-level details with more tokens. The complete input to the third image is "List all the objects on the desk. The objects on the desk include a computer monitor, a keyboard, a mouse, a cell phone, and a pair of headphones".LLaVA-1.5 while maintaining the same or even better performance; and when using 16 tokens, we achieve an 8x speed-up with a performance drop of only 2.4 points.Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?Answer: [NA] Answer: [Yes] Justification: We provide our code along with scripts to supplemental material . All the code and models will be open sourced. Guidelines:• The answer NA means that paper does not include experiments requiring code. Justification: We discuss ethics concern and broader impact in Appendix A.Guidelines:• The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.• If the authors answer No, they should explain the special circumstances that require a deviation from the Code of Ethics. • The authors should make sure to preserve anonymity (e.g., if there is a special consideration due to laws or regulations in their jurisdiction). 
analyses: 
 To better understand the meaning of visual tokens and to systematically study the number of tokens required by different vision-language tasks, we investigate two key questions: (1) How does the focus of the model change with varying numbers of visual tokens? ( §4.1); and (2) How do different numbers of visual tokens impact various tasks? ( §4.2)To explore what visual information each token encodes, we utilize Grad-CAM (Selvaraju et al., 2017) to visualize the focus of visual tokens. As illustrated in Figure 4 , we qualitatively analyze the results of using 8, 16, 64, and 256 tokens.We observe that the model's focus changes with the number of tokens used. When using a few tokens (e.g., 8), the model accurately concentrates on global visual concepts related to the question. As the number of tokens increases (e.g., 256), the model not only attends to the relevant objects but also delves into localized details. For example, in the third image, with 8 tokens, the model focuses on the monitor. With 16 tokens, it includes both the monitor and the mouse. With 64 tokens, it highlights the monitor and keyboard. Finally, with 256 tokens, the model encompasses several objects, including the monitor, keyboard, and cell phone. In the examples from the first and second images, our model effectively focuses on the man ironing behind the car and the two cats, even with only 8 tokens. The impressive qualitative results, especially those using only a few tokens, demonstrate the effectiveness of our approach and the strong capabilities obtained despite using a minimal number of tokens.When using varying numbers of visual tokens during inference, we observe that the model's performance change varies across different tasks. Tasks requiring a large number of visual tokens. Tasks that require fine-grained visual understanding and deep reasoning across multiple areas of the image naturally demand a higher number of visual tokens for optimal performance. When the number of visual tokens decreases, the encoded image information is reduced, leading to performance degradation. This trend is evident in tasks such as VQAv2, GQA, VizWiz, MMBench, LLaVA-Bench, and MM-Vet. As illustrated in Appendix Figure 8 , the performance on these tasks gradually declines as the number of visual tokens decreases from 256, with a more rapid decline observed when the tokens are further reduced.Tasks robust to visual token reduction. In contrast, for several benchmarks primarily targeting the visual perception skills of models, performance remains consistent when gradually reducing the number of visual tokens until a threshold is reached. Beyond this threshold, performance drops significantly (see Figure 5 and Appendix Figure 8 ). This "turning point" is observed in benchmarks such as MME Cognition, MME Perception, POPE, and MMMU.For instance, in MME-Cognition (see Figure 6 ), tasks involving commonsense reasoning, code reasoning, and numerical calculation can be performed effectively with as few as 16 visual tokens, allowing the model to focus on the relevant image sections. Similar results are seen in other tasks, like the hallucination question "Is there a car in the image?" from POPE. However, once the "turning point" is reached, further reducing the number of visual tokens prevents the model from attending to the correct objects, leading to a sharp decline in performance.Another notable observation comes from ScienceQA and MMMU, which contain subject-specific questions from school curricula. The model's performance on these tasks remains robust despite a decrease in visual tokens, achieving scores of 65.0 and 32.5, respectively, with only 2 tokens. This suggests that the reasoning required for academic questions is primarily conducted by the language model (LLM); even with minimal visual hints, the LLM can interpret the image content and perform the reasoning tasks effectively.When are fewer visual tokens better? As shown above, MQT-LLAVA with 16 tokens can achieve better performance on ScienceQA compared to MQT-LLAVA with 144 tokens. To understand why fewer tokens may benefit this task, we qualitatively analyze instances where MQT-LLAVA succeeded with 16 visual tokens, but failed with 144. We show a representative example in Figure 7 . MQT-LLAVA with 16 visual tokens attends to all three objects, allowing it to understand their mutual relationship and answer the question correctly. On the other hand, with 144 visual tokens, MQT-LLAVA focuses on various portions of the image and attend to each object independently. This discourages the model from reasoning with the common attributes among the three objects, thus predicting the wrong answer. In summary, fewer visual tokens seems to be preferable when fine-grained visual understanding is not required.However, it should be noted that using fewer tokens is not always better in this case. As shown in Figure 7 , MQT-LLAVA with 144 tokens precisely identified state of Virginia on the map and answered the question correctly. Whereas 16 tokens concentrated on another region which potentially confused its final prediction, lacking the abilities of distinguishing local details of the geographic shape on the map.We ablate several design choices across 11 benchmarks in Table 2 . Each ablation independently modifies our best variant, MQT-LLAVA, to create new variants. (i) linear vs. log-based token number selection. We replace our linear growth elastic tokens, i.e., m ∈ {2, 4, 6, . . . , 252, 254, 256} to the log-based approach of MRL, i.e., m ∈ {2, 4, 8, 16, . . . , 128, 256}. This results in an average accuracy of 57.3%, 2.1% lower than MQT-LLAVA, validating our hypothesis that gradually compressing the visual tokens helps the model perform better than log-based choices. (ii) query transformer architecture. As mentioned in §2, we choose to first perform cross-attention between query tokens and visual features, then project the learned visual tokens to the LLM. We call this technique "attention then projection". The alternative variant is "projection then attention", which achieves lowest average performance, with a score of 56.6%. This suggests that directly applying the attention mechanism helps preserve the rich grid features, making them better projected to the LLM. (iii) first-stage pretraining with query transformer. As mentioned in §3.1, we choose to apply our elastic training paradigm only during the second stage. Experimental results demonstrate that adopting elastic training during the first stage leads average performance dropped by 2.7%. We hypothesize that the first stage aims to align the randomly initialized query tokens with vision-language awareness.Therefore, it is important to train all 256 tokens with this prior knowledge before reducing the number of tokens in the second stage.5 Related Work 2023) have successfully showcased promising results on a wide variety of vision-language perception and reasoning tasks. Recent works further expand the capabilities of LVLMs to region-level image understanding (Huang et al., 2023; Peng et al., 2023; Chen et al., 2023; Cai et al., 2024a) , video understanding (Zhang et al., 2023; Li et al., 2023b; Jin et al., 2024) and 3D understanding (Hong et al., 2023; Szot et al., 2024) . These models mostly consist of a vision encoder and an LLM aligned by a vision-language connector module, which can be an MLP (Liu et al., 2023b; Hu et al., 2023) , Q-Former or queries through cross attentions (Dai et al., 2023; Bai et al., 2023; Ye et al., 2023) , or Resampler (Alayrac et al., 2022) . The number of visual tokens output by these modules can be very large, especially for higher image resolutions, multiple frames in video tasks, and multiple images for in-context learning. In this paper, we propose a training paradigm that can dynamically choose a number of visual tokens that adapts to variable computation costs at inference time. 
efficient vision transformers: 
 Reducing the computational cost of LVLMs at deployment is an active area of research. Several works, e.g., TinyLLaVA (Zhou et al., 2024) and LLaVA-Phi (Zhu et al., 2024) , reduce the size of the LLM backbone by replacing it with a smaller one, e.g., Phi-2 (Javaheripi et al., 2023) . MobileVLM (Chu et al., 2023) and MobileVLM-v2 (Chu et al., 2024) focus on a compact architecture design and training paradigms specifically for mobile usage. In these cases, the computation reductions come from reducing the size of either the LLM or vision encoder backbones, whereas our method focuses on increasing LVLM efficiency by dynamically reducing the number of visual tokens.A long-standing issue with vision transformers is that the attention mechanism introduces computational complexity that scales quadratically with the input tokens. Vision Longformer (Zhang et al., 2021) adopts sparse attention (Kitaev et al., 2020) to speed up vision transformers for larger inputs.Other works design various strategies to retain the most informative tokens and reduce the number of visual tokens at the inference stage (Fayyaz et al., 2021; Liang et al., 2022; Rao et al., 2021; Bolya et al., 2023; Yin et al., 2022) . Most similar to our work is SparseFormer (Gao et al., 2024) which employs cross-attention to learn sparse representations of both latent tokens and RoI descriptor tokens. In this work, we use the most simple query transformer architecture to actively control the number of visual tokens and explore the impact of reducing the number of visual tokens in LVLMs.Concurrent work (Cai et al., 2024b ) also studies matryoshka-style visual tokens, and designs 5 scales of pooling layers to control the granularity of images. Different from their work, we introduce a query transformer to extract visual tokens, enabling a more flexible choice of any number of visual tokens under a predefined maximum. Their work corroborates our findings by demonstrating robust performance and efficient use of a minimal number of visual tokens. 
conclusion: 
 In this work, we present MQT-LLAVA, a single vision-language model that enables elastic inference on various downstream tasks and computation resources. We demonstrate that our model achieves performance comparable to or better than training with a fixed number tokens. MQT-LLAVA matches the performance of LLaVA-1.5 across 11 benchmarks using less than half the number of visual tokens, and outperforms LLaVA-1.5 in 6 out of 11 benchmarks. We achieve an 8x less TFLOPs when reducing to 16 tokens while only sacrificing the performance on MMBench by 2.4 points. We hope our exploration of the trade-off between the accuracy and computational cost caused by the number of visual tokens will facilitate future research to achieve the best of both worlds. 
broader impacts: 
 Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?Answer: [Yes] Justification: We discuss broader impact in Appendix A.Guidelines:• The answer NA means that there is no societal impact of the work performed.• If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact. • Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.• The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is a direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that an improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that a generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster. • The authors should consider possible harms that could arise when the technology is being used as intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology. • If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML). 
safeguards: 
 Question: Does the paper describe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?Answer: [Yes] Justification: We require users to follow the guidelines such as LLM (Llama and Vicuna)'s guidelines when release the model.Guidelines:• The answer NA means that the paper poses no such risks.• Released models that have a high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters. • Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images. • We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make a best faith effort.12. Licenses for existing assets Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?Answer: [Yes] Justification: We properly credited the original owners and followed their license.Guidelines:• The answer NA means that the paper does not use existing assets.• The authors should cite the original paper that produced the code package or dataset.• The authors should state which version of the asset is used and, if possible, include a URL. • The name of the license (e.g., CC-BY 4.0) should be included for each asset.• For scraped data from a particular source (e.g., website), the copyright and terms of service of that source should be provided. • If assets are released, the license, copyright information, and terms of use in the package should be provided. For popular datasets, paperswithcode.com/datasets has curated licenses for some datasets. Their licensing guide can help determine the license of a dataset. • For existing datasets that are re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided. 
 
 Appendix: The deployment and release of MQT-LLAVA carry both potential benefits and risks. These considerations include visual aspects as well as common issues found in existing LLMs like Alpaca and Vicuna. Since MQT-LLAVA is built on LLaMA, Vicuna, and CLIP, it inherits certain challenges associated with LLMs and vision encoders.Hallucination Similar to other LLMs, MQT-LLAVA might produce outputs that are not based on factual information or input data. This raises concerns about the accuracy of inferences, particularly in critical applications such as medical fields.Biases Biases present in the base models can be brought to MQT-LLAVA, stemming from both the vision encoder (CLIP) and the language decoder (LLaMA/Vicuna). This may result in biased outcomes or unfair representations of diverse content.Energy Consumption We followed LLaVA-1.5's datasets (smaller datasets compared to other methods) in training our model, which makes energy consumption is not a primary concern.Despite our remarkable performance, one limitation of our work is that the maximum number of tokens MQT-LLAVA can accommodate at inference time is 256. We leave the exploration of using larger numbers in inference than training time to future work.We include how number of visual tokens impact the different tasks differently across all 11 benchmarks in Figure 8 and Figure 9 We present the results of choosing a random number of visual tokens, 77 as shown in Table 3 , to demonstrate our flexibility in selecting any number of tokens during inference.To demonstrate that the visual tokens used for visualization in Figure 4 are not cherry-picked, we present all the first eight tokens in Figure 10 .Figure 8 : The number of visual tokens impact the different tasks differently across 11 benchmarks. We log scaled x-axis for readability. Figure 10 : Grad-CAM visualization from all the tokens in our model when inference with 8 tokens. Input: "How many cats are there in the image? Answer: 2".NeurIPS Paper ChecklistUnlike previous work(Bai et al.,2023;Ye et al., 2023) that first applies projection followed by attention, we empirically find that our "attention then projection" architecture performs better (c.f. §4.3). 


Decision: This paper unanimously receives positive rates thanks to the novel use of visual tokens and convincing experiments. Although overall reviews are positive, the rebuttal should be reflected in the final draft.

Reviews: review: paper_summary: paper_summary: This paper addresses the challenge of achieving flexibility in the number of visual tokens to suit different tasks and computational resources. Inspired by Matryoshka Representation Learning (MRL), the authors propose the Matryoshka Query Transformer (MQT), which allows for any number of visual tokens during inference. Experimental results demonstrate considerable performance across varied visual token lengths and show promising results even with extreme token numbers.
strength_weakness: strength_weakness: 1. Can you explain the choice of the MQT transformer design? Does it need to maintain sufficient capability for extracting information from varied tokens? For example, Q-Former has multiple layers to achieve effective information extraction.

2. The extraction process V=Q(Z,G) lacks correlation with textual information. In Fig 4, different text prompts should result in the same attention map due to this. Therefore, this design might be limited when dealing with complex images, where visual representations should maintain different semantics.

3. How does the model perform in text-rich scenarios, such as TextVQA? The sophisticated recognition ability might be significantly compromised in these cases.

4. Since the number of visual tokens affects performance, how does the model perform when the number of visual tokens exceeds the original 576 tokens?

5. It would be better to provide a comparison with the LLaVA baseline, using adaptive pooling or vanilla Q-former to adjust token numbers.
questions: questions: As in LLaVA-1.5, the 576 visual tokens do not introduce significant computational overhead. It would be more meaningful to evaluate this method on a more extensive case like LLaVA-NEXT.
limitations: limitations: Please consider answer the questions above.
paper_summary: paper_summary: The paper considers multimodal vision transformers, in which we have a stream of both visual and textual tokens. Current architecture typically assume a fixed number $m$ of visual tokens. In contrast, the proposed `MQT` aims to achieves a dynamic number of visual tokens. This is done by integrating a form of compression+dropout mechanism during training. Specifically, the image tokens are compressed to $M$ tokens using cross-attention. To vary the number of tokens, the model only needs to use only  $m < M$ queries from te cross-attention mechanism. At training time, this is done by randomly selecting $m$ and feeding the **first** $m$ queries.
strength_weakness: strength_weakness: - **Lack of comparison with dynamic baselines**: in the field of image-only ViT, there is a flourishing literature on dynamically reducing the number of tokens using merging or pruning or scale selection. I am not as familiar with this literature for multimodal models, but there does seem to be similar existing approaches: for instance, *LLaVA-PruMerge: Adaptive Token Reduction for Efficient Large Multimodal Models* is one, though it might be too recent for the authors to have considered it for the submission. However, more generally, having a simple off-the-shelf token merging/pruning approach for the ViT encoder would be interesting.

- **No improvement over baselines in the low FLOPs regime**. From Figure 1, it seems that the benefits of the proposed `MQT` decrease in the low FLOPs regime, when increasing the sparsity ratio of $M / m$.  This limits the usefulness of the methods in the realm of model efficiency, though it seems to be a useful and simple training strategy for the high FLOPs regime.
questions: questions: - How are the $m$ latent queries chosen at inference ? I am assuming $m$ is chosen as a hyperparameter, and the first $m$ queries are kept, however it is not very clear from the methods section

- Addressing the **low FLOPs regime**:   From Figure 1, it seems that the benefits of the proposed `MQT` decrease in the low FLOPs regime, when increasing the sparsity ratio of $M / m$. I think it is not necessarily a surprising result as increasing sparsity too much will severely impact accuracy. However, I am wondering how would the trade-off curves would look like when starting with a base `QRT` model with a smaller $M$ and lower sparsity ratio, i.e. with a model that would achieves same number of tokens/FLOPs, but where the relative sparsity ratio to achieve is less aggressive.
limitations: limitations: The paper has a limitation section but in my opinion it does not fully address some limitations of the paper (e.g. no improvement in performance in the low FLOPs regime)
paper_summary: paper_summary: Authors try to use Matryoshka mechanism to guide the learning process of LVLMs such as llava. 

Authors show a pretty good scaling curve with different amount of visual tokens.
strength_weakness: strength_weakness: 1. For a given image, which scale should I choose to achieve best trade off? Is there an answer?
2. Authors use more compute (2 epoch) to get the similar performance as LLaVA, is there any reason for this?
3. All designs are center upon the qformer. There is little study on this. How many layers do we need? Do we really need it?
questions: questions: 1. What are the performance under other datasets like TextVQA and ChartQA? Those are dataset different from the natrual image distribution. 
2. The visualizations are good but we don't have quantitive numbers to support your claim. Since those pictures can be cherry picked.
limitations: limitations: See comments above.
paper_summary: paper_summary: This paper introduces a new concept called Matryoshka Query Transformer (MQT) that brings in the concept of Matryoshka information packing to make visual tokens flexible and can be used in multimodal vision language models like LLaVA. The reduced tokens tend to reduce the quadratic complexity of the language model due to potential flexibility in the processed visual tokens. 

Owing to this, there are significant efficiency benefits all the while retaining accuracy across various benchmarks. The paper also showcases analysis and further discussion on the technique. 



----------------------------
The review will be short and does not reflect the time put in for the review or the quality of the paper. When the ideas are simple and clear -- I tend to write shorter reviews to the point.
strength_weakness: strength_weakness: Most of these are not weaknesses, but rather questions

1) How are you measuring TFLOPs? What does it include, the cost of vision encoder? LLM processing and generation? I might have missed this in the paper and any pointer would be great.
2) While I understand the when you are picking a fixed # token for MQT, you make that choice before in hand. What happens if you actually just obtain 256 tokens and take the first 8? What will be the performance in that case? I know softmax etc make huge differences, but would be good to see that. 
3) Any thoughts on sampling vs joint training?
4) Why is log-based (actually exponential) spacing so poor (not 3.5% as mentioned in the paper but 2.1% from the table) so much worse than fine granularity? This is a bit surprising to me because 256 token performance should not be affected because of the lower token counts.
5) In visualizations, what do you mean by one random token? 

I am looking forward to your answers on these things.
questions: questions: See above.
limitations: limitations: Yes

Comments: title: Thank you for the discussion
Reply to the question on correct granularity.
Reply to authors
Thank you for the discussion
comment: Dear reviewer, Thank you for your feedback. We are encouraged by your decision to raise the score to accept this paper. We agree that the efficiency and accuracy trade-off method can be investigated more in-depth. As you mentioned, we demonstrated elastic inference and compared it with independently trained baselines to support our "beneficial property". We will leave a more in-depth exploration of the efficiency and accuracy trade-off to future work. Thank you!
Dear authors,

Thanks for your response. I do agree that elastic inference is a beneficial property, which MQT seems to reach with a simple training strategy. Although I still think the efficiency/accuracy trade-off of the method could be investigated a bit more in-depth, in particular for different number of baseline tokens.  (see below).
However, since I do not have any strong negative left after the rebuttal, I will raise my score to weak accept. 

**Regarding 2/4:** My suggestion was rather on trying different sparsity ratio vs baseline number of tokens trade-off, to see how this affects the performance in the lower FLOPs regime: 
Taking Figure 1 as an example, it is my understanding that the blue curve is a model trained with 576 tokens (baseline) + MQT to enable elastic inference with a lower number of tokens. so for instance, to achieve 64 tokens at inference a sparsity ratio of ~10% is applied. But another way to attain 64 tokens at inference would to use a 256 token baseline model (*which seems to only be roughly $0.2$ point of accuracy below the 576 tokens one in Figure 1*) with a less drastic sparsity ratio of 25%, which may hurt accuracy less and lead to a better trade-off curve.
Dear reviewer, Thank you for your feedback. As demonstrated in the original Matryoshka Representation Learning (MRL) work (Kusupati et al., NeurIPS 2022), the model can learn the correct granularity even when the supervision signals come from the same source. As mentioned in MRL work, this coarse-to-fine granularity is achieved by explicit optimization of $O(log(d))$ lower-dimensional vectors in a nested fashion.

Similarly in our work, although the supervision signal remains consistent, our model is trained with nested varying numbers of visual tokens, allowing it to adapt to different settings. We acknowledge that the model may learn the correct granularity implicitly, even though we employed an explicit variation in the number of visual tokens. Thus, to support this, we provide both quantitative data (Table 1) and qualitative visualizations (Figures 4, 5, and 7) demonstrating that this implicit training is effective and that the learned patterns are intuitive.

We hope our response can address your question. Thank you.
Thank you authors for providing the rebuttal.
I am still doubtful about the visualizations and even overall the story line provided. The supervision signal under different granularities of visual tokens are from the same set of captions. How can we make sure this makes sense and model really learns the correct granularity? If so, then the assumption about the visualizations do not make sense.
Dear reviewer, Thank you for your promptly and insightful feedback. We appreciate your lean to accept this paper. 

1. Thanks. 
2. The results of inference with the first 8 tokens are in our paper's Table 1. We have tested your proposed scenarios and found to be exact the same numbers. 
3. Thanks. 
4. Thank you for your advice, we realized this confusion, too. We will revise this expression in line 249 and 258 to absolute scores. 
5. Thanks, we will definitely add this visualization to the main paper.
Thanks for the rebuttal.

1) Good to know about what was being measured in TFLOPs

2) Can you please share the results of the experiment of picking first 8 tokens?

3) Sounds good about joint vs sampled

4) I would prefer having aboslute numbers than relative to reduce confusion. But I am still suprised by the performance gap here. But I can see why it might happen.

5) Thanks for the information. Please add this to the main paper as promised. 

I shall look at other reviews, rebuttals and discssion before making a final decision, but I lean to accept this paper even in the current stage.