Kernel solver scales efficiently with new memory-saving preconditioner
BLOKS: Scaling Kernel Solvers with Sketched Block Preconditioning
Machine Learning
Summary
Kernel Ridge Regression is a powerful but usually slow method for learning from large data. To make it faster, people use tricks like iterative solving, sketching to reduce problem size, and preconditioning to speed up calculations. The authors introduce BLOKS, a way to combine these tricks with a new preconditioner that uses less memory and works well on GPUs. This allows solving big problems faster without losing accuracy.
What this means in practice
- •For machine learning engineers: Train kernel-based predictive models on large datasets with lower memory usage and faster convergence using BLOKS.
- •For data platform operators: Implement scalable approximate kernel solvers in data processing pipelines to handle increasing data sizes efficiently.
Authors
Giordano Vitale, Giacomo Meanti, Marco Rando, Blanca Ayuso de Dios, Lorenzo Rosasco
Abstract
Kernel Ridge Regression (KRR) is a theoretically sound algorithm for nonparametric learning, which is also often used to investigate properties of deep learning. However, the exact KRR formulation incurs time and memory costs preventing applicability to large datasets, motivating the development of approximate solvers. These algorithms rely on (i) iterative methods, to avoid the computational burden of direct solutions; (ii) sketching or batching, to reduce the problem size; and (iii) preconditioning, to accelerate convergence. While crucial for fast convergence, preconditioning incurs memory costs that can be managed by lowering the number of free parameters of the model, potentially at the expense of expressivity and hence accuracy. To address this challenge, we propose BLOKS, a Conjugate Gradient-based solver for Nyström-sketched KRR built around a novel preconditioner that combines a block Jacobi approximation with sketching. The BLOKS preconditioner has memory cost independent of both the dataset size and the number of parameters, allowing scaling to large models. We provide a memory-efficient, GPU-accelerated implementation and show that BLOKS achieves excellent accuracy while retaining favorable time/memory complexity. Indeed, on standard large-scale datasets used to benchmark kernel solvers, BLOKS achieves comparable or better performances than state-of-the-art methods.