Differentiable Kernel Ridge Regression for Deep Learning Pipelines

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.

• The paper introduces a modular Sparse Kernel (SK) framework that integrates kernel ridge regression into deep learning, enabling training-free transfer and nonlinear probing.
• It demonstrates efficient, localized computations via M-NN tessellations that reduce inference complexity from cubic to linear in dataset size.
• Empirical results show that SK modules enhance accuracy in computer vision and reinforcement learning tasks compared to traditional neural readouts.

Differentiable Kernel Ridge Regression for Deep Learning Pipelines

Introduction

The paper "Differentiable Kernel Ridge Regression for Deep Learning Pipelines" (2605.02313) presents a scalable, modular approach to integrating kernel ridge regression (KRR) into standard deep learning frameworks. Through the introduction of Sparse Kernels (SKs)—a localized, differentiable, and lazy variant of KRR—the authors address critical limitations of scalability and nonparametric modeling in deep learning workflows. The framework explicitly separates feature representations, target values, and evaluation points, each of which may be fixed or learned, thereby expanding the design possibilities for model architects. Empirical evaluations across computer vision and reinforcement learning tasks substantiate the effectiveness and versatility of SKs as both stand-alone predictors and as augmentations to neural architectures.

Modular Sparse Kernel Ridge Regression

The proposed SKs operate as local kernel methods relying on $M$-nearest neighbor (M-NN) tessellations, with each prediction reduced to solving a small $M \times M$ KRR system at inference time. This approach yields computational efficiency—linear in the dataset size and $\mathcal{O}(M^3)$ per prediction—sidestepping the cubic scaling bottleneck of conventional dense kernels for large $N$.

Kernel layers are implemented as differentiable PyTorch modules with explicit parameterizations for "features" (input representations), "targets" (outputs associated with anchor points), and "evaluation points" (locations at which predictions are made). This separation enables several training and inference regimes:

• Training-free transfer: Kernel readouts are applied atop frozen backbones, leveraging pretrained representations without additional parametric optimization.
• Nonlinear probing: SK layers provide expressive, nonlinear probes of intermediate model activations, exceeding the capacity of traditional linear probes.
• Hybrid and trainable kernel modules: Complex architectures allow kernel parameters (including evaluation points and targets) to be learned by gradient descent alongside neural components.

The integration of the SK mechanism within standard backends ensures end-to-end trainability, facilitating experimentation and deployment within existing deep learning pipelines.

Empirical Results

Transfer Learning Performance

On the CIFAR-10 transfer benchmark using a frozen ResNet-18 pretrained on ImageNet, the SK-based readout matches or exceeds the accuracy of both linear and MLP heads, without requiring any parameter training. Notably, strong test accuracy is achieved even with limited labeled samples, demonstrating the ability of SKs to effectively leverage the geometry of pretrained representations in a nonparametric fashion.

Figure 1: Test accuracy for different readout mechanisms highlights the SK readout’s ability to reach strong accuracy without parametric training.

Probing Intermediate Representations

Probing experiments on VGG-19 and vision transformer (ViT) architectures reveal that kernel-based nonlinear probes over intermediate layers often yield higher test accuracy than probes over the final classifier head. This effect is especially pronounced in ViTs, where over-specialization of the final head reduces the efficacy for transfer tasks. The results indicate that nonparametric kernel probes are more sensitive to representational richness and less vulnerable to information loss due to head specialization.

Figure 2: VGG-19 probing results demonstrate the predictive power of intermediate representations when read out by the SK probe.

Learning with Differentiable Kernel Probes

By enabling differentiation through the kernel layer, the authors support end-to-end training of probes at various depths. Experiments corroborate that maximal test accuracy is generally obtained when training kernel predictors on representations from intermediate layers, paralleling trends observed in probe evaluations. This highlights the practical utility of selectively positioning kernel modules for improved generalization and transfer.

Figure 3: VGG-19 learning with kernel-based predictors at different depths mirrors the probing results, with intermediate layers providing superior utility.

Reinforcement Learning Augmentation

Augmenting the Double DQN algorithm with kernel modules demonstrates improved sample efficiency and faster reward acquisition on the LunarLander environment. Kernel readouts attached to the Q-network (either as additive corrections or perturbations to neural layers) function as flexible, trainable modules within the RL loop. The kernel-augmented agent outpaces the baseline in acquiring higher rewards, despite incurring higher TD losses during exploration, suggesting more accurate policy learning.

Figure 4: Episode rewards in RL show kernel-augmented Double DQN agents achieve higher and faster reward acquisition.

Theoretical and Practical Implications

This work formalizes a practical instantiation of localized KRR within the interpolating regime and extends classical RKHS-based error estimates to the local, sparse setting. By exposing features, targets, and evaluation points as adjustable parameter sets, the framework generalizes the classical representer theorem for hybrid neural-kernel systems. Notably, the SK approach provides a flexible platform for:

• Training-free nonparametric regression/classification in deep pipelines
• Nonlinear geometric probing of learned feature spaces
• Hybrid architectures blending parametric and nonparametric components
• Efficient resource usage due to inference-time localization and lazy solves

Practitioners gain an expanded set of levers for model design, including the ability to regularize the structure and support of SK modules (e.g., with optimal-transport-based loss terms).

Limitations center on evaluation scale: benchmarks were limited to CIFAR-10 and LunarLander in a single-machine, CPU-only environment. Wall-clock performance and kernel parameter selection (e.g., neighborhood size $M$, kernel type) merit further investigation on large-scale workloads or alternative modalities (e.g., NLP, generative modeling). For applications requiring global function continuity, hierarchical or blended variants of SKs are discussed but not exhaustively benchmarked.

Future Directions

Several lines of inquiry are opened by this work:

• Systematic exploration of kernel and bandwidth selection strategies, including data-driven adaptive mechanisms.
• Integration and benchmarking of SKs on ImageNet-scale datasets, large vision transformers, foundation models, and diffusion models.
• Expansion into continuous, hybrid, and hierarchical SK constructions for globally smooth predictions and compatibility with downstream optimization.
• Study of regularization methods for models with learnable evaluation points to ensure stability and generalization.

Further research into automated selection of representational depth and kernel configuration is likely to enhance transfer and generalization performance in multimodal deep learning systems.

Conclusion

The differentiable SK framework provides a computationally viable, modular mechanism for incorporating KRR into deep learning architectures. Experiments demonstrate that SK-based modules frequently match or exceed neural network baselines in both transfer learning and RL applications, commonly with reduced or zero training requirements. The theoretical generality and empirical flexibility of SKs underscore their promise as core building blocks for future hybrid, scalable function approximators in deep learning pipelines.

---

Reference: "Differentiable Kernel Ridge Regression for Deep Learning Pipelines" (2605.02313)