TorchKM: A GPU-Oriented Library for Kernel Learning and Model Selection

Yikai Zhang; Gaoxiang Jia; Jie Ding; Boxiang Wang

cs.LGJun 4, 2026

TorchKM: A GPU-Oriented Library for Kernel Learning and Model Selection

Authors: Yikai Zhang, Gaoxiang Jia, Jie Ding, Boxiang Wang

Organizations: Department of Statistics and Actuarial Science, University of Iowa, Iowa City, IA 52242, USA · Individual Researcher · AIScientists, Inc. (MorphMind), 245 Main Street, Cambridge, MA 02142, USA

Abstract

TorchKM is an open-source library for kernel machines, including support vector machines, kernel logistic regression, and kernel quantile regression, with GPU acceleration. The library features a scikit-learn-style API and is designed to exploit GPU-friendly linear algebra, accelerating the full training and model-selection pipeline through intelligent reuse of matrix operations. Benchmarks show competitive predictive performance with substantial speedups over standard baselines. The efficiency and programmable design also make TorchKM a kernel-learning component for AI-driven workflows. Code and documentation are available at https://github.com/YikaiZhang95/torchkm, and the package can be easily installed via PyPI.

Explore similar work

Tensor Network Kernel Machines: A JAX Framework for Machine Learning and Nonlinear System Identification
Aug 7, 2026Albert Saiapin, Kim BatselierTensor Networks Neural Tangent Kernel
AMDKernelVault: Large-Scale Datasets and Agentic Training for AMD GPU Kernel Optimization
Sep 14, 2026Ji Liu, Saptarshi Majumder, Yiqing Huang +14Graphics Processing Unit Kernels Pytorch
FastKernels: Benchmarking GPU Kernel Generation in Production
May 22, 2026Gabriele Oliaro, Yichao Fu, May Jiang +5Graphics Processing Unit Kernels

Aug 7, 2026 · cs.MSJ/K move · Enter open · S save

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

Albert Saiapin, Kim Batselier

Delft Center for Systems and Control, TU Delft, Netherlands

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.

Tensor Networks Neural Tangent Kernel Machine Learning System Identification

Open paper Ask about it PDF

Abstract

Explore similar work

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

AMDKernelVault: Large-Scale Datasets and Agentic Training for AMD GPU Kernel Optimization

FastKernels: Benchmarking GPU Kernel Generation in Production