OpenReservoirComputing: GPU-Accelerated Reservoir Computing in JAX
This work addresses the need for efficient and scalable RC implementations for researchers and practitioners in machine learning, particularly for time-series analysis, but it is incremental as it primarily packages existing methods into a new library.
The authors tackled the problem of accelerating reservoir computing (RC) for tasks like time-series forecasting by developing OpenReservoirComputing, a Python library built on JAX and Equinox, which provides GPU acceleration, JIT compilation, and automatic vectorization to enable faster prototyping and more powerful reservoir architectures.
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.