Faster Kernel Ridge Regression Using Sketching and Preconditioning

TL;DR

Random-feature preconditioned PCG solves exact KRR, reaching high accuracy on up to one million examples in about one hour.

math.NA 🔴 Advanced 2016-11-10 10 views
Haim Avron Kenneth L. Clarkson David P. Woodruff
kernel ridge regression random features preconditioned CG TensorSketch sketching

Key Findings

Methodology

The paper uses Random Fourier Features or TensorSketch to build a low-rank feature matrix Z, not as the final kernel model but as a preconditioner. For (K+λI)c=y, PCG uses ZZᵀ+λI; Woodbury and a Cholesky factorization of ZᵀZ+λI reduce each application to λ^{-1}(x−UᵀUx), where U=L^{-T}Zᵀ.

Key Results

  • For the q-degree polynomial kernel, s≥4(2+3^q)s_λ(K)^2/δ yields, with probability at least 1−δ, (2/3)(ZZᵀ+λI)≼K+λI≼2(ZZᵀ+λI). The preconditioned condition number is at most 3, giving ⌈√3/2·ln(2/ε)⌉ PCG iterations.
  • Experiments cover Gaussian and polynomial kernels and datasets with up to 1,000,000 training examples. The distributed implementation solves a million-example system to relatively high accuracy in about one hour on an EC2 cluster; the paper releases libSkylark code.
  • Direct random-feature regression can require s close to n and may still have worse test error than exact KRR. Preconditioning uses features only to accelerate convergence, preserving the statistical capacity of the full kernel system.

Significance

This work reframes random features from a model approximation into a numerical preconditioning device. The final predictor remains the solution of the original KRR system, while a compact sketch improves its conditioning. That separation addresses a persistent tension between kernel accuracy and scalability, offering researchers and practitioners a route to high-accuracy nonlinear regression at very large sample sizes.

Technical Contribution

The main contribution is a sketch-to-precondition framework with a statistical-dimension analysis. For TensorSketch, a matrix-multiplication bound establishes spectral equivalence between the sketched and exact regularized systems, producing a constant condition number and logarithmic dependence on target accuracy. The implementation combines Woodbury, Cholesky, and GEMM-friendly application, plus multiple-level sketching and empirical preconditioner testing.

Novelty

Unlike Nyström, direct random-feature regression, and conventional sketch-and-solve methods, the sketch is not treated as the final kernel. It only constructs a preconditioner for iterative solution of the exact system. The paper also provides a TensorSketch theorem tied to statistical dimension, proposes multilevel sketching, and discusses practical tests for detecting weak preconditioners.

Limitations

  • The formal theorem is mainly for TensorSketch with polynomial kernels; an analogous approximate-multiplication property is not proved for Random Fourier Features, limiting rigorous guarantees for Gaussian kernels.
  • The algorithm explicitly forms K, so storage and Θ(n²) matrix-vector products dominate at very large n. Preprocessing also costs Θ(ns²), and the theoretical feature bound may be conservative.
  • Iteration guarantees assume exact arithmetic and high-probability events. Floating-point PCG therefore needs residual-based stopping and careful numerical monitoring.

Future Work

Important directions include proving guarantees for Random Fourier Features, replacing explicit K with fast or implicit kernel products, and estimating statistical dimension adaptively. Multilevel SRHT or Johnson–Lindenstrauss compression, adaptive feature budgets, finite-precision analysis, and heterogeneous or streaming implementations could extend the approach beyond the reported setting.

AI Executive Summary

Kernel ridge regression is conceptually simple: solve (K+λI)c=y and predict with f(x)=∑ic_ik(x_i,x). The difficulty is scale. Popular Gaussian kernels produce dense, often ill-conditioned n×n matrices; direct methods cost Θ(n³), while unpreconditioned iterative methods may require many iterations. Nyström and Random Features reduce computation, but their approximation becomes the final model and can lose accuracy—even when the number of features approaches n.

Avron, Clarkson, and Woodruff propose a different division of labor. A random feature map creates Z∈R^{n×s}, but ZZᵀ is not substituted for K. Instead, ZZᵀ+λI preconditions Preconditioned Conjugate Gradients (PCG) applied to the exact system. After factoring ZᵀZ+λI=L Lᵀ, the Woodbury identity gives efficient applications through U=L^{-T}Zᵀ. Random Fourier Features serve Gaussian kernels; TensorSketch serves polynomial kernels and can be computed in O(q(nnz(x)+s log s)).

