Skip to content
Preprint

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

Aug 2026 · 0 citations · 51 references
Computer Science Engineering

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.

View source

Similar papers

#artificial intelligence Preprint Aug 2026

Iterative tensor network transformations for element-wise evaluation of elementary and filtering functions

Tensor networks are powerful formats for compressing large-scale data. However, their application to general data processing has been limited by the difficulty of performing nonlinear operations. Here, we introduce iterative tensor network transformations (ITNTs), a general algorithmic framework for the element-wise evaluation of elementary and nonlinear filtering functions on data encoded as tensor trains (TTs), a class of tensor networks. Our approach operates entirely in the compressed domain, enabling efficient computation on exponentially large datasets while maintaining a controlled computational cost. We demonstrate its power in two key areas: (I) evaluating highly nonlinear elementary and filtering functions on a 3D reactive flow field, enabling high-fidelity reaction rate computation and region filtering, and (II) finding extrema in complex optimization problems, such as solving Max-SAT instances on spaces up to $2^{70}$ configurations. These results establish ITNT as a foundational tool that provides tensor network methods with the capability for general-purpose data science and large-scale optimization.

Xiao Wang, Tomohiro Hashizume, Pia Siegl et al. · 2 citations
Preprint Aug 2026

Accelerated Learning of High Dimensional Functions with a Tensor-Featured Training Network

In this work we present a method to accelerate the optimization of learning high dimensional functions using deep neural network (DNN). This optimization procedure introduces contextual features into the first layer of a DNN. The parameters of DNN are optimized via standard gradient descent while keeping the input-feature basis fixed. After optimization of the DNN parameters, the feature layer is provided a chance to update and change before DNN optimization resumes. The feature layer has two types of functions: those that can be evaluated quickly in a matrix-free way on the domain (i.e. rank-1 features) and more complex features that must first be decomposed using tensor network (TN) decomposition strategies (tensor features). In particular, we study the effect of adding features which distill pretrained DNN into TNs using a discretize and decompose strategy. To efficiently decompose high-dimensional functions constructed from discretized DNN, we leverage a randomized tensor decomposition strategy. Using randomization, we are able to reduce the storage cost of decomposing high dimensional functions by at least 8 orders of magnitude. Using this approach, we are able to efficiently train models between 5 and 40 dimensions.

Karl Pierce, Y. Khoo, Haizhao Yang · 0 citations
Preprint Aug 2026

Kernel Methods for Learning Operators with Multiple Inputs and Outputs

Learning mappings between infinite-dimensional objects is a central challenge in scientific machine learning. We introduce a general kernel-based encoder-decoder framework for operator learning that separates observation, representation, learning, and reconstruction. We develop this framework for multi-input, multi-output operator learning, where operators map between products of potentially distinct function spaces. Our approximation theory shows that, although the number of inputs and outputs can increase, the convergence rate is governed by the most challenging constituent approximation problem rather than the overall problem dimension. The framework leads to practical kernel methods with closed-form training and inference, combining mathematical tractability with computational efficiency. We further specialize the approach to multiple operator learning by introducing KernelMO, a family of kernel methods with complementary operator-valued and product-space formulations. Across five families of parametric partial differential equations, the proposed methods achieve competitive or state-of-the-art predictive accuracy while reducing training and inference costs relative to neural operator architectures and deep learning based models, offering an efficient and lightweight alternative.

Adrien Weihs, Chunyang Liao, Jingmin Sun et al. · 0 citations
Conference Jul 2026

DATT: Dimension-Augmented Tensor-Train Decomposition for Neural Network Compression

Traditional model reduction techniques for neural network compression usually leverage existing matrix/tensor decomposition methods to explore the redundancy among model parameters, which may overlook the local structures of the models. In this work, we propose a novel low-rank approximation algorithm, called Dimension-Augmented Tensor-Train (DATT) decomposition, which integrates dimension slicing with the Tensor-Train decomposition. DATT allows different partitions to be represented using distinct, non-shared spanning sets. This enables the decomposition to adapt more flexibly to localized structure, potentially achieving better compression. With the localized low-rank structures, DATT provides a more compact representation of neural network weights. We evaluated the performance of DATT with the other state-of-the-art low-rank approximation-based model compression algorithms, and demonstrated that DATT outperforms them in terms of parameter reduction and model accuracy. In particular, for compressing ResNet-20 on CIFAR-10, DATT achieves a 77.00% reduction in parameters with only a 0.37% drop in top-1 accuracy.

Yu-Chuan Tai, Cheng-Yu Sie, Che-Rung Lee · 0 citations
Book Open access Aug 2026

Interplay Between Classical Tensor Methods And Foundation Models

Tensor methods have played a central role in analyzing multi-dimensional data across a wide range of real-world applications. At the same time, these methods provide low-rank representations, which enable trustworthy and parsimonious machine learning systems. However, even though they hold great promise, tensor methods remain relatively underexplored in the context of modern neural architectures and foundation models. This workshop aims to provide a forum for advancing tensor methods and applying them to various applications at the intersection of data mining and modern machine learning. Supported by organizers and keynote speakers with broad expertise across machine learning, signal processing, and data mining, the workshop aims to foster an interactive environment for researchers to exchange ideas and build connections across communities.

Dawon Ahn, Yong-chan Park, Taehyung Kwon et al. · 0 citations
Jul 2026

Variational Inference and Density Estimation with Non-negative Tensor Train

Abstract. This work proposes an efficient numerical approach for compressing a high-dimensional discrete distribution function into a non-negative tensor train (NTT) format. The two settings we consider are variational inference and density estimation, whereby one has access to either the unnormalized analytic formula of the distribution or the samples generated from the distribution. In particular, the compression is done through a two-stage approach. In the first stage, we use existing subroutines to encode the distribution function in a tensor train format. In the second stage, we use an NTT ansatz to fit the obtained tensor train. For the NTT fitting procedure, we use a log barrier term to ensure the positivity of each tensor component and then utilize a second-order alternating minimization scheme to accelerate convergence. In practice, we observe that the proposed NTT fitting procedure exhibits drastically faster convergence than an alternative multiplicative update method that has been previously proposed. Through challenging numerical experiments, we show that our approach can accurately compress target distribution functions. Reproducibility of computational results. This paper has been awarded the “SIAM Reproducibility Badge: Code and Data Available” as a recognition that the authors have followed reproducibility principles valued by SISC and the scientific computing community. Code and data that allow readers to reproduce the results in this paper are available at https://github.com/Xun-Tang123/Non_Negative_Tensor_Train and in the supplementary materials ( Non_Negative_Tensor_Train-main.zip [51.4MB]), linked from the main article webpage. [Formula: see text]

Xun Tang, R. Dwaraknath, Lexing Ying · 0 citations