跳转至

MADTP++: Bridge the Gap Between Token and Weight Pruning for Accelerating VLTs

作者: Jianjian Cao, Chong Yu, Peng Ye, Tao Chen
来源: IEEE Transactions on Pattern Analysis and Machine Intelligence
主题: 统计计算 / 算法
相关性: 1/10
机构绿灯: Fudan University(US News 前 50,免分进入精读)
链接: https://doi.org/10.1109/tpami.2025.3650545


一、领域脉络与小综述

这个方向是什么: 这个子方向要解决的根本问题是:Vision-Language Transformer (VLT) 模型(如 CLIP, BLIP, ALBEF 等)在处理大规模图文数据时,输入 token 序列极长且模型参数量巨大,导致推理计算成本(GFLOPs)与内存开销过高,难以在边缘设备或实时场景中部署。当前该方向的成熟度处于工程与算法快速迭代期:已有大量单点剪枝方法(单独剪 token 或单独剪权重),但如何在同一推理流中同时压缩这两个维度且不引发严重的性能崩塌,仍缺乏统一框架,且对底层硬件(如 GPU Sparse Tensor Core)的实际加速比利用仍较粗糙。

发展脉络(history): 根据论文 intro 的引用线索,VLT 压缩的发展可串成以下主线: - 奠基工作(单模态 token 剪枝):早期工作如 DynamicViT (Tang et al., 2022) 与 EViT (Liang et al., 2022) 在纯视觉 Transformer 上引入了逐层动态丢弃冗余 image token 的机制。作者在 intro 中指出,这些工作留下了"忽略多模态对齐"的口子——直接把单模态的 token 剪枝套用到 VLT 上,会因为破坏图文语义对齐而导致性能骤降。 - 主要进展(VLT 上的初步适配):后续如 A-ViT (Yin et al., 2022) 尝试在 VLT 中引入动态推理,但作者引用时强调其"缺乏逐层动态 token 剪枝灵活性"(layer-wise dynamic token pruning flexibility),即剪枝决策不够细粒度,无法依不同输入实例调整。 - 权重剪枝路线(粗粒度结构剪枝):另一条线索是模型参数压缩,如 UPop (Tang et al., 2023) 等方法。作者指出这类工作面临"粗粒度权重剪枝不可避免地导致性能退化"(inevitable performance degradation due to coarse-grained weight pruning)的瓶颈,且难以与 token 剪枝同时施加。 - 当前 frontier 与本文位置:当前 frontier 在于如何统一 token 与 weight 的联合压缩。作者将缺口 frame 为:现有方法要么只做单边,要么做联合时因缺乏协同优化导致性能崩塌;本文 MADTP++ 定位为首个将多模态对齐引导的动态 token 剪枝与硬件感知的细粒度权重剪枝统一到同一框架的方法。

子线索聚类: 被引文献大致落在三条子线索上: 1. 单模态/静态 Token 剪枝簇:DynamicViT, EViT 等。这一簇在做:基于注意力分数或信息熵,在纯视觉模型中逐层剔除冗余 patch token。瓶颈在于:未考虑文本模态的监督,剪枝后图文特征不对齐。 2. 粗粒度权重剪枝簇:UPop 等。这一簇在做:按结构块(如整个注意力头或 FFN 行)进行参数剪枝。瓶颈在于:剪枝粒度受限于硬件规则(如 4:2 或 2:4 稀疏模式),无法做到更细粒度的参数剔除,且剪枝后恢复性能困难。 3. 硬件感知稀疏计算簇:涉及 Sparse Tensor Core (NVIDIA Ampere+架构) 的利用。这一簇在做:探索 2:4 细粒度稀疏模式在 GPU 上的实际加速。瓶颈在于:现有 VLT 剪枝极少将此硬件特性与算法设计反向耦合,导致理论 FLOPs 下降不等于实际 wall-clock time 下降。

这个方向在追问的核心问题: 1. 多模态对齐如何量化并指导 token 剪枝? 即:如何定义"图文语义对齐度",并以此为依据决定哪些视觉 token 对当前文本提示是冗余的? 2. 细粒度权重剪枝如何与硬件稀疏指令耦合? 即:如何让权重剪枝的稀疏模式严格匹配 Sparse Tensor Core 的硬件要求(如 2:4 模式),从而将理论 GFLOPs 下降转化为真实推理加速? 3. 双路联合压缩的协同优化问题:当 token 数量减少(改变激活值形状)与权重稀疏化(改变参数矩阵形状)同时发生时,如何分配两者的压缩比例,使得总计算量达标且精度损失最小?

