OpenReservoirComputing:基于 JAX 的 GPU 加速滞洛计算
机器学习
2026-03-17 v1
摘要
OpenReservoirComputing (ORC) 是一个基于 JAX (Bradbury 等,2018) 和 Equinox (Kidger and Garcia,2021) 编写的滞洛计算 (reservoir computing, RC) Python 库。JAX 是一个用于高性能数值计算的 Python 库,支持自动微分、即时编译 (JIT) 以及 GPU/TPU 加速,而 Equinox 是 JAX 的神经网络框架。RC 是一种机器学习方法,通过将低维序列或信号映射到高维动力学系统,然后训练一个简单的线性读取层,将高维动力学的结果映射回感兴趣的低维量来实现。RC 最常见的应用是时间序列预测,即预测信号的未来演化。RC 在这一任务上已取得最先进性能,尤其是应用于混沌动力学系统时。此外,RC 方法还可以用于分类和控制任务。ORC 提供用于构建自定义 RC 模型的模块化组件以及用于预测、分类和控制的内置模型。通过基于 JAX 和 Equinox,ORC 提供 GPU 加速、JIT 编译和自动向量化。这些能力使新模型的原型设计更快,并能够构建更大更强大的滞洛结构。端到端的可微性还使其能够无缝集成到由 Equinox 构建的其他深度学习模型中。
引用
@article{arxiv.2603.14802,
title = {OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX},
author = {Jan Williams and Dima Tretiak and Steven L. Brunton and J. Nathan Kutz and Krithika Manohar},
journal= {arXiv preprint arXiv:2603.14802},
year = {2026}
}