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

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, consistently improving over T-JEPA while remaining competitive with strong raw-feature baselines.

• The paper derives Var-JEPA as a generative JEPA trained with one ELBO, treating the predictor as a learned conditional prior and replacing heuristic anti-collapse methods with principled KL regularization.
• Experiments show the full ELBO prevents representation collapse and matches explicit SIGReg regularization in simulations, while reconstruction losses preserve predictive information in the learned latents.
• Var-T-JEPA improves stability over T-JEPA on tabular benchmarks and supports selective prediction, raising Adult accuracy from 0.852 to 0.921 when the most uncertain half of samples is discarded.

Overview and motivation

The Joint-Embedding Predictive Architecture (JEPA) has been positioned, since its introduction by LeCun, as a deliberately non-generative alternative to likelihood-based self-supervised learning: it predicts target embeddings from context embeddings in latent space rather than reconstructing observations. This paper by Gögl and Yau argues that this separation is "largely rhetorical rather than structural." The central claim is that the canonical JEPA design—coupled encoders with a context-to-target predictor—mirrors the variational posteriors and learned conditional priors of a coupled variational autoencoder (VAE), and that standard JEPA is best understood as a deterministic specialization of such a model in which anti-collapse regularization is supplied implicitly through architectural heuristics (EMA targets, auxiliary costs) rather than an explicit likelihood.

From this reinterpretation the authors derive Var-JEPA, a generative JEPA trained on a single Evidence Lower Bound (ELBO). The formulation yields representations without ad-hoc anti-collapse regularizers and provides principled per-sample uncertainty estimates in latent space. The framework is instantiated for heterogeneous tabular data as Var-T-JEPA and validated in a controlled simulation study and across seven downstream datasets.

The variational perspective on JEPA

The authors reverse-engineer a probabilistic model whose ELBO recovers the JEPA predictor as a learned latent-space conditional prior. The generative process is a directed acyclic graph over context observation $x$, context latent $s_x$, auxiliary latent $z$, target latent $s_y$, and target observation $y$, with joint factorization:

$$p_\theta(x,y,s_x,z,s_y) = p(s_x)\,p(z)\,p_\theta(x\mid s_x)\,p_\theta(s_y\mid s_x,z)\,p_\theta(y\mid s_y),$$

where $p(s_x)=p(z)=\mathcal{N}(0,I)$ and all conditionals are Gaussian networks. The amortized posterior factorizes as $$q_\phi(s_x\mid x)\,q_\phi(z\mid s_x)\,q_\phi(s_y\mid s_x,z,y)$$; notably, both $s_x$ and $z$ are inferred without access to the target, preserving the JEPA principle of preventing target information leakage into the predictive pathway.

Maximizing the resulting ELBO decomposes into five interpretable loss terms, each with its own weight $s_x$0: context reconstruction ($s_x$1), target generation ($s_x$2), KL regularization of $s_x$3 and $s_x$4 toward $s_x$5, and a KL term on $s_x$6 against the learned conditional prior $s_x$7. The last term combines a prediction component (how well $s_x$8 and $s_x$9 forecast $z$0) with an entropy component that keeps the target posterior sufficiently uncertain—the paper's principled substitute for collapse-prevention heuristics. Under this reading, the JEPA predictor is the conditional prior $z$1, and deterministic JEPA corresponds to fixing posterior covariances and dropping decoders while imposing distributional control indirectly via EMA or surrogate costs.

The paper also connects the framework to LeJEPA's SIGReg regularization (Balestriero et al., 11 Nov 2025). Via the standard "ELBO surgery" decomposition [Hoffman & Johnson, 2016], the per-sample KL terms for latents with fixed priors split into an aggregated-posterior mismatch plus an information-bottleneck term:

$z$2

Because the target-latent term regularizes toward a learned conditional prior rather than $z$3, it does not admit an analogous decomposition—an asymmetry the authors examine empirically rather than resolve theoretically.

Simulation study

A controlled synthetic setting with a known ground-truth generative process (mixture-structured $z$4, correlated $z$5, auxiliary factor $z$6, nonlinear frozen maps $z$7, Gaussian observation noise) isolates the effects of per-sample KL versus aggregated-distribution (SIGReg) regularization across ten ablation variants (A–J).

