Var-JEPA: A Variational Formulation of the Joint-Embedding Predictive Architecture - Bridging Predictive and Generative Self-Supervised Learning

Moritz Gögl; Christopher Yau

cs.LGMar 20, 2026

Var-JEPA: A Variational Formulation of the Joint-Embedding Predictive Architecture - Bridging Predictive and Generative Self-Supervised Learning

Authors: Moritz Gögl, Christopher Yau

Organizations: University of Oxford, UK · Health Data Research UK

Abstract

The Joint-Embedding Predictive Architecture (JEPA) is often seen as a non-generative alternative to likelihood-based self-supervised learning, emphasizing prediction in representation space rather than reconstruction in observation space. We argue that the resulting separation from probabilistic generative modeling is largely rhetorical rather than structural: the canonical JEPA design (coupled encoders with a context-to-target predictor) mirrors the variational posteriors and learned conditional priors obtained when variational inference is applied to a particular class of coupled latent-variable models, and standard JEPA can be viewed as a deterministic specialization in which regularization is imposed via architectural and training heuristics rather than an explicit likelihood. Building on this view, we derive the Variational JEPA (Var-JEPA), which makes the latent generative structure explicit by optimizing a single Evidence Lower Bound (ELBO). This yields meaningful representations without ad-hoc anti-collapse regularizers and allows principled uncertainty quantification in the latent space. We instantiate the framework for tabular data (Var-T-JEPA) and achieve strong representation learning and downstream performance, improving over T-JEPA across real-world tabular benchmarks while remaining competitive with strong raw-feature baselines.

Figures & tables

Figure 1: Comparison between JEPA (left) and Var-JEPA (right). Var-JEPA extends JEPA by replacing deterministic encoders and predictor with probabilistic distributions, and adding generative networks (decoders) to enable variational learning under a unified ELBO.

Exp.	Objective	$\text{Acc}(s)$	$\mathrm{KL}_{\text{agg}}\!\left(q(s)\\|\mathcal{N}(0,I)\right)$	$\mathrm{SIGReg}\text{-}\mathrm{MSE}(s)$	$\\|\mathrm{Cov}(s)-I\\|_{F}$	$\\|\mathbb{E}[s]\\|_{2}$
Context latent diagnostics ( $s_{x}$ )
(A)	ELBO (Var-JEPA)	0.996 $\pm$ 0.002	0.113 $\pm$ 0.037	4.0e-4 $\pm$ 1.0e-4	0.649 $\pm$ 0.127	0.082 $\pm$ 0.009
(B)	ELBO+SIGReg ( $\lambda_{s_{x}}\!=\!10,\lambda_{s_{y}}\!=\!0$ )	0.996 $\pm$ 0.002	0.108 $\pm$ 0.032	4.0e-4 $\pm$ 1.0e-4	0.639 $\pm$ 0.113	0.077 $\pm$ 0.018
(C)	ELBO+SIGReg ( $\lambda_{s_{x}}\!=\!0,\lambda_{s_{y}}\!=\!10$ )	0.996 $\pm$ 0.002	0.110 $\pm$ 0.033	4.0e-4 $\pm$ 1.0e-4	0.647 $\pm$ 0.118	0.077 $\pm$ 0.018
(D)	ELBO+SIGReg	0.996 $\pm$ 0.002	0.107 $\pm$ 0.030	3.9e-4 $\pm$ 1.0e-4	0.634 $\pm$ 0.108	0.078 $\pm$ 0.018
(E)	ELBO ( $\alpha^{\text{KL}}_{s_{x}}\!=\!0$ )	0.996 $\pm$ 0.002	8.374 $\pm$ 1.944	5.1e-2 $\pm$ 1.9e-2	7.197 $\pm$ 1.522	2.550 $\pm$ 0.607

Table 1: Simulation study: Var-JEPA and SIGReg ablations (mean

\pm

std over 5 runs). We report linear-probe accuracy (predicting the mixture component) and LeJEPA-link diagnostics for the aggregated latent distributions of

s_{x}

and

s_{y}

Figure 2: Epoch-wise distribution diagnostics for selected experiments. We show how aggregated KL divergences, SIGReg-MSE, isotropy metrics, and conditional-prior coupling evolve during training.

