Robust fine-tuning of zero-shot models

Abstract

We introduce new methods for 1) accelerating and 2) stabilizing training for large language-vision models. 1) Towards accelerating training, we introduce Switch-Back, 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) Towards stable training, we analyze loss spikes and find they consistently occur 1-8 iterations after the squared gradients become underestimated by their AdamW second moment estimator. As a result, we recommend an AdamW-Adafactor hybrid, which we refer to as StableAdamW because it avoids loss spikes when training a CLIP ViT-Huge model and outperforms gradient clipping.

Prompt Generation Networks for Efficient Adaptation of Frozen Vision Transformers

“…In contrast, Jia et al (2022) learns both additional inputs to multiple layers in a pretrained vision transformer and also finetunes a linear classifier on top. Another finetuning-based approach proposed in Wortsman et al (2022b) is to ensemble the weights between zero-shot and finetuned models and Zhang et al (2020) which trains additional networks that are fused via summation. Specific to vision-language models, Zhou et al (2022b) learn an adaptation network between CLIP's vision and text encoders.…”

Vision-Language Pre-training: Basics, Recent Advances, and Future Trends

“…• Robustness. Wortsman et al (2022) study robust fine-tuning of zero-shot models. Fang et al (2022a) report that data determines distributional robustness in CLIP.…”

Momentum-based Weight Interpolation of Strong Zero-Shot Models for Continual Learning

“…However, retaining a large collection of fine-tuned task expert models in the CL setting is memory intensive, impractical, and undesired. Instead, we show that we can simulate the empirical benefits highlighted in [41] through repeated momentum interpolation between our foundation model and a continuously fine-tuned variant. This allows us to avoid the drawbacks of pure fine-tuning, while both specializing on the new stream of tasks, and retaining the generalizability of our foundation model.…”

“…Their application to the CL problem set, which tackles a continuous distribution shift, stands to reason, with recent works showing notable benefits in the use of foundation models [40,21,31,42], particularly highlighting a reduction in catastrophic forgetting. Still, as learners are adapted to continuously shifting training distribution, even foundation models will suffer from forgetting through fine-tuning [41].…”

“…This slow model is excluded from the direct CL optimization process, and is only updated through linear momentum-interpolation with a task-adapted model copy (denoted as fast model). This is motivated by insights made in [41], who show that simple linear interpolation in weight space between the original zero-shot model and a variant fine-tuned to a task at hand allows for adaptation, while retaining better generalizability as compared to sole fine-tuning. However, retaining a large collection of fine-tuned task expert models in the CL setting is memory intensive, impractical, and undesired.…”