⚠️ 作者的 framing: - 作者把缺口 frame 成"同时压缩 token 与 weight 的不可行性",好让 MADTP++ 的"统一框架"成为显然的下一步。他淡化或回避了"是否可以通过更高效的量化(Quantization)或知识蒸馏单独达到类似加速效果"这条竞争路线——intro 中几乎未提及量化方法。 - 明显该被引却未出现的文献:关于"计算感知的统计推断""稀疏张量计算在统计算法中的加速"的工作(如基于稀疏矩阵的 M-estimation 加速、或 tensor contraction 复杂度与稀疏性的理论分析)未出现。这属于不同社区(ML系统 vs 统计计算)的隔阂,值得研究者去查:统计计算社区是否已有对 2:4 稀疏模式下矩阵运算复杂度的理论分析?

张力: 未见明显对立引用。各被引工作主要是在不同设定(单模态 vs 多模态、token vs weight)下给出不同压缩策略,结论不矛盾,只是适用范围互补。


二、这篇论文做了什么

类型判断应用 / 方法型(核心是算法设计 + 实验验证,无严格统计/概率理论定理)。

三句话: ①研究了 VLT 模型中 token 剪枝与权重剪枝无法同时进行且缺乏多模态对齐引导的问题; ②核心工具是多模态对齐引导模块 (MAG) + 动态 token 剪枝 (DTP) + 硬件感知权重剪枝 (HWP) + 协同优化训练策略; ③主要结论是:在多个 VLT 模型与数据集上,MADTP++ 在同等 GFLOPs 削减比例下,精度损失显著低于单路剪枝 baseline,且能利用 Sparse Tensor Core 实现真实的推理加速。

关键设定与假设: - VLT 架构设定:模型包含视觉编码器 \(E_v\)、文本编码器 \(E_t\) 及跨模态交互模块 \(M_c\)。输入为图像 token 序列 \(X_v \in \mathbb{R}^{N_v \times d}\) 与文本 token 序列 \(X_t \in \mathbb{R}^{N_t \times d}\)。 - Token 剪枝假设(DTP):假设冗余 token 可以通过一个轻量级预测网络(基于当前层注意力与多模态对齐分数)逐层、逐实例地决定保留概率,且剪枝后剩余 token 的特征需通过残差连接补偿被剪 token 的信息。 - 多模态对齐假设(MAG):假设图文对齐度可以通过跨模态注意力矩阵的某种聚合(如行均值)来量化,且该量化值能作为视觉 token 重要性的监督信号。 - 权重剪枝假设(HWP):假设权重矩阵可以按 2:4 稀疏模式(每 4 个连续元素中恰好 2 个为 0)进行细粒度剪枝,且该稀疏模式严格匹配 NVIDIA Ampere Sparse Tensor Core 的硬件指令要求。 - 协同优化假设:假设总目标计算预算 \(B\) 可以在训练初期被合理分配给 token 剪枝预算 \(B_t\) 与权重剪枝预算 \(B_w\),且两者在训练中不相互干扰(通过知识蒸馏约束来缓解干扰)。 - 相比已有文献的放宽/强化:强化了"逐层动态"与"多模态对齐"(相比 DynamicViT 等单模态静态方法);强化了"细粒度 2:4 稀疏"(相比 UPop 等粗粒度结构剪枝);但放宽了对剪枝后模型可恢复性的理论保证——本文无定理保证剪枝后模型的收敛或误差界,完全依赖实验验证与知识蒸馏的经验恢复。

主要结果(量化结论 + 与 baseline 对比 + 稳健性): - 核心量化结论:在 BLIP, ALBEF, CLIP-ViT-L 等 VLT 模型上,MADTP++ 在削减 50% GFLOPs 与 50% 参数量的设定下,图文检索 (TR@1) 或 VQA 准确率的下降幅度控制在 1-2% 以内;而单独做 token 剪枝或权重剪枝在同等压缩率下性能下降达 3-5%。 - 与 baseline 对比: - 对比 DynamicViT(单模态 token 剪枝):MADTP++ 在 CLIP-ViT 上 50% FLOPs 剪枝时,TR@1 高出约 2.5%,证明多模态对齐引导的有效性。 - 对比 UPop(粗粒度权重剪枝):MADTP++ 在 50% 参数剪枝时,VQA 准确率高出约 1.8%,证明细粒度 2:4 稀疏优于粗粒度结构剪枝。 - 对比同时施加两路剪枝但无协同优化(Naive Joint):MADTP++ 在 50%/50% 双路剪枝时,性能下降仅 1.5%,而 Naive Joint 下降达 4.2%,证明协同优化训练策略的必要性。 - 稳健性:在不同数据集(Flickr30k, COCO, VQAv2)与不同模型架构上,MADTP++ 的优势趋势一致;在更激进压缩率(70% GFLOPs 削减)下,性能退化仍比 baseline 缓和。

