⚡ electrical engineering

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

This paper introduces "tnkm," an open-source JAX-based Python library that enables the construction and training of expressive yet computationally efficient tensor network kernel machines for nonlinear system identification through a unified interface supporting diverse feature maps, architectures, and optimization strategies.

Original authors: Albert Saiapin, Kim Batselier

Published 2026-08-10

📖 3 min read☕ Coffee break read

CC BY 4.0

Original authors: Albert Saiapin, Kim Batselier

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a computer to predict the future, whether it's guessing the next note in a song, forecasting the weather, or controlling a robot arm. To do this, the computer needs to understand how different inputs (like wind speed or a joystick movement) turn into outputs (like a cloud forming or a wheel spinning). The tricky part is that the real world is rarely a straight line; it's full of curves, loops, and sudden jumps. This is the world of "nonlinear systems."

For decades, scientists have tried to build models that can handle these curves. Some models are like giant, messy libraries that can remember everything but take forever to find a book (they are accurate but slow). Others are like tiny, fast flashcards that are quick to read but might miss the big picture (they are fast but not very smart). The big challenge in machine learning is finding a way to have the best of both worlds: a model that is smart enough to understand complex curves but small and fast enough to run on a regular laptop without melting it. This is where a clever idea called "tensor networks" comes in. Think of a tensor network as a way to fold a massive, unwieldy map into a tiny, compact origami crane that still holds all the important directions.

In this paper, the authors, Albert Saiapin and Kim Batselier, introduce a new open-source software tool called tnkm (Tensor Network Kernel Machines). They built this tool to make it much easier for anyone to build and train these "origami" models using a powerful programming language called JAX. Before this, building these models was like trying to assemble a complex piece of furniture without a manual or the right tools; it was possible, but frustrating and limited to experts. The authors created a "Swiss Army knife" for these models, allowing users to mix and match different types of "feature maps" (the way the computer sees the data) with different "tensor networks" (the way the computer stores the rules) and "optimization methods" (the way the computer learns).

The paper doesn't just claim this tool works; they put it to the test. They trained these models on several famous "benchmark" problems, which are like standardized math puzzles used to see how good a computer is at learning. In one test involving the noise of an airplane wing, they found that their method could learn the pattern just as well as the heavy, slow methods, but it finished the job in about 1 second instead of taking minutes. In other tests involving fluid tanks and electric motors, the new tool matched the accuracy of complex black-box models but did so with a fraction of the computing power. The authors show that by using these compact, folded mathematical structures, you can get high-accuracy predictions without needing a supercomputer. They also compared two ways of teaching the model: one that solves the math step-by-step (called Alternating Least Squares) and one that uses a more general, trial-and-error approach (gradient-based optimization). They found that the step-by-step method was often much faster and more stable for these specific types of problems, though the trial-and-error method offers more flexibility for different kinds of tasks.

Ultimately, the paper suggests that this new framework makes it possible to build powerful, efficient models for things like robotics and process control without getting bogged down in computational complexity. It's a step toward making advanced machine learning accessible, reproducible, and practical for solving real-world problems that involve messy, nonlinear data.

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

Problem Statement

Developing nonlinear models that balance high expressive power with computational efficiency remains a significant challenge in machine learning and nonlinear system identification. Traditional approaches, such as kernel methods, Gaussian processes, and deep neural networks, often face scalability issues where model complexity or computational costs grow prohibitively as the number of parameters or training samples increases. Specifically, representing nonlinear mappings via explicit feature maps in high-dimensional spaces typically leads to an exponential explosion in the number of parameters, rendering direct optimization infeasible for high-dimensional problems.

Methodology

The paper introduces Tensor Network Kernel Machines (TNKM), a framework that addresses scalability by combining explicit nonlinear feature representations with compact, low-rank tensor-network parameterizations. The core mathematical formulation represents a nonlinear function $f(x)$ as a linear model in a high-dimensional feature space:
$f(x) = \phi(x)^\top w$
where $\phi(x)$ is a nonlinear feature map and $w$ is the parameter vector. To overcome the curse of dimensionality inherent in $w$ , the TNKM approach employs two key strategies:

