Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

KerJEPA: Kernel Discrepancies for Euclidean Self-Supervised Learning

Eric Zimmermann, Harley Wiltzer, Justin Szeto, David Alvarez-Melis, Lester Mackey

TL;DR

KerJEPA generalizes Euclidean self-supervised learning by introducing kernel-based regularizers that align learned embeddings with isotropic priors. It unifies and extends prior work (LeJEPA) through both sliced and unsliced discrepancy measures (MMD and KSD), offering flexible choices of kernels and priors to control embedding geometry and training stability. The framework provides theoretical connections (EP as MMD, spectral forms of KSD, and SKSD) and practical implementations (Sliced MMDReg, Sliced KSDReg, and generalized discrepancy regularization), and demonstrates competitive performance on a vision SSL benchmark. Overall, KerJEPA offers a scalable, principled toolbox for regularizing Euclidean representations in SSL with tunable priors and kernel choices to suit downstream tasks.

Abstract

Recent breakthroughs in self-supervised Joint-Embedding Predictive Architectures (JEPAs) have established that regularizing Euclidean representations toward isotropic Gaussian priors yields provable gains in training stability and downstream generalization. We introduce a new, flexible family of KerJEPAs, self-supervised learning algorithms with kernel-based regularizers. One instance of this family corresponds to the recently-introduced LeJEPA Epps-Pulley regularizer which approximates a sliced maximum mean discrepancy (MMD) with a Gaussian prior and Gaussian kernel. By expanding the class of viable kernels and priors and computing the closed-form high-dimensional limit of sliced MMDs, we develop alternative KerJEPAs with a number of favorable properties including improved training stability and design flexibility.

KerJEPA: Kernel Discrepancies for Euclidean Self-Supervised Learning

TL;DR

KerJEPA generalizes Euclidean self-supervised learning by introducing kernel-based regularizers that align learned embeddings with isotropic priors. It unifies and extends prior work (LeJEPA) through both sliced and unsliced discrepancy measures (MMD and KSD), offering flexible choices of kernels and priors to control embedding geometry and training stability. The framework provides theoretical connections (EP as MMD, spectral forms of KSD, and SKSD) and practical implementations (Sliced MMDReg, Sliced KSDReg, and generalized discrepancy regularization), and demonstrates competitive performance on a vision SSL benchmark. Overall, KerJEPA offers a scalable, principled toolbox for regularizing Euclidean representations in SSL with tunable priors and kernel choices to suit downstream tasks.

Abstract

Recent breakthroughs in self-supervised Joint-Embedding Predictive Architectures (JEPAs) have established that regularizing Euclidean representations toward isotropic Gaussian priors yields provable gains in training stability and downstream generalization. We introduce a new, flexible family of KerJEPAs, self-supervised learning algorithms with kernel-based regularizers. One instance of this family corresponds to the recently-introduced LeJEPA Epps-Pulley regularizer which approximates a sliced maximum mean discrepancy (MMD) with a Gaussian prior and Gaussian kernel. By expanding the class of viable kernels and priors and computing the closed-form high-dimensional limit of sliced MMDs, we develop alternative KerJEPAs with a number of favorable properties including improved training stability and design flexibility.
Paper Structure (29 sections, 22 theorems, 101 equations, 3 figures, 3 tables)

This paper contains 29 sections, 22 theorems, 101 equations, 3 figures, 3 tables.

Table of Contents

  1. Introduction
  2. Related Work
  3. Matching Distributions with Kernel Discrepancies
  4. Kernels, Reproducing Kernel Hilbert Spaces, and Distribution Comparison
  5. Maximum Mean Discrepancy
  6. Epps-Pulley via Maximum Mean Discrepancy
  7. Kernel Stein Discrepancy
  8. Spectral Representation of the Kernel Stein Discrepancy
  9. Sliced Integral Probability Metrics
  10. Tradeoffs with Sliced Metrics
  11. Sliced Kernel Stein Discrepancy
  12. Discrepancy Regularization in Self-Supervised Learning
  13. Sliced Discrepancy Regularization
  14. Generalized Discrepancy Regularization
  15. Experimental Results
  16. ...and 14 more sections

Key Result

Proposition 1

$\mathtt{EP}(P) = \mathtt{MMD}_{k_{\mathrm{gsn}}}^2(P, \mathcal{N}(0, \sigma^2))$.

Figures (3)

  • Figure 1: Pseudo-code for the finite-sliced discrepancy regularization losses for both Gaussian and Laplace priors using the Gaussian kernel. Integration is done as Gauss-Hermite quadrature approximation.
  • Figure 2: Impacts of slicing in various dimensions for the MMD regularizer on ImageNette, measuring test accuracy over various training horizons. Results indicate that insufficient slices slow convergence rates as a function of dimension and introduces instability in the early half of training compared to the analytically sliced counterpart.
  • Figure 3: Impacts of slicing in various dimensions for the KSD regularizer on ImageNette, measuring test accuracy over various training horizons. Results indicate that insufficient slices slow convergence rates as a function of dimension and introduces instability in the early half of training compared to the analytically sliced counterpart.

Theorems & Definitions (23)

  • Proposition 1
  • Proposition 2: lejepa, Theorem 6
  • Corollary 3: lejepa, Theorem 6
  • Proposition 3
  • Corollary 3
  • Theorem 4: Cramér-Wold
  • Definition 1: Sliced divergence
  • Theorem 5
  • Theorem 6
  • Theorem 7
  • ...and 13 more