Rectified LpJEPA: Sparse, Nonnegative JEPA

• The paper introduces rectified distribution matching regularization using a rectified generalized Gaussian prior to enforce sparsity and nonnegativity.
• It replaces traditional dense Gaussian priors with a sparse formulation that yields interpretable, maximum-entropy latent codes while avoiding representational collapse.
• Empirical results show competitive accuracy on benchmarks alongside flexible control over sparsity–performance trade-offs via sliced-2-Wasserstein matching.

Rectified LpJEPA is a self-supervised learning framework that extends Joint-Embedding Predictive Architectures (JEPAs) to produce sparse, nonnegative, and view-invariant representations by introducing a novel rectified distribution matching regularization. This approach addresses the limitations of Gaussian-based JEPA regularization by explicitly controlling sparsity through a rectified generalized Gaussian prior, resulting in feature distributions capable of representing efficient, interpretable, and robust latent codes while maintaining competitive performance across canonical vision benchmarks (Kuang et al., 1 Feb 2026).

1. Foundations of JEPA and Projection-Based Distribution Matching

JEPAs learn representations by predicting one randomized augmentation of an input from another within the latent space, entirely avoiding pixel-level reconstruction. Given augmentations $x, x' \sim \mathbb P_{x,x'}$ and a shared encoder $f_\theta$, JEPAs generate embeddings $z = f_\theta(x)$ and $z' = f_\theta(x')$. View invariance is promoted via a penalty $\|z-z'\|_2^2$, but this term alone causes representational collapse (embeddings become constant).

To mitigate collapse, recent JEPA variants, notably LeJEPA [1], employ distribution matching as a regularizer: they align all one-dimensional projections $c^\top z$ with a univariate reference distribution using random projection vectors $c \in \mathbb S^{d-1}$. The process leverages Cramér–Wold theorem guarantees, practically enforcing high-entropy, isotropic latent laws via sliced one-dimensional two-sample tests.

2. Motivation: Limitations of Gaussian Latent Targets in Sparsity

Traditional Gaussian priors in JEPA settings maximize entropy under quadratic (energy) constraint, resulting in latent embeddings that are dense and sign-symmetric. However, sparse nonnegative codes have proven important for efficiency, interpretability, and robustness—motivated by findings in neuroscience, compressed sensing, and ReLU-based deep models. Gaussian targets preclude exact zeros and nonnegativity, making them suboptimal for learning efficient sparse representations (Kuang et al., 1 Feb 2026).

3. Rectified Generalized Gaussian (RGG) Distribution: Theory and Parameterization

To directly encode sparsity and nonnegativity, Rectified LpJEPA introduces the rectified generalized Gaussian family $\mathcal{RGN}_p(\mu,\sigma)$. Each $X \sim \mathcal{RGN}_p$ is given by $X = \max\{0, G\}$, where $f_\theta$0 and the density of $f_\theta$1 is

$f_\theta$2

$f_\theta$3 is a mixture, with a point mass at zero (Dirac) and a truncated generalized Gaussian over $f_\theta$4. The explicit control of expected sparsity arises from the closed form of the expected $f_\theta$5 norm:

$f_\theta$6

where $f_\theta$7 is the regularized Gamma function. This formulation enables explicit, independent adjustment of both the moment ($f_\theta$8) and the sparsity ($f_\theta$9). The rectified case strictly generalizes Gaussians, which are recovered when $z = f_\theta(x)$0 and $z = f_\theta(x)$1.

4. Rectified Distribution Matching Regularization (RDMReg)

RDMReg introduces the RGG prior into the JEPA training objective via a sliced two-sample matching loss. The encoder and projection network output $z = f_\theta(x)$2 (imposing nonnegativity), and synthetic targets $z = f_\theta(x)$3. For each of $z = f_\theta(x)$4 random projections $z = f_\theta(x)$5, the slices $z = f_\theta(x)$6 and $z = f_\theta(x)$7 are compared batchwise using the empirical sliced-2-Wasserstein distance,

$z = f_\theta(x)$8

where $z = f_\theta(x)$9 denotes sorting. The overall loss combines view-invariance and distribution matching:

$z' = f_\theta(x')$0

This critically makes use of nonparametric sliced-matching, as the RGG is not closed under arbitrary projections and thus precludes closed-form parametric matching.

5. Theoretical Properties and Sparsity–Performance Trade-off