Method	AD $\uparrow$	CO $\uparrow$	EL $\uparrow$	CC $\uparrow$	BM $\uparrow$	MNIST $\uparrow$	SIM $\uparrow$
MLP	0.849 $\pm$ 1.4e -3	0.750 $\pm$ 3.7e -3	0.781 $\pm$ 3.6e -3	0.816 $\pm$ 2.5e -3	0.898 $\pm$ 1.4e -3	0.822 $\pm$ 1.1e -2	0.823 $\pm$ 7.5e -3
+T-JEPA	0.849 $\pm$ 4.4e -3	0.408 $\pm$ 2.3e -1	0.627 $\pm$ 4.1e -2	0.774 $\pm$ 0e -0	0.616 $\pm$ 3.7e -1	0.113 $\pm$ 6.2e -3	0.692 $\pm$ 3.5e -2
+Var-T-JEPA	0.852 $\pm$ 2.1e -3	0.679 $\pm$ 2.9e -2	0.790 $\pm$ 3.3e -3	0.818 $\pm$ 1.5e -3	0.900 $\pm$ 2.0e -3	0.822 $\pm$ 1.8e -2	0.695 $\pm$ 1.4e -2
+Var-T-JEPA (10%)	0.865 $\pm$ 2.0e -3	0.696 $\pm$ 2.8e -2	0.793 $\pm$ 2.9e -3	0.827 $\pm$ 1.4e -3	0.905 $\pm$ 2.2e -3	0.838 $\pm$ 1.6e -2	0.705 $\pm$ 1.3e -2
+Var-T-JEPA (20%)	0.883 $\pm$ 1.7e -3	0.697 $\pm$ 2.7e -2	0.795 $\pm$ 2.8e -3	0.835 $\pm$ 1.2e -3	0.908 $\pm$ 2.5e -3	0.856 $\pm$ 1.9e -2	0.706 $\pm$ 1.5e -2
+Var-T-JEPA (50%)	0.921 $\pm$ 2.1e -3	0.706 $\pm$ 1.4e -2	0.804 $\pm$ 3.9e -3	0.841 $\pm$ 7.0e -4	0.921 $\pm$ 1.4e -3	0.901 $\pm$ 2.1e -2	0.729 $\pm$ 1.5e -2

Table 2: Downstream performance comparison across tabular datasets. Results (mean

\pm

std over 5 downstream model runs; XGBoost single deterministic run) show test accuracy. AD=Adult, CO=Covertype, EL=Electricity, CC=Credit Card Default, BM=Bank Marketing. Light blue rows indicate selective evaluation (reduced coverage).

Figure 3: Uncertainty quantification on the MNIST (A) and SIM (B) datasets. Left: risk–coverage curve induced by abstaining on samples with highest latent uncertainty. Middle: standardized latent uncertainty versus simulated ambiguity. Right: ROC curve for detecting high-ambiguity samples from latent uncertainty.

Appendix figures & tables9 assets

Supplementary material from the paper’s appendix.

Appendix

Figure 4: Architectural overview of Var-T-JEPA for tabular data. The model implements the theoretical framework through specialized components: (1) the context encoder processes masked tabular features to produce variational context representations $q_{\phi}(s_{x}^{(j)}|x)$ ; (2) the auxiliary encoder infers predictive latents $q_{\phi}(z|s_{x})$ ; (3) the target posterior prepends context and auxiliary latents as special tokens before feature processing $q_{\phi}(s_{y_{k}}^{(j)}|s_{x},z,w;m_{\text{trg}}^{(k)})$ ; (4) the predictor generates target representations $p_{\theta}(s_{y_{k}}^{(j)}|s_{x},z;m_{\text{trg}}^{(k)})$ , and (5) the decoders reconstruct original tabular features from latent representations.

