arXiv:2603.14802cs.LG2026-03被引 1

用JAX加速的脉冲神经网络库,让时间序列预测更快更准。

OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX

  • 基于JAX和Equinox构建,支持GPU加速与自动微分
  • 可实现高维动态系统建模,对混沌系统预测效果好
  • 适合需要快速迭代的时序建模研究者使用

OpenReservoirComputing(ORC)是基于JAX(Bradbury et al. 2018)和Equinox(Kidger and Garcia 2021)开发的脉冲计算(RC)Python库。RC是一种通过将低维信号映射到高维动力系统并训练线性读出层来完成任务的机器学习方法,常用于时间序列预测,尤其在混沌系统上表现优异。该库提供模块化组件和内置模型,支持预测、分类与控制任务。依托JAX的自动微分、即时编译(JIT)及GPU/TPU加速能力,ORC实现端到端可微,支持高效向量化与大规模架构,便于与等效框架中的深度学习模型无缝集成。

原文摘要 · Abstract (English)

OpenReservoirComputing (ORC) is a Python library for reservoir computing (RC) written in JAX (Bradbury et al. 2018) and Equinox (Kidger and Garcia 2021). JAX is a Python library for high-performance numerical computing that enables automatic differentiation, just-in-time (JIT) compilation, and GPU/TPU acceleration, while Equinox is a neural network framework for JAX. RC is a form of machine learning that functions by lifting a low-dimensional sequence or signal into a high-dimensional dynamical system and training a simple, linear readout layer from the high-dimensional dynamics back to a lower-dimensional quantity of interest. The most common application of RC is time-series forecasting, where the goal is to predict a signal's future evolution. RC has achieved state-of-the-art performance on this task, particularly when applied to chaotic dynamical systems. In addition, RC approaches can be adapted to perform classification and control tasks. ORC provides both modular components for building custom RC models and built-in models for forecasting, classification, and control. By building on JAX and Equinox, ORC offers GPU acceleration, JIT compilation, and automatic vectorization. These capabilities make prototyping new models faster and enable larger and more powerful reservoir architectures. End-to-end differentiability also enables seamless integration with other deep learning models built with Equinox.

脉冲计算时间序列JAXGPU加速

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。