Track trials, findings, and clinical evidence as it lands.
Search papers, labs, and topics across Lattice.
We introduce BLOKS, a Conjugate Gradient solver for Nyström-sketched Kernel Ridge Regression that utilizes a novel sketched block Jacobi preconditioner. This preconditioner achieves memory costs independent of dataset and parameter size, enabling scalability to large models. BLOKS demonstrates comparable or superior accuracy to state-of-the-art methods on large-scale kernel solver benchmarks, with favorable time and memory complexity.
Kernel solvers can now scale to massive models without sacrificing accuracy, thanks to a memory-independent preconditioner that cleverly combines block Jacobi and sketching.
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\"om-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.