TokenFLEX: Unified VLM Training for Flexible Visual Tokens Inference

📄 arXiv: 2504.03154v1 📥 PDF

作者: Junshan Hu, Jialiang Mao, Zhikang Liu, Zhongpu Xia, Peng Jia, Xianpeng Lang

分类: cs.CV

发布日期: 2025-04-04

💡 一句话要点

TokenFLEX：提出一种统一的VLM训练框架，实现视觉tokens数量的灵活推理。

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 视觉语言模型 可变tokens 自适应学习 多模态融合 大语言模型 视觉理解 计算效率

📋 核心要点

  1. 现有VLM使用固定数量的视觉tokens，无法兼顾不同任务的复杂度和计算效率。
  2. TokenFLEX通过随机调整token数量的训练范式和自适应token投影器，实现视觉tokens数量的灵活调整。
  3. 实验表明，TokenFLEX在不同token数量下均优于固定token模型，并在多个视觉-语言任务上取得了显著提升。

📝 摘要（中文）

传统的视觉-语言模型(VLM)通常使用固定数量的视觉tokens，而忽略了任务复杂性。这种一刀切的策略效率低下：对于简单任务，过多的tokens导致不必要的计算开销；对于复杂任务，不足的tokens会损害细粒度的视觉理解。为了克服这些限制，我们提出了TokenFLEX，一种创新的、自适应的视觉-语言框架，可以将图像编码为可变数量的tokens，以便与大型语言模型(LLM)高效集成。我们的方法基于两项关键创新。首先，我们提出了一种新的训练范式，通过在训练期间随机调整token数量来提高不同数量视觉tokens的性能。其次，我们设计了一个轻量级的视觉token投影器，包含一个自适应池化层和SwiGLU，允许灵活地对视觉tokens进行下采样，并自适应地选择针对特定token数量定制的特征。综合实验表明，TokenFLEX始终优于其固定token的同类产品，在八个视觉-语言基准测试中，分别使用64、144和256个tokens时，平均性能分别提高了1.6%、1.0%和0.4%。这些结果突出了TokenFLEX的卓越灵活性，同时保持了高性能的视觉-语言理解。

🔬 方法详解

问题定义：现有视觉-语言模型（VLM）通常采用固定数量的视觉tokens来表示图像，这种“一刀切”的方法忽略了不同任务对视觉信息粒度的需求。对于简单的任务，使用过多的tokens会导致不必要的计算开销；而对于复杂的任务，固定的少量tokens又无法捕捉到足够的细粒度视觉信息，从而限制了模型的性能。因此，如何根据任务的复杂程度自适应地调整视觉tokens的数量，是当前VLM面临的一个重要挑战。

核心思路：TokenFLEX的核心思路是使VLM能够根据任务需求灵活地调整视觉tokens的数量。为了实现这一目标，论文提出了两种关键技术：一是随机token数量训练范式，二是自适应视觉token投影器。通过随机token数量训练，模型可以学习到在不同token数量下的鲁棒表示；而自适应视觉token投影器则负责将图像特征映射到不同数量的tokens，并选择与token数量相适应的特征。

技术框架：TokenFLEX的整体框架主要包含以下几个模块：1) 图像编码器：用于提取图像的初始视觉特征。2) 自适应视觉token投影器：将初始视觉特征投影到不同数量的视觉tokens。该模块包含一个自适应池化层和SwiGLU激活函数，用于灵活地进行下采样和特征选择。3) 大型语言模型（LLM）：接收视觉tokens和文本输入，进行多模态融合和推理。4) 训练模块：采用随机token数量训练范式，使模型能够适应不同数量的视觉tokens。

关键创新：TokenFLEX最重要的技术创新在于其能够根据任务需求自适应地调整视觉tokens的数量。与传统的固定token VLM相比，TokenFLEX能够更有效地利用计算资源，并在不同复杂度的任务上取得更好的性能。随机token数量训练范式和自适应视觉token投影器是实现这一目标的关键。

关键设计：在自适应视觉token投影器中，自适应池化层用于将图像特征下采样到目标token数量。SwiGLU激活函数则用于增强特征的表达能力。在训练过程中，采用随机token数量训练范式，即在每个训练batch中，随机选择一个token数量，并使用该数量的tokens进行训练。损失函数采用标准的交叉熵损失函数，用于优化模型的预测结果。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，TokenFLEX在多个视觉-语言基准测试中均优于其固定token的同类产品。具体来说，在使用64、144和256个tokens时，TokenFLEX的平均性能分别提高了1.6%、1.0%和0.4%。这些结果表明，TokenFLEX能够有效地利用不同数量的视觉tokens，并在不同复杂度的任务上取得更好的性能。

🎯 应用场景

TokenFLEX具有广泛的应用前景，例如智能问答、图像描述、视觉推理等。通过自适应地调整视觉tokens的数量，TokenFLEX可以更有效地处理不同复杂度的任务，提高模型的性能和效率。此外，TokenFLEX还可以应用于资源受限的场景，例如移动设备或嵌入式系统，通过减少tokens数量来降低计算开销。

📄 摘要（原文）

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.