This repo contains official Pytorch and JAX implementation of CLIPA in our paper: An Inverse Scaling Law for CLIP Training
Overview of the Inverse Scaling Law: larger image/text encoders
enable training with fewer image/text tokens while maintaining competitive performance
[Note] All of our CLIPA-v2 models were trained on TPU using our JAX codebase. We followed the same pre-training process as CLIPA-v1, but with a more efficient fine-tuning strategy. In order to replicate our results, we have provided the training configuration (e.g. the H-14 model in this folder here), along with the pre-trained weights, configuration, and logs, which can be found here.
[2023.6.16] We release CLIPA-v2. Compared to the prior best publicly available CLIP model, our CLIPA-v2 can be trained significantly faster and yields stronger performance. Our best model is H/14@336x336 on DataComp-1B with an accuracy of 81.8, and its estimated training cost is <$15k!
data | Schedule | GPU Hours | Estimated Cost | zero-shot IN-1K | model weight | |
---|---|---|---|---|---|---|
H/14 | LAION-2B | 12.8B@84 + 512M@224 + 128M@336 | 8640 | $13613 | 79.1 | PyTorch / JAX |
L/14 | DataCOMP-1B | 12.8B@84 + 512M@224 +128M@336 | 4520 | $7124 | 80.3 | PyTorch / JAX |
H/14 | DataCOMP-1B | 12.8B@84 + 512M@224 + +128M@336 | 8640 | $13613 | 81.8 | PyTorch / JAX |
Our CLIPA-v2โs GPU hour is estimated using an 8-A100 80GB GPU machine on Google Cloud. The corresponding training cost is estimated based on 80GB A100โs cloud pricing.
CLIP, the first foundation model that connects images and text, has enabled many recent breakthroughs in computer vision. However, its associated training cost is prohibitively high, imposing a significant barrier to its widespread exploration. In this paper, we present a surprising finding that there exists an inverse scaling law for CLIP training, whereby the larger the image/text encoders used, the shorter the sequence length of image/text tokens that can be applied in training. Moreover, we showcase that the strategy for reducing image/text token length plays a crucial role in determining the quality of this scaling law.
As a result of this finding, we are able to successfully train CLIP even by using academic resources. For example, on an A100 eight-GPU server, our CLIP models achieve zero-shot top-1 ImageNet accuracies of 63.2% in about 2 days, 67.8% in about 3 days, and 69.3% in about 4 days. By reducing the computation barrier associated with CLIP, we hope to inspire more research in this field, particularly from academics.
Our experiments are conducted on both GPUs and TPUs. Both the JAX and PyTorch implementation enable TPU training. But how to gain access and setup TPU machines? Check this brief doc. In a nutshell, you can access TPU machines on Google Cloud for free!
This project is under the Apache 2.0 License.
The jax repo is built on big vision, and the pytorch repo is built on OpenCLIP. We've also borrowed some code from TIMM and MAE. Many thanks to the awesome works from the open-source community!
We are also very grateful that this work is supported by a gift from Open Philanthropy, TPU Research Cloud (TRC) program, and Google Cloud Research Credits program.
@article{li2023clipa,
title={An Inverse Scaling Law for CLIP Training},
author={Xianhang Li and Zeyu Wang and Cihang Xie},
journal={arXiv preprint arXiv:2305.07017},
year={2023},
}
@article{li2023clipav2,
title={CLIPA-v2: Scaling CLIP Training with 81.1% Zero-shot ImageNet Accuracy within a $10,000 Budget; An Extra $4,000 Unlocks 81.8% Accuracy},
author={Xianhang Li and Zeyu Wang and Cihang Xie},
journal={arXiv preprint arXiv:2306.15658},
year={2023},
}
If you have any question, please feel free to raise an issue or contact us directily: Xianhang Li: [email protected]; Zeyu Wang: [email protected]