TokenFLEX: Unified VLM Training for Flexible Visual Tokens Inference

TL;DR

TokenFLEX tackles the inefficiency of fixed vision-token budgets in vision-language models by introducing a stochastic dynamic token training regime and a lightweight, token-adaptive projector that can produce a flexible number of vision tokens $N$ for integration with a Large Language Model. The projector uses adaptive average pooling to form $N=n^{2}$ tokens and SwiGLU-based feature weighting to preserve salient information across token budgets. Across eight OpenCompass benchmarks, TokenFLEX consistently outperforms fixed-token baselines and reduces visual-token usage by up to 28% while cutting training time by about 13%, demonstrating a practical path to efficient, flexible VLM deployment. The approach achieves competitive performance with state-of-the-art methods even when using far fewer vision tokens, illustrating the value of dynamic token counts for real-world multi-modal reasoning and task adaptability.

Abstract

Conventional Vision-Language Models(VLMs) typically utilize a fixed number of vision tokens, regardless of task complexity. This one-size-fits-all strategy introduces notable inefficiencies: using excessive tokens leads to unnecessary computational overhead in simpler tasks, whereas insufficient tokens compromise fine-grained visual comprehension in more complex contexts. To overcome these limitations, we present TokenFLEX, an innovative and adaptable vision-language framework that encodes images into a variable number of tokens for efficient integration with a Large Language Model (LLM). Our approach is underpinned by two pivotal innovations. Firstly, we present a novel training paradigm that enhances performance across varying numbers of vision tokens by stochastically modulating token counts during training. Secondly, we design a lightweight vision token projector incorporating an adaptive pooling layer and SwiGLU, allowing for flexible downsampling of vision tokens and adaptive selection of features tailored to specific token counts. Comprehensive experiments reveal that TokenFLEX consistently outperforms its fixed-token counterparts, achieving notable performance gains across various token counts enhancements of 1.6%, 1.0%, and 0.4% with 64, 144, and 256 tokens, respectively averaged over eight vision-language benchmarks. These results underscore TokenFLEX's remarkable flexibility while maintaining high-performance vision-language understanding.

Figures (3)

• Figure 1: Comparison of TokenFLEX with methods using fixed vision token counts. TokenFLEX is trained with a stochastic dynamic image token count, while other methods are trained with fixed token counts of 64, 144, and 256, respectively. The performance is measured as the average score across 8 multi-modal benchmarks on OpenCompass 2023opencompass. TokenFLEX consistently outperforms the fixed-token methods, particularly when using fewer tokens, such as 64.
• Figure 2: Overall architecture. TokenFLEX combines a Visiual Encoder and a Large Language Model connected by a dynamic token projector. (Left) During the training phase, each sample randomly selects an image token length; during inference, the image token length can be chosen on-demand based on task complexity and computational budget. (Right) Our light-weight projector acts as an adaptive filter, allowing it to selectively emphasize important image tokens.
• Figure 3: Ablation study on the dynamic token mechanism. The X-axis represents the number of vision tokens used during inference, while the Y-axis indicates benchmark performance. Three models are trained using a fixed number of vision tokens: 64, 144, and 256. Additionally, a model utilizing the dynamic token mechanism is trained with token counts of $\{64, 144, 256\}$. The $\star$ symbol denotes the number of vision tokens employed during training. The dynamic token mechanism achieved competitive results across various vision token counts.