Rectified Distribution Matching Reg.

• RDMReg is a regularization technique that induces explicit ℓ0-controlled sparsity by aligning non-negative JEPA features with a rectified generalized Gaussian prior.
• It employs a sliced Wasserstein loss to match learned representations to a tailored RGG distribution, ensuring maximum-entropy properties under ℓp constraints.
• The method enables tunable sparsity-performance trade-offs and outperforms traditional Gaussian regularizations by offering controlled sparsity with high information capacity.

Rectified Distribution Matching Regularization (RDMReg) is a regularization technique designed to induce explicit $\ell_0$-controlled sparsity while maintaining maximum-entropy properties in representations learned by Joint-Embedding Predictive Architectures (JEPAs). RDMReg aligns learned non-negative representations to samples from a Rectified Generalized Gaussian (RGG) prior using a sliced two-sample distribution-matching loss. This method enables precise tuning of expected sparsity and generalizes prior Gaussian-based regularizations by introducing rectification, thus producing representations that are sparse, non-negative, and information-rich under $\ell_p$ norm constraints (Kuang et al., 1 Feb 2026).

1. Mathematical Formulation

RDMReg centers on matching representations to the RGG distribution. For $p > 0$, mean $\mu \in \mathbb{R}$, and scale $\sigma > 0$, the univariate RGG, denoted $\mathrm{RGN}_p(\mu,\sigma)$, is defined as a mixture: a point mass at zero with weight $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$ (the CDF of the unrectified generalized Gaussian at $-\mu/\sigma$) and a truncated generalized Gaussian on $(0,\infty)$. The density for $X \sim \mathrm{RGN}_p(\mu,\sigma)$ with respect to $\ell_p$0 is: $\ell_p$1 For independent RGG coordinates $\ell_p$2, the expected $\ell_p$3 sparsity is $\ell_p$4, allowing direct control over the fraction of nonzero activations by tuning $\ell_p$5.

2. Distribution Matching Objective

Given encoder-projector $\ell_p$6, two augmented views $\ell_p$7, and non-negative representations $\ell_p$8, the RDMReg loss matches the feature distribution to the RGG prior using a sliced Wasserstein approach. For random projection $\ell_p$9:

• Project both learned and RGG target features: $p > 0$0.
• Compute sliced 2-Wasserstein:

$p > 0$1

where $p > 0$2 denotes sorting.

The regularization is the expectation over projections and view-pairs: $p > 0$3 The JEPA loss becomes

$p > 0$4

3. Maximum-Entropy and Sparsity Guarantees

Among all distributions on $p > 0$5 with fixed expected $p > 0$6 norm, the truncated product Generalized Gaussian achieves maximum differential entropy. For $p > 0$7, each coordinate is a positive-orthant truncated generalized Gaussian. The RGG, while not absolutely continuous (due to its atom at zero), retains this maximum-entropy property on its support.

The effective entropy is characterized via the Rényi information dimension. Theorem 3.8 provides that the information dimension of an RGG distributed vector is $p > 0$8, where $p > 0$9 denotes ambient dimension. The entropy decomposes into a Bernoulli ($\mu \in \mathbb{R}$0) term plus the entropy of the truncated-GG on the active coordinates, enabling explicit $\mu \in \mathbb{R}$1 control without sacrificing information content.

4. Algorithmic Implementation

The implementation of RDMReg within JEPA consists of the following steps:

1. Forward computation: Given batch size $\mu \in \mathbb{R}$2 of view pairs $\mu \in \mathbb{R}$3, compute $\mu \in \mathbb{R}$4, $\mu \in \mathbb{R}$5.
2. Invariance loss: $\mu \in \mathbb{R}$6.
3. Sample RGG targets: Draw $\mu \in \mathbb{R}$7 by sampling i.i.d. $\mu \in \mathbb{R}$8 then applying $\mu \in \mathbb{R}$9.
4. Sliced Wasserstein regularization: Sample projection directions $\sigma > 0$0 (for $\sigma > 0$1 projections). For each, compute $\sigma > 0$2 for both $\sigma > 0$3 and $\sigma > 0$4, average to get $\sigma > 0$5.
5. Backpropagation: Update via $\sigma > 0$6.

Typical runtime is determined by sorting operations, with $\sigma > 0$7 per slice and practical $\sigma > 0$8 in the few thousands even for $\sigma > 0$9.

5. Hyperparameterization and Sparsity–Performance Trade-off

RDMReg affords direct and interpretable control over sparsity and representational properties:

• Norm order $\mathrm{RGN}_p(\mu,\sigma)$0: $\mathrm{RGN}_p(\mu,\sigma)$1 (rectified Gaussian) yields dense codes; $\mathrm{RGN}_p(\mu,\sigma)$2 (rectified Laplace) produces sparser codes; $\mathrm{RGN}_p(\mu,\sigma)$3 further approximates $\mathrm{RGN}_p(\mu,\sigma)$4 sparsity.
• Mean shift $\mathrm{RGN}_p(\mu,\sigma)$5: Controls expected active fraction, corresponding to $\mathrm{RGN}_p(\mu,\sigma)$6 per dimension; empirically, $\mathrm{RGN}_p(\mu,\sigma)$7.
• Scale $\mathrm{RGN}_p(\mu,\sigma)$8: Two options: (i) $\mathrm{RGN}_p(\mu,\sigma)$9 for unit variance of pre-ReLU Gaussian, (ii) $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$0 for unit variance of RGG; (i) is recommended due to simplicity and empirical parity on performance frontiers.
• Number of slices $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$1: Effective performance with $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$2–$\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$3 largely independent of feature dimension.
• Trade-off: Performance remains stable up to approximately 95% zeros; higher sparsity induces sharp accuracy drops. Optimal Pareto fronts are attained with $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$4 and moderate $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$5 or $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$6.

6. Comparison to Gaussian-Based Regularizations

Prior approaches, notably LeJEPA (SIGReg), regularize toward isotropic Gaussian targets via similar sliced-Wasserstein losses, enforcing only $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$7 constraints and producing dense, non-negative activations without control over $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$8 sparsity.

RDMReg strictly generalizes these methods:

• If $\Phi_{\mathrm{GN}_p(0,1)}(-\mu/\sigma)$9 and rectification is absent, the method recovers the Gaussian case.
• Adding ReLU and RGG targets introduces an explicit $-\mu/\sigma$0 handle, preserving maximum entropy in the nonzero coordinates.

Empirical results demonstrate RDMReg learns representations with 5–10% density while matching or exceeding accuracy of dense baselines on linear-probe and transfer tasks, and yields more factorial codes with lower nHSIC. Qualitative probes indicate preservation of semantic content even at extreme sparsity levels.

7. Recommendations and Practical Use

Recommended settings for RDMReg in practice:

• Choose $-\mu/\sigma$1 (rectified Laplace) or intermediate $-\mu/\sigma$2 for moderate sparsity.
• Adjust $-\mu/\sigma$3 to control expected fraction of zeros (50–90%).
• Use $-\mu/\sigma$4 for simplicity.
• Set the number of projections $-\mu/\sigma$5–$-\mu/\sigma$6 for $-\mu/\sigma$7.
• Balance invariance and RDMReg losses with $-\mu/\sigma$8–$-\mu/\sigma$9.

RDMReg can be integrated as a “plug-in” into any JEPA variant, enabling explicit, controlled transitions between dense maximum-entropy and sparse maximum-entropy codes while preserving high information capacity under $(0,\infty)$0 constraints (Kuang et al., 1 Feb 2026).