PaperlayerView exact revision on arXiv

Source-linked AI summary

BLOKS: Scaling Kernel Solvers with Sketched Block Preconditioning

Giordano Vitale, Giacomo Meanti, Marco Rando, Blanca Ayuso de Dios, Lorenzo Rosasco

arXiv:2610.13382v1cs.LGstat.ML

TL;DR

Exact KRR's time and memory costs make it difficult to apply to large datasets. BLOKS combines Conjugate Gradient for Nyström-sketched KRR with a sketched block Jacobi preconditioner. The paper reports higher accuracy than state-of-the-art methods under the same memory budget, with competitive training time.

  • Problem

    Exact KRR requires O(n^3) time and O(n^2) memory, limiting its applicability as datasets grow.

    Source: §1 · p. 1

  • Method

    BLOKS solves Nyström-sketched KRR with Conjugate Gradient and a preconditioner combining block Jacobi approximation with sketching.

    Source: §1 · p. 2

  • Results

    Under the same memory budget, BLOKS achieves higher accuracy than state-of-the-art methods with competitive training time.

    Source: §6 · p. 11

  • Takeaways & Limitations

    BLOKS decouples model size from memory requirements and scales to 3.2M parameters with linear training time.

    Source: §6 · p. 10

    Reducing block size mitigates spatial complexity but causes significant temporal overhead from more frequent data swapping.

    Source: §2 · pp. 3–4

Abstract

from arXiv · show

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.

1 Introduction

Exact KRR is difficult to scale, and existing approximate solvers trade off per-iteration cost, convergence, and preconditioner memory. BLOKS combines Nyström sketching, Conjugate Gradient, and a sketched block Jacobi preconditioner to scale to large models while achieving comparable or better performance than state-of-the-art methods.

1 Introduction

  • Exact KRR requires solving an n×n system, costing O(n^3) time and O(n^2) memory, which becomes prohibitive as n grows.Source: §1 · p. 1
  • Existing approaches trade off computational cost and convergence: full-problem methods can retain O(n^2) iteration costs, while sketched methods face O(m^2) preconditioner memory.Reducing m can limit accuracy, while cheaper stochastic or block-coordinate solvers give up CG's fast convergence.Source: §1 · p. 2
  • BLOKS achieves comparable or better performance than state-of-the-art methods on standard large-scale kernel-solver datasets.Its memory-efficient, GPU-accelerated implementation scales Nyström KRR to m = 3.2M parameters.Source: §1 · p. 2
  • BLOKS is a CG-based solver for Nyström-sketched KRR with a preconditioner combining block Jacobi approximation and sketching.Its preconditioner's spatial complexity is independent of both dataset size n and model size m.Source: §1 · p. 2

2 Background and setting

The section sets out KRR and Nyström KRR, then explains how iterative methods and preconditioning address computational costs while introducing their own constraints. In particular, smaller blocks reduce memory demands but increase data-swapping overhead, and a full preconditioner can require quadratic memory in model size.

2 Background and setting

  • Exact KRR becomes infeasible at scale: direct Cholesky costs O(n^3) time and O(n^2) memory.Source: §2 · p. 3
  • Nyström KRR restricts predictors to combinations of m sampled centers, reducing the coefficient vector to m dimensions.The centers are sampled uniformly at random from the training data, with m ≤ n.Source: §2 · p. 3
  • Iterative methods avoid explicitly inverting H, but too few iterations may give a poor solution approximation; convergence depends on conditioning.Smaller blocks can mitigate spatial complexity, but frequent data swapping introduces significant temporal overhead.Source: §2 · pp. 3–4
  • Preconditioning transforms the system into an equivalent one with improved condition number, supporting faster convergence of gradient-based methods.Source: §2 · p. 4
  • A full unstructured approximation of H can require O(m^2) memory, leaving a preconditioner memory bottleneck even after sketching and iteration reduce other costs.Source: §2 · p. 4

3 Introducing BLOKS

BLOKS uses a sketched block-diagonal preconditioner over disjoint subsets of Nyström centers within Conjugate Gradient. Its block structure supports efficient action computation and memory use independent of the model size when blocks are processed sequentially.

