Decision SpikeFormer: Spike-Driven Transformer for Decision Making

📄 arXiv: 2504.03800v1 📥 PDF

作者: Wei Huang, Qinying Gu, Nanyang Ye

分类: cs.LG, cs.AI, cs.NE

发布日期: 2025-04-04

备注: This work has been accepted to CVPR 2025

🔗 代码/项目: PROJECT_PAGE

💡 一句话要点

提出Decision SpikeFormer，一种用于离线强化学习的脉冲驱动Transformer模型。

🎯 匹配领域: 支柱二：RL算法与架构 (RL & Architecture) 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 离线强化学习 脉冲神经网络 Transformer 自注意力机制 序列建模

📋 核心要点

  1. 现有基于ANN的离线强化学习方法计算和能量需求高，不适用于能量受限的嵌入式AI应用。
  2. DSFormer通过时间脉冲自注意力和位置脉冲自注意力捕捉序列建模中的时间与位置依赖性，并使用PTBN保持SNN的脉冲特性。
  3. 在D4RL基准测试中，DSFormer优于SNN和ANN，节能78.4%，同时保持了竞争性的性能。

📝 摘要（中文）

离线强化学习(RL)允许仅基于预先收集的数据进行策略训练，避免了直接的环境交互，这对于能量受限的嵌入式AI应用至关重要。虽然基于人工神经网络(ANN)的方法在离线RL中表现良好，但它们的高计算和能量需求促使人们探索更有效的替代方案。脉冲神经网络(SNNs)由于其低功耗而显示出在这类任务中的潜力。本文介绍DSFormer，这是第一个旨在通过序列建模解决离线RL的脉冲驱动Transformer模型。与现有侧重于视觉任务空间维度的SNN Transformer不同，我们在DSFormer中开发了时间脉冲自注意力(TSSA)和位置脉冲自注意力(PSSA)，以捕捉RL序列建模必不可少的时间和位置依赖性。此外，我们提出了渐进阈值依赖批量归一化(PTBN)，它结合了LayerNorm和BatchNorm的优点，以保持时间依赖性，同时保持SNN的脉冲特性。在D4RL基准测试中的全面结果表明，DSFormer优于SNN和ANN，实现了78.4%的节能，突出了DSFormer不仅在能源效率方面，而且在竞争性能方面的优势。代码和模型已公开。

🔬 方法详解

问题定义：论文旨在解决离线强化学习中，现有基于人工神经网络（ANN）的方法功耗过高的问题。这些方法虽然性能良好，但在能量受限的嵌入式AI应用中难以部署。现有的脉冲神经网络（SNN）在强化学习中的应用不足，尤其是在序列建模方面，缺乏有效的时序信息处理机制。

核心思路：论文的核心思路是利用SNN的低功耗特性，并结合Transformer模型的序列建模能力，设计一种新型的脉冲驱动Transformer模型（DSFormer）。通过专门设计的脉冲自注意力机制，DSFormer能够有效地捕捉强化学习任务中的时间依赖性和位置依赖性，从而在保证性能的同时降低功耗。

技术框架：DSFormer的整体架构基于Transformer模型，但所有计算都由脉冲神经网络执行。主要包含以下模块：输入编码层（将状态、动作等信息编码为脉冲序列）、时间脉冲自注意力（TSSA）层、位置脉冲自注意力（PSSA）层、前馈神经网络层（也是脉冲神经网络）、以及输出解码层（将脉冲序列解码为动作）。此外，还使用了渐进阈值依赖批量归一化（PTBN）来稳定训练过程并保持SNN的脉冲特性。

关键创新：论文的关键创新在于以下几点：1) 提出了时间脉冲自注意力（TSSA）和位置脉冲自注意力（PSSA），这是专门为脉冲神经网络设计的自注意力机制，能够有效地捕捉时间依赖性和位置依赖性。2) 提出了渐进阈值依赖批量归一化（PTBN），它结合了LayerNorm和BatchNorm的优点，能够在保持SNN脉冲特性的同时稳定训练过程。3) 将SNN和Transformer结合，用于离线强化学习，是该领域的首次尝试。

关键设计：TSSA和PSSA的关键在于如何将传统的自注意力机制转换为脉冲神经网络可以执行的形式。这涉及到脉冲的编码方式、注意力权重的计算方式、以及信息的聚合方式。PTBN的关键在于如何根据脉冲的阈值动态调整归一化的参数，以避免破坏脉冲的稀疏性。损失函数采用标准的强化学习损失函数，例如Q-learning或策略梯度相关的损失函数。网络结构方面，Transformer的层数、隐藏层维度、以及注意力头的数量都是重要的超参数。

🖼️ 关键图片

fig_0
fig_1
fig_2

📊 实验亮点

实验结果表明，DSFormer在D4RL基准测试中取得了显著的性能提升，并且相比于传统的ANN方法，实现了78.4%的节能。这表明DSFormer不仅具有竞争性的性能，而且在能源效率方面具有显著的优势。此外，DSFormer也优于其他SNN方法，证明了其在序列建模方面的有效性。

🎯 应用场景

DSFormer在能量受限的嵌入式AI应用中具有广泛的应用前景，例如机器人控制、自动驾驶、智能家居等。其低功耗特性使得它能够在资源有限的设备上运行复杂的强化学习算法，从而实现更智能、更高效的决策。此外，该研究也为SNN在强化学习领域的应用提供了新的思路和方法。

📄 摘要（原文）

Offline reinforcement learning (RL) enables policy training solely on pre-collected data, avoiding direct environment interaction - a crucial benefit for energy-constrained embodied AI applications. Although Artificial Neural Networks (ANN)-based methods perform well in offline RL, their high computational and energy demands motivate exploration of more efficient alternatives. Spiking Neural Networks (SNNs) show promise for such tasks, given their low power consumption. In this work, we introduce DSFormer, the first spike-driven transformer model designed to tackle offline RL via sequence modeling. Unlike existing SNN transformers focused on spatial dimensions for vision tasks, we develop Temporal Spiking Self-Attention (TSSA) and Positional Spiking Self-Attention (PSSA) in DSFormer to capture the temporal and positional dependencies essential for sequence modeling in RL. Additionally, we propose Progressive Threshold-dependent Batch Normalization (PTBN), which combines the benefits of LayerNorm and BatchNorm to preserve temporal dependencies while maintaining the spiking nature of SNNs. Comprehensive results in the D4RL benchmark show DSFormer's superiority over both SNN and ANN counterparts, achieving 78.4% energy savings, highlighting DSFormer's advantages not only in energy efficiency but also in competitive performance. Code and models are public at https://wei-nijuan.github.io/DecisionSpikeFormer.