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

AI summary

Overview

Research area: Self-supervised representation learning, specifically Joint-Embedding Predictive Architectures (JEPAs), with a focus on distribution-matching regularization and sparse representation learning.

Technical level: Advanced. The paper draws heavily on probability theory (maximum-entropy characterizations, truncated and rectified Generalized Gaussian distributions), information theory (Rényi information dimension), and optimal transport (sliced Wasserstein distances).

Scope: A one-sentence summary: the paper introduces a family of rectified Generalized Gaussian target distributions and a sliced two-sample distribution-matching loss (RDMReg) that equips JEPAs with controllable sparsity and maximum-entropy representations.

What This Paper Is About

JEPAs learn view-invariant features without reconstruction, but invariance alone admits degenerate solutions such as complete or dimensional collapse. Existing fixes regularize representations toward isotropic Gaussian targets, which are dense by construction and therefore cannot express sparsity — a property widely associated with efficient and interpretable codes. The paper's goal is a principled target distribution, and a matching loss, that yield sparse, non-negative, high-entropy representations while preserving task-relevant information and competitive downstream accuracy.

Key Contributions

1. Rectified Generalized Gaussian (RGG) distributions. The authors introduce the RGG family, relate its truncated continuous component to maximum-entropy distributions under expected ℓ_p norm and support constraints, and show that explicit rectification induces analytically controllable ℓ_0 sparsity. Extending the Generalized Gaussian to its rectified variant for p ≠ 2 is stated as novel.

2. Rectified LpJEPA with RDMReg. They propose Rectified LpJEPA, a JEPA architecture equipped with Rectified Distribution Matching Regularization — a two-sample sliced distribution-matching loss that aligns projected feature marginals to an RGG target. Rectified LpJEPA strictly generalizes LeJEPA, which corresponds to the dense (non-rectified, p = 2) special case.

3. Necessity of two-sample matching. Because the RGG family is not closed under linear projections (unlike the Gaussian, which is closed as a stable distribution), the authors argue that projection-based matching requires nonparametric two-sample hypothesis testing rather than a closed-form one-sample comparison. They instantiate this with the sliced 2-Wasserstein distance.

4. Empirical validation. They demonstrate controllable sparsity, favorable sparsity–performance tradeoffs, improved statistical independence (nHSIC), higher d(ξ)-dimensional entropy, dataset-adaptive sparsity, and competitive accuracy on image classification benchmarks.

Main Findings

• Rectification is necessary. In CIFAR-100 validation experiments, models without rectification reach competitive accuracy but produce dense representations with no zero entries. Rectifying the features while using a Gaussian target for those rectified features causes substantial performance drops. The setting that rectifies both features and target (RGN_p | z⁺) attains the best accuracy–sparsity tradeoff.

• Sparsity is controllable and predictable. Proposition 3.5 gives the expected ℓ_0 norm of an RGG vector as d · Φ_{GN_p(0,1)}(μ/σ). Pretrained Rectified LpJEPA features closely follow these theoretical predictions as μ and p vary.

• Favorable Pareto frontier. Across varying μ and p, performance drops sharply only once more than ∼95% of entries are zero, indicating sparsity can be pushed far before accuracy degrades.

• Competitive ImageNet-100 linear probe accuracy (Table 1). Rectified LpJEPA with RGN₂.₀(1.0, σ_GN) reaches the highest encoder Acc1 at 85.08 with projector Acc1 80.54, L1 sparsity 0.6278, L0 sparsity 0.8668. RGN₂.₀(0, σ_GN) also reaches 85.08 encoder Acc1 with projector 80.00 and much lower sparsity (0.3412 L1, 0.7298 L0). RGN₁.₀(0, σ_GN) reaches 84.72 / 80.40 with 0.2726 L1 and 0.6940 L0; RGN₁.₀(0.25, σ_GN) reaches 84.98 / 80.76 with 0.3745 L1 and 0.7437 L0.

• Highly sparse settings trade accuracy. RGN₂.₀(−2.5, σ_GN) gives encoder Acc1 82.02 and projector Acc1 67.82 with L1 0.0137 and L0 0.0224; RGN₁.₀(−3.0, σ_GN) gives 82.72 / 71.88 with L1 0.0058 and L0 0.0098.

• Baseline comparison. Dense baselines: VICReg (84.18 / 78.88, L1 0.7954, L0 1.0000), SimCLR (83.44 / 77.90, L1 0.6338, L0 1.0000), LeJEPA (84.80 / 79.52, L1 0.6365, L0 1.0000). Sparse baselines: NVICReg-ReLU (84.48 / 77.74, L1 0.5207, L0 0.7117), NCL-ReLU (82.58 / 76.88, L1 0.0037, L0 0.0085), NVICReg-RepReLU (84.20 / 78.18, L1 0.4965, L0 0.7549), NCL-RepReLU (82.76 / 76.70, L1 0.0024, L0 0.0048). Notably, NCL variants achieve extreme sparsity only at substantially lower projector accuracy.

• Greater statistical independence. Rectified LpJEPA achieves smaller normalized Hilbert-Schmidt Independence Criterion (nHSIC) values than VICReg or NVICReg, which penalize only second-order statistics.

• Higher entropy and sparsity–entropy tradeoff. Averaged univariate d(ξ)-dimensional entropy of Rectified LpJEPA features shows the expected sparsity–entropy tradeoff as μ and p vary.

• Dataset-adaptive sparsity. Rectified LpJEPA exhibits the highest variation in ℓ_1 and ℓ_0 sparsity across downstream datasets, measured by relative mean absolute deviations from the median.

