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

• URL: http://arxiv.org/abs/2608.07043v1
• Date: Fri, 07 Aug 2026 09:52:05 GMT
• Title: Tensor Network Kernel Machines: A JAX Framework for Machine Learning and Nonlinear System Identification
• Abstract summary: "tnkm" is an open-source Python library for constructing and training network kernel machines.<n>We show that "tnkm" achieves competitive prediction accuracy while retaining compact parameterizations and efficient training.
• Score: 6.101839518775971
• License: http://creativecommons.org/licenses/by/4.0/
• 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.

Related papers

• Implementation of tangent linear and adjoint models for neural networks based on a compiler library tool [14.856725081014424]
  This paper presents TorchNWP, a compilation library tool for the efficient coupling of artificial intelligence components and numerical models.<n>It aims to address the issues of poor cross-language compatibility, insufficient coupling flexibility, and low data transfer efficiency.<n>It can be efficiently integrated into numerical weather prediction models such as CMA-GFS and MCV.
  arXiv  Detail & Related papers  (2026-03-17T13:11:22Z)
• Elastic ViTs from Pretrained Models without Retraining [74.5386166956142]
  Vision foundation models achieve remarkable performance but are only available in a limited set of pre-determined sizes.<n>We introduce SnapViT: Single-shot network approximation for pruned Vision Transformers.<n>Our approach efficiently combines gradient information with cross-network structure correlations, approximated via an evolutionary algorithm.
  arXiv  Detail & Related papers  (2025-10-20T16:15:03Z)
• Deep Hierarchical Learning with Nested Subspace Networks [53.71337604556311]
  We propose Nested Subspace Networks (NSNs) for large neural networks.<n>NSNs enable a single model to be dynamically and granularly adjusted across a continuous spectrum of compute budgets.<n>We show that NSNs can be surgically applied to pre-trained LLMs and unlock a smooth and predictable compute-performance frontier.
  arXiv  Detail & Related papers  (2025-09-22T15:13:14Z)
• A projection-based framework for gradient-free and parallel learning [50.96641619247761]
  We introduce PJAX, a JAX-based software framework that enables this paradigm.<n>PJAX composes projection operators for elementary operations, automatically deriving the solution operators for the feasibility problems.<n>We train diverse architectures (MLPs, CNNs, RNNs) on standard benchmarks using PJAX, demonstrating its generality.
  arXiv  Detail & Related papers  (2025-06-06T08:44:56Z)
• Manifold meta-learning for reduced-complexity neural system identification [1.0276024900942875]
  We propose a meta-learning framework that discovers a low-dimensional manifold.<n>This manifold is learned from a meta-dataset of input-output sequences generated by a class of related dynamical systems.<n>Unlike bilevel meta-learning approaches, our method employs an auxiliary neural network to map datasets directly onto the learned manifold.
  arXiv  Detail & Related papers  (2025-04-16T06:49:56Z)
• Superpose Task-specific Features for Model Merging [33.761816717454025]
  We introduce a novel perspective on model merging by leveraging the fundamental mechanisms of neural network representation.<n>We propose a method that superposes task-specific features from individual models into a merged model.<n>Our approach specifically targets linear transformation matrices, which are crucial for feature activation and extraction in deep networks.
  arXiv  Detail & Related papers  (2025-02-15T07:05:55Z)
• Data Summarization via Bilevel Optimization [48.89977988203108]
  A simple yet powerful approach is to operate on small subsets of data. In this work, we propose a generic coreset framework that formulates the coreset selection as a cardinality-constrained bilevel optimization problem.
  arXiv  Detail & Related papers  (2021-09-26T09:08:38Z)
• Neural Network Compression Framework for fast model inference [59.65531492759006]
  We present a new framework for neural networks compression with fine-tuning, which we called Neural Network Compression Framework (NNCF) It leverages recent advances of various network compression methods and implements some of them, such as sparsity, quantization, and binarization. The framework can be used within the training samples, which are supplied with it, or as a standalone package that can be seamlessly integrated into the existing training code.
  arXiv  Detail & Related papers  (2020-02-20T11:24:01Z)
• Model Fusion via Optimal Transport [64.13185244219353]
  We present a layer-wise model fusion algorithm for neural networks. We show that this can successfully yield "one-shot" knowledge transfer between neural networks trained on heterogeneous non-i.i.d. data.
  arXiv  Detail & Related papers  (2019-10-12T22:07:15Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.