SparseVLMs: 高效视觉语言模型推理的视觉Token稀疏化

一、项目概述

项目	内容
标题	SparseVLM: Visual Token Sparsification for Efficient Vision-Language Model Inference
作者	Yuan Zhang, Chun-Kai Fan, Junpeng Ma, Wenzhao Zheng, Tao Huang, Kuan Cheng, Denis Gudovskiy, Tomoyuki Okuno, Yohei Nakata, Kurt Keutzer 等
机构	北京大学 (Peking University)、加州大学伯克利分校 (UC Berkeley)、松下控股 (Panasonic Holdings Corporation)
论文	arXiv:2410.04417
代码	GitHub: Gumpest/SparseVLMs
会议	ICML 2025 (International Conference on Machine Learning)
许可	Apache License 2.0
语言	Python
Star	265
Fork	23

SparseVLM+ (V2.0) 扩展信息：

项目	内容
标题	SparseVLM+: Visual Token Sparsification with Improved Text-Visual Attention Pattern
作者	Yuan Zhang, Junpeng Ma, Qizhe Zhang, Chun-Kai Fan, Wenzhao Zheng, Kuan Cheng, Jiwen Lu, Shanghang Zhang
机构	北京大学 (Peking University)、清华大学 (Tsinghua University)
发布	2025年12月11日

二、核心思想

问题定义

在视觉语言模型 (VLMs) 中，视觉Token通常占据大量的计算开销，尽管其信息密度相比文本Token更加稀疏。以 LLaVA-v1.5-7B 为例，一张 336x336 的图像经过 CLIP 编码器后会产生 576 个视觉Token，这些Token在 Transformer 解码器的每一层中都需要参与自注意力计算，带来了巨大的计算开销。

现有的视觉Token压缩方法主要存在两个问题：

文本无关的视觉稀疏化：如 FastV 等方法在选择视觉Token时完全忽略了语言Token的引导，这与多模态模型的”根据问题理解图像”的核心范式相矛盾。
修改图像编码器或投影器：这类方法需要重新训练模型组件，灵活性较差。

解决方案概述

SparseVLM 提出了一个核心洞察：视觉Token应该根据问题提示 (question prompt) 自适应地进行稀疏化。当模型面对不同问题时，应该关注图像的不同部分（例如前景或背景）。SparseVLM 通过利用 Transformer 解码器层中已经计算好的文本-视觉注意力权重来指导视觉Token的选择，无需额外的可学习参数，完全在推理阶段实现。

SparseVLM+ (V2.0) 进一步改进了文本-视觉注意力模式，通过三项关键创新提升了稀疏化质量：

注意力重力校正 (Attention Gravity Correction)：消除 RoPE 位置编码引入的注意力偏差
文本-视觉优先头选择 (Priority Heads Selection)：选择对文本-视觉关系最敏感的注意力头
取消Token合并，保留位置信息：简化流程并保留精确的位置编码

三、技术架构

整体框架

SparseVLM 基于 LLaVA-v1.5-7B 架构构建，核心组件包括：

输入图像 (336x336) --> CLIP Vision Encoder --> 576 个视觉Token (hidden_size=4096)
                                                    |
                                                    v
输入文本 --> Tokenizer --> 文本Token嵌入 -----> 拼接 --> LLaMA Decoder Layers
                                                          |
                                                          +-- Layer 2: 第一次稀疏化 (576 -> ~300/238/...)
                                                          +-- Layer 6: 第二次稀疏化 (~300 -> ~200/108/...)
                                                          +-- Layer 15: 第三次稀疏化 (~200 -> ~110/60/...)
                                                          |
                                                          v
                                                    最终保留的Token --> LM Head --> 输出

核心稀疏化流程

稀疏化发生在 LLaMA 解码器的特定层（Layer 2、6、15），每层执行以下步骤：

Step 1: 文本Rater选择 (Select Text Raters)

在进入解码器层之前，利用视觉Token和文本Token之间的交叉注意力关系选择”文本评分者”：