Tensor-Product Feature Maps: The feature map $\phi(x)$ is constructed as a tensor product of one-dimensional maps acting on individual input variables, naturally arising from product kernels.
Low-Rank Tensor Decomposition: Instead of storing the full parameter vector $w$ , the parameters are represented using low-rank tensor decompositions, specifically Canonical Polyadic (CP) and Tensor-Train (TT) formats. This reduces the parameter complexity from exponential to linear (or near-linear) with respect to the input dimension.

The authors developed tnkm, an open-source Python library built on JAX, to implement this framework. The library adopts a modular architecture separating three primary components:

Feature Maps: Supports polynomial, Fourier, and Volterra (memory-based) features.
Tensor-Network Parameterizations: Implements CP and TT decompositions via dedicated model classes (e.g., CPKM, TTKM).
Optimization Algorithms: Provides two complementary strategies:
- Alternating Least Squares (ALS): Exploits the multilinear structure for fast, deterministic convergence on least-squares problems.
- Gradient-Based Optimization: Utilizes JAX's automatic differentiation and the Optax library (e.g., Adam, SGD) for flexibility with arbitrary differentiable loss functions.

Key Contributions

The paper outlines four main contributions:

Unified JAX Framework: A minimal-dependency, open-source library (tnkm) that unifies different feature maps, tensor-network backbones (CP and TT), and optimization strategies under a single interface.
Scalable Training Implementations: Integration of structure-aware ALS optimization alongside general gradient-based methods, enabling application to both regression and broader learning tasks.
Modular Software Architecture: A design allowing the systematic swapping of components (feature maps, model classes, optimizers) and providing utilities for data preprocessing, simulation, and evaluation specific to nonlinear system identification.
Comprehensive Experimental Evaluation: A rigorous benchmarking of the framework on nonlinear regression and system identification tasks, demonstrating competitive accuracy with efficient training times.

Experimental Results

The authors evaluated tnkm on two fronts:

1. Optimization Method Comparison (Airfoil Self-Noise Dataset):

ALS vs. Adam: ALS converged rapidly to a validation Mean Squared Error (MSE) of 0.102, closely matching the Kernel Ridge Regression (KRR) baseline (0.15), in just 1.16 seconds. In contrast, Adam required significantly more time (3.25s for mini-batch) and epochs to reach a higher error (0.145).
Conclusion: ALS offers superior efficiency for least-squares TNKM problems, while gradient-based methods provide necessary flexibility for non-least-squares objectives.

2. Nonlinear System Identification Benchmarks:
The framework was tested on three standard benchmarks: Coupled Electric Drives, Cascaded Tanks, and Silverbox.

Coupled Electric Drives: TNKM achieved an RMSE comparable to the best-performing methods (GPNARX and dynoNet) but with a training time of < 0.1 seconds, whereas competitors required seconds to minutes.
Cascaded Tanks: TNKM provided competitive accuracy (RMSE 0.347) among black-box approaches, significantly outperforming methods like SUBNET and PNARX in terms of training speed (< 0.1s vs. 1-60 min).
Silverbox: TNKM achieved performance comparable to leading black-box methods (GPNARX, SUBNET) with training times of 1–60 seconds. While it did not match the accuracy of physics-informed approaches (NL-LFR), it demonstrated effective long-horizon prediction stability.

Significance and Claims

The paper claims that tnkm facilitates the reproducible development and application of tensor-network-based learning methods. By decoupling feature maps, tensor parameterizations, and optimization algorithms, the library enables systematic exploration of TNKM configurations without altering the underlying workflow.

The authors position TNKM as an intermediate approach between classical kernel methods and deep neural networks. Like kernel methods, TNKM uses explicit feature representations, offering greater control and the ability to incorporate prior knowledge, but it replaces the full parameter vector with a compact tensor network to ensure scalability. Unlike neural networks, it does not learn features directly from data, which the authors acknowledge as a limitation requiring manual feature selection.

The significance of the work lies in providing a practical, extensible software foundation that bridges the gap between theoretical tensor network methods and real-world nonlinear system identification, demonstrating that these models can achieve high predictive accuracy while maintaining low computational costs.