TU Delft Releases JAX Library for Tensor Network Kernel Machines
TU Delft researchers Albert Saiapin and Kim Batselier released tnkm, an open-source JAX library that combines nonlinear feature maps with compact tensor-network parameterizations. Their August 7 paper reports support for alternating least squares and gradient-based training; on one Airfoil benchmark, ALS reached 0.102 validation MSE in 1.161 seconds.
TU Delft researchers Albert Saiapin and Kim Batselier released tnkm, an open-source Python library for building tensor network kernel machines with JAX. The accompanying preprint was submitted on August 7 and the public repository supplies the implementation, documentation and reproducible experiment code.
Compressing a very large feature space
Tensor network kernel machines combine explicit nonlinear feature maps with low-rank representations of the model coefficients. The goal is to retain expressive interactions without storing every coefficient in the full tensor-product feature space.
The paper gives a concrete scale example: with 10 input dimensions, 20 features per dimension and rank 10, the uncompressed coefficient vector would contain roughly 20^10, or 10 trillion, entries. A canonical polyadic representation needs about 2,000 parameters instead. The library supports both canonical polyadic and tensor-train backbones, along with polynomial, Fourier and memory-based Volterra features.
The software separates feature maps, tensor representations and optimizers behind reusable interfaces. It offers alternating least squares for multilinear least-squares problems and JAX/Optax gradient methods for more general differentiable objectives. The repository documents installation from PyPI or source and includes experiment code.
What the reported benchmarks show
Saiapin and Batselier compared alternating least squares with Adam on the Airfoil Self-Noise regression dataset using the same rank-10 model and a 90/10 train-validation split. They report validation mean-squared error of 0.102 after 1.161 seconds for alternating least squares, compared with 0.145 after 3.251 seconds for mini-batch Adam and 0.183 after 1.675 seconds for full-batch Adam. A kernel-ridge reference scored 0.15.
The paper also evaluates the framework on three nonlinear system-identification benchmarks: Coupled Electric Drives, Cascaded Tanks and Silverbox. Those experiments are useful evidence that the package goes beyond a toy API, but they remain author-run results from a new preprint rather than independent validation.
For practitioners, the immediate value is a reproducible JAX implementation for experimenting with compact nonlinear models and explicit rank control. The evidence does not establish that tnkm will beat neural networks or kernel methods generally; performance depends on the feature map, rank, optimizer and problem structure.
Key Points
- 1The August 7 preprint introduces tnkm, an open-source JAX library for tensor network kernel machines.
- 2The package supports canonical polyadic and tensor-train representations, several nonlinear feature maps, and both alternating least squares and gradient optimizers.
- 3In the authors' Airfoil benchmark, alternating least squares reached 0.102 validation MSE in 1.161 seconds under the reported setup.
- 4The benchmark evidence is reproducible from the public repository but remains author-reported and problem-specific.
Scoring Rationale
The release provides a reproducible JAX framework and measured results for a specialized class of compact nonlinear models. Its practical reach is narrower than a general-purpose ML framework, and all reported benchmarks come from the authors' preprint.
Sources
Primary source and supporting public references used for this report.
Practice interview problems based on real data
1,625 SQL & Python problems across 15 industry datasets — the exact type of data you work with.
Try 250 free problems