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

From Alignment to Prediction: A Study of Self-Supervised Learning and Predictive Representation Learning
Apr 15, 2026Mintu Dutta, Ritesh Vyas, Mohendra RoySelf-Supervised Learning Predictive Representation Learning
Department of Information and Communication Technology, School of Technology, Pandit Deendayal Energy University, Gandhinagar, Gujarat 382007, India
UR-JEPA: Uniform Rectifiability as a Regularizer for Joint-Embedding Predictive Architectures
May 31, 2026Triet M. LeJoint-Embedding Predictive Architecture Geometric Representation Learning
Drive vs. Decay: On the Training Dynamics of Joint-Embedding Predictive Architectures
Oct 1, 2026José Lucas De Melo Costa, Seong Woo Ahn, Fabrice Popineau +2Joint-Embedding Predictive Architecture Representation Collapse
Université Paris-Saclay, CNRS, CentraleSupélec, LISN, Orsay, France

Apr 15, 2026J/K move · Enter open · S save

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

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

Open paper AskPDF arXiv

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.

Abstract

Explore similar work

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

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

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