中文

SpikeRL:用于深度脉冲强化学习的可扩展且节能的框架

机器学习 2025-02-26 v1 人工智能 神经与进化计算

摘要

在 AI 革命的时代,对大规模数据驱动 AI 系统的大规模投资需要高性能计算,消耗巨大的能源和资源。这一趋势在不牺牲可扩展性或性能的情况下优化可持续性提出了新挑战。在传统冯·诺依曼架构的节能替代方案中,神经形态计算及其脉冲神经网络由于其固有的能效而成为极具前景的选择。然而,在诸如复杂连续控制任务等某些实际应用场景中,SNNs 通常缺乏传统人工神经网络所具备的性能优化。研究人员通过将 SNNs 与深度强化学习相结合来解决这一问题,但可扩展性仍未被探索。在本文中,我们扩展了之前关于 SpikeRL 的工作,这是一个用于基于 DeepRL 的 SNNs 的可扩展且节能的连续控制框架。在我们最初的 SpikeRL 框架实现中,我们依赖于来自种群编码脉冲演员网络方法的种群编码来构建我们的 SNN 模型,并通过 mpi4py 利用消息传递接口(MPI)实现了分布式训练。同时,通过使用混合精度进行参数更新来进一步优化我们的模型训练。在我们新的 SpikeRL 框架中,我们实现了自带种群编码的 DeepRL-SNN 组件,并使用带有 NCCL 后端的 PyTorch Distributed 包进行分布式训练,同时仍通过混合精度训练进行优化。我们新的 SpikeRL 实现比最先进的 DeepRL-SNN 方法快 4.26 倍,能效高 2.25 倍。我们提出的 SpikeRL 框架展示了在实际应用中的复杂连续控制任务的真正可扩展和可持续的解决方案。

关键词

引用

@article{arxiv.2502.17496,
  title  = {SpikeRL: A Scalable and Energy-efficient Framework for Deep Spiking Reinforcement Learning},
  author = {Tokey Tahmid and Mark Gates and Piotr Luszczek and Catherine D. Schuman},
  journal= {arXiv preprint arXiv:2502.17496},
  year   = {2025}
}