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

TL;DR

Rectified LpJEPA targets representations that are both sparse and information-rich by aligning JEPA features to Rectified Generalized Gaussian distributions. The method combines a Cramér-Wold-based two-sample distribution-matching objective (RDMReg) with a rectified target that induces explicit ℓ0 sparsity while preserving maximum entropy under ℓp constraints. The authors show that rectification disrupts closure under linear projections, necessitating slice-based hypothesis tests (SWD) and yielding a form of Non-Negative VCReg that promotes reduced higher-order dependencies. Empirically, Rectified LpJEPA achieves controllable sparsity, favorable sparsity–performance tradeoffs, and competitive transfer performance on image classification benchmarks, while also revealing interpretable sparsity patterns and strong independence properties. This work provides a principled, parameterized bias toward sparse, non-negative representations that maintain information content for self-supervised JEPA frameworks.

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.

Figures (19)

• Figure 1: Rectified LpJEPA. (a) Two views $(x,x')$ of the same underlying data are embedded and rectified to obtain $\operatorname{ReLU}(\mathbf{z})$ and $\operatorname{ReLU}(\mathbf{z}') \in \mathbb{R}^d$. Rectified LpJEPA minimizes the $\ell_2$ distance between rectified features while regularizing the $d$-dimensional rectified feature distribution towards a product of i.i.d. Rectified Gaussian distributions $\operatorname{ReLU}(\mathcal{N}(\mu,\sigma^2))$ using RDMReg. As a result, each coordinate of the learned representation aligns towards a Rectified Gaussian distribution (CDF shown above), a special case of the Rectified Generalized Gaussian family $\mathcal{RGN}_p(\mu,\sigma)$ when $p=2$. In the absence of rectification on both the features and the target distribution, Rectified LpJEPA reduces to isotropic Gaussian regularization as in LeJEPA balestriero2025lejepaprovablescalableselfsupervised. (b) Samples from $2$-dimensional Gaussian $\mathcal{N}(\mathbf{0},\mathbf{I})$ and Rectified Gaussian $\operatorname{ReLU}(\mathcal{N}(\mathbf{0},\mathbf{I}))$ are drawn and projected along a certain direction $\mathbf{c}$. As opposed to Gaussian which is closed under linear combinations, the projected marginals of the Rectified Gaussian distribution no longer fall in the same family, motivating the necessity of using two-sample distribution-matching losses.
• Figure 2: Rectified Laplace ($p=1$) and Rectified Gaussian ($p=2$) as special cases of Rectified Generalized Gaussian distributions. Assume $\mu=0$ and $\sigma=1$. For any $p>0$, the Truncated Generalized Gaussian $\prod_{i=1}^{d}\mathcal{TGN}_p$ over the support $(0,\infty)^d$ is the maximum differential entropy distribution under a fixed expected $\ell_p$-norm constraint. For $p\in\{1,2\}$, $\prod_{i=1}^{d}\mathcal{TGN}_p$ further admits a radial--angular decomposition $\mathbf{x}=r\cdot\mathbf{u}$ with $r\perp\!\!\!\perp\mathbf{u}$, where $\mathbf{u}$ is uniformly distributed with respect to the surface measure on the unit $\ell_p$-sphere confined to the positive orthant and $r^p$ follows a Gamma distribution. Rectified Laplace and Rectified Gaussian arise via coordinate-wise mixing of the corresponding truncated distributions with a Dirac measure at zero, yielding a distribution with expected $\ell_0$-norm guarantees, where $\Phi_{\mathcal{L}}$ and $\Phi_{\mathcal{N}}$ denote the cumulative distribution functions of the standard Laplace and standard Gaussian distributions respectively.
• Figure 3: Rectified LpJEPA achieves controllable sparsity and favorable sparsity-performance tradeoffs under proper parameterizations. (a) We report CIFAR-100 validation accuracy and the $\ell_0$ sparsity metric $1-(1/d)\cdot\mathbb{E}[\|\mathbf{x}\|_0]$ for four settings where we match non-rectified features $\mathbf{z}$ or rectified features $\mathbf{z}^{+}:=\operatorname{ReLU}(\mathbf{z})$ to either Rectified Generalized Gaussian $\mathcal{RGN}_p$ or conventional Generalized Gaussian $\mathcal{GN}_p$. Rectified LpJEPA $(\mathcal{RGN}_p\mid\mathbf{z}^{+})$ achieves the best sparsity-performance tradeoffs compared to other settings. (b) We compare the normalized $\ell_0$ norm of pretrained Rectified LpJEPA features against the theoretical predictions of Proposition \ref{['proposition:expectedl0normforrecgengauss']} as $\mu$ varies. Empirical sparsity closely follows the predicted behavior across different values of $\mu$ and $p$. (c) We plot the Pareto frontier of sparsity versus accuracy across varying values of $\mu$ and $p$. Performance drops sharply only when more than $\sim\!95\%$ of entries are zero.
• Figure 4: Rectified LpJEPA empirically achieves higher-entropy, more independent features with dataset-adaptive sparsity. (a) The averaged univariate $d(\xi)$-dimensional entropy of the Rectified LpJEPA features are computed against the $\ell_1$ sparsity metric $1-(1/D)\cdot\mathbb{E}[\|\mathbf{z}\|_1^2/\|\mathbf{z}\|_2^2]$ across varying $\mu$ and $p$. Overall, we observe the expected behavior of sparsity-entropy tradeoff (b) We evaluate the normalized Hilbert-Schmidt independence Criterion (nHSIC) for LpJEPA, Rectified LpJEPA, and other baselines. Rectified LpJEPA achieves smaller nHSIC values compared to VICReg or NVICReg that only penalizes second-order statistics. (c) The relative mean absolute deviations (MAD) away from the median of the $\ell_1$ and $\ell_0$ sparsity metrics are computed over different methods. Rectified LpJEPA exhibits the highest variations of sparsity for different downstream dataset. Additional visualizations can be found in \ref{['fig:total_transfer_sparsity']}.
• Figure 5: The Probability Density Function of Generalized Gaussian $\mathcal{GN}_p$, Truncated Generalized Gaussian $\mathcal{TGN}_p$, and Rectified Generalized Gaussian $\mathcal{RGN}_p$ across varying $p$ with fixed $\mu=-0.5$ and $\sigma=1$.$\Phi_{\mathcal{GN}_p(0,1)}$ is the CDF of the Generalized Gaussian $\mathcal{GN}_{p}(0,1)$. (a) The case when $p=0.5$. (b) When $p=1$, we obtain Laplace, Truncated Laplace, and Rectified Laplace. (c) For $p=2$, we have Gaussian, Truncated Gaussian, and Rectified Gaussian.

Theorems & Definitions (51)

• Definition 3.1: Generalized Gaussian Distribution
• Definition 3.2: Truncated Generalized Gaussian Distribution
• Proposition 3.3: Maximum Entropy Characterizations of Multivariate Truncated Generalized Gaussian Distributions
• proof
• Definition 3.4: Rectified Generalized Gaussian
• Proposition 3.5: Sparsity
• proof
• Theorem 3.6: Rényi Information Dimension Characterizations of Multivariate Rectified Generalized Gaussian Distributions
• proof
• Remark 2.1