• Projection design. Using eigenvectors of the empirical feature covariance as projection vectors leads to faster convergence to optimal performance. A small, dimension-independent number of projections N is reported as sufficient for strong empirical performance.

Methodology in Plain English

The authors start from the idea that a good representation should be invariant across views of the same input while its coordinates follow a well-chosen probability distribution that prevents collapse. Their target distribution is built in three steps. First, they take the Generalized Gaussian distribution, whose shape parameter p controls tail behavior (p = 1 gives Laplace, p = 2 gives Gaussian), and truncate it to the positive half-line; this truncated form is the maximum-entropy distribution under a fixed expected ℓ_p norm. Second, they mix that truncated continuous part with a point mass at exactly zero, producing the Rectified Generalized Gaussian. The size of that point mass is determined by μ and σ, and it directly sets how many coordinates are exactly zero — giving an explicit, closed-form handle on sparsity.

To make a network's features follow this target, they use the Cramér–Wold device, which says two random vectors match in distribution if and only if all their one-dimensional linear projections match. This turns one hard high-dimensional matching problem into many easy one-dimensional problems. Because the rectified Generalized Gaussian family is not preserved under projection, they cannot compare a projected sample to a known closed-form marginal; instead they draw samples from the target, project them the same way, and compare the two projected empirical distributions using a sliced 2-Wasserstein distance based on sorting. The training objective is this distribution-matching term applied to both views, plus an ℓ_2 invariance term pulling the two views' rectified features together.

The network is the standard encoder-plus-projector design (ResNet or ViT encoder feeding an MLP), with a ReLU applied at the output so features are non-negative. Linear probes are evaluated both on the rectified projector output and on the encoder features. Sparsity is measured with two metrics: m_{ℓ_1}(x) = (1/D)·E[‖x‖₁²/‖x‖₂²], which ranges from 1/D for extremely sparse vectors to 1 for dense ones, and m_{ℓ_0}(x) = (1/D)·E[‖x‖₀], the fraction of nonzero entries. The target scale is set by σ_GN = Γ(1/p)^{1/2}/(p^{1/p}·Γ(3/p)^{1/2}), which makes the pre-rectification variance 1; an alternative σ_RGN that makes the post-rectification variance 1 is found by bisection search.

Why This Matters

Impact on research. The paper reframes collapse prevention in JEPAs as a target-distribution design problem rather than just a variance or covariance constraint. By showing that Gaussian targets are a dense special case of a much larger rectified family, it opens a principled axis — the choice of p, μ, and σ — for trading off sparsity against maximum entropy. It also connects distribution matching to second-order methods: the authors show a conditional link between RDMReg and Non-Negative VCReg (if projected marginals match the RGG target along covariance eigenvectors, the centered covariance is isotropic), which helps explain why eigenvector projections speed up convergence. The information-theoretic treatment via Rényi information dimension provides tools for reasoning about entropy when part of the distribution is a point mass at zero and differential entropy is ill-defined.

Real-world applications (implied by the properties studied; the paper itself evaluates image classification benchmarks):

• Efficient storage and transmission of learned features, where a large fraction of exactly-zero coordinates can reduce memory and bandwidth costs.
• Compressed sensing and signal processing pipelines, where sparse, non-negative codes underpin recovery and reconstruction methods the paper cites.
• Parts-based, interpretable decompositions, in the spirit of non-negative matrix factorization, where non-negativity and sparsity both aid interpretability.
• On-device or resource-constrained inference, where representations with many zero entries can be exploited by sparse compute paths.

Industry relevance. Standard self-supervised pipelines already rely on variance–covariance or distribution-matching regularizers, so RDMReg is a drop-in style change to an existing loss rather than a new architecture. The paper reports that a small, dimension-independent number of projections suffices, and that linear probe accuracy on par with dense baselines is achievable while retaining substantial sparsity — properties attractive to teams that care about both representation quality and inference or storage cost.

Future Directions

1. Scaling beyond the reported benchmarks. The empirical evidence presented here centers on CIFAR-100 and ImageNet-100 with ResNet- or ViT-based encoders; whether the sparsity–accuracy frontier holds at larger data scales and with larger backbones is not established in the provided content.

2. Lifting the finite-N approximation. The authors acknowledge that matching with a finite number of projections remains an approximation to the population Cramér–Wold criterion. A tighter theory of how many projections are needed, and how the choice of projection distribution affects matching, remains open.

3. Exploring p < 1 and alternative matching losses. The maximum-entropy characterization is stated to hold for 0 < p < 1, giving a continuous spectrum of sparse distributions, and the loss is instantiated only with the sliced 2-Wasserstein distance. Other two-sample test statistics and the nonconvex p < 1 regime are natural extensions.

4. Extending beyond images. The distributional argument is modality-agnostic in principle, so applying RDMReg to domains where sparse, non-negative codes are natural — such as audio, neural recordings, or other sensory streams — is untested here.

Target Audience

This paper is best suited to machine learning researchers and graduate students working on self-supervised learning, representation learning, and JEPA-style architectures, particularly those interested in the theoretical design of target distributions and collapse-prevention losses. It also appeals to readers with a background in probability, information theory, or sparse coding who want to see maximum-entropy and ℓ_0 arguments operationalized inside a deep learning objective. Practitioners seeking a drop-in regularizer for sparse features will find the empirical tables useful, though the derivations assume comfort with the underlying measure-theoretic and distributional machinery.

Authors’ 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 its continuous truncated component admits a maximum-entropy characterization under expected $\ell_p$ norm and support 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, showing that RDMReg can enforce sparsity while preserving task-relevant information.