Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

Enhancing JEPAs with Spatial Conditioning: Robust and Efficient Representation Learning

Etai Littwin, Vimal Thilak, Anand Gopalakrishnan

TL;DR

This work condition the target encoder and context encoder modules in IJEPA with positions of context and target windows respectively, and shows performance gains on several image classification benchmark datasets, improved robustness to context window size and sample-efficiency during pretraining.

Abstract

Image-based Joint-Embedding Predictive Architecture (IJEPA) offers an attractive alternative to Masked Autoencoder (MAE) for representation learning using the Masked Image Modeling framework. IJEPA drives representations to capture useful semantic information by predicting in latent rather than input space. However, IJEPA relies on carefully designed context and target windows to avoid representational collapse. The encoder modules in IJEPA cannot adaptively modulate the type of predicted and/or target features based on the feasibility of the masked prediction task as they are not given sufficient information of both context and targets. Based on the intuition that in natural images, information has a strong spatial bias with spatially local regions being highly predictive of one another compared to distant ones. We condition the target encoder and context encoder modules in IJEPA with positions of context and target windows respectively. Our "conditional" encoders show performance gains on several image classification benchmark datasets, improved robustness to context window size and sample-efficiency during pretraining.

Enhancing JEPAs with Spatial Conditioning: Robust and Efficient Representation Learning

TL;DR

This work condition the target encoder and context encoder modules in IJEPA with positions of context and target windows respectively, and shows performance gains on several image classification benchmark datasets, improved robustness to context window size and sample-efficiency during pretraining.

Abstract

Image-based Joint-Embedding Predictive Architecture (IJEPA) offers an attractive alternative to Masked Autoencoder (MAE) for representation learning using the Masked Image Modeling framework. IJEPA drives representations to capture useful semantic information by predicting in latent rather than input space. However, IJEPA relies on carefully designed context and target windows to avoid representational collapse. The encoder modules in IJEPA cannot adaptively modulate the type of predicted and/or target features based on the feasibility of the masked prediction task as they are not given sufficient information of both context and targets. Based on the intuition that in natural images, information has a strong spatial bias with spatially local regions being highly predictive of one another compared to distant ones. We condition the target encoder and context encoder modules in IJEPA with positions of context and target windows respectively. Our "conditional" encoders show performance gains on several image classification benchmark datasets, improved robustness to context window size and sample-efficiency during pretraining.

Paper Structure

This paper contains 13 sections, 4 figures, 8 tables.

Table of Contents

  1. Introduction
  2. Method
  3. IJEPA
  4. EC-IJEPA
  5. Results
  6. Conclusion
  7. Experimental Details
  8. Architecture Details.
  9. Pretraining Details.
  10. Evaluation on ImageNet-1k
  11. Evaluation on out-of-distribution (OOD) datasets
  12. Additional Results
  13. Average Pooling Ablation.

Figures (4)

  • Figure 1: Conditioning the Context and Target Encoders in IJEPA with positions of the target (blue box) and context windows (red box) respectively. Patches marked with X indicate positional information while those with solid color fill indicate feature information is extracted at those locations.
  • Figure 2: Ablation on ranges of context window scale used for pretraining.
  • Figure 3: Classification performance on ImageNet-1k measured during pretraining cycle in IJEPA (blue) and EC-IJEPA (orange) at two encoder sizes (left: ViT-L/16 and right: ViT-H/14).
  • Figure 4: Linear probing accuracy on Imagenet-1k dataset w.r.t kernel size and stride.