Method	AD $\uparrow$	CO $\uparrow$	EL $\uparrow$	CC $\uparrow$	BM $\uparrow$	MNIST $\uparrow$	SIM $\uparrow$
MLP	0.845 $\pm$ 3.7e -3	0.747 $\pm$ 5.0e -3	0.780 $\pm$ 3.5e -3	0.792 $\pm$ 6.9e -3	0.887 $\pm$ 4.8e -3	0.822 $\pm$ 1.1e -2	0.821 $\pm$ 9.0e -3
+T-JEPA	0.842 $\pm$ 9.0e -3	0.274 $\pm$ 1.5e -1	0.545 $\pm$ 7.4e -2	0.675 $\pm$ 0e -0	0.583 $\pm$ 3.9e -1	0.023 $\pm$ 2.3e -3	0.662 $\pm$ 3.3e -2
+Var-T-JEPA	0.845 $\pm$ 3.3e -3	0.657 $\pm$ 3.8e -2	0.789 $\pm$ 2.2e -3	0.797 $\pm$ 4.1e -3	0.889 $\pm$ 3.1e -3	0.823 $\pm$ 1.7e -2	0.659 $\pm$ 3.3e -2
+Var-T-JEPA (10%)	0.859 $\pm$ 3.0e -3	0.678 $\pm$ 3.4e -2	0.792 $\pm$ 2.0e -3	0.792 $\pm$ 5.6e -3	0.893 $\pm$ 3.5e -3	0.840 $\pm$ 1.5e -2	0.671 $\pm$ 3.1e -2
+Var-T-JEPA (20%)	0.877 $\pm$ 2.4e -3	0.680 $\pm$ 3.0e -2	0.794 $\pm$ 2.3e -3	0.792 $\pm$ 4.3e -3	0.895 $\pm$ 4.0e -3	0.857 $\pm$ 1.8e -2	0.671 $\pm$ 3.2e -2
+Var-T-JEPA (50%)	0.915 $\pm$ 2.6e -3	0.692 $\pm$ 1.7e -2	0.803 $\pm$ 3.4e -3	0.780 $\pm$ 1.9e -3	0.907 $\pm$ 3.9e -3	0.901 $\pm$ 2.0e -2	0.691 $\pm$ 3.7e -2

Appendix

Table 3: Downstream performance comparison across tabular datasets using macro F1-score. Results (mean

\pm

std over 5 downstream model runs; XGBoost single deterministic run) show test macro F1-score. AD=Adult, CO=Covertype, EL=Electricity, CC=Credit Card, BM=Bank Marketing. Light blue rows indicate selective evaluation (reduced coverage).

	Shallow	Low pred.-FF	Narrow	High pred.-FF	Deep	Low LR	Batch size 256
$d$ (hidden dim.)	64	64	32	64	64	64	64
$L$ (layers)	4	8	8	8	16	8	8
$\mathrm{ff}_{p}$ (pred.)	256	128	128	512	256	256	256
$\eta$ (LR)	$10^{-3}$	$10^{-3}$	$10^{-3}$	$10^{-3}$	$10^{-3}$	$5\times 10^{-4}$	$10^{-3}$
$B$ (batch)	512	512	512	512	512	512	256
MLP + Var-T-JEPA	0.850 $\pm$ 2.1e -3	0.852 $\pm$ 2.1e -3	0.850 $\pm$ 1.6e -3	0.852 $\pm$ 2.1e -3	0.832 $\pm$ 1.7e -3	0.849 $\pm$ 3.2e -3	0.851 $\pm$ 2.7e -3

Appendix

Table 5: Adult sensitivity analysis for Var-T-JEPA architecture and optimization settings. Columns correspond to model/optimizer ablations varying

(d,L,\mathrm{ff}_{p},\eta,B)

. We report downstream test accuracy (mean

\pm

std over 5 downstream runs) for MLP and ResNet probes, including selective evaluation (light blue rows).

Dataset	Method	Spearman $\uparrow$	AUC $\uparrow$	AURC $\downarrow$
SIM	Var-T-JEPA (latent uncertainty)	0.640	0.805	0.209
	Probe entropy	0.339	0.692	0.272
	Probe MSP	0.313	0.668	0.267
	MC dropout	0.560	0.599	0.013
	Deep ensemble	0.586	0.612	0.011
MNIST	Var-T-JEPA (latent uncertainty)	0.949	0.986	0.123

Appendix

Table 6: Uncertainty quantification: Var-T-JEPA’s native latent uncertainty versus four post-hoc baselines on the same frozen embeddings, for the controlled SIM and MNIST datasets. Metrics: Spearman correlation with the ground-truth simulated ambiguity score, AUC for detecting high-ambiguity samples, and area under the risk–coverage curve (AURC).

Dataset	Samples	Features	(num/cat)	Classes	Task	Domain
Adult (AD) ( Kohavi, 1996 )	48,842	14	(6/8)	2	Classification	Census
Covertype (CO) ( Blackard, 1998 )	10,000	54	(54/0)	7	Classification	Forestry
Electricity (EL) ( Harries, 1999 )	45,312	8	(8/0)	2	Classification	Time series
Credit Card Default (CC) ( Yeh, 2009 )	30,000	23	(23/0)	2	Classification	Finance
Bank Marketing (BM) ( Moro et al., 2014 )	45,211	16	(16/0)	2	Classification	Marketing
MNIST ( LeCun et al., 1998 )	10,000	784	(784/0)	10	Classification	Vision

