PaperlayerView exact revision on arXiv

Source-linked AI summary

Combined Scaling for Zero-shot Transfer Learning

Hieu Pham, Zihang Dai, Golnaz Ghiasi, Kenji Kawaguchi, Hanxiao Liu, Adams Wei Yu, Jiahui Yu, Yi-Ting Chen, Minh-Thang Luong, Yonghui Wu, Mingxing Tan, Quoc V. Le

arXiv:2111.10050v3cs.LGcs.CLcs.CV

TL;DR

Existing image-text contrastive models remained behind supervised systems, motivating stronger zero-shot transfer without labeled ImageNet examples. BASIC jointly scales data, model, and contrastive batch size, using memory-saving methods and theoretical analysis; it reaches 85.7% ImageNet top-1 accuracy and improves robustness benchmarks.

  • Problem

    CLIP and ALIGN remain less competitive than supervised and semi-supervised models when labeled data is available, leaving a performance gap for zero-shot transfer.

  • Method

    BASIC jointly scales dataset size, model size, and contrastive batch size, using gradient checkpointing and model-parallel methods to address accelerator memory limits.

  • Results

    85.7% top-1 accuracy on ImageNet without labeled ImageNet examples surpasses CLIP and ALIGN by 9.3%, while BASIC also improves robustness and outperforms CLIP on 13 of 17 benchmarks.

  • Takeaways & Limitations

    Combined scaling boosts zero-shot transfer accuracy and robustness, while larger contrastive batches are associated with smaller generalization gaps.

  • Takeaways & Limitations

    BASIC performs poorly on underrepresented domains such as MNIST and PCam, and zero-shot accuracy can be hurt by poorly chosen or adversarial prompts.

Abstract

from arXiv · show

We present a combined scaling method - named BASIC - that achieves 85.7% top-1 accuracy on the ImageNet ILSVRC-2012 validation set without learning from any labeled ImageNet example. This accuracy surpasses best published similar models - CLIP and ALIGN - by 9.3%. Our BASIC model also shows significant improvements in robustness benchmarks. For instance, on 5 test sets with natural distribution shifts such as ImageNet-{A,R,V2,Sketch} and ObjectNet, our model achieves 84.3% top-1 average accuracy, only a small drop from its original ImageNet accuracy. To achieve these results, we scale up the contrastive learning framework of CLIP and ALIGN in three dimensions: data size, model size, and batch size. Our dataset has 6.6B noisy image-text pairs, which is 4x larger than ALIGN, and 16x larger than CLIP. Our largest model has 3B weights, which is 3.75x larger in parameters and 8x larger in FLOPs than ALIGN and CLIP. Finally, our batch size is 65536 which is 2x more than CLIP and 4x more than ALIGN. We encountered two main challenges with the scaling rules of BASIC. First, the main challenge with implementing the combined scaling rules of BASIC is the limited memory of accelerators, such as GPUs and TPUs. To overcome the memory limit, we propose two simple methods which make use of gradient checkpointing and model parallelism. Second, while increasing the dataset size and the model size has been the defacto method to improve the performance of deep learning models like BASIC, the effect of a large contrastive batch size on such contrastive-trained image-text models is not well-understood. To shed light on the benefits of large contrastive batch sizes, we develop a theoretical framework which shows that larger contrastive batch sizes lead to smaller generalization gaps for image-text models such as BASIC.

1. Introduction

