This repository is the official implementation of StoP.

Given a partial image of a dog, can you precisely determine the location of its tail? Existing Masked Image Modeling (MIM) models like MAE and I-JEPA predict tokens deterministically and do not model location uncertainties (a), we propose to predict the target (masked tokens) in stochastic positions (StoP) which prevents overfitting to locations features. StoP leads to improved MIM performance on downstream tasks, including linear probing on ImageNet (b).

Please follow the installation instruction from the I-JEPA repo.

The commands for pretraining I-JEPA StoP. In the original setting, all models were trained using 4 V100 GPU nodes. ViT-B/L were trained with float32 while ViT-H is trained on half precision (float16).

Training command:

torchrun --nnodes=4 --nproc-per-node=8 --node_rank=<node_rank 0-3> --master_addr=<master_addr> --master_port=<master_port> --backend=nccl main.py --fname configs/pretrain/vit-b16.yaml

• Prior to running, set the path to imagenet (image_folder, root_folder) in the config file.
• to train ViT-L/ViT-H change vit-b16 to vit-l16/vit-h16.
• We --node_rank as 0 on the first node. On other nodes, run the same command with --node_rank=1,...,3 respectively. --master_addr is set as the ip of the node 0.

python logistic_eval.py \
  --subset-path imagenet_subsets1/1percent.txt \
  --root-path /path/to/datasets --image-folder imagenet_folder/ \
  --device cuda:0 \
  --pretrained /path/to/checkpoint/folder \
  --fname checkpoint_name.pth.tar \
  --model-name deit_base \
  --patch-size 16 \
  --penalty l2 \
  --lambd 0.0025

To perform linear probing on ImageNet, you can follow instruction from VISSL. Alternatiely, we provide a bash script to convert checkpoints to vissl format and launch experiments on 8 V100 machines each with 8 gpus on SLURM:

bash bash/in1k_eval_vissl.sh <output_dir> <checkpoint_path> <dataset_root> <arch>

• arch is vitb, vitl or vitb.
• checkpoint path is the full path to the checkpoint saved durning training
• dataset_root is the path to imagenet folder.

• I-JEPA was trained using bfloat16 which is supported in new nvidia gpus (e.g, A100, H100) and beneficial to stabilize training. Here we used older gpus hence used float32 for the smaller models ViT-B/L and float16 for ViT-H.
• For ViT-H training, we deviated from the cosine LR schedule after 250 epochs and continued training with fixed low learning rate to push performance.
• StoP relies on a previous internal implementation of I-JEPA which utilized additional image augmentations compared to the official repo. The main difference is the use of hflip, gaussian blur and random grayscale during training. In the ablation experiments (Figures 3-4, Tables 5-6 & 8) we compare different positional embeddings using this same set of augmentations.

The codebase relies on the implementation of I-JEPA.

If you found this code helpful, feel free to cite our work:

@inproceedings{barstochastic,
  title={Stochastic positional embeddings improve masked image modeling},
  author={Bar, Amir and Bordes, Florian and Shocher, Assaf and Assran, Mido and Vincent, Pascal and Ballas, Nicolas and Darrell, Trevor and Globerson, Amir and LeCun, Yann},
  booktitle={Forty-first International Conference on Machine Learning}
}```