Reproducible scaling laws for contrastive language-image learning

In plain words

A model that learns to match pictures with captions can be made bigger, shown more data, or trained for longer. This paper asks how much each of those helps, and whether the gains follow a predictable curve. The authors trained many versions of such a model on large public collections of image-caption pairs, then tested each version on image classification, image search and further training steps on top of it.

The gains were smooth and followed power-law curves, so a curve fitted to smaller runs can be used to forecast bigger ones, with caution. The curves depended on which collection the pairs came from: models trained on the public data improved faster at image search, and models trained on the original CLIP's private data improved faster at classifying images.

TL;DR

Open CLIP-style models trained on LAION data follow power-law scaling in model size, data size and samples seen, on zero-shot classification, retrieval, linear probing and fine-tuning §4. The strongest reported zero-shot result, ViT-H/14 on LAION-2B after 34B samples, reaches 78.0% ImageNet top-1, which the authors report tab. 1. The fitted curves predict higher numbers at larger scales, but those are extrapolations. Use the training recipe in §3.2 with the per-run settings in tab. 18, and note the fit and code caveats under Paper vs code.

Problem

Large contrastive image-text models were trained on private data, or with multi-stage recipes that make it hard to tell what scale alone buys §2. Earlier scaling-law studies covered language models or vision-only models §2. Open datasets were too small until LAION-400M and LAION-5B were released §2. The authors want a systematic scaling study with open data and open code, varying model size, data size and samples seen, with power-law fits for each downstream task §1.

Method

The study varies three scales: the visual encoder (model), the pre-training data (LAION subset) and the number of samples seen (training length) §3.2. Each combination is one run with its own learning-rate schedule §3.2. The text encoder is scaled with the visual encoder §3.2 (text tower widths in tab. 25).

Data. LAION-400M and LAION-5B are used §3.1. The English subset of LAION-5B, called LAION-2B here, has 2.32 billion pairs §3.1. LAION-80M is an 80M subset of LAION-400M §3.2.

Model and loss. Models are OpenCLIP dual encoders with ViT image towers (B/32, B/16, L/14, H/14, g/14) §3.2 tab. 25. Training uses the InfoNCE contrastive loss, following the original CLIP procedure §3.2. The direction of the loss and the temperature are not stated.

Optimiser and schedule. AdamW with β1 = 0.9, β2 = 0.98 and weight decay 0.2 §3.2. Each samples-seen scale gets its own cosine annealing schedule, so short runs are not read off a longer run's schedule §3.2. Peak learning rate, warmup steps and batch size are set per run tab. 18. Mixed precision is used; bfloat16 replaces float16 for the larger models §3.2.

Distributed data parallelism. Training uses PyTorch DDP across up to 1,520 A100 GPUs (text) on JUWELS Booster, partly on Stability AI hardware §3.2. Contrastive negatives span the global batch, so every worker gathers all embeddings app. "Sharding contrastive loss". The gather must be differentiable app. "Sharding contrastive loss".

# One pre-training step, data-parallel over workers  [§3.2]
# Sharding details: [app. "Sharding contrastive loss"]
# Input:  local images (n, 3, H, W), local texts (n, L)
# Output: InfoNCE loss; gradients flow through all gathered embeddings

def training_step(local_images, local_texts):
    # Encode the local shard with both towers  [§3.2]
    local_image_emb = image_tower(local_images)     # (n, d)
    local_text_emb = text_tower(local_texts)        # (n, d)
    # Gather every worker's embeddings; the gather must be differentiable
    # [app. "Sharding contrastive loss"]
    all_image_emb = differentiable_all_gather(local_image_emb)  # (N, d)
    all_text_emb = differentiable_all_gather(local_text_emb)    # (N, d)
    # Each worker scores its own n pairs against all N pairs
    # [app. "Sharding contrastive loss"]
    logits_image = local_image_emb @ all_text_emb.T   # (n, N)
    logits_text = local_text_emb @ all_image_emb.T    # (n, N)
    # InfoNCE: the matching pair is the positive  [§3.2]
    # (embedding normalisation and temperature in training: not stated)
    loss = info_nce(logits_image, local_positive_index)
    return loss

Zero-shot classification. Each class embedding is the mean of its prompt embeddings, then L2-normalised. An image gets the class with the largest cosine similarity to its normalised embedding. Top-1 accuracy is reported §4.1. The prompts come from prior work and are not listed here §4.1.

# Zero-shot classification  [§4.1]
# Input:  class names, prompt templates per class, one image
# Output: predicted class index

def build_class_embeddings(classes, templates):
    class_embs = []
    for each class in classes:
        # Average the text embeddings of that class's prompts  [§4.1]
        prompt_embs = text_tower(prompts_for(class, templates))
        class_embs.append(mean(prompt_embs))
        # ...then L2-normalise the average  [§4.1]
    return normalize(stack(class_embs))              # (C, d)

def classify(image, class_embs):
    image_emb = normalize(image_tower(image))        # (d,)
    # Largest cosine similarity wins  [§4.1]
    return argmax(class_embs @ image_emb)

