Differentiable Kernel Ridge Regression for Deep Learning Pipelines

TL;DR

This paper introduces differentiable, scalable kernel ridge regression modules called Sparse Kernels that can be integrated into deep learning pipelines, enabling new transfer, probing, and hybrid modeling capabilities.

Contribution

The authors present Sparse Kernels, a differentiable, local, and lazy kernel ridge regression variant that can be integrated into deep learning models as modular layers with trainable or fixed parameters.

Findings

SK modules match neural readouts with less training

SK modules improve performance when added to existing models

Kernel methods can be integrated with deep learning effectively

Abstract

Deep neural networks dominate modern machine learning, while alternative function approximators remain comparatively underexplored at scale. In this work, we revisit kernel methods as drop-in components for standard deep learning pipelines. We introduce \emph{Sparse Kernels} (SKs), a differentiable, localized, and lazy variant of kernel ridge regression (KRR) that defers training to inference time and reduces to the solution of small local systems. We integrate SKs into PyTorch as modular layers that preserve end-to-end trainability, and we show that they expose three distinct sets of parameters -- feature representations, target values, and evaluation points -- each of which can be fixed or learned. This decomposition broadens the design space available to practitioners, enabling, in particular, training-free transfer, nonlinear probing, and hybrid kernel-neural models. Across convolutional networks, vision transformers, and reinforcement learning, SK-based modules serve two complementary roles: in some settings, they match the performance of trained neural readouts with substantially less training; in others, they augment existing models and improve their performance when used as additional components. Our results suggest that kernel methods, once made scalable and differentiable, can be readily integrated with deep learning rather than treated as a separate paradigm.