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

2025-01-14

AI Summary

To address the low computational efficiency of path signatures for long sequences, poor multi-backend compatibility, and high memory overhead for large-scale signatures in deep learning, this paper introduces the first native Keras 3–compatible, high-performance path signature library. Methodologically, it replaces conventional recursive implementations with a novel tensor-reshaping strategy—eliminating C++ dependencies—and designs a unified, multi-backend (TensorFlow/PyTorch/JAX) API alongside a higher-order GPU-accelerated batching algorithm. Experimental results demonstrate a 55% reduction in training time, 5–10× speedup in direct signature computation, stable CPU performance, and excellent scalability with respect to both sequence length and signature order. This work provides a lightweight, general-purpose, and high-performance foundational toolkit for signature-based deep learning models.

Technology Category

• Machine Learning: Matrix & Tensor Methods
• Computer Vision: Learning & Optimization for CV
• Search and Optimization: Learning to Search

Application Category

• Graph Algorithms and Modeling for the Web: Algorithms and analysis for heterogeneous, signed, attributed, multi-relational, temporal, higher-order, and annotated Web-related graphs
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Web learning to rank, online learning, and counterfactual learning for ranking

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, extit{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.

Problem

• Deep Learning
• Signature Computation
• Memory Optimization

Innovation

• Path Signatures
• GPU Acceleration
• Keras Integration