Loss-Guided Auxiliary Agents for Overcoming Mode Collapse in GFlowNets¶
会议: AAAI 2026
arXiv: 2505.15251
代码: 有
领域: 生成式流网络 / 探索
关键词: GFlowNet, 模式坍缩, 辅助Agent, 损失引导, 多样性采样
一句话总结¶
提出 LGGFN（Loss-Guided GFlowNets），用辅助 GFlowNet 的探索直接由主 GFlowNet 的训练损失驱动——辅助 Agent 的奖励 = 原始奖励 + λ·主模型损失，优先采样主模型理解不足的区域，在网格/序列/贝叶斯结构学习任务上分别发现 40× 更多唯一模式、99% 探索误差降低。
研究背景与动机¶
-
领域现状：GFlowNet 设计用于从奖励函数按比例采样多模态分布（而非仅找最优点），理论上避免模式坍缩。但实践中 on-policy 训练仍面临模式坍缩——被早期发现的高奖励模式吸引。
-
现有痛点：
- 现有探索技术依赖启发式新奇信号（如状态计数、RND）——与模型实际学习状态脱节
- 新奇信号可能引导探索到无关区域（新但不重要）
-
缺乏利用模型自身训练信号来指导探索的机制
-
核心矛盾：GFlowNet 需要广泛探索来学习完整的分布，但 on-policy 采样被高奖励区域主导——需要一种机制将采样重定向到模型尚未学好的区域。
-
本文要解决什么？ 用主模型的训练损失作为辅助 Agent 的探索信号，实现"哪里不会去哪里"。
-
切入角度：训练损失高的区域 = 主模型理解不足的区域 → 辅助 Agent 优先探索这些区域 → 将样本反馈给主模型训练。
-
核心 idea 一句话：辅助 GFlowNet 奖励 = 主模型奖励 + 主模型损失 → 定向探索薄弱区域。
方法详解¶
整体框架¶
双 GFlowNet 架构：主 Agent（学习目标分布）+ 辅助 Agent（学习"主 Agent 的损失加权分布"）。主 Agent 的训练数据 = on-policy 采样 ∪ 辅助 Agent 采样。辅助 Agent 定期更新其奖励（反映主 Agent 当前的薄弱点）。
关键设计¶
- 损失引导辅助奖励:
- 做什么：让辅助 Agent 优先探索主 Agent 损失高的区域
- 核心思路：\(R_{aux}(x) = R(x) + \lambda \cdot L_{main}(x)\)，\(L_{main}\) 是主 GFlowNet 在轨迹 \(x\) 上的训练损失
- 设计动机：损失高 = 主模型不确定/理解不足 → 最需要更多数据的区域。比新奇信号更直接——直接测量"学得不好"而非"去过没去过" - 实现细节：损失可以是基于轨迹的 TB 损失、基于转换的 FM 损失或子轨迹损失，方法与具体训练目标无关；\(\lambda\) 保持 \(R_{aux}\) 与 \(R_{main}\) 同量级，防止训练不稳定
- 混合训练策略:
- 做什么：主 Agent 从两个来源学习
- 核心思路：训练 batch = α · on-policy 样本 + (1-α) · 辅助 Agent 样本。α 控制探索 vs 利用
-
设计动机：纯辅助采样可能偏离奖励分布太远；混合保持分布学习的稳定性
-
神经网络泛化利用:
- 做什么：利用神经网络的泛化能力使损失信号在相邻状态间"扩散"
- 核心思路：主模型在某个轨迹上损失高 → 通过网络泛化，相邻轨迹的损失也会较高 → 辅助 Agent 的探索不是点对点而是区域性的
- 设计动机：这是比状态计数更高效的探索——一次有信息量的探索覆盖整个"不确定区域"
损失函数 / 训练策略¶
- 主 Agent：标准 GFlowNet 训练损失（TB/DB/SubTB）
- 辅助 Agent：TB 损失，奖励 = R + λL_main
实验关键数据¶
主实验¶
| 任务 | LGGFN | 最佳基线 | 改善 |
|---|---|---|---|
| Hypergrid 128×128 L1↓ | 0.83±0.21 | 0.92±0.36 (AdaTeachers) | 10% |
| 序列生成 唯一模式数↑ | 40× 更多 | 基线 | 40× |
| 序列生成 探索误差↓ | 99% 降低 | 基线 | 99% |
| 贝叶斯结构学习 | +10% | 基线 | 10% |
消融¶
| 配置 | 效果 |
|---|---|
| 仅 on-policy（无辅助） | 严重模式坍缩 |
| 随机探索（辅助无引导） | 中等改善 |
| 新奇性引导 | 改善但不稳定 |
| 损失引导（LGGFN） | 最优且稳定 |
关键发现¶
- 损失信号 > 新奇性信号：损失直接测量"学得不好"，比"去没去过"更精确
- 40× 更多模式（序列生成）：证明 LGGFN 真正解决了模式坍缩而非仅改善近似精度
- 辅助 Agent 的"损失递减"：随着主模型改善，辅助 Agent 的探索方向自然转变——自适应机制
亮点与洞察¶
- 用训练损失直接驱动探索的思路简洁但强大——"哪里不会去哪里"是最自然的探索策略
- 对任何需要多样性采样的生成模型（如药物设计、材料搜索）都有迁移价值
- GFlowNet + 辅助 Agent 的双模型框架可推广到其他 RL 场景
局限性 / 可改进方向¶
- 辅助 Agent 的 λ 超参数需要调节
- 双 GFlowNet 增加了计算和内存开销
- 仅在离散空间任务上验证
相关工作与启发¶
- vs Adaptive Teachers (Jain et al.)：用教师分布引导。LGGFN 直接用损失信号更精准
- vs RND / Curiosity：新奇性驱动探索侧重于状态访问频率，LGGFN 的损失驱动更精准地指向模型实际没学好的区域
- vs Diverse RL：多样性 RL 通常用内在奖励或信息论正则化，而 LGGFN 通过辅助 Agent 的显式引导更可控且可解释
- 损失引导探索可推广到 RLHF 中的多样性生成，也可用于对抗样本挖掘
评分¶
- 新颖性: ⭐⭐⭐⭐ 损失引导探索概念简洁有力，用辅助 Agent 的高损失区域引导主 Agent 探索是非常自然的解耦设计
- 实验充分度: ⭐⭐⭐⭐ 三类任务（超图、分子、序列生成）、多基线、消融完整，跨域验证说服力强
- 写作质量: ⭐⭐⭐⭐ 动机链清晰，从 mode collapse 问题到损失引导方案的推导逻辑性好
- 价值: ⭐⭐⭐⭐ 对多模态采样和分子设计有直接应用价值，是 GFlowNet 社区解决 mode collapse 的重要改进
补充说明¶
- 损失引导的辅助 Agent 设计思路不仅限于 GFlowNet，可推广到任何需要多样性探索的生成模型场景