# 计算视觉-文本注意力矩阵
v_t = hidden_states[:, v_token_start: text_token_start, :]  # 视觉Token
t_t = hidden_states[:, text_token_start:, :]                # 文本Token
m_v_t = v_t @ t_t.transpose(1, 2)  # [B, 576, text_len]
m_v_t = m_v_t.softmax(2).mean(1)   # [B, text_len]

# 选择高于平均注意力的文本Token作为"评分者"
t_token_idx = torch.where(m_v_t > m_v_t.mean())

Step 2: 注意力后处理与Top-K选择

在每个稀疏化层的注意力计算完成后，利用多头注意力权重对视觉Token进行评分：

def attn_postprocess_topk(self_attn_weights, v_token_start, v_token_num,
                          text_token_start, t_token_idx, layer_idx):
    # 对所有注意力头取平均
    self_attn_weights = self_attn_weights.mean(1)  # [B, L[Q], L[K]]

    # 提取文本评分者对视觉Token的注意力
    t_token_idx = t_token_idx[1] + text_token_start
    relation_vis_text = self_attn_weights[:, t_token_idx,
                         v_token_start: v_token_start+v_token_num]  # [B, text_raters, v_tokens]
    relation_vis_text = relation_vis_text.mean(1)  # [B, v_tokens]

    # Top-K 选择保留的视觉Token
    sparse_token_list = sparse_token_dict[RETAIN_TOKN]  # 根据保留数量配置
    _, indices = torch.topk(relation_vis_text,
                           min(sparse_token_list[layer_dict[layer_idx]], v_token_num - 1))
    mask[0][indices] = 1  # 选中的Token标记为1

    return mask, s_flag, relation_vis_text

Step 3: Token合并 (V1.0 模式)

对于未被选中的视觉Token，V1.0 模式会进行聚类合并：

def cluster_and_merge(x, cluster_num):
    B, N, C = x.shape
    # 计算Token间距离矩阵
    x1 = ein.rearrange(x, "b l r -> b l () r")
    x2 = ein.rearrange(x, "b l r -> b () l r")
    distance = (x1 - x2).norm(dim=-1, p=2)
    dist_matrix = distance / (C ** 0.5)

    # 基于局部密度选择聚类中心
    dist_nearest, index_nearest = torch.topk(dist_matrix, k=cluster_num,
                                             dim=-1, largest=False)
    density = (-(dist_nearest ** 2).mean(dim=-1)).exp()

    # 距离指示器：密度高的Token倾向于成为中心
    mask = density[:, None, :] > density[:, :, None]
    dist, index_parent = (dist_matrix * mask + dist_max * (1 - mask)).min(dim=-1)

    # 聚类分数 = 距离 * 密度
    score = dist * density
    _, index_down = torch.topk(score, k=cluster_num, dim=-1)

    # 分配Token到最近的中心并合并
    # ... (加权平均合并)

    return x_merged  # [B, cluster_num, C]

Step 4: 注意力掩码更新

使用 softmax_with_policy 函数在注意力计算中屏蔽被移除的Token：

def softmax_with_policy(attn, policy, eps=1e-6):
    # policy: [B, N, 1]，1表示保留，0表示移除
    attn_policy = policy.reshape(B, 1, 1, N)
    eye = torch.eye(N, ...).view(1, 1, N, N)
    attn_policy = attn_policy + (1.0 - attn_policy) * eye  # 对角线保持为1

    # 创建-inf掩码屏蔽被移除的Token
    policy_bias.masked_fill_(attn_policy.logical_not(), float("-inf"))
    attn += policy_bias
    attn = torch.softmax(attn, dim=-1)
    return attn

SparseVLM+ (V2.0) 的额外创新

创新1: 注意力重力校正 (Attention Gravity Correction)

RoPE (旋转位置编码) 会引入与内容无关的位置偏差，导致注意力权重不能准确反映Token间的语义相关性。V2.0 通过计算纯位置诱导的注意力来消除这种偏差：

