# Tensor Network Kernel Machines: A JAX Framework for Machine Learning and Nonlinear System Identification

arXiv: https://arxiv.org/abs/2608.07043 · HTML: https://mlsanity.com/paper/2608.07043/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 24 of its 51 references that are in the map.
- First posted on arXiv 7 Aug 2026.

## Builds on (most cited first, 24 of 24)

- [Adam: A Method for Stochastic Optimization](https://mlsanity.com/paper/1412.6980/) (2014)
- [Generating Sequences With Recurrent Neural Networks](https://mlsanity.com/paper/1308.0850/) (2013)
- [An Introduction to Convolutional Neural Networks](https://mlsanity.com/paper/1511.08458/) (2015)
- [Hyperparameter Optimization: Foundations, Algorithms, Best Practices and Open Challenges](https://mlsanity.com/paper/2107.05847/) (2021)
- [Tensor Ring Decomposition](https://mlsanity.com/paper/1606.05535/) (2016)
- [TensorLy: Tensor Learning in Python](https://mlsanity.com/paper/1610.09555/) (2016)
- [Hilbert Space Methods for Reduced-Rank Gaussian Process Regression](https://mlsanity.com/paper/1401.5508/) (2014)
- [Variational Fourier features for Gaussian processes](https://mlsanity.com/paper/1611.06740/) (2016)
- [Nonlinear System Identification: A User-Oriented Roadmap](https://mlsanity.com/paper/1902.00683/) (2019)
- [Tensor Networks for Dimensionality Reduction and Large-Scale Optimizations. Part 2 Applications and Future Perspectives](https://mlsanity.com/paper/1708.09165/) (2017)
- [Nonlinear state-space identification using deep encoder networks](https://mlsanity.com/paper/2012.07697/) (2020)
- [dynoNet: a neural network architecture for learning dynamical systems](https://mlsanity.com/paper/2006.02250/) (2020)
- [Parallelized Tensor Train Learning of Polynomial Classifiers](https://mlsanity.com/paper/1612.06505/) (2016)
- [Large-Scale Learning with Fourier Features and Tensor Decompositions](https://mlsanity.com/paper/2109.01545/) (2021)
- [Continuous-time identification of dynamic state-space models by deep subspace encoding](https://mlsanity.com/paper/2204.09405/) (2022)
- [Quantized Fourier and Polynomial Features for more Expressive Tensor Network Models](https://mlsanity.com/paper/2309.05436/) (2023)
- [tn4ml: Tensor Network Training and Customization for Machine Learning](https://mlsanity.com/paper/2502.13090/) (2025)
- [A Tensor Network Kalman filter with an application in recursive MIMO Volterra system identification](https://mlsanity.com/paper/1610.05434/) (2016)
- [Interpretable Bayesian Tensor Network Kernel Machines with Automatic Rank and Feature Selection](https://mlsanity.com/paper/2507.11136/) (2025)
- [Tensor Network Based Feature Learning Model](https://mlsanity.com/paper/2512.02547/) (2025)
- [Data-driven augmentation of first-principles models under constraint-free well-posedness and stability guarantees](https://mlsanity.com/paper/2604.11421/) (2026)
- [Inference and Learning of Nonlinear LFR State-Space Models](https://mlsanity.com/paper/2503.14409/) (2025)
- [Laplace Approximation For Tensor Train Kernel Machines In System Identification](https://mlsanity.com/paper/2512.02532/) (2025)
- [Laplace Approximation for Bayesian Tensor Network Kernel Machines](https://mlsanity.com/paper/2604.26673/) (2026)

## Abstract (arXiv)

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.