3 Introducing BLOKS

  • BLOKS partitions Nyström centers into disjoint blocks and builds a sketched block-diagonal preconditioner from each block’s kernel matrix.Each block matrix is defined by evaluating the kernel on points in that block.Source: §3 · pp. 4–5
  • The preconditioner acts on a vector through q decoupled linear systems, one for each block.The input vector is partitioned into q blocks of size m0, matching the Nyström-center partition.Source: §3 · pp. 4–5
  • Sequential block processing costs O(m0^2) memory independent of m; plain Block Jacobi instead requires unstable, inefficient Gram products costing O(nm0^2).Sequential processing is used when the block size prevents storing all blocks simultaneously.Source: §3 · p. 5
  • The sketched block action avoids forming the kernel Gram product, while left PCG computes each block only once per iteration.The bilateral formulation requires a second recomputation when only one block can be stored.Source: §3 · p. 6
  • Computing the block factors requires storing only one m0 × m0 matrix, reused for in-place Cholesky decompositions.The upper and lower triangular portions hold the successive factors.Source: §3 · p. 6

4 Related works

Related kernel solvers either address full KRR or a sketched approximation, using different combinations of iterative optimization and preconditioning. BLOKS is positioned among sketching-based methods that retain CG while reducing preconditioner memory.

Related-work scope

  • Data-, input-, and feature-space partitioning are complementary scalability strategies; this review focuses on methods directly solving full KRR or a sketched approximation.The paper separates related solvers according to which of those two problems they address.Source: §4 · p. 7

Full KRR methods

  • Full KRR PCG methods use randomized preconditioners, but each iteration costs O(n2), limiting them to modest-size datasets.Other full-problem approaches include stochastic optimization and block-coordinate descent.Source: §4 · p. 7

Sketching-based methods

  • Sketching-based methods include SGD variants and PCG solvers such as Falkon and Krill; the latter require O(m2) preconditioner memory.BLOKS instead uses O(m0^2) memory independent of m and n, and recovers Falkon when m0 = m.Source: §4 · p. 7

5 Experiments

Across large-scale datasets, BLOKS reaches state-of-the-art kernel-method accuracy with competitive training times, while medium-scale results often show it among the most accurate methods. Experiments also show how model size, inference cost, and block size affect the accuracy–efficiency trade-off.

Parameter scaling

  • On HIGGS, BLOKS scales to 3.2M parameters with fixed memory, while Falkon runs out of memory beyond 100k parameters; BLOKS training time scales linearly and test AUC improves consistently.The fixed memory level is essentially one m0 × m0 matrix, with m0 = 100k in this experiment.Source: §5 · p. 8, Figure 2 · p. 9, Table 2 · p. 9

Inference times

  • BLOKS is consistently the second-fastest inference method after Falkon and has the best overall trade-off between inference time and average test ranking across datasets.Inference time increases linearly with the number of parameters; EigenPro 4 and Joker use batching that slows inference, while ASkotch depends on training-set size.Source: §5 · p. 8, Figure 3 · p. 10

Overall comparisons

  • BLOKS reaches state-of-the-art kernel-method accuracy on the large-scale datasets while maintaining competitive training times except against Falkon; on medium-scale datasets it is often most accurate while Falkon is often fastest.Hyperparameters were tuned independently, subject to reasonable training-time constraints; those constraints could not be met for some ASkotch and EigenPro 4 runs.Source: Table 1 · p. 8, §5 · p. 9, Table 3 · p. 10

Block size ablation

  • Increasing block size m0 lowers test error but increases training time on Airlines and COMET MC, with diminishing returns especially on Airlines.When blocks no longer fit in memory together, recomputation makes preconditioning cost more significant; increasing m improves model capacity, whereas increasing m0 improves preconditioner accuracy.Source: §5 · pp. 9–10, Figure 4 · p. 10

6 Conclusions

BLOKS is a PCG solver for Nyström-sketched KRR with a sketched block Jacobi preconditioner whose memory cost is independent of dataset size and model size. The paper reports scaling to 3.2M parameters with linear training time and consistently improving accuracy, and identifies several extensions for future work.

Conclusion

  • BLOKS uses a sketched block Jacobi preconditioner with memory cost O(m0^2), independent of dataset size n and model size m.This decouples model size from memory while retaining PCG's fast, step-size-free convergence.Source: §6 · p. 10
  • BLOKS scaled to 3.2M parameters with linear training time and consistently improving accuracy, achieving higher accuracy than state-of-the-art methods under the same memory budget with competitive training time.The paper provides a GPU-accelerated implementation.Source: §6 · p. 11
  • Future directions include other loss functions, alternative block partitions, and coarse corrections that capture interblock interactions.The paper notes that a logistic BLOKS variant would enable larger-model kernel-based new-physics searches and goodness-of-fit tests.Source: §6 · p. 11