if V2_0:
    # 计算纯RoPE诱导的注意力偏差
    query_states = hidden_states.new_ones(bs, seq, num_heads, head_dim).transpose(1, 2)
    key_states = hidden_states.new_ones(bs, seq, num_heads, head_dim).transpose(1, 2)
    cos, sin = decoder_layer.self_attn.rotary_emb(key_states, seq_len=position_ids.max().item() + 1)
    query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)
    rope_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(head_dim)

    # 消除位置偏差：注意力权重 / RoPE权重
    attn_logits = attn_logits / rope_weights

创新2: 文本-视觉优先头选择 (Priority Heads Selection)

不同的注意力头对文本-视觉关系的建模能力不同。V2.0 选择对文本-视觉交互最敏感的注意力头（Top-14/28）：

def select_attn_head_by_sum(self_attn_weights, t_token_idx, v_token_start, text_token_start):
    # 提取每个头的文本->视觉注意力
    each_head_text_to_visual_attn = self_attn_weights[0][:, t_token_idx,
                                     v_token_start: text_token_start]
    # 按头求和，选择注意力最集中的头
    sum_attn_per_head = each_head_text_to_visual_attn.sum((1, 2))  # [28]
    select_attn_head_idx = sum_attn_per_head.topk(14)[1]  # 选择Top-14个头

    return self_attn_weights[:, select_attn_head_idx, :, :]

创新3: 位置ID保留与取消合并

V2.0 取消了Token合并操作（s_flag = False），改为直接选择保留的Token，并保留精确的位置编码：

if V2_0:
    s_flag = False  # 取消合并
    # 保留被选中Token的位置ID
    position_ids = batch_index_select(position_ids, select_token_idx)
else:
    # V1.0: 截断位置ID
    position_ids = position_ids[:, :len(select_token_idx[0])]

核心公式

注意力权重计算（带Policy掩码）

$\text{Attn}(Q, K, V, \pi) = \text{softmax}\left(\frac{QK^T}{\sqrt{d_k}} + M_\pi\right) V$

其中 $M_\pi$ 是基于 policy $\pi$ 的掩码矩阵： $M_\pi[i,j] = \begin{cases} 0 & \text{if } \pi[j] = 1 \text{ or } i = j \\ -\infty & \text{otherwise} \end{cases}$

RoPE重力校正 (V2.0)

$\hat{A} = \frac{A}{A_{\text{RoPE}}}$

其中 $A$ 是原始注意力权重， $A_{\text{RoPE}}$ 是仅由RoPE位置编码诱导的注意力权重。

文本-视觉注意力评分

$s_i = \frac{1}{|T_r|} \sum_{t \in T_r} A[t, v_i]$

其中 $T_r$ 是被选中的文本Rater集合， $v_i$ 是第 $i$ 个视觉Token， $A$ 是注意力权重矩阵。

Token合并中的密度聚类

密度计算： $\rho_i = \exp\left(-\frac{1}{k} \sum_{j \in \mathcal{N}_k(i)} d_{ij}^2\right)$

聚类分数： $\text{score}_i = d_i \cdot \rho_i$

其中 $d_i$ 是到最近更高密度点的距离， $\rho_i$ 是局部密度。

模型组件

组件	说明	关键参数
CLIP Vision Tower	视觉编码器，提取图像特征	ViT-L/14, image_size=336, patch_size=14
MM Projector	视觉-语言对齐投影器	2层MLP: 4096 -> 4096 (hidden_size)
LLaMA Decoder	语言模型骨干	32层, hidden_size=4096, 32头
稀疏化层	执行Token选择的层	Layer 2, 6, 15 (pruning_loc)
LlamaDynamicvitAttention	带Policy掩码的注意力层	支持标准/SDPA/FlashAttention2
LlamaDynamicvitDecoderLayer	带稀疏化支持的解码器层	包含标准FFN + Policy注意力

训练流程

SparseVLM 的一个关键优势是无需额外训练。它是一个纯推理阶段的方法，直接应用于已有的预训练VLM（如 LLaVA-v1.5-7B）。训练流程仅涉及原始LLaVA的训练：

预训练阶段：训练 MM Projector（视觉-语言对齐）
- 冻结 CLIP 和 LLaMA，只训练投影器
- 使用 CC3M 等图文对数据
指令微调阶段：端到端微调
- 冻结 CLIP，训练投影器和 LLaMA
- 使用 LLaVA-Instruct-150K 等指令数据