BASIC scales image-text contrastive learning across data, model, and batch size to narrow the gap between zero-shot and supervised image classification. It addresses accelerator memory limits and reports stronger ImageNet and robustness performance than CLIP and ALIGN.

  • 6.6B noisy image-text pairs, a 3B-parameter CoAtNet, and a 65536-example contrastive batch implement BASIC’s three-dimensional scaling.The dataset is almost 4x larger than ALIGN’s, while the model has 3.75x more weights and 8x more FLOPs than the largest CLIP and ALIGN models.
  • Gradient checkpointing, micro-batch pipelining, gradient accumulation, and model parallelism address the accelerator-memory bottleneck created by larger models and batches.SPMD distributes selected layer weights across devices, while rematerialization saves memory by recomputing discarded forward-pass values.
  • Larger contrastive batches are theoretically linked to smaller generalization gaps in image-text models.
  • Table 1 reports BASIC, CLIP, and ALIGN top-1 accuracy on ImageNet and robustness sets, with BASIC exceeding the baselines by 10.1 percentage points on average.
  • 85.7% top-1 accuracy on ImageNet is achieved by the largest BASIC model without labeled ImageNet examples, surpassing CLIP and ALIGN by 9.3%.
  • 83.7% average top-1 accuracy across five natural distribution-shift test sets shows improved robustness with only a small drop from ImageNet accuracy.

2. Related Work

Related work spans weakly supervised and self-supervised visual pretraining, image-text modeling, zero-shot learning distinctions, and scaling data, models, and batches. BASIC builds on image-text contrastive learning while pushing all three scaling dimensions.

  • Image-text models learn visual representations from paired images and text, with classification from textual descriptions rather than caption generation as the paper’s focus.
  • Earlier zero-shot learning commonly means generalizing to classes or tasks without training data, whereas BASIC follows the distinct zero-shot transfer-learning setting.
  • Zero-shot transfer learning permits supervised information during pretraining but uses no supervised examples during transfer and skips fine-tuning.
  • Scaling data, model size, and batch size has been used to improve vision-model pretraining, especially for contrastive losses.

3. Background on Image-text Contrastive Learning and Zero-shot Transfer Learning

Image-text contrastive learning maps images and text into a shared latent space, bringing semantically matched pairs together and separating mismatched pairs. Zero-shot classification uses these learned mappings with textual descriptions.

  • Image and text encoders F and G map inputs to D-dimensional unit-sphere embeddings.
  • The contrastive objective makes matched image-text pairs nearby in latent space and mismatched pairs distant.
  • The similarity matrix A is converted into row- and column-based softmax-like losses, with temperature τ controlling distribution sharpness.

4. Batch Size Scaling with Pipelining and Gradient Accumulation

Large contrastive batches and networks create an accelerator-memory bottleneck, while naive gradient accumulation cannot directly compute contrastive losses. The proposed rematerialization-based procedure enables exact microbatch gradients with lower peak memory.

  • Memory is the crucial bottleneck hindering simultaneous scaling of contrastive model size and batch size.
  • Naive gradient accumulation fails because each contrastive loss requires the full similarity matrix A, and storing cumulative gradients consumes parameter-scale memory.
  • Rematerialization computes the full similarity matrix while discarding intermediate hidden states, then recomputes them during backpropagation for each microbatch.
  • Θ(M · max {Mem(F), Mem(G)}) peak memory enables exact microbatch gradients instead of Θ(B · (Mem(F) + Mem(G))).
  • The method can use different microbatch sizes for the image and text networks, supporting more efficient computation when their memory requirements differ.
  • Approximating optimizer moment updates introduces one source of inexactness, while the square-of-sum versus sum-of-squares mismatch creates another.

5. Batch Size Scaling with the Single-Program Multiple-Data (SPMD) Scheme

SPMD addresses accelerator-memory bottlenecks for large contrastive models by sharding weights across cores, while rematerialization saves activation memory. Compared with pipelining and gradient accumulation, it provides exact computations and faster training steps, at the cost of greater memory use and manual design.

  • 5. Batch Size Scaling: SPMD avoids the inexact gradients of microbatch pipelining, but requires manual designs that make it less generic for arbitrarily large contrastive batches.Pipelining can differ from global-batch computation when networks depend on batch statistics, especially with small microbatches.
  • 5.1 Weight Sharding: SPMD distributes encoder weight tensors and their optimizer moments across accelerator cores, combining shards only when complete tensors are needed.This reduces replicated weight storage while supporting computations such as sharded convolution.
  • 5.2 Rematerialization: Rematerialization discards selected forward-pass values and recomputes them during backpropagation, preserving accelerator memory.The strategy keeps weight-involving layers and rematerializes activation and normalization layers.
  • 5.2 Rematerialization: More than half of encoder activation values can be removed after each forward pass while only lightweight computations are repeated during backpropagation.This benefit follows from rematerializing layers without weights, such as activations and normalization layers.
  • 5.3 Comparison with Pipelining and Gradient Accumulation: SPMD is faster than pipelining and gradient accumulation at matched settings, although it has a larger device-memory footprint.In the largest reported setting, its backward pass was more than 1.2 seconds faster, about 10% of total step time.