The central theorem states that for a q-degree polynomial kernel, s≥4(2+3^q)s_λ(K)^2/δ yields a preconditioned condition number at most 3 with probability at least 1−δ, so ε-accuracy requires ⌈√3/2·ln(2/ε)⌉ iterations. Distributed experiments cover Gaussian and polynomial kernels at up to one million training examples; a million-example system reaches relatively high accuracy in about an hour on an EC2 cluster. The approach preserves exact KRR capacity, but explicit K and Θ(n²) kernel products remain decisive limitations.

Deep Analysis

Background

KRR performs regularized linear regression in an RKHS and solves (K+λI)c=y. Dense kernels make direct factorization Θ(n³), while ill-conditioning slows CG. Nyström methods and Rahimi–Recht random features improve scalability, but their low-rank approximation also defines the final predictor. Prior sketching work distinguished sketch-and-solve from sketch-to-precondition; this paper imports the latter idea into kernel learning and targets high-accuracy exact KRR.

Core Problem

The goal is to accelerate the original kernel system without replacing it by a crude approximation. The preconditioner must be much smaller than K, cheap to factor and apply, parallelizable, and spectrally close enough to reduce the condition number. This is difficult because K is dense, λ may be small, and the useful rank depends on the regularized spectrum rather than simply on n.

Innovation

  • �� Use random features only to construct ZZᵀ+λI.
  • �� Relate feature count to statistical dimension s_λ(K).
  • �� Prove a condition-number guarantee for TensorSketch polynomial kernels.
  • �� Implement Woodbury application with GEMM-friendly operations.
  • �� Introduce multilevel sketching and practical preconditioner testing.
  • �� Preserve the exact KRR solution rather than accepting low-rank model bias.

Methodology

  • �� Input: training pairs, kernel k, regularization λ, feature count s, and tolerance ε.
  • �� Feature generation: sample Random Fourier Features for Gaussian kernels; use TensorSketch for k(x,z)=(xᵀz)^q.
  • �� Construction: stack z_i=ϕ(x_i)ᵀ into Z and factor ZᵀZ+λI=L Lᵀ.
  • �� Preconditioner: set U=L^{-T}Zᵀ and apply (ZZᵀ+λI)^{-1}x=λ^{-1}(x−UᵀUx).
  • �� Solve: run PCG on the unmodified K+λI until a residual criterion is met.
  • �� Analysis: use the TensorSketch multiplication lemma and a λ-LQ factorization to prove spectral sandwiching.
  • �� Extension: compose TensorSketch with SRHT or Johnson–Lindenstrauss transforms for multilevel compression.

Experiments

The study evaluates Gaussian and polynomial kernels in a distributed-memory implementation, comparing exact-kernel solution behavior with random-feature approximation and preconditioned solving. Reported scales reach 1,000,000 training examples, with costs including feature construction, ZᵀZ factorization, PCG iterations, and kernel matrix-vector products. The supplied paper text does not identify named public datasets or provide a complete numerical error table, so no unreported benchmark score should be inferred.

Results

Theorem 1 gives a constant preconditioned condition number, at most 3, under the stated TensorSketch feature bound; convergence then depends logarithmically on 1/ε. In engineering experiments, the method handles Gaussian and polynomial kernels and reaches relatively high accuracy for one million examples in roughly one hour on EC2 resources. Compared with direct random features, it avoids making approximation error part of the final model.

Applications

The method suits large nonlinear regression workloads where exact-kernel capacity matters: scientific surrogate modeling, engineering response prediction, and high-dimensional continuous-output learning. Users need distributed linear algebra and either explicit kernel storage or a compatible fast kernel-product engine. Fast Gauss Transform, tree codes, or hierarchical matrices could remove the explicit-K bottleneck.

Limitations & Outlook

The strongest theory covers TensorSketch and polynomial kernels, not Random Fourier Features for Gaussian kernels. Explicit K causes Θ(n²) storage or multiplication costs, while preprocessing costs Θ(ns²); these limits become severe beyond the reported scale. Statistical dimension is nontrivial to estimate and the theorem may overstate the required s. Future work should combine implicit kernel products, adaptive multilevel sketches, robust finite-precision analysis, and adaptive stopping.

Plain Language Accessible to non-experts

Imagine a huge factory where every order affects every other order. The exact factory rules produce the best result, but planning all interactions is painfully slow. Random features are a small rough map of the factory. The map is not accurate enough to replace the factory, yet it quickly reveals where traffic jams are likely to occur. This paper uses the rough map only to build express lanes. The full factory still makes the final decisions, so the answer keeps the accuracy of the original system. Each time the planner gets stuck, the express lanes suggest a better direction; consequently, far fewer corrections are needed. For polynomial relationships, TensorSketch creates the small map efficiently rather than listing every possible combination. Experiments reached one million training orders and solved the planning problem to relatively high accuracy in about an hour on a distributed EC2 cluster. The catch is that the factory still has to remember many pairwise interactions, so storage and repeated checking can remain expensive.