Rectified LpJEPA achieves maximum-entropy latent distributions up to rescaling under fixed sparsity and moment constraints by invoking Rényi’s information dimension $z' = f_\theta(x')$1 and dimension-appropriate entropies. For $z' = f_\theta(x')$2,

$z' = f_\theta(x')$3

and $z' = f_\theta(x')$4 maximizes entropy under an $z' = f_\theta(x')$5 moment. The mechanism produces a provable trade-off: decreasing $z' = f_\theta(x')$6 or $z' = f_\theta(x')$7 increases sparsity (decreases $z' = f_\theta(x')$8), yielding lower information dimension and adjusted $z' = f_\theta(x')$9-entropy. In the fully dense (non-rectified) limit, this reduces to the classical Gaussian maximum entropy principle. Sliced matching via Wasserstein distance is necessary, as RGG distributions are not closed under linear combinations.

6. Architecture and Training Procedures

The Rectified LpJEPA implementation comprises:

• Encoder $\|z-z'\|_2^2$0 (ResNet-50 or ViT-Small)
• Projector $\|z-z'\|_2^2$1, a 2–3-layer MLP (hidden dim 2048, output $\|z-z'\|_2^2$2)
• ReLU nonlinearity enforcing nonnegativity and sparsity: $\|z-z'\|_2^2$3
• A sampling block drawing random directions $\|z-z'\|_2^2$4 and optionally using empirical covariance eigenvectors for accelerated decorrelation

Hyperparameterization outlines:

• Pretraining on ImageNet-100, 1000 epochs, batch 128, LARS or AdamW optimizer
• $\|z-z'\|_2^2$5, $\|z-z'\|_2^2$6
• $\|z-z'\|_2^2$7 projections for sliced Wasserstein distance
• Generalized Gaussian scale $\|z-z'\|_2^2$8 fixes pre-ReLU variance; $\|z-z'\|_2^2$9 (post-ReLU variance) as alternative
• Standard image augmentations: crop, flip, color jitter, blur, solarize
• Single-GPU, mixed-precision, training duration 2–3 days for 1000 epochs

Ablation confirms the necessity of rectifying both feature and target distributions, and shows that mixing eigenvector and random projections accelerates decorrelation and convergence.

7. Empirical Performance and Ablations

Rectified LpJEPA demonstrates high sparsity–performance flexibility and competitive accuracy across benchmarks:

• On ImageNet-100 (ResNet-50, linear probe): encoder top-1 ≈ 84–85 %; projector top-1 ≈ 80 %; $c^\top z$0 sparsity efficiently tunable between 2 % and 0 % zeros, yielding a full sparsity–performance Pareto frontier
• In contrast, LeJEPA (Gaussian) yields no exact zeros (always dense)
• On CIFAR-100, empirical $c^\top z$1 closely tracks the theoretical $c^\top z$2; accuracy degrades only when sparsity exceeds 95 % zeros
• HSIC independence metrics show Rectified LpJEPA achieves lower higher-order statistical dependence than VICReg or NVICReg
• Ablations reveal Rectified LpJEPA outperforms or matches VICReg/LeJEPA for both small/large projector dimension, and both $c^\top z$3 and $c^\top z$4 deliver similar trade-offs
• Direct RDMReg matching significantly outperforms continuous-mapping approaches; joint rectification is essential for true sparsity
• Transfer learning across downstream classification tasks retains competitive accuracy, with independent control of sparsity levels

8. Implications, Limitations, and Future Directions

Rectified LpJEPA generalizes isotropic Gaussian priors to a broad class of sparse, nonnegative distributions, maintaining strong collapse-prevention and predictive-view invariance properties central to JEPA. The explicit control of sparsity under fixed moment constraint produces a continuum of operational points for sparsity–performance trade-offs, enabling learning of factorial, maximum-entropy, and interpretable sparse codes.

Proposed future research directions include:

• Tighter theoretical bounds on sliced-matching sample complexity for RGG
• Exploration of alternative two-sample tests (e.g., MMD, energy distance)
• Extensions to structured sparsity (group, top-$c^\top z$5)
• Application to domains demanding structured sparse codes (e.g., RL, multimodal)
• Meta-learning of hyperparameters $c^\top z$6 for dataset-adaptive priors
• Investigating connections between RGG-induced sparsity and neural coding in biological networks

A plausible implication is that the mechanisms encoded in Rectified LpJEPA could inform both artificial and theoretical models where efficient, sparse, high-entropy representations are critical (Kuang et al., 1 Feb 2026).