Rectified LpJEPA: Joint-Embedding Predictive Architectures with Sparse and Maximum-Entropy Representations

TL;DR

This paper introduces Rectified LpJEPA, a novel joint-embedding predictive architecture that enforces sparsity and maximum-entropy representations using a new regularization method, improving efficiency and performance in image classification tasks.

Contribution

It proposes RDMReg, a sliced distribution-matching loss that aligns representations to a Rectified Generalized Gaussian distribution, enabling explicit sparsity control in JEPAs.

Findings

Learns sparse, non-negative representations with favorable trade-offs.

Achieves competitive performance on image classification benchmarks.

Effectively enforces sparsity while maintaining task-relevant information.

Abstract

Joint-Embedding Predictive Architectures (JEPA) learn view-invariant representations and admit projection-based distribution matching for collapse prevention. Existing approaches regularize representations towards isotropic Gaussian distributions, but inherently favor dense representations and fail to capture the key property of sparsity observed in efficient representations. We introduce Rectified Distribution Matching Regularization (RDMReg), a sliced two-sample distribution-matching loss that aligns representations to a Rectified Generalized Gaussian (RGG) distribution. RGG enables explicit control over expected $\ell_0$ norm through rectification, while preserving maximum-entropy up to rescaling under expected $\ell_p$ norm constraints. Equipping JEPAs with RDMReg yields Rectified LpJEPA, which strictly generalizes prior Gaussian-based JEPAs. Empirically, Rectified LpJEPA learns sparse, non-negative representations with favorable sparsity-performance trade-offs and competitive downstream performance on image classification benchmarks, demonstrating that RDMReg effectively enforces sparsity while preserving task-relevant information.