ELI14 Explained like you're 14

Suppose you are choosing the best gear for a million game characters, and every character affects every other character. Checking every interaction directly would be ridiculously slow. Random features make a small scouting map. It is too rough to decide the final gear, but it quickly shows the important routes. The paper uses that map as a shortcut while still checking answers with the complete game rules. PCG is the repeated fixer: it guesses, checks, and improves. Because the shortcut points it in a useful direction, it needs many fewer rounds.

For polynomial-style relationships, TensorSketch builds the scouting map quickly using clever hashing and fast calculations. The theory says that when the map is large enough compared with the problem’s genuinely important complexity, the shortcut can keep the number of fixing rounds small—about ⌈√3/2·ln(2/ε)⌉ in the theorem.

The experiments are impressive: Gaussian and polynomial kernels were tested, with up to one million training examples, and a distributed EC2 setup reached relatively high accuracy in around an hour. Best of both worlds, right? Speed comes from the small map, while accuracy comes from the full rules!

But there is a boss battle left: the full interaction table can still be enormous. The next upgrade would avoid storing that table and prove stronger guarantees for Gaussian-kernel features.

Glossary

Kernel Ridge Regression

A regularized regression method that measures similarity through a kernel rather than explicit raw features. Its dual coefficients solve (K+λI)c=y.

The exact learning problem accelerated by the paper.

Preconditioned Conjugate Gradients

An iterative solver for symmetric positive-definite systems that transforms the geometry using a preconditioner. Better conditioning usually means faster convergence.

PCG solves the original KRR system.

Random Fourier Features

For shift-invariant kernels, sampled frequencies and phases create cosine features whose expected inner product equals the kernel. They provide a finite-dimensional approximation.

Used to build Gaussian-kernel preconditioners.

TensorSketch

A randomized hash-and-sign transform that compresses tensor-product features, with FFT-based computation. It avoids explicitly materializing the d^q-dimensional polynomial feature vector.

Used for polynomial kernels and the main theorem.

Statistical dimension

s_λ(K)=Tr((K+λI)^{-1}K), the effective number of spectral degrees of freedom after regularization. It can be far smaller than n when eigenvalues decay quickly.

Controls the theoretical feature budget.

Sketch-to-precondition

A strategy that sketches a matrix only to construct a preconditioner, then iteratively solves the unsketched original problem. This differs from sketch-and-solve approximation.

The paper’s central algorithmic paradigm.

Open Questions Unanswered questions from this research

  • 1 No proof establishes the required approximate-multiplication property for Random Fourier Features, so the Gaussian-kernel theory remains weaker than the TensorSketch result.
  • 2 Explicit K and Θ(n²) products constrain further scaling; combining this preconditioner with implicit fast kernel multiplication remains unresolved.

Applications

Immediate Applications

Large-scale nonlinear regression

Researchers can generate Random Fourier Features or TensorSketch features, build the preconditioner, and run PCG on exact KRR. This is useful when model fidelity matters more than accepting low-rank approximation bias. Users should tune λ and s and monitor residuals.

Scientific and engineering prediction

Distributed teams can use the libSkylark implementation for high-dimensional continuous-response problems. Before deployment, they should verify memory for K, feature-construction cost, kernel-product throughput, and numerical stability under the chosen regularization.

Long-term Vision

Implicit exact-kernel platforms

Combining random-feature preconditioning with Fast Gauss Transform, tree codes, or hierarchical matrices could eliminate explicit K storage. Such systems might extend high-accuracy exact-KRR training beyond the reported million-example scale, provided kernel products remain sufficiently accurate and fast.

Abstract

Kernel Ridge Regression (KRR) is a simple yet powerful technique for non-parametric regression whose computation amounts to solving a linear system. This system is usually dense and highly ill-conditioned. In addition, the dimensions of the matrix are the same as the number of data points, so direct methods are unrealistic for large-scale datasets. In this paper, we propose a preconditioning technique for accelerating the solution of the aforementioned linear system. The preconditioner is based on random feature maps, such as random Fourier features, which have recently emerged as a powerful technique for speeding up and scaling the training of kernel-based methods, such as kernel ridge regression, by resorting to approximations. However, random feature maps only provide crude approximations to the kernel function, so delivering state-of-the-art results by directly solving the approximated system requires the number of random features to be very large. We show that random feature maps can be much more effective in forming preconditioners, since under certain conditions a not-too-large number of random features is sufficient to yield an effective preconditioner. We empirically evaluate our method and show it is highly effective for datasets of up to one million training examples.

math.NA cs.DS cs.LG