Add config for ViT-XX-gap-norep

[juntang-zhuang]

Add configs for 'ViT-S_32-gap-norep', 'ViT-S_16-gap-norep', 'ViT-B_32-gap-norep', 'ViT-B_16-gap-norep', which are compatible with GSAM checkpoints.

ViT-S_32-gap-norep: https://colab.research.google.com/drive/1menAIRCDxwx6PfUKupFerKOP5Fzm1APp?usp=sharing
ViT-S_16-gap-norep: https://colab.research.google.com/drive/1sAMdbu4GIsMuNCRDyNGgMHxxyBEvXI1h?usp=sharing
ViT-B_32-gap-norep: https://colab.research.google.com/drive/12SGIVL2iigLbdGBQDTZS8aOAfLzwW0KV?usp=sharing
ViT-B_16-gap-norep: https://colab.research.google.com/drive/1sJ4_Rgs_7LK3hfyMnhlQMnnC_nOOwvBH?usp=sharing

[andsteing]

Thanks for updating the configs!

I pulled the following CIFAR10 accuracies (after 100 steps) from your Colabs for reference:

model	acc
gsam-S/32	95.10
gsam-S/16	96.05
gsam-B/32	96.21
gsam-B/16	96.46
B/16	97.62

Please also fix the tests to make the Github build action pass - you'll need to update the expected params for the configs you have added here:

vision_transformer/vit_jax/models_test.py

Lines 25 to 48 in 40c0b78

	MODEL_SIZES = {
	'LiT-B16B': 195_871_489,
	'LiT-L16L': 638_443_521,
	'LiT-L16S': 331_140_353,
	'LiT-L16Ti': 311_913_089.,
	'Mixer-B_16': 59_880_472,
	'Mixer-B_32': 60_293_428,
	'Mixer-L_16': 208_196_168,
	'R+ViT-Ti_16': 6_337_704,
	'R26+ViT-B_32': 101_383_976,
	'R26+ViT-S_32': 36_431_912,
	'R50+ViT-B_16': 98_659_112,
	'R50+ViT-L_32': 328_994_856,
	'ViT-B_8': 86_576_872,
	'ViT-B_16': 86_567_656,
	'ViT-B_32': 88_224_232,
	'ViT-H_14': 632_045_800,
	'ViT-L_16': 304_326_632,
	'ViT-L_32': 306_535_400,
	'ViT-S_16': 22_050_664,
	'ViT-S_32': 22_878_952,
	'ViT-Ti_16': 5_717_416,
	'testing': 21_390,
	}

[juntang-zhuang]

@andsteing Thanks a lot for feedback! I added the expected param count in the test file, and the local test is all passed.

For the finetune results, GSAM results are pre-trained on ImageNet1k, but for ViT models I only found checkpoints pretrained on ImageNet21k, I wonder if there are finetuning results for ViT trained I1k with pure AdamW?

[andsteing]

@juntang-zhuang

No, we only published the imagenet-21k checkpoints because they produce strictly better results.

Did you try GSAM pre-training on imagenet-21k as well? It would be interesting to see how much the performance can be improved in that setting, maybe even combining with the data augmentation from the AugReg paper, to see if the improvements are additive.

[juntang-zhuang]

@andsteing We did not experiment on ImageNet-21k due to time limits. I'm currently not affiliated with Google, would you or @tingliu be available to try GSAM with more data or augmentation? If time does not permit, we might provide ImageNet checkpoints only.