For self supervised learning, we use the same setup as [1]} and build on their codebasehttps://github.com/Sara-Ahmed/SiT. We evaluate the proposed model in two settings. In the first setting we pretrain the model and evaluate the learned representation using linear probing on same dataset. We consider both the CIFAR10 dataset and the STL10 dataset for this setting. In the second setting we measure transfer across datasets. Here, we pretrain the model on STL10 dataset and then evaluate the learned representation using linear probing on the CIFAR10 dataset. We use images of size \(224 \times 224\) . We use the same augmentation policy as [1]}. For the baseline, we use a 12 layered Vision Transformer with 12 heads and embedding dimension 768. We use an FFN dimension of 3072. For the proposed model, we use a 12-layered model with \(R = 2\) . We use 5 temporal latent bottleneck state vectors. We use a patch size of 16 for both the models. For SiT + TLB, we use a chunk length of 20. Similar to [1]}, we use the Adam optimizer with batch size 72, momentum 0.9, weight decay 0.05 and learning rate of 0.0005. We train the models for 5 days on 1 RTX8000 GPU completing 400 pretraining epochs. For models pretrained on the CIFAR10 dataset, we perform linear probing training for 500 epochs. For models pretrained on the STL10 dataset, we perform linear probing training for 300 epochs.
