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

2026-08-07

AI Summary

This work addresses the challenge of balancing expressive power and computational efficiency in nonlinear models by proposing and open-sourcing “tnkm,” a JAX-based Python library that unifies nonlinear feature mappings with low-rank tensor network architectures. The framework offers the first scalable and modular implementation of tensor network kernel machines, enabling flexible composition of feature maps, network topologies, and optimization strategies—including alternating least squares and gradient-based methods. Experimental results demonstrate that the approach achieves competitive predictive accuracy on multiple nonlinear benchmark tasks while substantially reducing model parameter count and improving training efficiency. By providing a reproducible and high-performance platform, this work advances research in efficient nonlinear modeling.

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.

Problem

• nonlinear system identification
• tensor network
• kernel machines
• machine learning
• software framework

Innovation

• Tensor Network Kernel Machines
• JAX framework
• nonlinear system identification
• low-rank parameterization
• alternating least squares