6. Theoretical Insight on the Role of Contrastive Batch Size

The paper analyzes how contrastive batch size affects generalization from image-text training examples to unseen data. Its theoretical results show that increasing both sample count and contrastive batch size reduces the generalization gap under stated assumptions and for general model classes.

  • Generalization Setup: The analysis compares normalized empirical training loss with normalized testing loss to study generalization from paired training data to unseen inputs.The testing-time prediction uses image-text representation similarities against candidate text inputs.
  • Theoretical Results: A small contrastive batch can leave a large generalization gap even when the training sample count is large.The paper contrasts this with standard supervised classification, where increasing the sample count can be sufficient to reduce the gap.
  • General Models: Theorem 2 extends the decreasing-gap result to general function classes F and G, with additional terms representing model complexity.These terms include R_m(F_k) + R_m(G_k) and a complexity term for G.
  • General Models: For standard deep neural networks, the model-complexity terms typically scale as O(1/√m) and O(1/√B), supporting the use of large contrastive batches.This theoretical insight motivates BASIC’s batch-size scaling design.

7. Data and Model Scaling

BASIC scales image-text contrastive learning through a substantially expanded dataset and larger encoder models. The design uses ALIGN+JFT data, favors image-encoder capacity, and defines Small, Medium, and Large model variants.

  • Data Scaling: The training corpus expands ALIGN’s 1.7B weakly aligned image-text pairs with 5B JFT-derived pairs, producing roughly four times more data.JFT classes are converted into text sequences and combined with ALIGN instances.
  • Data Scaling: A tokenizer trained on 200M sampled sentences uses a 32K-piece vocabulary and discards sequences longer than 64 tokens.The paper reports that adapting this tokenizer and filtering improved ImageNet top-1 accuracy by more than 1% in preliminary experiments.
  • Model Scaling: At the same computational budget, the authors find it more beneficial to scale the image encoder than the text encoder.They therefore use the largest CoAtNet architecture for images and a simple Transformer for text.
  • Model Scaling: BASIC uses CoAtNet image encoders with convolutional layers followed by attention layers and averages top-layer text representations across sequence steps.This differs from ALIGN’s use of a final [CLS] representation for text.
  • Model Scaling: Three model sizes, BASIC-S, BASIC-M, and BASIC-L, are selected by extrapolating scaling findings from small models to larger models.The paper reports their architectures and computational costs separately.

8. Pretraining and Finetuning

The paper evaluates pretraining followed by contrastive text-encoder training as a memory-efficient alternative, but finds that excluding noisy image-text data can cause failures on some tasks. Its best results use a hybrid procedure with later joint finetuning.

  • Pretraining: Pretraining fixes a pretrained image encoder while contrastive learning trains the text encoder, reducing peak memory usage.The procedure avoids computing gradients for both encoders simultaneously.
  • Limitations: The pretraining-finetuning scheme can fail on tasks requiring skills underrepresented in the labeled pretraining data, such as MNIST digit recognition.The paper attributes this to the image encoder’s lack of exposure to noisy image-text data containing digit examples.
  • Hybrid Training: The best experimental results use a hybrid procedure combining image pretraining, contrastive text-encoder training, and joint finetuning of both encoders.Gradient accumulation is used during the final finetuning stage when needed.

