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

TL;DR

This work targets efficient computation of path signatures in deep learning by introducing Keras Sig, a pure Python implementation built on Keras 3 that runs across PyTorch, JAX, and TensorFlow backends. The core method reformulates signature computation to exploit GPU parallelism via cumulative sums and parallel tensor operations, trading memory usage for substantial speedups. Benchmarking shows up to 5–10× improvements in direct signature computation on GPU (and up to ~55% faster training for models incorporating signatures) relative to existing libraries, with backend-specific trade-offs: JAX/XLA often leads, while CPU performance aligns more closely across methods. The design emphasizes automatic backend selection, maintainability via high-level tensor ops, and practical guidance on memory constraints and compilation behavior, suggesting a viable path for broader adoption of signature-based tools in deep learning workflows.

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.

Figures (12)

• Figure 1: Average signature computation time on GPU as a function of batch size (Seq. Length=100, Depth=4).
• Figure 2: Average signature computation time on GPU as a function of sequence length (Batch Size=128, Depth=4).
• Figure 3: Average signature computation time on GPU as a function of signature depth (Batch Size=128, Seq. Length=100).
• Figure 4: Average signature computation time on CPU as a function of batch size (Seq. Length=100, Depth=4).
• Figure 5: Average signature computation time on CPU as a function of sequence length (Batch Size=128, Depth=4).