Probabilistic image-based joint embedding predictive architecture

In the field of self-supervised learning, the Image-based Joint-Embedding Predictive Architecture (I-JEPA) has shown impressive results in capturing semantic information from images without the need for strong augmentations between views. I-JEPA learns to predict the representations of multiple masked target blocks within the same image, using the representation of a single unmasked context block. To achieve this, I-JEPA reconstructs the masked patches in the representation space rather than at the pixel level, as done in previous works. However, accurately reconstructing abstract representations of patches is a challenging task, as pointwise matching focuses on exact reconstruction of individual target representations rather than the preservation of their local geometric structure. In this work, we present a probabilistic matching approach that can be applied to I-JEPA’s learning process to encourage the output of semantically richer representations. Specifically, we propose a probabilistic formulation of I-JEPA, which involves efficiently modeling and matching the probability distributions of the missing patches between the target block representations and the network’s predictions based solely on the context block. Through rigorous evaluations, we demonstrate that the proposed methodology provides consistent improvements over I-JEPA in downstream classification and transfer learning tasks.