尧图网站建设 
1. 项目概述当ViT遇上动态令牌合并DORA如何用强化学习“瘦身”模型最近在优化视觉Transformer模型时大家可能都遇到过这个头疼的问题模型精度和推理速度就像天平的两端难以兼得。尤其是在处理高分辨率图像时ViT模型需要将图像分割成大量的“令牌”进行处理计算量和内存开销呈平方级增长直接拖慢了整个系统的响应。传统的令牌剪枝或合并方法比如Token Merging虽然能减少令牌数量但大多是静态或启发式的往往“一刀切”在处理不同复杂度的图像内容时要么合并过度损失细节要么合并不足效率低下。这让我一直在想有没有一种更智能、更自适应的方式能根据每张图片的“内容复杂度”动态决定合并多少、合并哪些令牌这就是DORADynamic Online Reinforcement Agent for Token Merging要解决的核心问题。它不是一个全新的模型架构而是一个精巧的“插件式”智能体专门为ViT模型设计。其核心思想非常巧妙将令牌合并这个决策过程建模为一个序列决策问题并引入强化学习来训练一个智能体让它在线地、根据当前图像特征动态地决定每一步最优的合并策略。简单来说DORA就像一个坐在ViT模型里的“调度员”它实时观察经过处理的令牌特征然后判断“这一批令牌里哪些信息是冗余的、可以合并的合并多少才能既保住关键特征又最大化计算效率” 这个决策不再是固定的而是每张图、每一层都可能不同。对于需要部署ViT模型到资源受限环境如移动端、边缘设备的工程师或者任何关心模型效率的研究者来说DORA提供了一种全新的思路。它不改变ViT的主干网络而是通过一个轻量级的智能体模块实现模型推理过程的动态优化在几乎不损失精度的情况下显著提升吞吐量。最近社区里关于“Agentic RL”智能体化强化学习和高效微调技术如LoRA的讨论很多DORA可以看作是这两个趋势在模型推理优化上的一个有趣交汇用类似智能体的方式Agentic解决一个具体的优化问题并且其“动态性”与LoRA的“低秩适配”思想有异曲同工之妙都是在不改变原模型大规模参数的前提下通过添加小型可学习模块来实现特定目标。2. DORA核心设计思路将合并决策转化为智能体游戏要理解DORA我们得先拆解一下它面对的问题本质。在标准的ViT中一张图片被切成N个图像块每个块经过线性投影变成一个令牌。这些令牌会依次通过多个Transformer编码器层。Token Merging类方法的基本操作是在某一层计算所有令牌之间的相似度然后将最相似的一对令牌合并成一个用它们的加权平均来代表从而将令牌数量从N减少到N-1。这个过程可以重复进行直到达到预设的令牌保留率。传统方法的局限性在于这个“合并谁”以及“合并多少”的决策要么是基于一个全局固定的相似度阈值要么是根据一个预设的压缩率。但不同的图像内容差异巨大一张纯色背景的图片很多令牌高度相似可以大胆合并而一张细节丰富的医学影像过度合并则会丢失关键病理特征。一个静态策略无法适应这种变化。DORA的突破性思路在于它不再将合并决策视为一个基于规则的计算问题而是视为一个顺序决策过程。我们来想象这样一个场景环境当前Transformer层的所有令牌特征。智能体DORA模块它能够观察环境令牌特征。动作在每一步智能体需要做出两个决策(a) 是否在此刻执行合并(b) 如果合并选择哪一对令牌进行合并状态随着合并动作的执行令牌集合的特征和数量发生变化环境状态也随之改变。奖励智能体的最终目标是平衡效率和精度。因此奖励信号被设计为在推理结束时如果模型预测精度高则给予正奖励同时每合并一个令牌即减少计算量也给予一个小的正奖励。如果合并导致精度严重下降则给予负奖励。通过这样的建模训练DORA就变成了训练一个强化学习智能体让它学会在长期的“合并游戏”中最大化累积奖励。这意味着它必须学会评估当前令牌集合的“可合并性”并做出对未来最终分类精度最有利的即时决策。2.1 为什么选择强化学习而非其他方法你可能会问为什么不用一个简单的神经网络直接预测合并方案或者用可微分的软合并这里涉及到几个关键考量决策的离散性与序列性合并是一个典型的离散决策合并/不合并合并哪一对并且当前决策会影响后续决策可用的令牌集合。强化学习特别是基于策略梯度的方法非常擅长处理这类序列化的离散决策问题。稀疏延迟奖励我们最关心的模型最终精度是在所有合并决策完成后才得到的。这是一个典型的延迟且稀疏的奖励信号。强化学习算法如PPO、A2C本身就是为优化这类长期回报而设计的。在线适应训练好的DORA智能体在推理时是根据当前输入实时做出决策的实现了真正的“动态在线”无需任何额外的元数据或预处理。注意DORA的设计隐含了一个重要假设即存在一个“策略网络”它能够从令牌特征中提取出用于决策的抽象表示。这个网络通常非常小可能只有几层MLP以确保其引入的计算开销远小于它通过合并令牌所节省的开销。2.2 DORA与ToMe、DiffRate等方法的本质区别为了更清楚我们可以将DORA与之前的主流方法做个对比特性传统Token Merging (ToMe)可学习Token Merging (DiffRate)DORA (本文方法)决策依据固定的余弦相似度阈值可学习的评分网络预测每个令牌的保留分数强化学习智能体基于当前状态序列决策合并粒度逐层固定比例或阈值逐层可学习保留率逐样本、逐步骤动态决策优化目标启发式相似度最高端到端训练最小化任务损失最大化长期奖励精度效率动态性静态静态每层参数固定高度动态每张图决策不同额外开销几乎为零需要训练评分网络需要训练智能体推理时需运行策略网络可以看出DORA的核心优势在于其决策的上下文感知能力和灵活性。它不像DiffRate那样为每一层学习一个固定的“压缩率”而是为每一个输入实例学习一套“压缩策略”。3. DORA架构详解与实操要点理解了核心思想我们深入到DORA的具体实现架构。一个完整的DORA系统包含三个核心组件状态编码器、策略网络和价值网络。下面我们逐一拆解并说明在实操中需要注意的关键点。3.1 状态编码器如何让智能体“看懂”令牌智能体做出决策的前提是感知环境。在DORA中环境状态就是当前层的所有令牌特征。然而直接将所有令牌的原始特征例如768维向量输入给智能体是不现实的因为令牌数量N可能很大且是变化的。因此需要一个状态编码器来提取固定维度的全局状态表示。一个常见且有效的设计是对当前所有令牌特征进行全局平均池化得到一个汇总全局信息的向量。同时计算令牌特征的某些统计量如标准差、最大值等以反映特征的离散程度间接反映内容复杂度。将这些汇总信息拼接起来再通过一个小型MLP输出一个固定维度的状态向量s_t。实操要点编码器必须轻量状态编码器的计算量必须严格控制通常就是一层或两层线性层。它的目的是提供决策的“摘要”而非精细描述。包含序列信息对于更复杂的场景可以考虑引入轻量级的序列建模如一个单层的Transformer编码器或LSTM来处理令牌间的相对关系。但这会显著增加开销需要仔细权衡。归一化很重要输入状态编码器的特征应先进行层归一化确保训练稳定性。3.2 策略网络与动作空间设计策略网络π(a_t | s_t)接收状态向量s_t输出在当前状态下采取各个动作的概率分布。DORA的动作空间设计是其精妙之处。动作通常被设计为离散动作动作0不执行合并直接进入下一层或结束当前层的合并过程。动作1到K执行合并并指定合并哪一对令牌。这里的K是候选令牌对的数量。直接对所有可能的令牌对共N*(N-1)/2对进行建模是不现实的。因此需要一种高效的采样机制。一种实践方案是基于令牌特征的余弦相似度快速预筛选出最相似的Top-M个令牌对例如M10或20。策略网络的输出对应这M个候选对外加一个“不合并”动作。这样动作空间大小就是M1变得可管理。策略网络本身可以是一个简单的MLP输入状态s_t输出M1维的logits经过softmax后得到动作概率。实操心得探索与利用的平衡在强化学习训练初期需要鼓励智能体探索不同的合并策略。可以通过在策略网络的输出概率上添加熵正则项来实现防止智能体过早收敛到次优策略。动作掩码对于无效动作例如当令牌数已降到下限时“合并”动作应被禁止需要使用动作掩码将无效动作的概率置零并重新归一化。训练不稳定性策略梯度方法容易训练不稳定。务必使用像PPO近端策略优化这类更稳定的算法并仔细调整学习率、折扣因子等超参数。3.3 价值网络与奖励函数设计价值网络V(s_t)用于估计当前状态s_t的长期价值即从该状态开始遵循当前策略所能获得的期望累积奖励。它在计算优势函数用于更新策略时至关重要。奖励函数r_t的设计是引导智能体学习的关键。一个典型的设计是稀疏的终端奖励在整个推理过程结束时即图像经过所有Transformer层并得到分类结果后计算任务损失如交叉熵损失的负值作为最终奖励R_Terminal -α * Loss。α是一个缩放系数。稠密的步进奖励为了鼓励合并在每一步执行了“合并”动作后立即给予一个小的正奖励r_step β。β是一个很小的正数如0.01或0.001用于提供即时反馈引导智能体倾向于合并。总奖励最终的累积奖励是步进奖励之和加上终端奖励。注意事项奖励缩放终端奖励基于损失和步进奖励的量级可能相差很大。需要仔细调整α和β使两者处于同一数量级否则智能体可能只会优化其中一个目标。折扣因子γ用于计算累积奖励时对未来奖励的折扣。在DORA中由于合并决策的长期影响对最终精度的影响至关重要γ通常应设置得比较高如0.99让智能体更有远见。价值网络训练价值网络通过最小化时序差分误差来训练。它是独立于策略网络的一个回归头通常与策略网络共享状态编码器后的底层特征以提升效率。4. DORA训练与部署全流程实操纸上得来终觉浅我们来一步步拆解DORA从训练到部署的完整过程。这里以在ImageNet数据集上为一个预训练的ViT-B/16模型添加DORA为例。4.1 第一阶段环境搭建与模型准备首先你需要一个预训练好的ViT模型作为“环境”。假设我们使用timm库。import torch import torch.nn as nn import timm # 1. 加载预训练ViT-B/16作为主干网络 backbone timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes0) # 先不要分类头 backbone.eval() # 初始阶段冻结主干网络参数 # 2. 准备你的DORA智能体模块包含状态编码器、策略网络、价值网络 class DORAAgent(nn.Module): def __init__(self, token_dim, state_dim, num_actions): super().__init__() self.state_encoder nn.Sequential( nn.Linear(token_dim * 2, 256), # 输入全局平均标准差 nn.ReLU(), nn.Linear(256, state_dim) ) self.policy_net nn.Linear(state_dim, num_actions) self.value_net nn.Linear(state_dim, 1) # ... 其他初始化 def forward(self, token_features): # 计算全局状态 global_mean token_features.mean(dim1) global_std token_features.std(dim1) state torch.cat([global_mean, global_std], dim-1) encoded_state self.state_encoder(state) # 获取动作概率和状态价值 action_logits self.policy_net(encoded_state) state_value self.value_net(encoded_state) return action_logits, state_value # 初始化智能体 agent DORAAgent(token_dim768, state_dim128, num_actions21) # 假设20个候选对1个“不合并”关键步骤解析初始阶段我们通常冻结主干网络只训练DORA智能体。这是因为同时训练所有参数极其困难且容易破坏预训练模型已经学到的强大特征。num_actions21对应了20个最相似的候选令牌对加上“不合并”动作。这个数量需要根据你的计算预算和效果进行权衡。4.2 第二阶段强化学习训练循环这是最核心也是最复杂的部分。我们需要在一个循环中让智能体与环境ViT主干交互收集数据然后更新策略。import torch.optim as optim from torch.distributions import Categorical # 假设我们使用PPO算法 def compute_ppo_loss(old_logits, new_logits, actions, advantages, returns, epsilon0.2): old_dist Categorical(logitsold_logits) new_dist Categorical(logitsnew_logits) # 计算重要性采样比率 ratio torch.exp(new_dist.log_prob(actions) - old_dist.log_prob(actions.detach())) # PPO裁剪目标 surr1 ratio * advantages surr2 torch.clamp(ratio, 1-epsilon, 1epsilon) * advantages policy_loss -torch.min(surr1, surr2).mean() return policy_loss # 训练循环伪代码框架 optimizer optim.Adam(agent.parameters(), lr1e-4) for epoch in range(num_epochs): for batch_images, batch_labels in dataloader: # 1. 前向传播交互 tokens backbone.patch_embed(batch_images) # ... 添加位置编码等 trajectory [] # 存储 (state, action, reward, logits, value) current_tokens tokens for layer_idx in range(num_layers_to_apply_dora): # 智能体观察状态并做出决策 action_logits, state_value agent(current_tokens) action_dist Categorical(logitsaction_logits) action action_dist.sample() # 执行动作合并令牌或跳过 if action.item() ! 0: # 非“不合并”动作 # 根据动作索引找到对应的令牌对(i,j)并合并 i, j find_token_pair_to_merge(current_tokens, action) new_token (current_tokens[:, i] current_tokens[:, j]) / 2 # 更新令牌序列 # ... 实现合并逻辑current_tokens数量减1 step_reward merge_reward_bonus # 小正奖励 else: step_reward 0.0 # 存储数据 trajectory.append({ state: current_tokens.detach(), action: action, logits: action_logits.detach(), value: state_value.detach(), reward: step_reward }) # 应用当前层的Transformer计算简化表示 current_tokens backbone.blocks[layer_idx](current_tokens) # 2. 计算终端奖励 final_output backbone.norm(current_tokens) cls_output final_output[:, 0] # 接上一个临时的分类头计算损失 loss criterion(cls_output, batch_labels) terminal_reward -loss.item() * terminal_reward_scale # 将终端奖励加到最后一步的奖励上 trajectory[-1][reward] terminal_reward # 3. 计算优势函数和回报 # ... 使用GAE(Generalized Advantage Estimation)等方法计算 # 4. 更新智能体参数PPO更新 optimizer.zero_grad() # 使用收集的轨迹数据重新前向计算新的logits和value # 计算策略损失和价值损失 policy_loss compute_ppo_loss(old_logits, new_logits, ...) value_loss nn.MSELoss()(new_values, returns) total_loss policy_loss 0.5 * value_loss # 加权和 total_loss.backward() optimizer.step()实操现场记录与避坑指南训练速度慢这是RL训练的通病。一个批次的数据只能产生一条轨迹一个决策序列。为了稳定训练你需要收集大量轨迹才能进行一次有效的更新。这要求数据加载和交互过程必须高效。考虑使用向量化环境或重放缓冲区。奖励稀疏问题终端奖励非常稀疏。虽然步进奖励提供了即时反馈但智能体可能很难将最终的精度下降归因于早期某个错误的合并决策。课程学习是一个有效技巧先从简单的任务如合并率很低开始训练逐步增加任务难度允许更高的合并率。梯度流确保在交互阶段对主干网络的令牌特征执行.detach()防止主干网络的梯度通过交互过程回传除非你在后期进行联合微调。智能体的梯度应只通过其自身的网络计算。超参数敏感PPO的epsilon裁剪范围、学习率、折扣因子gamma、GAE的lambda等超参数对训练结果影响巨大。建议使用超参数优化工具如Optuna进行系统调优。4.3 第三阶段推理部署与性能评估训练完成后DORA智能体就可以用于动态加速推理了。# 推理模式下的DORA应用 torch.no_grad() def inference_with_dora(model, agent, image): tokens model.patch_embed(image) # ... 位置编码 for layer in model.blocks: # 在每一层或指定层前调用智能体决策 action_logits, _ agent(tokens) action torch.argmax(action_logits, dim-1) # 贪婪策略选择最大概率动作 if action ! 0: i, j find_token_pair_to_merge(tokens, action) # 执行合并 tokens merge_tokens(tokens, i, j) # 通过当前Transformer层 tokens layer(tokens) output model.norm(tokens) return model.head(output[:, 0])性能评估指标精度在测试集如ImageNet val上的Top-1/Top-5准确率。对比原始模型和使用DORA后的模型。加速比测量平均每张图片的推理时间或FPS。由于DORA动态合并令牌不同图片的加速效果不同需要统计平均值。令牌保留率统计在整个测试集上DORA平均将原始令牌数量减少了多少百分比。这是一个重要的效率指标。FLOPs/内存占用理论计算量FLOPs和峰值内存占用的减少比例。这能更直接地反映计算效率的提升。提示在部署时DORA智能体本身也会引入少量计算开销状态编码和策略前向传播。务必确保这个开销远小于因令牌合并而节省的Transformer层计算量。通常智能体网络设计得足够小例如参数量在万级别这个条件是满足的。5. 常见问题、调优技巧与扩展思考在实际操作中你肯定会遇到各种各样的问题。下面是我在复现和实验过程中遇到的一些典型情况及其解决方法以及一些让DORA工作得更好的技巧。5.1 训练不稳定与收敛困难这是强化学习项目最常见的问题。现象奖励曲线剧烈震荡没有上升趋势策略熵迅速降为零智能体停止探索价值估计爆炸。排查与解决检查奖励尺度终端奖励基于分类损失通常绝对值较大而步进奖励很小。如果两者尺度差异过大智能体可能只关注终端奖励而忽略效率。尝试将终端奖励归一化例如除以一个基线损失值如原始模型的损失。调整折扣因子γ如果智能体表现得非常短视只追求即时合并奖励尝试降低gamma如从0.99降到0.9。如果智能体难以学到长期影响尝试提高gamma。使用更稳定的算法确保你正确实现了PPO的裁剪和目标函数。可以考虑使用像Stable-Baselines3这样的成熟RL库来验证你的训练循环。增加批量大小策略梯度方法方差较大。增加用于计算梯度估计的轨迹数量批量大小可以显著稳定训练。引入基线在计算优势函数时使用价值网络V(s)作为基线是标准做法。确保你的价值网络训练得足够好能够提供准确的状态价值估计。5.2 智能体学不到有效策略现象训练后智能体要么永远选择“不合并”要么随机合并导致精度或效率没有提升。排查与解决奖励设计问题这是最可能的原因。如果“不合并”动作没有惩罚而合并动作可能带来风险精度下降智能体自然会倾向于保守。可以尝试给“不合并”动作一个微小的负奖励如-0.001鼓励其探索合并。或者调整终端奖励的权重让效率提升带来的正奖励更加突出。探索不足在训练初期策略网络的输出概率应尽可能均匀。可以通过设置较高的熵正则化系数来强制探索。随着训练进行再逐渐减小该系数。动作空间太大如果候选令牌对M设置得太大比如50智能体可能难以学习。尝试减小M比如5或10让决策变得更简单。也可以考虑对动作空间进行分层设计先决策是否合并再决策合并哪一对。5.3 如何将DORA应用到其他视觉任务或模型DORA的思想并不局限于图像分类和ViT。目标检测/分割对于DETR或Mask2Former这类基于Transformer的检测/分割模型令牌通常对应图像特征图上的位置或对象查询。DORA同样可以学习动态合并这些令牌或查询。奖励函数需要相应调整例如结合检测的mAP或分割的mIoU作为终端奖励。视频理解视频Transformer需要处理时空令牌计算量巨大。DORA可以扩展为同时合并空间和时间维度上相似的令牌。状态编码器需要能同时捕捉时空特征。多模态模型对于像CLIP这样的模型DORA可以分别应用于图像编码器和文本编码器或者设计一个跨模态的合并策略合并掉那些跨模态对齐后冗余的信息。与高效微调结合这是一个非常有趣的方向。想象一下在应用LoRA对ViT进行微调的同时也训练一个DORA智能体来动态管理令牌。LoRA负责高效地适应新任务DORA负责高效地执行推理两者结合可以实现“训练和推理双高效”。5.4 高级调优技巧分层策略不必在所有层都部署DORA。浅层特征通常包含更多低级细节合并需谨慎深层特征更抽象冗余度可能更高。可以设计一个“元控制器”动态决定在哪一层激活DORA或者为不同层使用不同复杂度的智能体。知识蒸馏使用原始未加速的ViT作为教师模型来蒸馏训练集成了DORA的模型。这可以为DORA智能体提供更丰富的监督信号尤其是在终端奖励稀疏的情况下。离线训练可以先使用一个简单的启发式策略如ToMe在大量数据上运行收集“状态-动作-奖励”轨迹然后用这些离线数据对DORA智能体进行预训练离线强化学习再进行在线微调。这可以大大加快训练收敛速度。硬件感知奖励在奖励函数中除了考虑FLOPs还可以加入对目标硬件如特定型号的GPU、NPU实际延迟的估计作为奖励的一部分让DORA学习到的策略对特定部署平台更友好。DORA为我们打开了一扇门将强化学习这种序列决策工具应用于模型推理过程中的自适应优化。它背后的思想——让模型学会如何更高效地运行自己——具有很大的潜力。虽然当前的实现和训练有一定复杂性但其带来的动态性、自适应性和可观的效率提升使得它在追求极致性能的场景下非常具有吸引力。在实际项目中不妨从一个简化版本开始比如固定只在最后几层应用使用较小的动作空间验证其收益再逐步增加复杂性。这个过程中对强化学习训练技巧的掌握和对视觉模型的理解两者缺一不可。 DORA：基于强化学习的ViT动态令牌合并，实现模型推理自适应加速
1. 项目概述当ViT遇上动态令牌合并DORA如何用强化学习“瘦身”模型最近在优化视觉Transformer模型时大家可能都遇到过这个头疼的问题模型精度和推理速度就像天平的两端难以兼得。尤其是在处理高分辨率图像时ViT模型需要将图像分割成大量的“令牌”进行处理计算量和内存开销呈平方级增长直接拖慢了整个系统的响应。传统的令牌剪枝或合并方法比如Token Merging虽然能减少令牌数量但大多是静态或启发式的往往“一刀切”在处理不同复杂度的图像内容时要么合并过度损失细节要么合并不足效率低下。这让我一直在想有没有一种更智能、更自适应的方式能根据每张图片的“内容复杂度”动态决定合并多少、合并哪些令牌这就是DORADynamic Online Reinforcement Agent for Token Merging要解决的核心问题。它不是一个全新的模型架构而是一个精巧的“插件式”智能体专门为ViT模型设计。其核心思想非常巧妙将令牌合并这个决策过程建模为一个序列决策问题并引入强化学习来训练一个智能体让它在线地、根据当前图像特征动态地决定每一步最优的合并策略。简单来说DORA就像一个坐在ViT模型里的“调度员”它实时观察经过处理的令牌特征然后判断“这一批令牌里哪些信息是冗余的、可以合并的合并多少才能既保住关键特征又最大化计算效率” 这个决策不再是固定的而是每张图、每一层都可能不同。对于需要部署ViT模型到资源受限环境如移动端、边缘设备的工程师或者任何关心模型效率的研究者来说DORA提供了一种全新的思路。它不改变ViT的主干网络而是通过一个轻量级的智能体模块实现模型推理过程的动态优化在几乎不损失精度的情况下显著提升吞吐量。最近社区里关于“Agentic RL”智能体化强化学习和高效微调技术如LoRA的讨论很多DORA可以看作是这两个趋势在模型推理优化上的一个有趣交汇用类似智能体的方式Agentic解决一个具体的优化问题并且其“动态性”与LoRA的“低秩适配”思想有异曲同工之妙都是在不改变原模型大规模参数的前提下通过添加小型可学习模块来实现特定目标。2. DORA核心设计思路将合并决策转化为智能体游戏要理解DORA我们得先拆解一下它面对的问题本质。在标准的ViT中一张图片被切成N个图像块每个块经过线性投影变成一个令牌。这些令牌会依次通过多个Transformer编码器层。Token Merging类方法的基本操作是在某一层计算所有令牌之间的相似度然后将最相似的一对令牌合并成一个用它们的加权平均来代表从而将令牌数量从N减少到N-1。这个过程可以重复进行直到达到预设的令牌保留率。传统方法的局限性在于这个“合并谁”以及“合并多少”的决策要么是基于一个全局固定的相似度阈值要么是根据一个预设的压缩率。但不同的图像内容差异巨大一张纯色背景的图片很多令牌高度相似可以大胆合并而一张细节丰富的医学影像过度合并则会丢失关键病理特征。一个静态策略无法适应这种变化。DORA的突破性思路在于它不再将合并决策视为一个基于规则的计算问题而是视为一个顺序决策过程。我们来想象这样一个场景环境当前Transformer层的所有令牌特征。智能体DORA模块它能够观察环境令牌特征。动作在每一步智能体需要做出两个决策(a) 是否在此刻执行合并(b) 如果合并选择哪一对令牌进行合并状态随着合并动作的执行令牌集合的特征和数量发生变化环境状态也随之改变。奖励智能体的最终目标是平衡效率和精度。因此奖励信号被设计为在推理结束时如果模型预测精度高则给予正奖励同时每合并一个令牌即减少计算量也给予一个小的正奖励。如果合并导致精度严重下降则给予负奖励。通过这样的建模训练DORA就变成了训练一个强化学习智能体让它学会在长期的“合并游戏”中最大化累积奖励。这意味着它必须学会评估当前令牌集合的“可合并性”并做出对未来最终分类精度最有利的即时决策。2.1 为什么选择强化学习而非其他方法你可能会问为什么不用一个简单的神经网络直接预测合并方案或者用可微分的软合并这里涉及到几个关键考量决策的离散性与序列性合并是一个典型的离散决策合并/不合并合并哪一对并且当前决策会影响后续决策可用的令牌集合。强化学习特别是基于策略梯度的方法非常擅长处理这类序列化的离散决策问题。稀疏延迟奖励我们最关心的模型最终精度是在所有合并决策完成后才得到的。这是一个典型的延迟且稀疏的奖励信号。强化学习算法如PPO、A2C本身就是为优化这类长期回报而设计的。在线适应训练好的DORA智能体在推理时是根据当前输入实时做出决策的实现了真正的“动态在线”无需任何额外的元数据或预处理。注意DORA的设计隐含了一个重要假设即存在一个“策略网络”它能够从令牌特征中提取出用于决策的抽象表示。这个网络通常非常小可能只有几层MLP以确保其引入的计算开销远小于它通过合并令牌所节省的开销。2.2 DORA与ToMe、DiffRate等方法的本质区别为了更清楚我们可以将DORA与之前的主流方法做个对比特性传统Token Merging (ToMe)可学习Token Merging (DiffRate)DORA (本文方法)决策依据固定的余弦相似度阈值可学习的评分网络预测每个令牌的保留分数强化学习智能体基于当前状态序列决策合并粒度逐层固定比例或阈值逐层可学习保留率逐样本、逐步骤动态决策优化目标启发式相似度最高端到端训练最小化任务损失最大化长期奖励精度效率动态性静态静态每层参数固定高度动态每张图决策不同额外开销几乎为零需要训练评分网络需要训练智能体推理时需运行策略网络可以看出DORA的核心优势在于其决策的上下文感知能力和灵活性。它不像DiffRate那样为每一层学习一个固定的“压缩率”而是为每一个输入实例学习一套“压缩策略”。3. DORA架构详解与实操要点理解了核心思想我们深入到DORA的具体实现架构。一个完整的DORA系统包含三个核心组件状态编码器、策略网络和价值网络。下面我们逐一拆解并说明在实操中需要注意的关键点。3.1 状态编码器如何让智能体“看懂”令牌智能体做出决策的前提是感知环境。在DORA中环境状态就是当前层的所有令牌特征。然而直接将所有令牌的原始特征例如768维向量输入给智能体是不现实的因为令牌数量N可能很大且是变化的。因此需要一个状态编码器来提取固定维度的全局状态表示。一个常见且有效的设计是对当前所有令牌特征进行全局平均池化得到一个汇总全局信息的向量。同时计算令牌特征的某些统计量如标准差、最大值等以反映特征的离散程度间接反映内容复杂度。将这些汇总信息拼接起来再通过一个小型MLP输出一个固定维度的状态向量s_t。实操要点编码器必须轻量状态编码器的计算量必须严格控制通常就是一层或两层线性层。它的目的是提供决策的“摘要”而非精细描述。包含序列信息对于更复杂的场景可以考虑引入轻量级的序列建模如一个单层的Transformer编码器或LSTM来处理令牌间的相对关系。但这会显著增加开销需要仔细权衡。归一化很重要输入状态编码器的特征应先进行层归一化确保训练稳定性。3.2 策略网络与动作空间设计策略网络π(a_t | s_t)接收状态向量s_t输出在当前状态下采取各个动作的概率分布。DORA的动作空间设计是其精妙之处。动作通常被设计为离散动作动作0不执行合并直接进入下一层或结束当前层的合并过程。动作1到K执行合并并指定合并哪一对令牌。这里的K是候选令牌对的数量。直接对所有可能的令牌对共N*(N-1)/2对进行建模是不现实的。因此需要一种高效的采样机制。一种实践方案是基于令牌特征的余弦相似度快速预筛选出最相似的Top-M个令牌对例如M10或20。策略网络的输出对应这M个候选对外加一个“不合并”动作。这样动作空间大小就是M1变得可管理。策略网络本身可以是一个简单的MLP输入状态s_t输出M1维的logits经过softmax后得到动作概率。实操心得探索与利用的平衡在强化学习训练初期需要鼓励智能体探索不同的合并策略。可以通过在策略网络的输出概率上添加熵正则项来实现防止智能体过早收敛到次优策略。动作掩码对于无效动作例如当令牌数已降到下限时“合并”动作应被禁止需要使用动作掩码将无效动作的概率置零并重新归一化。训练不稳定性策略梯度方法容易训练不稳定。务必使用像PPO近端策略优化这类更稳定的算法并仔细调整学习率、折扣因子等超参数。3.3 价值网络与奖励函数设计价值网络V(s_t)用于估计当前状态s_t的长期价值即从该状态开始遵循当前策略所能获得的期望累积奖励。它在计算优势函数用于更新策略时至关重要。奖励函数r_t的设计是引导智能体学习的关键。一个典型的设计是稀疏的终端奖励在整个推理过程结束时即图像经过所有Transformer层并得到分类结果后计算任务损失如交叉熵损失的负值作为最终奖励R_Terminal -α * Loss。α是一个缩放系数。稠密的步进奖励为了鼓励合并在每一步执行了“合并”动作后立即给予一个小的正奖励r_step β。β是一个很小的正数如0.01或0.001用于提供即时反馈引导智能体倾向于合并。总奖励最终的累积奖励是步进奖励之和加上终端奖励。注意事项奖励缩放终端奖励基于损失和步进奖励的量级可能相差很大。需要仔细调整α和β使两者处于同一数量级否则智能体可能只会优化其中一个目标。折扣因子γ用于计算累积奖励时对未来奖励的折扣。在DORA中由于合并决策的长期影响对最终精度的影响至关重要γ通常应设置得比较高如0.99让智能体更有远见。价值网络训练价值网络通过最小化时序差分误差来训练。它是独立于策略网络的一个回归头通常与策略网络共享状态编码器后的底层特征以提升效率。4. DORA训练与部署全流程实操纸上得来终觉浅我们来一步步拆解DORA从训练到部署的完整过程。这里以在ImageNet数据集上为一个预训练的ViT-B/16模型添加DORA为例。4.1 第一阶段环境搭建与模型准备首先你需要一个预训练好的ViT模型作为“环境”。假设我们使用timm库。import torch import torch.nn as nn import timm # 1. 加载预训练ViT-B/16作为主干网络 backbone timm.create_model(vit_base_patch16_224, pretrainedTrue, num_classes0) # 先不要分类头 backbone.eval() # 初始阶段冻结主干网络参数 # 2. 准备你的DORA智能体模块包含状态编码器、策略网络、价值网络 class DORAAgent(nn.Module): def __init__(self, token_dim, state_dim, num_actions): super().__init__() self.state_encoder nn.Sequential( nn.Linear(token_dim * 2, 256), # 输入全局平均标准差 nn.ReLU(), nn.Linear(256, state_dim) ) self.policy_net nn.Linear(state_dim, num_actions) self.value_net nn.Linear(state_dim, 1) # ... 其他初始化 def forward(self, token_features): # 计算全局状态 global_mean token_features.mean(dim1) global_std token_features.std(dim1) state torch.cat([global_mean, global_std], dim-1) encoded_state self.state_encoder(state) # 获取动作概率和状态价值 action_logits self.policy_net(encoded_state) state_value self.value_net(encoded_state) return action_logits, state_value # 初始化智能体 agent DORAAgent(token_dim768, state_dim128, num_actions21) # 假设20个候选对1个“不合并”关键步骤解析初始阶段我们通常冻结主干网络只训练DORA智能体。这是因为同时训练所有参数极其困难且容易破坏预训练模型已经学到的强大特征。num_actions21对应了20个最相似的候选令牌对加上“不合并”动作。这个数量需要根据你的计算预算和效果进行权衡。4.2 第二阶段强化学习训练循环这是最核心也是最复杂的部分。我们需要在一个循环中让智能体与环境ViT主干交互收集数据然后更新策略。import torch.optim as optim from torch.distributions import Categorical # 假设我们使用PPO算法 def compute_ppo_loss(old_logits, new_logits, actions, advantages, returns, epsilon0.2): old_dist Categorical(logitsold_logits) new_dist Categorical(logitsnew_logits) # 计算重要性采样比率 ratio torch.exp(new_dist.log_prob(actions) - old_dist.log_prob(actions.detach())) # PPO裁剪目标 surr1 ratio * advantages surr2 torch.clamp(ratio, 1-epsilon, 1epsilon) * advantages policy_loss -torch.min(surr1, surr2).mean() return policy_loss # 训练循环伪代码框架 optimizer optim.Adam(agent.parameters(), lr1e-4) for epoch in range(num_epochs): for batch_images, batch_labels in dataloader: # 1. 前向传播交互 tokens backbone.patch_embed(batch_images) # ... 添加位置编码等 trajectory [] # 存储 (state, action, reward, logits, value) current_tokens tokens for layer_idx in range(num_layers_to_apply_dora): # 智能体观察状态并做出决策 action_logits, state_value agent(current_tokens) action_dist Categorical(logitsaction_logits) action action_dist.sample() # 执行动作合并令牌或跳过 if action.item() ! 0: # 非“不合并”动作 # 根据动作索引找到对应的令牌对(i,j)并合并 i, j find_token_pair_to_merge(current_tokens, action) new_token (current_tokens[:, i] current_tokens[:, j]) / 2 # 更新令牌序列 # ... 实现合并逻辑current_tokens数量减1 step_reward merge_reward_bonus # 小正奖励 else: step_reward 0.0 # 存储数据 trajectory.append({ state: current_tokens.detach(), action: action, logits: action_logits.detach(), value: state_value.detach(), reward: step_reward }) # 应用当前层的Transformer计算简化表示 current_tokens backbone.blocks[layer_idx](current_tokens) # 2. 计算终端奖励 final_output backbone.norm(current_tokens) cls_output final_output[:, 0] # 接上一个临时的分类头计算损失 loss criterion(cls_output, batch_labels) terminal_reward -loss.item() * terminal_reward_scale # 将终端奖励加到最后一步的奖励上 trajectory[-1][reward] terminal_reward # 3. 计算优势函数和回报 # ... 使用GAE(Generalized Advantage Estimation)等方法计算 # 4. 更新智能体参数PPO更新 optimizer.zero_grad() # 使用收集的轨迹数据重新前向计算新的logits和value # 计算策略损失和价值损失 policy_loss compute_ppo_loss(old_logits, new_logits, ...) value_loss nn.MSELoss()(new_values, returns) total_loss policy_loss 0.5 * value_loss # 加权和 total_loss.backward() optimizer.step()实操现场记录与避坑指南训练速度慢这是RL训练的通病。一个批次的数据只能产生一条轨迹一个决策序列。为了稳定训练你需要收集大量轨迹才能进行一次有效的更新。这要求数据加载和交互过程必须高效。考虑使用向量化环境或重放缓冲区。奖励稀疏问题终端奖励非常稀疏。虽然步进奖励提供了即时反馈但智能体可能很难将最终的精度下降归因于早期某个错误的合并决策。课程学习是一个有效技巧先从简单的任务如合并率很低开始训练逐步增加任务难度允许更高的合并率。梯度流确保在交互阶段对主干网络的令牌特征执行.detach()防止主干网络的梯度通过交互过程回传除非你在后期进行联合微调。智能体的梯度应只通过其自身的网络计算。超参数敏感PPO的epsilon裁剪范围、学习率、折扣因子gamma、GAE的lambda等超参数对训练结果影响巨大。建议使用超参数优化工具如Optuna进行系统调优。4.3 第三阶段推理部署与性能评估训练完成后DORA智能体就可以用于动态加速推理了。# 推理模式下的DORA应用 torch.no_grad() def inference_with_dora(model, agent, image): tokens model.patch_embed(image) # ... 位置编码 for layer in model.blocks: # 在每一层或指定层前调用智能体决策 action_logits, _ agent(tokens) action torch.argmax(action_logits, dim-1) # 贪婪策略选择最大概率动作 if action ! 0: i, j find_token_pair_to_merge(tokens, action) # 执行合并 tokens merge_tokens(tokens, i, j) # 通过当前Transformer层 tokens layer(tokens) output model.norm(tokens) return model.head(output[:, 0])性能评估指标精度在测试集如ImageNet val上的Top-1/Top-5准确率。对比原始模型和使用DORA后的模型。加速比测量平均每张图片的推理时间或FPS。由于DORA动态合并令牌不同图片的加速效果不同需要统计平均值。令牌保留率统计在整个测试集上DORA平均将原始令牌数量减少了多少百分比。这是一个重要的效率指标。FLOPs/内存占用理论计算量FLOPs和峰值内存占用的减少比例。这能更直接地反映计算效率的提升。提示在部署时DORA智能体本身也会引入少量计算开销状态编码和策略前向传播。务必确保这个开销远小于因令牌合并而节省的Transformer层计算量。通常智能体网络设计得足够小例如参数量在万级别这个条件是满足的。5. 常见问题、调优技巧与扩展思考在实际操作中你肯定会遇到各种各样的问题。下面是我在复现和实验过程中遇到的一些典型情况及其解决方法以及一些让DORA工作得更好的技巧。5.1 训练不稳定与收敛困难这是强化学习项目最常见的问题。现象奖励曲线剧烈震荡没有上升趋势策略熵迅速降为零智能体停止探索价值估计爆炸。排查与解决检查奖励尺度终端奖励基于分类损失通常绝对值较大而步进奖励很小。如果两者尺度差异过大智能体可能只关注终端奖励而忽略效率。尝试将终端奖励归一化例如除以一个基线损失值如原始模型的损失。调整折扣因子γ如果智能体表现得非常短视只追求即时合并奖励尝试降低gamma如从0.99降到0.9。如果智能体难以学到长期影响尝试提高gamma。使用更稳定的算法确保你正确实现了PPO的裁剪和目标函数。可以考虑使用像Stable-Baselines3这样的成熟RL库来验证你的训练循环。增加批量大小策略梯度方法方差较大。增加用于计算梯度估计的轨迹数量批量大小可以显著稳定训练。引入基线在计算优势函数时使用价值网络V(s)作为基线是标准做法。确保你的价值网络训练得足够好能够提供准确的状态价值估计。5.2 智能体学不到有效策略现象训练后智能体要么永远选择“不合并”要么随机合并导致精度或效率没有提升。排查与解决奖励设计问题这是最可能的原因。如果“不合并”动作没有惩罚而合并动作可能带来风险精度下降智能体自然会倾向于保守。可以尝试给“不合并”动作一个微小的负奖励如-0.001鼓励其探索合并。或者调整终端奖励的权重让效率提升带来的正奖励更加突出。探索不足在训练初期策略网络的输出概率应尽可能均匀。可以通过设置较高的熵正则化系数来强制探索。随着训练进行再逐渐减小该系数。动作空间太大如果候选令牌对M设置得太大比如50智能体可能难以学习。尝试减小M比如5或10让决策变得更简单。也可以考虑对动作空间进行分层设计先决策是否合并再决策合并哪一对。5.3 如何将DORA应用到其他视觉任务或模型DORA的思想并不局限于图像分类和ViT。目标检测/分割对于DETR或Mask2Former这类基于Transformer的检测/分割模型令牌通常对应图像特征图上的位置或对象查询。DORA同样可以学习动态合并这些令牌或查询。奖励函数需要相应调整例如结合检测的mAP或分割的mIoU作为终端奖励。视频理解视频Transformer需要处理时空令牌计算量巨大。DORA可以扩展为同时合并空间和时间维度上相似的令牌。状态编码器需要能同时捕捉时空特征。多模态模型对于像CLIP这样的模型DORA可以分别应用于图像编码器和文本编码器或者设计一个跨模态的合并策略合并掉那些跨模态对齐后冗余的信息。与高效微调结合这是一个非常有趣的方向。想象一下在应用LoRA对ViT进行微调的同时也训练一个DORA智能体来动态管理令牌。LoRA负责高效地适应新任务DORA负责高效地执行推理两者结合可以实现“训练和推理双高效”。5.4 高级调优技巧分层策略不必在所有层都部署DORA。浅层特征通常包含更多低级细节合并需谨慎深层特征更抽象冗余度可能更高。可以设计一个“元控制器”动态决定在哪一层激活DORA或者为不同层使用不同复杂度的智能体。知识蒸馏使用原始未加速的ViT作为教师模型来蒸馏训练集成了DORA的模型。这可以为DORA智能体提供更丰富的监督信号尤其是在终端奖励稀疏的情况下。离线训练可以先使用一个简单的启发式策略如ToMe在大量数据上运行收集“状态-动作-奖励”轨迹然后用这些离线数据对DORA智能体进行预训练离线强化学习再进行在线微调。这可以大大加快训练收敛速度。硬件感知奖励在奖励函数中除了考虑FLOPs还可以加入对目标硬件如特定型号的GPU、NPU实际延迟的估计作为奖励的一部分让DORA学习到的策略对特定部署平台更友好。DORA为我们打开了一扇门将强化学习这种序列决策工具应用于模型推理过程中的自适应优化。它背后的思想——让模型学会如何更高效地运行自己——具有很大的潜力。虽然当前的实现和训练有一定复杂性但其带来的动态性、自适应性和可观的效率提升使得它在追求极致性能的场景下非常具有吸引力。在实际项目中不妨从一个简化版本开始比如固定只在最后几层应用使用较小的动作空间验证其收益再逐步增加复杂性。这个过程中对强化学习训练技巧的掌握和对视觉模型的理解两者缺一不可。