9. Experiments

The experiments evaluate BASIC across image-classification and robustness benchmarks, finding strong zero-shot transfer and higher effective robustness than ImageNet-trained models. Adding labeled ImageNet data improves ImageNet accuracy but generally reduces robustness.

  • Image Classification Benchmarks: BASIC and CLIP are compared on 17 image-classification benchmarks using consistent textual prompts.The evaluation includes models with similar computational budgets.
  • Image Classification Benchmarks: BASIC models outperform comparable CLIP models on 13 of 17 datasets.BASIC-L reaches 97.9% mean per-class recall on Oxford IIIT Pets without seeing training images from that dataset.
  • Image Classification Benchmarks: BASIC models perform poorly on EuroSAT, MNIST, and PCam, with MNIST reaching at most 40.3% accuracy.The paper links failures on MNIST and PCam to relatively few corresponding examples in ALIGN and JFT.
  • Robustness Benchmarks: Zero-shot models have higher effective robustness than ImageNet-trained models at the same ImageNet accuracy.The comparison uses five robustness benchmarks derived from natural distribution shifts.
  • Robustness Benchmarks: Training BASIC on more labeled ImageNet data increases ImageNet accuracy, but gains generally do not transfer to robustness benchmarks.ImageNet-V2 is an exception, remaining stable for BASIC-L and slightly increasing for BASIC-M.
  • Robustness Benchmarks: A 3% ImageNet accuracy gain coincides with an 8.3% accuracy drop on ImageNet-R in the extreme case.The authors caution that their experiment shows correlation rather than causal attribution to ImageNet.

10. Ablation Study

The ablation study examines batch size, data scaling, model scaling, and pretraining. Larger batches improve accuracy beyond what additional training steps achieve, while model scaling and combined pretraining choices provide further gains.

  • The Importance of Batch Size Scaling: Larger batch sizes yield higher ImageNet accuracy even when models see the same number of training examples.The controlled comparison varies batch size and training steps while keeping the number of examples fixed.
  • The Importance of Batch Size Scaling: Large batch sizes continue benefiting BASIC through 32768 and 65536, beyond the 8192 saturation observed for SimCLR.The paper attributes the continued benefit to the larger dataset and model sizes used by BASIC.
  • Data Scaling, Model Scaling, and Pretraining: Expanding the contrastive dataset from ALIGN to ALIGN+JFT improves BASIC-S by 5.3% and BASIC-M by 5.8%.Enlarging the model from BASIC-S to BASIC-M produces an 11.7% accuracy gain.
  • Data Scaling, Model Scaling, and Pretraining: JFT-pretrained image encoders with text encoders trained on ALIGN perform similarly to models trained from scratch on ALIGN+JFT.This suggests the training loss choice has a smaller effect than the datasets in these settings.
  • Data Scaling, Model Scaling, and Pretraining: Training text encoders on ALIGN+JFT instead of ALIGN adds 1% for BASIC-S and 1.8% for BASIC-L.The authors speculate that cleaner JFT class-name prompts may improve textual representations, but call for further study.
  • Data Scaling, Model Scaling, and Pretraining: Continuing low-learning-rate training of both encoders adds 1.4% for BASIC-S, 0.6% for BASIC-M, and 0.4% for BASIC-L.This phase follows JFT image-encoder pretraining and ALIGN+JFT text-encoder training.

11. Limitations

BASIC models have important limitations in domain coverage, prompt dependence, and training cost. Their failures are most pronounced when evaluation data or prompts differ from the conditions represented during training.

  • Data coverage: BASIC models perform poorly on test sets underrepresented in their training datasets, including MNIST and PCam.The paper attributes these failures to few handwritten-digit and lymph-node images in ALIGN and JFT.
  • Data coverage: BASIC performs better on EuroSAT than on MNIST and PCam, but remains less accurate than CLIP because its satellite class names are not descriptive enough.
  • Prompt dependence: Zero-shot transfer requires prompt engineering because poorly chosen or adversarial prompts can flip model predictions.The paper uses CLIP prompts for comparability and leaves prompt engineering outside its scope.
  • Compute cost: Combined scaling has substantial hardware and training-time costs, although using the resulting models without fine-tuning avoids fine-tuning expense.The paper suggests larger accelerator memory could reduce extra re-materialization steps.