Three findings stand out. First, the KL terms alone achieve distributional properties for $z$8 comparable to explicit SIGReg: variant (A), the full ELBO, attains aggregated KL of 0.113 ± 0.037 and linear-probe accuracy of 0.996 ± 0.002, essentially matching variants augmented with SIGReg. Second, removing KL terms causes severe distributional collapse—variant (I), with all KL terms removed, shows aggregated KL of 10.465 ± 2.207 for $z$9 and 20.276 ± 11.331 for $s_y$0—while adding SIGReg to the no-KL objective (variant J) only partially compensates (aggregated KL of 2.115 and 2.215 respectively), leaving the full ELBO optimal. Third, removing reconstruction terms entirely (variant G) collapses probe accuracy to near chance (0.571 for $s_y$1), demonstrating that reconstruction is what forces latents to encode predictive information; SIGReg alone cannot restore this (variant H reaches only 0.834). For $s_y$2, the aggregated distribution deviates from isotropy under the full ELBO, which the authors argue is the theoretically correct behavior given regularization toward a learned conditional prior rather than $s_y$3.

Downstream evaluation on tabular data

Var-T-JEPA adapts the framework to tabular data following T-JEPA's feature-level masking strategy [thimonier2025_tjepa]: features are tokenized into transformer sequences, Gaussian posteriors are inferred for context, auxiliary, and target latents, and the predictor operates purely in latent space with learnable mask tokens. Decoders use feature-specific heads with Gaussian likelihoods for numerical and categorical likelihoods for categorical features. Training uses separate linear KL-annealing schedules to avoid posterior collapse.

Downstream evaluation covers five real-world tabular datasets (Adult, Covertype, Electricity, Credit Card Default, Bank Marketing), MNIST treated as tabular features with controlled input corruption, and a synthetic dataset with prototype-mixing ambiguity. Six predictor families (MLP, DCNv2, ResNet, AutoInt, FT-Transformer, XGBoost) are trained either on raw features or on learned embeddings.

The results support two claims. First, Var-T-JEPA consistently improves over T-JEPA and avoids its failure modes: T-JEPA exhibits representation collapse on several dataset–predictor combinations (e.g., accuracy of 0.408 ± 0.231 on Covertype with MLP probes, 0.113 on MNIST with MLP probes, versus 0.679 and 0.822 for Var-T-JEPA), whereas Var-T-JEPA remains stable and competitive with raw-feature baselines. Second, the learned posterior uncertainties enable effective selective prediction: discarding the most uncertain 50% of test samples raises Adult accuracy from 0.852 to 0.921 (MLP probe) and Bank Marketing from 0.900 to 0.921, a consistent coverage–accuracy trade-off across all predictor families. On the semi-synthetic datasets, latent uncertainty correlates positively with the ground-truth ambiguity score and discriminates high-ambiguity samples with ROC curves substantially above chance, indicating the uncertainty signal is calibrated to known corruption structure rather than an artifact. Sensitivity analyses over loss weights, architecture width/depth, learning rate, and batch size show downstream performance is broadly robust, though very deep encoders (16 layers) degrade accuracy.

Limitations and open questions

The paper is explicit about scope. All empirical validation is on tabular and small-scale semi-synthetic data; scaling to vision and video is proposed but not demonstrated, and the appendix outlines the transfer to ViT-style patch tokenization only at the level of a design sketch. The evaluation intentionally validates the conceptual contribution rather than benchmarking comprehensively against the broader self-supervised learning landscape. Two further caveats bear directly on the headline claims. First, the claim that the ELBO "naturally prevents collapse" rests on simulation evidence and tabular experiments; whether per-sample KL regularization suffices at scale, where SIGReg was designed for computational efficiency, remains untested. Second, the theoretical link between Var-JEPA and LeJEPA is incomplete: the target-latent KL does not decompose into an aggregated-posterior term, so the relationship between conditional-prior regularization and minimax-optimal isotropic embeddings is asserted diagnostically rather than proven. Finally, downstream embedding generation conditions the target posterior on the full feature vector $s_y$4—a deviation from the general test-time procedure where targets may be unavailable—so the reported accuracies measure the representation-learning regime, not conditional generation quality.

Conclusion

This paper reframes JEPA not as a departure from probabilistic generative modeling but as a deterministic point within the latent-variable model design space, and derives Var-JEPA as the explicit variational counterpart in which the predictor is a learned conditional prior inside a single ELBO. Empirically, the formulation removes the need for heuristic anti-collapse mechanisms, matches explicit distributional regularizers on fixed-prior latents, and delivers usable per-sample uncertainty for selective prediction on tabular tasks. The principal open questions are whether these properties persist at vision/video scale and how the conditional-prior regularization of the target latent relates formally to aggregated-distribution optimality.