AI use statement

The authors report using generative AI tools for language work and utility scripts, while taking responsibility for the final content.

AI use statement

  • Generative AI tools were used for grammar and language polishing, text editing, and utility scripts for plotting and launching experiments.Source: AI use statement · p. 11
  • The authors take responsibility for the final content, including text, claims, and artifacts produced with AI assistance.Source: AI use statement · p. 11

A Further comparison with the literature

Kernel-solver methods combine iteration, sketching, and preconditioning in different ways, with memory or per-iteration costs limiting several approaches. The comparison covers methods for full KRR and Nyström-sketches, including their scaling bottlenecks.

FALKON and Krill

  • FALKON and Krill have preconditioner memory bottlenecks: FALKON uses O(m^2) memory and often runs out of memory near m = 10^5, while Krill inherits the same bottleneck.The survey also describes RPCholesky for full KRR when data are modest and eigenvalues decay sufficiently rapidly.Source: A Further comparison with the literature · p. 15

EigenPro

  • EigenPro addresses full-gradient and leading-eigenpair costs with stochastic updates and approximate eigenpairs, but EigenPro 3 retains an O(m^2) memory bottleneck and EigenPro 4 delays projections to reduce time.EigenPro 2’s Nyström approximation can require memory Ω(m0^2).Source: A Further comparison with the literature · pp. 15–16

Joker

  • Joker uses block coordinate descent with trust-region updates for KRR and proposes Inexact Joker with m Random Fourier Features to approximate the kernel matrix.For more general losses, it also proposes a truncated CG-Steihaug heuristic for box constraints.Source: A Further comparison with the literature · p. 16

Skotch and ASkotch

  • Skotch and ASkotch approximate randomized Newton updates with sketch-and-project methods, using a rank-m0 Nyström approximation and Nesterov acceleration for ASkotch.Their default sampling parameter is m = n/100, making the full MNWT solve at each step infeasible at large scale.Source: A Further comparison with the literature · pp. 16–17

B Block preconditioning strategies and comparisons

The three block-preconditioning strategies have nearly indistinguishable test performance, though their training-time costs differ. BLOKS sketches block Jacobi preconditioning to reduce block-computation burdens, and its HIGGS test curve avoids the overfitting seen with the other methods.

Sketched block Jacobi preconditioning

  • BLOKS sketches each block Jacobi block by sampling rows, avoiding the instability and O(nm0^2) Gram-product cost of forming and inverting the unsketched blocks.Inverting and stacking the sketched blocks diagonally yields the BLOKS preconditioner.Source: B Block preconditioning strategies and… · pp. 17–18

Block Jacobi–Falkon

  • The Block Jacobi–Falkon variant block-diagonalizes the Falkon preconditioner, but remains intractable for very large m because its Gram product is unstable and costs O(mm0^2).The approach sketches these blocks as well.Source: B Block preconditioning strategies and… · p. 18

Block preconditioning strategies and comparisons

  • Across experiments, the three strategies have almost indistinguishable test performance, while Block Jacobi is best on training error, Block Jacobi–Falkon next, and BLOKS slower, at longer training times.The authors suggest the test-time parity may reflect additional regularization from block approximation.Source: B Block preconditioning strategies and… · pp. 18, 20

Experimental evaluation

  • The comparison evaluates true relative error, preconditioner-norm residual, and test error; the small-scale experiments permit exact solves and report means over five-fold cross-validation.The relative residual is measurable without the exact solution, while test error is the quantity of greatest machine-learning interest.Source: B Block preconditioning strategies and… · p. 19, Table 6 · p. 20

Subsampled AIRLINES

  • On AIRLINES, test performances are indistinguishable and no method overfits within 40 iterations; the kernel bandwidth is σ = 0.5 because a larger value makes Block Jacobi blocks numerically singular.Block Jacobi remains more than an order of magnitude slower; Block Jacobi–Falkon becomes nearly as fast as BLOKS.Source: B Block preconditioning strategies and… · p. 20, Table 6 · p. 20

Subsampled HIGGS

Subsampled AIRLINES

