强化注意力学习

Reinforced Attention Learning

arXiv: 2602.04884v1

论文信息

标题: Reinforced Attention Learning

作者: Bangzheng Li, Jianmo Ni, Chen Qu, et al.

发布日期: 2026-02-04

arXiv ID: 2602.04884v1

PDF 链接: 下载 PDF

3 分钟速览

  • 研究问题:传统强化学习通过优化输出 token 概率来提升多模态大语言模型(MLLM)的推理能力,但对真实的视觉感知任务提升有限,甚至可能损害基础感知性能。本论文试图解决如何更直接地优化 MLLM 的内部视觉信息分配过程。
  • 核心方法:提出强化注意力学习(RAL)框架,将模型内部的注意力分布视为一种策略,利用策略梯度方法进行优化。该方法不再试图教会模型 “说什么”,而是教会它 “看哪里”,通过强化高奖励的注意力模式来提升视觉基础能力。
  • 关键结果:RAL 在图像和视频的多个视觉问答基准上,一致性地超越了传统的 GRPO 强化学习基线。例如,在需要细粒度视觉搜索的 V* Bench 上,RAL 相对 GRPO 有**+5.8**分的显著提升(见表 3)。
  • 主要局限:论文在实验部分说明,此方法需要在训练时提取并存储模型的注意力权重,计算开销较大,且目前仅在 Qwen-2.5-VL 系列模型上验证,尚未在其他模型架构上进行测试。
  • 适合读者:对多模态大模型、强化学习、模型内部机制可解释性及知识蒸馏感兴趣的机器学习研究员、算法工程师和研究生。

论文背景和研究动机

大型语言模型(LLM)在数学、编程等复杂推理领域的成功,很大程度上归功于其后训练阶段使用的强化学习(RL)。主流范式,如近端策略优化(PPO)和分组相对策略优化(GRPO),其核心是优化模型输出 token 的概率分布,以最大化一个基于答案正确性或人类偏好的奖励函数。这种策略可以被理解为,让模型学会 “说什么” 才能获得高分。

当研究者们试图将这一成功范式迁移到多模态大语言模型(MLLM)上时,却遇到了一个根本性问题。为了提升模型在视觉问答(VQA)任务上的表现,现有方法会鼓励模型在给出最终答案前,先生成一个冗长的文本化思考过程。然而,论文作者的实验结果显示,对于以图像或视频感知为核心的任务,这种 “显式思考” 带来的增益微乎其微,有时甚至会损害模型原本的感知能力。

作者将这一局限归因于优化目标的错位。他们指出,在典型的 MLLM 架构中,准确回答一个细粒度的视觉问题,依赖于模型的注意力机制能够精准地从密集的视觉信息中,定位并关注与任务最相关的部分。这是一个内部的 “信息分配” 过程。标准的 RLHF 优化的是结果(生成的 token),而非这个内部过程。因此,即便一个回答获得了高奖励,被强化的也只是其外在的文本表现形式,而非其内在的、可泛化的视觉推理路径。这为论文的核心创新——直接优化注意力分布——提供了明确的研究动机。

核心方法和技术细节

针对上述问题,论文提出了强化注意力学习(RAL)框架。其核心思想是重构 MLLM 后训练的策略目标,从优化 “下一个 token 的预测概率” 转向优化模型内部的 “注意力分布”。

注意力策略的构建

RAL 首先将模型的注意力模式形式化为一个策略。对于一个由提示词和生成回复组成的序列 S=(x1,…,xT)S = (x_1, \dots, x_T),对于每个生成 token tt,模型会计算它对之前所有位置 ii 的原始注意力权重 αt,i\alpha_{t,i}。RAL 取最终 Transformer 层所有注意力头的平均值,并为每个生成 token tt 定义一个因果注意力分布策略 pθtp_{\theta}^{t}:

pθt(i)=αt,i∑j=1t−1αt,j,∀i∈[1,t−1]p_{\theta}^{t}(i) = \frac{\alpha_{t,i}}{\sum_{j=1}^{t-1} \alpha_{t,j}}, \quad \forall i \in [1, t-1]

这个分布清晰地表征了在生成第 tt 个词时,模型从其历史上下文中 “汲取” 信息的比例,它包含了模型对原始指令、视觉输入以及自己已生成的思考过程的综合关注。

优势加权注意力散度

有了注意力策略,下一步是形成一个可优化的目标函数。作者借鉴了 PPO/GRPO 中的重要性采样和优势函数思想,定义了注意力强化学习的损失函数:

LAttnRL=Et[At⋅D(pθt∥poldt)]L_{\text{AttnRL}} = \mathbb{E}_{t} \left[ A_t \cdot D(p_{\theta}^{t} \parallel p_{\text{old}}^{t}) \right]

其中,AtA_t 代表序列级别的优势函数。如果一次生成的整体奖励高(At>0A_t > 0),梯度更新会最小化当前注意力分布 pθtp_{\theta}^{t} 与产生高奖励的旧注意力分布 poldtp_{\text{old}}^{t} 之间的差异,从而 “鼓励” 这种成功的注意力模式。反之,若奖励低(At<0A_t < 0),则会增大这种差异,推动模型远离失败的注意力模式。为保证训练的稳定性,论文采用 Jensen-Shannon 散度(JSD)作为 D(⋅)D(\cdot)。

整体优化与知识蒸馏扩展

最终的训练目标是将这个注意力损失与标准的 GRPO 损失 LRL\mathcal{L}_{\text{RL}} 相结合,形成 Ltotal=LRL+λattnLAttnRL\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{RL}} + \lambda_{\text{attn}} L_{\text{AttnRL}},实现输出质量和内部推理过程的双重优化。

