K-JEPA: From Point to Distributional Prediction by Kernelized Joint-Embedding Predictive Architectures

Joint-Embedding Predictive Architectures (JEPA) learn representations by predicting latent target features from a visible context. Although effective, standard JEPA objectives treat prediction as a point-estimation problem: the population-optimal predictor recovers a single point summary of the conditional target distribution. When a visible context $x_c$ admits multiple plausible masked or future targets, this point estimate retains insufficient distributional information to identify the underlying conditional target distribution. We propose K-JEPA, a novel, simple, kernelized extension that applies a characteristic-kernel transformation to JEPA latent targets while preserving the architecture and prediction setup. Rather than predicting only the latent target $z_t$, K-JEPA predicts its kernel representation $\varphi (z_t)$. We prove this changes the population target from a point summary of the latent target distribution to its conditional kernel mean embedding, which uniquely identifies the corresponding conditional target distribution when the kernel is characteristic. Empirically, we evaluate this distributional behavior in controlled settings and introduce a systematic target-retrieval comparison for direct evaluation of JEPA target prediction across heterogeneous formulations and modalities. We complement this evaluation with probabilistic scoring, calibration, selective prediction, prediction-set efficiency, and downstream transfer analyses. Our results characterize when K-JEPA improves target prediction and how these changes relate to downstream performance and reliability. Code is made available to the reviewers during the discussion process.