SparseVLM 在推理时直接加载预训练好的 LLaVA-v1.5-7B 权重，通过修改解码器的前向传播来实现Token稀疏化。

四、核心创新

创新点	说明	理论/实验依据
文本引导的视觉稀疏化	利用文本-视觉注意力权重选择与问题最相关的视觉Token，而非文本无关的选择	动机示例：同一图像面对不同问题时应关注不同区域
多层级渐进稀疏化	在3个不同深度的解码器层（Layer 2/6/15）逐步减少视觉Token数量	浅层Token信息冗余度高可大幅裁剪，深层需保留更多语义
无需额外训练	完全在推理阶段实现，不引入额外可学习参数	直接利用已有注意力权重进行Token评分
注意力重力校正 (V2.0)	消除RoPE位置编码引入的注意力偏差，获得更准确的语义相关性度量	RoPE对远距离Token施加系统性偏差，除法校正可消除
优先注意力头选择 (V2.0)	选择对文本-视觉交互最敏感的Top-14/28注意力头进行稀疏化决策	不同头对跨模态关系的建模能力差异显著
位置ID精确保留 (V2.0)	使用batch_index_select保留被选中Token的原始位置编码	保持位置信息的连续性和准确性
灵活的Token保留配置	支持192/128/96/64四种Token保留数量	用户可根据效率-精度权衡选择不同配置
Token聚类合并 (V1.0)	对未选中的Token进行密度聚类合并，保留信息多样性	基于TCFormer的密度聚类算法

五、代码实现分析

项目结构

SparseVLMs/
├── llava/                              # 核心模型代码
│   ├── model/
│   │   ├── language_model/
│   │   │   ├── score.py                # Token评分与注意力后处理
│   │   │   ├── sparse_llava_llama.py   # SparseVLM主模型类
│   │   │   ├── modelling_sparse_llama.py # 稀疏注意力机制核心实现
│   │   │   ├── utils.py                # 工具函数(聚类合并、注意力计算等)
│   │   │   ├── llava_llama.py          # 原始LLaVA LLaMA模型
│   │   │   └── llava_mistral.py        # Mistral变体
│   │   ├── llava_arch.py              # VLM架构基类(含prepare_sparse_inputs)
│   │   ├── multimodal_encoder/        # 视觉编码器
│   │   └── multimodal_projector/      # 视觉-语言投影器
│   ├── train/
│   │   ├── sparse_train.py            # 稀疏化训练脚本
│   │   ├── sparse_llava_trainer.py    # 稀疏化Trainer
│   │   └── train.py                   # 标准训练脚本
│   ├── eval/                          # 评估脚本
│   │   ├── model_vqa_loader.py        # 批量VQA评估
│   │   └── ...                        # 各benchmark评估
│   ├── constants.py                   # 常量定义
│   └── conversation.py                # 对话模板
├── scripts/
│   └── v1_5/eval/                     # 各benchmark评估脚本
│       ├── mme.sh
│       ├── textvqa.sh
│       ├── mmbench.sh
│       └── ...
├── pyproject.toml                     # 项目配置
└── README.md

关键文件详解

1. `score.py` - Token评分核心

这是SparseVLM的核心评分逻辑，定义了：

Token保留配置：4种保留数量（192/128/96/64），每种对应3层各自的保留数
V1.0 vs V2.0 配置差异：通过环境变量 USE_VERSION 切换
attn_postprocess_topk：基于注意力权重的Top-K Token选择
select_attn_head_by_sum：注意力头选择（V2.0）

# Token保留配置示例 (V2.0)
sparse_token_list_192 = [300, 200, 118]  # Layer2保留300, Layer6保留200, Layer15保留118
sparse_token_list_128 = [238, 108, 60]
sparse_token_list_96  = [246, 54, 28]
sparse_token_list_64  = [66, 34, 20]

layer_dict = {2:0, 6:1, 15:2}  # 层索引映射

2. `modelling_sparse_llama.py` - 稀疏解码器核心