此外,论文还将这一思想扩展到了On-Policy 注意力蒸馏。不同于传统的知识蒸馏只让学生模型模仿教师模型的输出概率,作者提出让学生模型在自己的推理轨迹上,直接模仿教师的注意力分布 pϕtp_{\phi}^{t},损失函数为 LAttnDistill=Eτ∼πθ[∑tJSD(pθt∥pϕt)]\mathcal{L}_{\text{AttnDistill}} = \mathbb{E}_{\tau \sim \pi_{\theta}}[\sum_{t} \text{JSD}(p_{\theta}^{t} \parallel p_{\phi}^{t})]。这是一种更深层次的结构性知识迁移,旨在赋予学生模型与教师相似的 “视觉信息筛选” 习惯,而不仅仅是得出相同的结论。

创新点和贡献

该论文的首要和核心创新在于视角的转换。它将一个长期被看作是模型内部不可干预的 “机制”——注意力,提升为一个可以直接优化的 “第一公民” 策略目标。这打破了强化学习在 LLM 后训练中只关心输出 token 的范式,开辟了一条以过程为导向的优化路径。

第二个贡献是方法论的统一与泛化。RAL 的框架具有很强的延伸性。论文成功将其应用于 On-Policy 蒸馏,证明了 “在哪里看” 的知识可以像 “预测什么” 的知识一样被有效传递。这种将注意力作为一种通用策略的方法,为解决多模态对齐、减少幻觉等问题提供了一种直接且优雅的方案。

第三个贡献在于其实证发现。论文通过RAL-zero的消融实验(一种完全移除显式文本思考过程的配置)表明,即使没有长篇大论的推理文字,仅通过优化注意力分布,模型也能在多项视频理解任务上超越带有思考过程的 GRPO 基线(见表 2, NExT-QA +5.5, VideoMME +3.1)。这一发现深刻地揭示了,多模态大模型中很多视觉理解能力可能已经被编码在内部,只是缺乏正确的 “信息检索” 路径来激发,而 RAL 正是提供了这样一条路径。

实验结果分析

论文在多项图像和长视频的 VQA 基准上进行了严格的实验,核心结论如下:

一致性的性能超越:无论是与基础模型 Qwen-2.5-VL-7B 还是与 GRPO 基线相比,RAL 都展现了强大且一致的性能优势。在图像任务中,它以全胜的战绩超越了 GRPO(见表 3),尤其在需要从高分辨率图像中寻找目标的 V*基准上取得了大幅领先。在长视频基准上,RAL 在 7 个数据集中的 6 项超越了 GRPO(见表 2),特别是在需要强大时序推理的 NExT-QA 上表现优异。上述结果表明,优化注意力分布比优化纯文本输出提供了更稳定、更泛化的学习信号。

与视觉信息量正相关的性能增益:论文进一步通过在不同图像分辨率(512, 1024, 2048 tokens)和视频帧数(32, 64, 128 帧)下进行实验,揭示了一个重要规律(见图 3)。RAL 相对于 GRPO 的优势,随着输入包含的视觉信息量(分辨率或帧率)增加而扩大。例如在 V*基准上,当图像 token 从 512 增至 2048 时,RAL 对 GRPO 的领先优势从+1.6 分扩大到+6.3 分。这表明 RAL 尤其擅长在信息密度大、需要精挑细选的环境中运作,其能力不仅没有饱和,反而随着信息粒度的细化而增强,展现出极佳的可扩展性。

基于注意力的蒸馏显著优于传统方式:实验还证明了 On-Policy 注意力蒸馏的价值。在大多数基准上,将注意力分布从 32B 教师模型蒸馏给 7B 学生模型,相比仅蒸馏输出概率,带来了显著的额外提升(在 V*上带来+4.2 分增益,在 VideoMME 上带来+2.6 分增益,见表 2 和表 3)。这直接证明了注意力分布中包含着文本输出概率所无法完全捕获的深层次、结构化的推理知识。

实践建议

对于有兴趣落地和应用此项工作的工程师与研究者,以下建议可供参考:

  1. 架构适应性改造:RAL 的实现需要提取训练时每个 token 的注意力权重。这在标准实现(如 Flash Attention)中为节约显存是默认关闭的。因此在工程上,你需要将模型中的注意力层替换为默认的 “eager” 模式实现,这会显著增加内存占用和训练时间(论文在 8 张 H100 集群上训练需约 120 小时)。建议引入梯度检查点或采用最新的支持在 Flash Attention 中高效输出注意力权重的技术,以在性能与资源间取得平衡。

  2. 超参数搜索策略:RAL 引入了新的超参数 λattn\lambda_{\text{attn}} 来控制注意力损失的强度(论文搜索范围为{0.5, 1, 5}, 见表 1)。与标准的语言生成损失不同,注意力损失是在一个分布上计算的,其量级可能差异巨大。建议在实际应用中,以输出奖励的方差为参考,先对注意力散度项进行归一化处理,再进行加权组合,这可以使超参数在不同任务和模型间更具通用性。

  3. 作为泛化能力增强器而非基础能力开创者:RAL-zero 实验表明,即使在无文本推理链时,该方法也能带来增益。这提示 RAL 可以作为一种强大的 “后训练插件”,用于微调那些已经预训练好的模型,以提升其泛化能力和抗遗忘能力。尤其对于容易产生幻觉的细粒度视频分析或多轮视觉对话任务,仅用少量数据进行 RAL 微调,可能比使用大量数据的 SFT 更有效。在实践中,可以尝试将 RAL 作为业务流水线上 SFT 之后的一个标准化增强步骤,专门用于提升模型应对高分辨率、长序列输入时的鲁棒性和精确度。