Tensor Network Kernel Machines: A JAX Framework for Machine Learning and Nonlinear System Identification
Albert Saiapin, Kim Batselier
cs.MS, cs.LG, cs.SY, eess.SY
Submitted: 2026-08-07
Updated: 2026-08-10
Comments: 10 pages, 6 figures, 4 tables, 1 listing. Code available at: https://github.com/AlbMLpy/tnkm
Code: https://github.com/AlbMLpy/tnkm
License: http://creativecommons.org/licenses/by/4.0/
The gist: Developing nonlinear models that are both expressive and computationally efficient remains a challenge in machine learning and nonlinear system identification.
Terminology
Abstract
Developing nonlinear models that are both expressive and computationally efficient remains a challenge in machine learning and nonlinear system identification. Tensor network kernel machines (TNKM) address this challenge by combining nonlinear feature representations with compact low-rank tensor-network parameterizations. However, practical and extensible software frameworks for developing TNKM models remain limited. In this work, we introduce "tnkm", an open-source Python library for constructing and training TNKM models using JAX. The library provides a unified interface for combining different feature maps, tensor-network architectures, and optimization strategies, including alternating least squares and gradient-based methods. We demonstrate the capabilities of "tnkm" on nonlinear benchmark problems, showing that the implemented models achieve competitive prediction accuracy while retaining compact parameterizations and efficient training. The proposed framework facilitates reproducible development and application of tensor-network-based learning methods.
Sources
- Hyperparameter Optimization: Foundations, Algorithms, Best Practices and Open Challenges
- Data-driven augmentation of first-principles models under constraint-free well-posedness and stability guarantees
- Interpretable Bayesian Tensor Network Kernel Machines with Automatic Rank and Feature Selection
- An Introduction to Convolutional Neural Networks
- tn4ml: Tensor Network Training and Customization for Machine Learning
- Laplace Approximation For Tensor Train Kernel Machines In System Identification
- Laplace Approximation for Bayesian Tensor Network Kernel Machines
- Tensor Ring Decomposition
Related papers
- Scen-Opt: A Scenario Optimization Toolbox for Data-Driven Convex Programming
- The Pauli Lightcone: Information-Theoretic Error Mitigation Beyond the Autocorrelation
- FP8 is All You Need (Part 2): Full-FP64 3-D FFT on FP8-Generation Tensor CoresThe Integer-Epilogue Wall and the Minimal Hardware That Would Remove It
- GRADSOLVE: fast exact gradients for ODE ensembles on GPUs
- Learning to Optimize by Differentiable Programming