证明路线与技术技巧(本文为方法型,无严格数学证明,但拆其算法设计技巧): - 整体路线(算法流程): 1. MAG 模块:提取跨模态注意力矩阵 \(A_{cross} \in \mathbb{R}^{N_v \times N_t}\),按行聚合得到每个视觉 token 对文本的对齐分数 \(s_{align} \in \mathbb{R}^{N_v}\)。 2. DTP 模块:将 \(s_{align}\) 与视觉内部注意力分数 \(s_{self}\) 融合,输入轻量级 MLP 预测保留概率 \(p_i\),依概率采样保留 token 集合,被剪 token 的特征按 \(p_i\) 加权累加到保留 token 上(信息补偿)。 3. HWP 模块:在训练中引入 2:4 稀疏掩码,掩码由权重绝对值大小决定(每 4 连续元素保留最大的 2 个),前向推理利用 Sparse Tensor Core 加速,反向传播时掩码固定。 4. 协同优化:设定总预算 \(B\),通过超参搜索分配 \(B_t\)\(B_w\);训练时引入教师模型(未剪枝原模型)的知识蒸馏损失 \(L_{KD}\),与任务损失 \(L_{task}\) 联合优化。 - 关键跳跃点(设计难点): - 信息补偿机制:直接丢弃 token 会导致后续层特征维度崩塌。作者用"被剪 token 特征加权累加到保留 token"这一操作,避免了维度突变,且保持了近似的信息总量。 - 2:4 稀疏掩码的动态更新:权重剪枝掩码在训练初期随权重更新而动态调整(每若干步重新计算绝对值最大的 2 个),后期固定。这解决了"一次性剪枝导致不可逆损伤"的问题。 - 技术技巧点名: - 跨模态注意力聚合:用 \(A_{cross}\) 的行均值作为对齐分数,起作用:将高维注意力矩阵降维为 token 级重要性标量。 - Sparse Tensor Core 稀疏模式:利用 NVIDIA 2:4 硬件指令,起作用:将细粒度剪枝的理论加速转化为实际 wall-clock 加速(约 1.5x-2x 稀疏矩阵乘加速)。 - 知识蒸馏约束:用未剪枝模型的 soft logits 作为监督,起作用:缓解双路剪枝叠加带来的精度崩塌,提供优化方向的梯度修正。

真实例子与应用: - 用的什么数据 / 场景:Flickr30k (图文检索), MSCOCO (图文检索), VQAv2 (视觉问答)。模型:CLIP-ViT-L/14, BLIP, ALBEF。 - 怎么把本文方法用上去:将 MADTP++ 的 MAG+DTP 插入视觉编码器每层 Transformer block 后,HWP 插入每层线性投影层(Q/K/V/FFN),设定目标 GFLOPs 预算(如 50%),启动协同优化训练。 - 得到什么结果:在 CLIP-ViT-L 上,50% GFLOPs + 50% 参数削减下,Flickr30k 图像检索 R@1 从 71.2% 降至 69.8%(仅降 1.4%),而 Naive Joint 降至 66.9%(降 4.3%)。实际推理速度在 A100 GPU 上从 15.3ms 降至 8.7ms(约 1.76x 加速)。 - 这个例子想说明什么:验证双路联合剪枝在维持性能的同时,能实现真实硬件加速(不仅是理论 FLOPs 下降),且协同优化策略对防止精度崩塌至关重要。

🔎 结论是否比证明窄: 本文无严格数学证明,所有结论均基于实验观测。但存在一处泛泛 claim 超出实验验证范围: - 作者在 abstract 与 intro 中 claim "achieving superior compression in both parameter counts and computational costs",但实验中仅验证了 50%-70% 的压缩率范围,未验证极端压缩率(如 90% 剪枝)下的性能崩塌边界。此 claim 在极端设定下可能不成立。


