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

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

Problem Statement
Kernel machines, including Support Vector Machines (SVMs), kernel logistic regression, and kernel quantile regression, are fundamental algorithms known for strong predictive performance and rigorous theoretical foundations. However, their adoption in modern practice is often limited by high computational costs, particularly during the model selection phase. Traditional workflows treat training and model selection (e.g., hyperparameter tuning via cross-validation) as separate stages. This necessitates repeated model fitting across a grid of tuning parameters and cross-validation folds. For kernel methods, which involve solving linear systems with $O(n^3)$ complexity, this repeated refitting creates a prohibitive computational bottleneck. While hardware acceleration via Graphics Processing Units (GPUs) has been explored (e.g., ThunderSVM), simply porting existing CPU solvers to GPUs often yields only modest improvements because the repeated matrix operations across tuning parameters and folds still dominate the computational cost.

Methodology
TorchKM addresses these limitations through an "algorithm–hardware co-design" approach. Rather than directly porting a CPU solver to the GPU, the library is built from the ground up to exploit GPU-friendly linear algebra. The core methodology relies on two key algorithmic innovations that allow for the exact solution of the training and tuning pipeline without approximation:

1. Exact Cross-Validation Formula: To avoid the computational expense of retraining models for each cross-validation fold, TorchKM utilizes a reformulation where the kernel matrix remains unchanged. Instead of refitting, the method modifies the response vector for each fold. This allows the reuse of the same kernel matrix and its associated decompositions, significantly reducing the overhead of cross-validation.
2. Spectral Algorithm: To handle the tuning of regularization parameters (the solution path), TorchKM employs a spectral approach. It computes a single eigendecomposition of the kernel matrix ($K = U\Lambda U^\top$). Subsequent updates for different tuning parameters are performed via matrix-vector multiplications ($O(n^2)$) rather than repeated matrix inversions ($O(n^3)$). This structure is inherently parallelizable and well-suited for GPU architecture.

These algorithms are integrated into a unified pipeline that supports a "scikit-learn-style" API. The library supports various methods including SVM, Distance-Weighted Discrimination (DWD), kernel logistic regression, and kernel quantile regression. It also offers a Nyström approximation for large-scale datasets where exact solutions are infeasible, and includes Platt scaling for probability calibration.

Key Contributions

• Integrated Pipeline: Unlike libraries that separate training and tuning, TorchKM integrates both processes into a single algorithmic step, eliminating the need for repeated external refitting.
• Exact Solutions: The library provides exact solutions (not approximations) for the full training-and-tuning pipeline, distinguishing it from methods that rely on heuristic approximations for speed.
• Broad Method Support: Beyond SVMs, the library implements kernel logistic regression, kernel quantile regression, and DWD, all sharing a consistent API.
• Open-Source Implementation: The library is released under the MIT license, available via PyPI and GitHub, and includes comprehensive documentation, tutorials, and benchmarks.

Results
Benchmarks presented in the paper demonstrate that TorchKM achieves competitive predictive performance with substantial speedups over standard baselines:

• Comparison with scikit-learn: TorchKM significantly outperforms CPU-based implementations. In a scenario with $n=20,000$ samples and $p=1,000$ features, scikit-learn failed to complete the training-and-tuning task within an 8-hour limit, whereas TorchKM completed the task in 129.3 seconds.
• Comparison with ThunderSVM: While ThunderSVM accelerates single-model fitting on GPUs, TorchKM further improves performance by accelerating the entire tuning pipeline. In the same $n=20,000, p=1,000$ setting, TorchKM (129.3s) was faster than ThunderSVM (580.8s).
• Scalability: The paper notes that for $n=10,000$, TorchKM is more than two orders of magnitude faster than standard CPU implementations and nearly two orders of magnitude faster than standard GPU implementations that do not utilize the integrated spectral and cross-validation strategies.
• Calibration: The library produces well-calibrated probability estimates, with an Expected Calibration Error (ECE) of 0.044 and a Brier score of 0.130 in tested scenarios.

Significance
The paper posits that the primary barrier to the widespread use of kernel machines is the computational cost of hyperparameter tuning. By redesigning the algorithm to shift the cubic-cost burden into matrix-vector operations that leverage GPU parallelism, TorchKM makes the full training-and-tuning pipeline feasible for large-scale structured data. The authors argue that this approach allows kernel methods to achieve their full predictive potential by enabling thorough model selection that was previously computationally prohibitive. The library represents a step toward making exact kernel learning practical for modern datasets through the synergy of algorithmic innovation and hardware acceleration.