Stable and low-precision training for large-scale vision-language models

By Mitchell Wortsman, Tim Dettmers, Luke S. Zettlemoyer, Ari S. Morcos, Ali Farhadi, Ludwig Schmidt. 2023

94 citations recorded by the source corpus.

We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models. 1) For acceleration, we introduce SwitchBack, a linear layer for int8 quantized training which provides a speed-up of 13-25% while matching the performance of bfloat16 training within 0.1 percentage points for the 1B parameter CLIP ViT-Huge -- the largest int8 training to date. Our main focus is int8 as GPU support for float8 is rare, though we also analyze float8 training through simulation. While SwitchBack proves effective for float8, we show that standard techniques are also successful if the network is trained and initialized so that large feature magnitudes are discouraged, which we accomplish via layer-scale initialized with zeros. 2) For stability, we analyze loss spikes and find they consistently occur 1-8 iterations after the squared gradients become under-estimated by their AdamW second moment estimator. As a result, we recommend an AdamW-Adafactor hybrid which avoids loss spikes when training a CLIP ViT-Huge model and outperforms gradient clipping at the scales we test.

Subfields: Artificial Intelligence, Computer Vision and Pattern Recognition, Multimodal Learning, Optimization & Training

How researchers see this paper

The researchers whose trusted networks rank it highest, strongest advocate first. Every researcher gets a different order.

In Artificial Intelligence

1. Luke S. Zettlemoyer
2. Ari Holtzman
3. Tim Dettmers
4. Artidoro Pagnoni
5. Yejin Choi
6. Mohit Bansal
7. Noah A. Smith
8. Michael Lewis

In Computer Vision and Pattern Recognition

1. Tim Dettmers
2. Ari Holtzman
3. Luke S. Zettlemoyer
4. Artidoro Pagnoni
5. Hannaneh Hajishirzi
6. Yejin Choi
7. Percy Liang
8. Mohit Bansal

In Multimodal Learning

1. Ari Holtzman
2. Luke S. Zettlemoyer

In Optimization & Training

1. Luke S. Zettlemoyer
2. Tim Dettmers
3. Ari Holtzman
4. Torsten Hoefler
5. Tal Ben-Nun
6. Waseem Rawat
7. Zenghui Wang

Citations in the corpus

• QLoRA: Efficient Finetuning of Quantized LLMs (2023, background): “Besides our work, SwitchBack layers [66] is the only work that studies backpropagation through quantized weights at a scale beyond 1B parameters.”
• SpQR: A Sparse-Quantized Representation for Near-Lossless LLM Weight Compression (2023, unknown): “Recommendations are presented in”