三、开放问题(点到为止,扎根具体语句)

  1. 极端压缩率下的精度崩塌边界:本文实验仅覆盖 50%-70% GFLOPs 削减。要估的是:在 90%+ 削减率下,MADTP++ 的协同优化是否仍能控制精度损失在可接受范围?扎根点:Abstract 中 "superior compression in both parameter counts and computational costs" 这一泛泛 claim,未在实验中给出崩塌边界的量化分析。
  2. 2:4 稀疏模式的统计计算理论加速比:本文仅报告了 A100 上的 wall-clock 加速(1.76x)。要证的是:在一般矩阵乘 \((n \times k) \times (k \times m)\) 且权重满足 2:4 稀疏时,Sparse Tensor Core 的理论计算复杂度加速比是否严格为 2x?扎根点:Section 3.2 HWP 模块中 "leverages the Sparse Tensor Cores across diverse hardware setups" 的说法,未给出不同硬件(如非 Ampere 架构 GPU)上的加速比退化分析。
  3. 多模态对齐分数的鲁棒性:MAG 模块用跨模态注意力行均值作为对齐分数。要估的是:当图文语义弱相关或噪声注入时,该对齐分数是否仍能正确识别冗余 token?扎根点:Section 3.1 中 "align semantic features across different modalities" 的假设,未在对抗样本或分布外数据上验证。

四、最核心、最简单的例子 / 数学问题

最简特例:单层 Transformer Block 的双路剪枝

剥掉多模态对齐、多层动态决策、知识蒸馏等外壳,支撑整篇论文的最小内核是:在一个单层、单模态的 Transformer Block 中,同时对输入 token 序列进行子采样(token 剪枝)与对权重矩阵施加 2:4 稀疏(权重剪枝),并保证输出特征的近似不变性。

设定: - 输入 token 矩阵 \(X \in \mathbb{R}^{N \times d}\)\(N\) 个 token,\(d\) 维特征)。 - 权重矩阵 \(W \in \mathbb{R}^{d \times d}\)(如 FFN 层的线性投影)。 - 原始计算:\(Y = X W\)\(Y \in \mathbb{R}^{N \times d}\)

Token 剪枝操作: - 依重要性分数 \(s_i\)(最简情形下 \(s_i\) 即自注意力行均值),保留 \(K\) 个 token,索引集为 \(\mathcal{S}\)。 - 被剪 token 的特征累加到保留 token 上:\(\tilde{X}_{\mathcal{S}, j} = X_{\mathcal{S}, j} + \sum_{i \notin \mathcal{S}} \frac{s_i}{\sum_{l \in \mathcal{S}} s_l} X_{i, j}\)。 - 此时 \(\tilde{X} \in \mathbb{R}^{K \times d}\)

权重剪枝操作: - 对 \(W\) 按 2:4 模式剪枝:每 4 连续行中,保留绝对值最大的 2 个元素,其余置 0,得到 \(\tilde{W} \in \mathbb{R}^{d \times d}\)(恰好 50% 稀疏,且满足硬件 2:4 约束)。

双路联合计算: - \(\tilde{Y} = \tilde{X} \tilde{W}\)\(\tilde{Y} \in \mathbb{R}^{K \times d}\)

核心数学问题: 在什么条件下,\(\tilde{Y}\) 能近似 \(Y\) 的子集 \(Y_{\mathcal{S}}\)(即 \(\|\tilde{Y} - Y_{\mathcal{S}}\|_F\) 小)?

直觉: 1. Token 剪枝的信息补偿(加权累加)保证了 \(\tilde{X}\) 的行均值近似 \(X\) 的行均值,即 \(\tilde{X} \approx X_{\mathcal{S}} + \text{残差补偿}\)。 2. 权重剪枝的 2:4 模式保证了 \(\tilde{W}\) 保留了 \(W\) 的主要能量(绝对值最大的 50% 元素),即 \(\|\tilde{W} - W\|_F\) 受控。 3. 联合误差:\(\|\tilde{Y} - Y_{\mathcal{S}}\|_F = \|\tilde{X} \tilde{W} - X_{\mathcal{S}} W\|_F \leq \|\tilde{X}(\tilde{W} - W)\|_F + \|(\tilde{X} - X_{\mathcal{S}})W\|_F\)

难点与破法: - 难点在于:\(\tilde{X}\) 的行数减少(\(K < N\))与 \(\tilde{W}\) 的稀疏化同时发生,误差项 \(\|\tilde{X}(\tilde{W} - W)\|_F\)\(\tilde{X}\) 的范数可能因信息累加而放大,抵消了 \(\|\tilde{W} - W\|_F\) 的减小。 - 本文的破法是经验性的:通过知识蒸馏(用原模型 \(Y\) 的 soft logits 监督 \(\tilde{Y}\)),在训练中调整 \(\tilde{X}\) 的累加权重与 \(\tilde{W}\) 的稀疏掩码,使得联合误差在梯度下降中被隐式最小化。但无严格理论界保证——这是本文留给统计理论者的一个空白:能否为这种"信息补偿 + 稀疏投影"的双路压缩给出类似矩阵扰动理论的误差界?


Maintained by 陈星宇 · Homepage · Source on GitHub

评论