BLOKS: Breaking the Memory Bottleneck in Kernel Methods with Sketched Block Preconditioning

The Scalability Problem with Kernel Methods

Kernel Ridge Regression (KRR) is a theoretically well-grounded algorithm in nonparametric learning, widely used to study the intrinsic properties of deep learning. Its mathematical formulation is elegant and interpretable — but at a steep cost: exact solutions require time and memory that scale sharply with dataset size, making it impractical for large-scale applications.

This tension has driven extensive research into approximate solvers: how do we preserve KRR's theoretical advantages while making it viable on real-world large datasets? That is precisely the core problem targeted by BLOKS, a recent research contribution.

Three Pillars of Approximate Solvers

To make KRR scalable, existing methods typically combine three approaches:

• Iterative methods: Avoid the heavy computational burden of direct solvers by progressively approximating the solution instead of computing it exactly;
• Sketching / batching: Reduce problem size through dimensionality reduction or block decomposition;
• Preconditioning: Accelerate convergence of the iterative process.

Preconditioning is critical to convergence speed — but it also creates its own headache. Preconditioner matrices tend to consume large amounts of memory. The common workaround is to reduce the number of model parameters to control memory usage, but this sacrifices model expressiveness and ultimately hurts accuracy.

A Hard-to-Reconcile Trade-off

In short, researchers have long faced a vicious cycle: fast convergence requires preconditioning, preconditioning costs memory, saving memory means cutting parameters, and cutting parameters degrades accuracy. The value of BLOKS lies precisely in its attempt to escape this loop.

Nyström sketching is one of the most widely used dimensionality reduction techniques in kernel methods. The core idea is to randomly select $m \ll n$ "landmark points" from $n$ training samples and use the low-rank matrix generated by these landmarks to approximate the full $n \times n$ kernel matrix. This reduces the primary storage and computation bottleneck from $O(n^2)$ to $O(nm)$, making kernel methods feasible on large datasets. The quality of the Nyström approximation depends on $m$: the larger $m$ is, the more accurate the approximation — but memory and compute costs grow accordingly. BLOKS operates within the Nyström framework, and its preconditioner design must decouple memory usage without disrupting this low-rank structure.

The Core Idea Behind BLOKS

BLOKS is a Conjugate Gradient (CG)-based solver for Nyström-sketched KRR. Its innovation centers on a novel preconditioner: combining block Jacobi approximation with sketching techniques.

The key property of this design is that the preconditioner's memory cost is simultaneously independent of both the dataset size and the number of parameters. This means that regardless of how large the data is or how many parameters the model has, the memory footprint of the preconditioner does not grow. For scenarios that require scaling to large models, this is a structural breakthrough — it directly untangles the memory–accuracy trade-off described earlier.

Why Block Jacobi Plus Sketching Is the Key

Block Jacobi approximation is essentially a block-wise simplification of the full preconditioner matrix, retaining local information critical to convergence. Sketching further compresses the dimensions that need to be processed. Together, they maintain the acceleration effect that preconditioning provides while keeping memory costs decoupled from problem scale. This decoupling is the source of BLOKS's scalability.

Conjugate Gradient (CG) is a classical iterative algorithm for solving symmetric positive definite linear systems $Ax = b$. For kernel ridge regression, exact solution requires directly inverting an $n \times n$ kernel matrix at $O(n^3)$ time complexity. CG instead iteratively approximates the solution by constructing mutually "conjugate" search directions, requiring only one matrix-vector product per step. The number of iterations needed for convergence is closely tied to the matrix's condition number — the larger the condition number, the slower the convergence. This is where preconditioners come in: they apply an equivalent transformation to the original equation that drastically reduces the condition number of the transformed matrix, allowing CG to reach the same accuracy in far fewer iterations.

Block Jacobi preconditioning is a generalization of classic Jacobi iteration: Jacobi methods construct a preconditioner from the diagonal elements of the original matrix, while block Jacobi partitions the matrix into diagonal blocks and inverts each small block to form the preconditioner. This preserves local coupling information within each block, capturing more of the matrix's structure than a pure diagonal preconditioner, and typically delivers faster convergence. The trade-off is that each diagonal block can itself be large, with non-negligible storage and inversion costs. BLOKS's innovation is to apply sketch compression again to each diagonal block, capping the effective size of each block at a fixed upper bound. This decouples the total memory footprint of the preconditioner from both the data size $n$ and the parameter count $m$, making it depend only on a pre-specified block size hyperparameter.

Implementation and Performance

The authors provide a memory-efficient, GPU-accelerated implementation. This matters on a practical level — many theoretically elegant algorithms fail to gain traction due to lack of efficient implementations, and BLOKS completes the engineering dimension alongside its algorithmic contributions.

On standard large-scale datasets used for benchmarking kernel solvers, BLOKS achieves performance on par with or better than current state-of-the-art methods, while maintaining favorable time and memory complexity. In other words, it does not trade speed or memory for accuracy — it achieves a good balance across all three dimensions simultaneously.

Implications for Research and Applications

For researchers who use kernel methods as analytical tools for studying deep learning, BLOKS makes large-scale experiments more feasible. For practical modeling tasks, it lowers the barrier to applying KRR on large datasets. Kernel methods have long been squeezed out of mainstream large-scale applications due to scalability concerns; work like this is steadily closing the gap in engineering usability between kernel methods and modern deep learning approaches.

Summary

BLOKS's contribution can be summarized in a single sentence: by introducing a novel preconditioner whose memory cost is decoupled from problem scale, it advances the scalability of kernel ridge regression. It neither sidesteps the inherent computational challenges of kernel methods nor hides the problem by simply cutting parameters — instead, it addresses the critical bottleneck at the preconditioning stage, delivering a solution that balances accuracy, speed, and memory.

Note that this article is based on the paper's abstract; specific experimental results, comparison baselines, and hyperparameter details should be referenced from the full paper.