实现了 LlamaDynamicvitModel 和 LlamaDynamicvitForCausalLM，核心修改包括：

LlamaDynamicvitModel.__init__：初始化稀疏化配置
- pruning_loc = [2, 6, 15]：稀疏化发生的层
- init_token_total_shape：初始Token总数
- generate_process_count：生成过程中的Token增量计数
LlamaDynamicvitModel.forward：核心前向传播，包含稀疏化逻辑
- 在每个稀疏化层执行：注意力计算 -> Token评分 -> Top-K选择 -> 可选合并 -> 更新序列
- FLOPs统计：4 * n * d^2 + 2 * n^2 * d + 3 * n * d * m
LlamaDynamicvitAttention：支持Policy掩码的注意力层
- policy=None时使用标准注意力
- policy!=None时使用 softmax_with_policy 进行掩码注意力

3. `utils.py` - 工具函数

包含稀疏化所需的核心工具：

softmax_with_policy：带Policy掩码的softmax
scaled_dot_product_attention_with_policy：带Policy的缩放点积注意力
cluster_and_merge：基于密度的Token聚类合并
batch_index_select：批量索引选择
index_points：索引采样

4. `llava_arch.py` - 稀疏化输入准备

prepare_sparse_inputs_labels_for_multimodal 方法负责：

图像编码：[1, 3, 336, 336] -> [1, 576, 4096]
记录 pre_prompt_length：图像Token前的文本长度
记录 token_length_list：每个样本的有效Token长度
记录 image_shape：视觉Token数量（576）

环境变量配置

变量	说明	默认值	可选值
`RETAIN_TOKN`	稀疏化后保留的视觉Token数量	192	192, 128, 96, 64
`USE_VERSION`	算法版本	1_0	1_0, 2_0

依赖环境

Python >= 3.8
torch == 2.1.2
torchvision == 0.16.2
transformers == 4.37.2
flash_attn == 2.3.3
einops == 0.6.1
timm == 0.6.13
deepspeed == 0.12.6 (训练)

六、实验结果

基准测试

SparseVLM 在多个多模态基准上进行了评估，基线模型为 LLaVA-v1.5-7B。以下为基于代码配置的Token保留策略：

Token保留配置与效率

保留Token数	Layer 2	Layer 6	Layer 15	平均等效Token	估计加速比
192 (V2.0)	300	200	118	~206	~1.5x
128 (V2.0)	238	108	60	~135	~2.0x
96 (V2.0)	246	54	28	~109	~2.5x
64 (V2.0)	66	34	20	~40	~4.0x

V1.0 vs V2.0 配置对比

保留数	V1.0 Layer配置 [2,6,15]	V2.0 Layer配置 [2,6,15]	差异说明
192	[300, 200, 110]	[300, 200, 118]	V2.0深层保留更多
128	[303, 110, 36]	[238, 108, 60]	V2.0浅层减少、深层增加
96	[238, 48, 26]	[246, 54, 28]	V2.0整体略增
64	[66, 30, 17]	[66, 34, 20]	V2.0后两层增加

评估的基准

根据评估脚本，SparseVLM 在以下基准上进行了测试：

基准	类型	评估脚本
MME	综合感知与认知	`mme.sh`
TextVQA	文本识别VQA	`textvqa.sh`
MMBench	综合多模态基准	`mmbench.sh`
ScienceQA	科学问答	`sqa.sh`
GQA	场景图推理	`gqa.sh`
POPE	幻觉评估	`pope.sh`
VizWiz	视障辅助VQA	`vizwiz.sh`
VQAv2	通用VQA	`vqav2.sh`
SEED-Bench	综合评估	`seed.sh`
MM-Vet	复杂推理	`mmvet.sh`
LLaVA-Bench	对话质量	`llavabench.sh`
Q-Bench	低层视觉感知	`qbench.sh`

关键性能指标

根据论文（ICML 2025 接收）和项目页面的信息：

SparseVLM (V1.0) 核心结果：

在保留 192 个视觉Token时，与原始 LLaVA-v1.5-7B 性能相当，同时减少约 50% 的视觉Token
在 MME、TextVQA、ScienceQA 等基准上保持竞争力
在保留 64 个Token时（原始576的11%），仍能维持合理的性能