Appendix

Table 7: Dataset characteristics used in experimental evaluation.

Figure 5: Visualization of MNIST image corruption under different values of the ambiguity score $u_{i}$ .

Figure 6: Architectural overview of the T-JEPA pipeline by Thimonier et al. ( 2025 ) extended by learnable auxiliary tokens.

Dataset	batch	lr	warm	ctx	ctxM	trg	trgM	K	d	L	H	ff	drop	Lp	Hp	ffp	dropp	$\alpha^{\text{KL}}_{s_{x}}$	$\alpha^{\text{KL}}_{z}$	$\alpha^{\text{KL}}_{s_{y}}$	$T_{s_{x}}^{\text{anneal}}$	$T_{z}^{\text{anneal}}$	$T_{s_{y}}^{\text{anneal}}$	$\alpha^{\text{rec}}$	$\alpha^{\text{gen}}$
AD	512	0.001	10	0.1	0.3	0.1	0.6	4	64	8	4	256	0.001	16	4	256	0.002	1e-04	1e-06	1e-05	15	15	15	0.1	1
CO	512	5e-04	12	0.15	0.35	0.15	0.6	4	64	8	8	64	0.0015	16	4	256	0.002	1e-06	1e-05	1e-06	60	60	20	0.001	0.5
EL	512	5e-04	10	0.15	0.6	0.15	0.9	4	64	6	4	64	0.001	16	4	128	0.002	1e-06	1e-05	1e-06	40	40	15	0.001	0.25
CC	512	0.001	10	0.1	0.3	0.1	0.6	4	64	8	4	256	0.001	16	4	256	0.002	1e-04	1e-06	1e-05	15	15	15	0.1	1
BM	512	0.001	10	0.1	0.3	0.1	0.6	4	64	8	4	256	0.001	16	4	256	0.002	1e-04	1e-06	1e-05	15	15	15	0.1	1
MNIST	256	0.001	10	0.15	0.5	0.15	0.8	4	64	8	4	128	0.002	32	4	256	0.002	1e-06	1e-06	1e-06	100	100	100	0.001	0.1

Appendix

Table 8: Var-T-JEPA hyperparameters used for downstream experiments.

Dataset	batch	lr	warm	ctx	ctxM	trg	trgM	K	d	L	H	ff	drop	Lp	Hp	ffp	dropp	ema
AD	512	0.001	10	0.1	0.3	0.1	0.6	4	64	8	4	256	0.001	16	4	256	0.002	0.996
CO	512	5e-04	12	0.15	0.35	0.15	0.6	4	64	8	8	64	0.0015	16	4	256	0.002	0.996
EL	512	5e-04	10	0.15	0.6	0.15	0.9	4	64	6	4	64	0.001	16	4	128	0.002	0.996
CC	512	0.001	10	0.1	0.3	0.1	0.6	4	64	8	4	256	0.001	16	4	256	0.002	0.996
BM	512	0.001	10	0.1	0.3	0.1	0.6	4	64	8	4	256	0.001	16	4	256	0.002	0.996
MNIST	256	0.001	10	0.15	0.5	0.15	0.8	4	64	8	4	128	0.002	32	4	256	0.002	0.996

Appendix

Table 9: T-JEPA hyperparameters used for downstream experiments.

Explore similar work

Apr 15, 2026

From Alignment to Prediction: A Study of Self-Supervised Learning and Predictive Representation Learning

Self-supervised learning has emerged as a major technique for the task of learning from unlabeled data, where the current methods mostly revolve around alignment of representations and input recon struction. Although such approaches have demonstrated excellent performance in practice, their scope remains mostly confined to learning from observed data and does not provide much help in terms of a learning structure that is predictive of the data distribution. In this paper, we study some of the recent developments in the realm of self-supervised learning. We define a new category called Predictive Representation Learning (PRL), which revolves around the latent prediction of unobserved components of data based on the observation. We propose a common taxonomy that classifies PRL along with alignment and reconstruction-based learning approaches. Furthermore, we argue that Joint-Embedding Predictive Architecture(JEPA) can be considered as an exemplary member of this new paradigm. We further discuss theoretical perspectives and open challenges, highlighting predictive representation learning as a promising direction for future self-supervised learning research. In this study, we implemented Bootstrap Your Own Latent (BYOL), Masked Autoencoders (MAE), and Image-JEPA (I-JEPA) for comparative analysis. The results indicate that MAE achieves perfect similarity of 1.00, but exhibits relatively weak robustness of 0.55. In contrast, BYOL and I-JEPA attain accuracies of 0.98 and 0.95, with robustness scores of 0.75 and 0.78, respectively.