12. Conclusion

The paper argues that zero-shot transfer must narrow its accuracy gap with supervised models and presents BASIC’s combined scaling as a way to improve accuracy and robustness. It also identifies memory-efficient training as necessary for implementing this scaling.

  • Motivation: Zero-shot transfer models need to bridge the accuracy gap to supervised and semi-supervised models before becoming practical for real-world applications.
  • Contribution: BASIC combines scaling of data size, model size, and batch size, and the paper reports improvements in final accuracy and robustness.
  • Contribution: Gradient accumulation based on re-materialization addresses the memory limit created by combined scaling.
  • Model selection: Table 5 reports the final model sizes selected for compute-to-performance efficiency.

B. Hyperparameters and other implementation details

This section documents training and evaluation details, including hyperparameters, regularization choices, computational requirements, and robustness analyses. It also illustrates BASIC’s zero-shot predictions on varied examples and reports model-size specifications.

  • Experimental setup: The paper provides experiment hyperparameters in Table 6 and evaluation-dataset details in Table 7, whose results appear in Tables 1 and 3.
  • Regularization: BASIC uses no regularization beyond decoupled weight decay, because stochastic depth and dropout substantially reduce ImageNet top-1 accuracy.
  • Regularization: Re-materialization requires consistent forward passes, while dropout-like noise keeps training loss relatively large and creates optimization difficulty.
  • Robustness: Finetuning converged BASIC checkpoints on more ImageNet labeled data leads to worse robustness results on five ImageNet-derived benchmarks.
  • Robustness: BASIC-M robustness on ImageNet-V2 improves with ImageNet accuracy, whereas BASIC-L robustness does not; the paper suspects larger capacity picks up more spurious ImageNet patterns.
  • Robustness: On ObjectNet, BASIC’s improvement over ALIGN and CLIP is 6.6%, lower than its improvement on most other benchmarks, possibly because shared class names and prompts are less descriptive.
  • Compute: BASIC-L uses TPUv4 MegaCore memory near its limit, with 30.1GB used out of 32GB.
  • Qualitative examples: Figure 7 shows BASIC-L aligning unseen images with natural-language sequences, including eggs, cosplayed Pokémon, and dogs in different weather.

G.1 The benchmarks where BASIC fails

BASIC performs poorly on EuroSAT, MNIST, and Patch Camelyon, revealing limitations in zero-shot transfer for specialized or text-sensitive tasks. The failures also show that prompts, class names, counting, negation, and spatial relations can strongly affect predictions.

  • Low-performing benchmarks: 51.0%, 40.3%, and 59.6% are BASIC-L’s accuracies on EuroSAT, MNIST, and Patch Camelyon, respectively.These are identified as benchmarks where BASIC models have particularly low performance.
  • Patch Camelyon: 59.6% for BASIC-L and 63.0% for CLIP on Patch Camelyon are far below practical deployment standards.Because PCam is binary classification, both results are only slightly above random guessing.
  • EuroSAT: 51.0% to 55.7% shows that changing EuroSAT class names and prompts improves BASIC-L accuracy.The alternative words and phrases have essentially the same meaning to humans.
  • MNIST: 40.3% remains BASIC-L’s MNIST accuracy across digit-only, text-only, and combined digit-and-text class names.The visibly different confusion matrices suggest that class-name strings influence matching even when accuracy does not improve.
  • Failure modes: Prompt engineering affects desired outputs, while BASIC-L struggles with object counting, negation, and left-right distinctions on unseen images.The failure cases also show that prompts can provide sufficient context for some desired outputs.
Loading 2111.10050v3…