Retrieval. Image and text embeddings are compared by cosine similarity. Recall@5 is computed on MS-COCO and Flickr30K with the Karpathy test splits §4.2.

Linear probe. The frozen image features are cached once, with no data augmentation tab. 5. A softmax linear classifier is trained with Adam and batch size 256 §4.3. The learning rate and epoch count are chosen from a sweep §4.3 (grid in Hyperparameters).

Fine-tuning. Two regimes: a single model fine-tuned on ImageNet with timm §4.4, and one model jointly fine-tuned on eight datasets where zero-shot models are weak §4.4. The joint setup trains only the vision encoder and keeps the zero-shot text-derived heads frozen §4.4.

Scaling-law fit. Compute is measured in GMAC: per-sample multiply-adds times samples seen zeroshot_results.csv (see Paper vs code). Error is one minus top-1 accuracy figures.ipynb. The fit is a line in log-log space, on the best (lowest-error) points in each of five compute bins §4.1 app. "Predictions derived from scaling laws". Predictions use the fitted line app. "Predictions derived from scaling laws".

# Power-law fit on the lower envelope of runs  [§4.1]
# Details: [app. "Predictions derived from scaling laws"]
# Input:  per run: compute C (GMAC, log10 scale), error E
# Output: slope a and intercept b, so that E ~= 10**b * C**a

def fit_power_law(compute, error):
    log_c = log10(compute)
    log_e = log10(error)
    # Five equal-width bins over log compute  [code: figures.ipynb#L66-L91]
    edges = linspace(min(log_c), max(log_c) + 0.1, 6)
    centres = []
    minima = []
    for each bin with at least one run:
        # Keep the best (lowest-error) run in the bin
        # [app. "Predictions derived from scaling laws"]
        centres.append(mean(edges of this bin))
        minima.append(min(log_e in this bin))
    # Least-squares line through the bin centres  [code: figures.ipynb#L66-L91]
    a, b = least_squares_line(centres, minima)
    return a, b

def predict(compute, a, b):
    # Extrapolate with the fitted power law
    # [app. "Predictions derived from scaling laws"]
    return 10 ** b * compute ** a

Hyperparameters

All pre-training runs (29 models). Source: tab. 18 and cross-checked against trained_models_info.csv trained_models_info.csv. "BS" is the per-GPU batch with the global batch in parentheses. Warmup is in steps. Samples seen are nominal labels. Time is wall-clock hours.

Total cost reported: 1,058,318 GPU-hours and 370.41 MWh tab. 18.

Other settings:

ImageNet fine-tune, direct on ImageNet-1k tab. 10:

Intermediate fine-tune on ImageNet-12k tab. 11. Values that differ from the table above:

Architectures tab. 25:

Parameter counts are the paper's. The repository's arch_info.csv gives the same GMAC values arch_info.csv arch_info.csv arch_info.csv arch_info.csv arch_info.csv.

Results

Author-reported. Setup: OpenCLIP models trained on LAION subsets, evaluated at 224 px; baseline is OpenAI CLIP trained on the private WIT-400M tab. 1.

Table 1 does not state the sample count for the L/14 row. The matching value in tab. 19 (75.20) is the 34B-sample run.

Fitted power-law exponents, as reported §4.1 §4.2. Error is one minus accuracy, so a more negative exponent means faster improvement.

Extrapolations (predicted, not measured) tab. 4 §4.1:

Measured in the same table: ViT-H/14 at 34B 77.97 and ViT-g/14 at 13B 76.66 tab. 4.

Other results, as reported:

• Linear probes on ImageNet, CIFAR-100 and VTAB improve with scale, in the few-shot and full-data regimes §4.3. Per-model values are in tab. 5.
• Fine-tuned OpenCLIP models reach 82.6 to 88.5% ImageNet top-1, with the best from ViT-H/14 at 336 px after ImageNet-12k intermediate fine-tuning §4.4 tab. 9.
• Joint fine-tuning on eight tasks improves with scale §4.4 fig. 5.
• Zero-shot ImageNet robustness moves with ImageNet accuracy §5.

Paper vs code

The repository at commit 70e66ab4dd4872dfe7d20fa1c35c42355879690f does not contain training code. It holds result CSVs, two notebooks, a model download script and a one-line requirements.txt download_models.py. The paper says OpenCLIP is the training code §3.2; OpenCLIP was not checked here.

Discrepancies:

• Power-law fit. The paper fits both curves on the Pareto frontier of available models fig. 1 §4.1. The notebook bins the OpenCLIP points into five compute bins and keeps the lowest error per bin [code: figures.ipynb#L66-L91] figures.ipynb. For OpenAI CLIP it runs an ordinary least-squares line over all points, with no binning [code: figures.ipynb#L301]. The two exponents therefore come from different fits.
• Batch size. The text gives global batch sizes of "86-88K" §3.2. The run table has 79K (H/14 and B/32 at 2B, 34B), 64K (g/14), 33K (B/16 at 400M, 13B) and 32K (several B/32 runs) tab. 18 trained_models_info.csv. The text summary covers only part of the runs.
• GPU count. The text says "up to 1520" A100 GPUs §3.2. The largest run in tab. 18 uses 1024 GPUs, and no run lists 1520.
• Learning rate and warmup per run. All 29 rows of tab. 18 agree with trained_models_info.csv on peak LR, warmup steps, per-GPU batch and GPU count trained_models_info.csv. Example: H/14, 34B samples, LR 5e-4, warmup 2000, 824 GPUs, per-GPU batch 96 [code: trained_models_info.csv#L26] tab. 18.
• Zero-shot value. The H/14 ImageNet top-1 is 0.77972 in the repository's results, rounding to the 77.97 in tab. 19 zeroshot_results.csv.
• Compute measure. The repository's gmacs_total equals GMAC per sample times samples seen. For H/14: 190.97 × 34,725,395,968 = 6.6315 × 10^12 zeroshot_results.csv. This is inferred from one row. The per-run sample count in the CSV, 34.73e9, is the actual count, while the "34B" label is nominal.
• Error vs accuracy. The text's E is not defined in words. The code's E is one minus top-1 figures.ipynb. The negative exponents fit this reading.

Pitfalls

• Scaling-law forecasts are extrapolations. The 68B-sample predictions and the ViT-G/14 number are fitted, not measured tab. 4. The authors say extrapolation has limits §5.
• The fit uses few points. The paper does not state a bin count. The notebook uses five bins over log compute and drops empty bins [code: figures.ipynb#L66-L91]. A small change in the bin edges can move the line.
• Samples seen vs epochs. The paper says 13B samples is 32 full epochs of LAION-400M tab. 12 tab. 13. So samples seen is roughly epochs times dataset size. The 34B label is nominal (see Paper vs code).
• Warmup shape and final LR are not stated. Warmup is given in steps tab. 18, and the cosine schedule runs per samples-seen scale §3.2. The warmup shape (linear or other) and the minimum LR at the end of the cosine are not stated in the active text.
• Loss details not stated. The direction of InfoNCE (image-to-text, text-to-image or both), the temperature or logit scale, and whether embeddings are normalised during training are not stated §3.2.
• Differentiable gather is required. The plain all_gather does not pass gradients. The differentiable variant is required app. "Sharding contrastive loss". The paper's memory example: a full N×N similarity matrix for N = 32768 in FP32 takes 4 GB, and sharded n×N matrices with local batch 256 take 64 MB app. "Sharding contrastive loss".
• Loss spikes at scale. Mixed-precision float16 caused loss spikes for L/14, H/14 and g/14. Switching to bfloat16 fixed them. Float32 with TF32 also worked. Smaller models can become unstable at large learning rate and batch size app. "Training instabilities" §3.2. Extra normalisation layers, scaled cosine attention and gradient clipping did not help app. "Training instabilities" §3.2 footnote.
• Resampling changes little. Fixed-length "virtual" steps with resampling, instead of epoch-wise shuffling, changed ImageNet accuracy by 0.1 to 0.3% app. "Control experiments" tab. 15. Whether all large runs used it is not stated.
• Batch-size effect is small but not zero. Changing batch size changed zero-shot ImageNet by 0.2 to 0.5% app. "Control experiments" tab. 12.
• LAION-400M is not a subset of LAION-2B. The two were built by the same procedure but separately app. "Datasets employed in experiments". A random 400M subset of LAION-2B gave 63.56 vs 63.37 for LAION-400M on ViT-B/32 at 13B samples tab. 13.
• Linear-probe selection split is not stated. The paper says the best result of the sweep is reported §4.3. Whether selection used a validation split or the test set is not stated. Use a held-out split.
• Zero-shot prompts are not in the paper. The prompt sets come from prior work §4.1. An implementation must take them from that work.
• Duplicate check is partial. Downstream images were checked against LAION-400M with pHash only §4.1 app. "Datasets employed in experiments". Near-duplicates may remain §5.
• Tuning is limited. Full hyperparameter tuning was not possible at large scale §5. Control experiments cover a few settings at early stages §5.
• Costs. 1,058,318 GPU-hours and 370.41 MWh tab. 18. The largest single runs are L/14 at 34B on LAION-2B (319 h) and H/14 at 34B (279 h) tab. 18.
• Version. Metadata says arXiv v2 (updated 2024-07-13). This card follows that version. The CVPR 2023 version was not compared.
• Table and figure numbers. They were counted from the LaTeX source, not from the PDF. Check against the PDF before citing externally.

Lineage

• Builds on: contrastive language-image pre-training (CLIP, Radford et al.) §2, whose recipe is followed with a larger batch §3.2. Earlier scaling laws for language and vision models §2 (no cards yet).
• Data: LAION-400M [/papers/2111.02114] and LAION-5B [/papers/2210.08402] §3.1.
• Compared against: OpenAI CLIP models (ViT-B/32, B/16, L/14 on WIT-400M) §4.1.
• Superseded by: not stated in this wiki's content.