Mintu Dutta, Ritesh Vyas, Mohendra Roy

Department of Information and Communication Technology, School of Technology, Pandit Deendayal Energy University, Gandhinagar, Gujarat 382007, India

Self-Supervised Learning Predictive Representation Learning Visual Representation Learning Joint-Embedding Predictive Architecture

Ask PDF

May 31, 2026

UR-JEPA: Uniform Rectifiability as a Regularizer for Joint-Embedding Predictive Architectures

A central difficulty in training Joint-Embedding Predictive Architectures (JEPAs) is preventing representation collapse. LeJEPA addresses this by enforcing an isotropic Gaussian target on the embeddings via Sketched Isotropic Gaussian Regularization (SIGReg). This target is in tension with the manifold hypothesis, which expects embeddings to concentrate on a low-dimensional subset of the ambient space. We propose UR-JEPA, which targets a uniformly

n

-rectifiable measure of local tangent dimension

n

at small scales, realized through a Gaussian-kernel smoothed Carleson-type square function

\mathcal{L}^{\text{CGLT}}

, with a complementary Jones

β

-number formulation. On Inet10, UR-JEPA(

\mathcal{L}^{\text{CGLT}}

) attains

0.9141 \pm 0.0014

for a

+0.83

,pp gain over LeJEPA(

\mathcal{L}^{\text{SIGReg}}

) with

\sim 30\%

lower seed standard deviation; on matched-recipe Galaxy10~SDSS, a single-seed ImageNet-

100

run, and a

3

-seed EuroSAT remote-sensing run, the two methods lie in the same peak-accuracy band at convergence, with UR-JEPA retaining its lower-seed-variance signature. On EuroSAT the in-domain pair is competitive at

96.0

96.1\%

with large remote-sensing foundation-model transfer at a

25\times

smaller backbone. The distinction is geometric: direct visualization of the projector output distribution shows that on all four datasets UR--JEPA(

\mathcal{L}^{\text{CGLT}}

) produces a global PCA spectrum with a

4

5

order-of-magnitude drop at index

\sim 20

25

out of

D = 32

, while LeJEPA's spectrum is near-flat (top-to-bottom ratio at most

3.6

). Per-dimension marginals are simultaneously near-Gaussian for both methods (mean Shapiro-Wilk

W \in [0.992, 0.996]

) as a Diaconis-Freedman consequence. At matched accuracy the two regularizers therefore yield structurally distinct projected representations.

Triet M. Le

Joint-Embedding Predictive Architecture Geometric Representation Learning Self-Supervised Visual Representation Learning

Ask PDF

Oct 1, 2026

Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures

Joint-Embedding Predictive Architectures (JEPAs) are prone to representation collapse, typically mitigated through empirical heuristics. We develop an early-training stability theory that unifies these heuristics. Linearising the coupled JEPA gradient flow around the trivial fixed point reveals two competing effects: a driving force (

γ

) and a decay effect (

σ

). Under approximate spectral decoupling, a per-mode stability ratio

μ_i = γ_i / σ_i

factorises into independent data-side and predictor-side terms and the count of unstable modes tracks the rank of representations that can emerge. The framework predicts a phase boundary, which we confirm empirically across more than 800 Tabular-JEPA configurations. It also unifies predictor scaling, masking ratio, and EMA as distinct mechanisms for shifting

μ

. Guided by this analysis, we introduce ResidualPred, a transformer predictor whose attention is biased toward the identity at initialisation; it improves both effective rank and downstream accuracy on tabular benchmarks and in I-JEPA pretraining on CIFAR-10, CIFAR-100, STL-10, and ImageNet. Our framework connects empirical collapse-avoidance heuristics to an explicit dynamical picture, yielding theory-driven stabilizers. Code is available at https://github.com/jose-melo/drive-vs-decay.

José Lucas De Melo Costa, Seong Woo Ahn, Fabrice Popineau +2

Université Paris-Saclay, CNRS, CentraleSupélec, LISN, Orsay, France

Joint-Embedding Predictive Architecture Representation Collapse Self-Supervised Pre-Training

Ask PDF