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

2026-06-04

AI Summary

This work addresses the computational inefficiency of traditional kernel methods when applied to large-scale datasets, particularly their limited compatibility with GPU acceleration and high costs in training and hyperparameter selection. To overcome these limitations, we introduce TorchKM, an open-source library built on PyTorch that unifies the full pipeline of kernel methods—including support vector machines, kernel logistic regression, and kernel quantile regression—within a GPU-accelerated framework for the first time. By leveraging GPU-friendly matrix operations and offering a scikit-learn–style interface, TorchKM achieves substantial computational speedups while preserving predictive accuracy. Empirical evaluations demonstrate significant acceleration over standard baseline implementations. The library is publicly available and supports one-click installation via PyPI.

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 together with substantial speedups over standard baselines. Code and documentation are available at https://github.com/YikaiZhang95/torchkm, and the package can be easily installed via PyPI.

Problem

• kernel learning
• model selection
• GPU acceleration
• support vector machines
• kernel methods

Innovation

• GPU acceleration
• kernel machines
• model selection
• matrix operation reuse
• scikit-learn API