Text-Conditional JEPA for Learning Semantically Rich Visual Representations

Image-based Joint-Embedding Predictive Architecture (I-JEPA) offers a promising approach to visual self-supervised learning through masked feature prediction. However with the inherent visual uncertainty at masked positions, feature prediction remains challenging and may fail to learn semantic representations. In this work, we propose Text-Conditional JEPA (TC-JEPA) that uses image captions to reduce the prediction uncertainty. Specifically, we modulate the predicted patch features using a fine-grained text conditioner that computes sparse cross-attention over input text tokens. With such conditioning, patch features become predictable as a function of text, thus are more semantically meaningful. We show TC-JEPA improves downstream performance and training stability, with promising scaling properties. TC-JEPA also offers a new vision-language pretraining paradigm based on feature prediction only, outperforming contrastive methods on diverse tasks, especially those requiring fine-grained visual understanding and reasoning.

Modern AI systems can learn to understand images without human labels by predicting missing parts of an image from the surrounding visual context. However, this task is difficult because many different image contents may plausibly fit a masked region, making it hard for the model to learn meaningful visual concepts. In this work, we introduce TC-JEPA, a method that uses image captions to guide the prediction process and reduce this uncertainty. The model selectively connects words in the caption to specific image regions, helping it learn more semantically meaningful visual features. As a result, TC-JEPA improves both training stability and performance on a range of downstream tasks. The method also scales effectively to larger models and datasets. More broadly, this work presents a new way to train vision-language AI systems using feature prediction alone, rather than the contrastive learning methods commonly used today.