SparseVLM+ (V2.0) 改进：

在所有Token保留配置下均优于V1.0
通过注意力重力校正和优先头选择，稀疏化决策更加准确
取消合并后推理效率进一步提升

效率分析

代码中内置了效率统计：

# FLOPs计算
self.all_FLOPs += 4 * n * (d**2) + 2 * (n**2) * d + 3 * n * d * m
# 其中 n=序列长度, d=hidden_size(4096), m=intermediate_size

# 输出信息
loggerinfo.info(f"Equal Tokens: {avg_tokens}, Prefill Time: {time_ms}ms, TFLOPs: {tflops}")

Equal Tokens：所有层平均的等效Token数
Prefill Time：预填充阶段的CUDA时间
TFLOPs：总浮点运算量（TFLOPs）

七、总结

核心贡献

提出文本引导的视觉Token稀疏化范式：首次将问题提示 (question prompt) 引入视觉Token选择过程，使稀疏化与多模态理解范式一致。
无需训练的高效推理方法：直接利用解码器中已有的注意力权重进行Token评分，无需额外的可学习参数或微调。
多层级渐进稀疏化策略：在不同深度的解码器层（Layer 2/6/15）逐步减少视觉Token，浅层大幅裁剪冗余，深层精细保留语义。
SparseVLM+ 的注意力模式改进：
- RoPE注意力重力校正，消除位置编码偏差
- 文本-视觉优先注意力头选择
- 位置ID精确保留
灵活的效率-精度权衡：支持4种Token保留配置（192/128/96/64），用户可根据需求灵活选择。

技术影响

VLM效率优化新方向：SparseVLM 证明了利用文本-视觉注意力进行Token选择的有效性，为后续工作提供了新的研究范式。
即插即用的通用性：方法不依赖特定模型架构，理论上可应用于任何基于Transformer的VLM。
实际部署价值：对于资源受限的边缘设备部署VLM具有重要意义。
视频VLM扩展：项目已开源 VideoLLaVA 的稀疏化实现（video分支），证明方法的通用性。

局限性

推理阶段额外开销：注意力权重的后处理和Token选择引入了一定的计算开销，尤其在Token数较少时占比增大。
层间依赖：稀疏化决策在每层独立进行，未考虑跨层的全局最优选择。
文本Rater选择的启发式：使用平均注意力阈值选择文本Rater，可能存在次优选择。
V1.0的合并操作：Token聚类合并在某些情况下可能引入信息损失。
基座模型限制：当前实现主要基于 LLaVA-v1.5-7B，对其他VLM架构的适配需要额外工作。

八、参考资源

论文与代码

资源	链接
SparseVLM 论文	arXiv:2410.04417
SparseVLM+ 论文	ResearchGate
GitHub 仓库	Gumpest/SparseVLMs
项目主页	SparseVLMs Project Page
V1.5 分支	SparseVLMs v1.5
VideoLLaVA 分支	SparseVLMs video

引用

@inproceedings{zhang2024sparsevlm,
  title={SparseVLM: Visual Token Sparsification for Efficient Vision-Language Model Inference},
  author={Zhang, Yuan and Fan, Chun-Kai and Ma, Junpeng and Zheng, Wenzhao and Huang, Tao and Cheng, Kuan and Gudovskiy, Denis and Okuno, Tomoyuki and Nakata, Yohei and Keutzer, Kurt and others},
  booktitle={International Conference on Machine Learning},
  year={2025}
}

@article{zhangsparsevlm+,
  title={SparseVLM+: Visual Token Sparsification with Improved Text-Visual Attention Pattern},
  author={Zhang, Yuan and Ma, Junpeng and Zhang, Qizhe and Fan, Chun-Kai and Zheng, Wenzhao and Cheng, Kuan and Lu, Jiwen and Zhang, Shanghang}
}

致谢项目

项目	说明
LLaVA	基础VLM架构
TCFormer	Token聚类合并算法来源
MiniGemini	多模态参考
VideoLLaVA	视频VLM基线