Keras Sig: Efficient Path Signature Computation on GPU in Keras 3

TL;DR

Keras Sig is a GPU-accelerated library built in Keras 3 that significantly speeds up path signature computations for deep learning, achieving up to 10-fold improvements over existing methods while ensuring compatibility and scalability.

Contribution

The paper introduces Keras Sig, a novel GPU-optimized library for path signature computation that leverages high-level tensor operations for improved performance and compatibility.

Findings

Reduces signature computation time by up to 55%

Achieves 5 to 10-fold speedups over existing methods

Maintains scalable performance across various hardware and signature parameters

Abstract

In this paper we introduce Keras Sig a high-performance pythonic library designed to compute path signature for deep learning applications. Entirely built in Keras 3, \textit{Keras Sig} leverages the seamless integration with the mostly used deep learning backends such as PyTorch, JAX and TensorFlow. Inspired by Kidger and Lyons (2021),we proposed a novel approach reshaping signature calculations to leverage GPU parallelism. This adjustment allows us to reduce the training time by 55\% and 5 to 10-fold improvements in direct signature computation compared to existing methods, while maintaining similar CPU performance. Relying on high-level tensor operations instead of low-level C++ code, Keras Sig significantly reduces the versioning and compatibility issues commonly encountered in deep learning libraries, while delivering superior or comparable performance across various hardware configurations. We demonstrate through extensive benchmarking that our approach scales efficiently with the length of input sequences and maintains competitive performance across various signature parameters, though bounded by memory constraints for very large signature dimensions.