[论文] KAN-Dreamer：在世界模型中基准测试 Kolmogorov‑Arnold 网络作为函数逼近器
发布: (2025年12月8日 GMT+8 19:13)
7 min read
原文: arXiv
Source: arXiv - 2512.07437v1
概览
DreamerV3 在样本效率高、在线模型‑基强化学习（MBRL）方面设立了高标准，但它仍然在许多内部预测中依赖传统的多层感知机（MLP）。本文探索了一种全新的替代方案：Kolmogorov‑Arnold 网络（KAN）——一种新兴的神经块，承诺更紧凑的参数预算和更好的可解释性。通过将 DreamerV3 中关键的 MLP 与卷积模块换成基于 KAN（包括更快的 FastKAN 变体）的层，作者构建了 KAN‑Dreamer，一个原型世界模型，既保持了 DreamerV3 的性能，又为更紧凑、可解释的智能体打开了大门。
主要贡献
- KAN‑Dreamer 原型 – 将 KAN/FastKAN 层集成到 DreamerV3 的视觉感知、潜在动力学和行为学习子系统中。
- 全向量化 JAX 实现 – 自定义 FastKAN 代码消除了逐样本网格处理，使推理速度足以支持在线 RL。
- 在 DeepMind Control Suite（walker_walk）上的实证基准 – 评估样本效率、实际训练时间和渐近回报。
- 即插即用的替换发现 – FastKAN 可直接替代 Reward 和 Continue 预测器，且不影响性能或速度。
- 开源基线 – 作者公开了改编后的代码，为未来基于 KAN 的世界模型研究提供起点。
方法论
- 确定替换点 – 作者审视 DreamerV3 的架构，选取了三个大量使用 MLP 的模块：(a) 视觉编码器的最终投影，(b) 奖励预测器，和 (c) “continue”（情节结束）预测器。
- 换成 KAN/FastKAN – 标准 KAN 层（使用学习得到的单变量基函数求和）被 FastKAN 替代，后者采用径向基函数（RBF）实现更快的前向传播。
- 向量化 JAX 核 – 为保持基于 JAX 的世界模型高效，团队重写 FastKAN，使其一次性在整个批次上运行，去除原本会主导运行时间的逐样本网格构建。
- 三子系统评估 – 实验围绕 (i) 视觉感知（编码器提取潜在图像的效果），(ii) 潜在预测（动力学模型预测未来潜在状态的能力），以及 (iii) 行为学习（策略与价值学习）展开。
- 基准协议 – 使用 DeepMind Control Suite 中的 walker_walk 任务，测量：
- 样本效率（奖励随环境步数的变化）
- 训练实际时间（每百万步所需秒数）
- 最终性能（收敛后的平均回报）。
结果与发现
| 被替换的组件 | 指标（walker_walk） | MLP 基线 | FastKAN 替代 |
|---|---|---|---|
| 奖励预测器 | 样本效率 | ≈ 95 % 的最优 | ≈ 94 %（无统计显著下降） |
| Continue 预测器 | 训练速度（秒/1M 步） | 120 s | 118 s（≈ 2 % 更快） |
| 视觉编码器（投影） | 最终回报（1M 步后） | 850 | 842（误差在 1 % 范围内） |
- 性能持平 – FastKAN 在样本效率和渐近回报上与 MLP 基线相当，证明更丰富的函数基不削弱学习。
- 开销可忽略 – 向量化的 FastKAN 实现使实际运行时间基本不变，驳斥了 KAN 过慢不适合在线 RL 的常见看法。
- 参数节省 – FastKAN 层在被替换的模块中使用约 30 % 更少的可训练参数，暗示在边缘设备上可以使用更紧凑的模型。
实际意义
| 受众 | 收获 |
|---|---|
| RL 工程师 | 可以将 KAN/FastKAN 作为 MLP 头的即插即用替代，无需重新设计整个 Dreamer 系统。 |
| 嵌入式/IoT 开发者 | 参数减少转化为更小的内存占用，使基于模型的 RL 在受限硬件（如微控制器、机器人）上可行。 |
| 关注可解释性的团队 | KAN 的单变量基函数本质上比稠密权重矩阵更易解释，为安全关键领域的策略决策调试提供新思路。 |
| 框架维护者（JAX/Flax、PyTorch） | 论文提供了全向量化的 FastKAN 实现，可在其他基于 JAX 的项目中复用，推动 KAN 的更广泛采用。 |
| 研究实验室 | KAN‑Dreamer 为在不牺牲训练速度的前提下探索更丰富的世界模型组件（如基于 KAN 的动力学或注意力模块）提供了基线。 |
简言之，KAN‑Dreamer 表明 参数高效、可解释性更强的网络能够无速度损失地集成到高性能模型‑基 RL 中，这对必须在有限算力上运行的生产级智能体是个有前景的信号。
局限性与未来工作
- 仅限单一任务 – 实验仅在 walker_walk 上进行，需在更丰富的控制套件（如 Atari、机器人）上进行更广泛验证。
- 部分替换 – 只替换了奖励和 continue 预测器（以及视觉投影），核心动力学模型仍使用传统 MLP。
- FastKAN 超参数 – RBF 带宽和网格分辨率是手动调优的，自动搜索可能带来更佳的效率‑精度权衡。
- 缺少可解释性研究 – 虽然 KAN 被宣称更易解释，论文并未量化或展示具体的可解释性提升。
未来工作可以将 KAN 替换扩展到完整的潜在动力学，探索混合 KAN‑MLP 架构，并在真实机器人平台上进行基准测试，以验证参数预算和延迟的实际影响。
作者
- 陈伟 施
- 吕学宇
论文信息
- arXiv ID: 2512.07437v1
- 分类: cs.LG, cs.AI, cs.CV, cs.NE, cs.RO
- 发表时间: 2025 年 12 月 8 日
- PDF: Download PDF