C Bilateral vs. left preconditioning

BLOKS supports both bilateral and left preconditioning, but left PCG is more practical when its block preconditioner is recomputed. The comparison with FALKON and KRILL highlights how the preconditioner structure and sampling affect which implementation benefits apply.

C Bilateral vs. left preconditioning

  • Bilateral PCG applies the symmetric factors of M around H; left PCG uses the action of M as a whole and can be expressed through a change of variables or inner product.The paper describes bilateral preconditioning as mathematically equivalent to left preconditioning, while noting that the former requires an explicit symmetric decomposition.Source: C Bilateral vs. left preconditioning · p. 21
  • Left PCG materializes each non-overlapping preconditioner block once per iteration, whereas bilateral PCG forms the blocks twice around multiplication by H.When M is explicitly stored, the two variants have the same exact computational cost.Source: C Bilateral vs. left preconditioning · pp. 21–22

C.1 Comparison with FALKON and KRILL

  • FALKON is the q = 1 special case of BLOKS, and its single block can be stored, so bilateral PCG is more convenient there.For FALKON, bilateral PCG avoids a multiplication by the kernel matrix that the left variant requires.Source: §C.1 · p. 22
  • KRILL's refined sampling does not retain FALKON's double in-place Cholesky technique or cheaper preconditioned-matrix action.KRILL uses a sparse random sign embedding rather than simple row sampling.Source: §C.1 · pp. 22–23

D.1 Datasets

The appendix describes the datasets, evaluation metrics, and parameter-selection procedures used for the experiments. Most datasets use five-fold cross-validation, while MillionSongs uses a predefined test set without cross-validation.

D.1 Datasets

  • Except for MillionSongs, results are reported as means and standard deviations over five folds, with four-fifths of each dataset used for training and one-fifth for testing.MillionSongs instead uses a predefined test set and omits cross-validation to avoid potentially misleading high test performance from the “producer effect.”Source: §D.1 · pp. 23–24
  • The experiments cover datasets including Taxi (1.1 × 10^9 points), Cifar10-5M (6 × 10^6), and HIGGS (1.1 × 10^7), with regression and classification metrics selected by dataset.Taxi uses RMSE, Cifar10-5M uses classification error, and HIGGS uses one minus AUC.Source: §D.1 · p. 23

D.2 Algorithms

  • Competitor algorithms use their authors’ default parameters, while BLOKS parameters and each dataset’s selected settings are documented alongside the experiments.For Joker, both KRR and SVM are trained, and the better test-error model is reported.Source: §D.2 · p. 25
  • Kernel, bandwidth, and regularization are selected by grid search, with Gaussian and Laplacian kernels used for Joker and EigenPro 4 because their libraries lack Matérn support.For other settings, the bandwidth grid is centered on prior reported values or the median heuristic; regularization grids differ for Joker and EigenPro 4.Source: §D.2 · p. 25

D.3 Experiments

The appendix specifies the hardware and reporting setup for the experiments and provides BLOKS pseudocode. The implementation processes blocks sequentially and can recompute them when they do not fit in memory together.

D.3 Experiments

  • Section 5 experiments use single-precision arithmetic on one Leonardo cluster node with an NVIDIA Ampere A100 GPU; all but MillionSongs use five-fold cross-validation.Appendix B experiments instead use double precision on a laptop.Source: §D.3 · pp. 25–26
  • Tables 8 and 9 report training and test times and errors for datasets and algorithms, and for BLOKS model sizes on HIGGS, respectively.Table 8 reports five-split means and standard deviations except for MillionSongs; Table 9 reports mean and standard deviation for HIGGS AUC.Source: §D.3 · p. 26, Table 8 · p. 27, Table 9 · p. 27

E Complete BLOKS pseudocode

  • The pseudocode takes data, labels, regularization, model and block sizes, kernel, and iteration count as input, and returns an approximate Nyström coefficient vector.It starts from zero; blocks can be recomputed each iteration or precomputed if they fit simultaneously, and are processed sequentially in the described single-GPU setup.Source: E Complete BLOKS pseudocode · pp. 27–28
  • BLOKS achieves training times competitive with the state of the art.The pseudocode discussion notes that block solves could be distributed across multiple GPUs, so reported training times are conservative.Source: E Complete BLOKS pseudocode · pp. 28–29
Loading 2610.13382v1…