This section is about how CLIP is trained. CLIP (Radford et al. 2021), short for Contrastive Language-Image Pre-training, is a pair of networks from OpenAI. One turns a photograph into a vector and the other turns a caption into a vector. Training makes a photo and its own caption point the same way, and makes mismatched pairs point apart. Its recipe is still the template for every model of this kind, so knowing it in detail is what lets you build, debug, and budget one.

Two features make CLIP's training unlike ordinary training. Its loss ties every example in a batch to every other example. And its batch of 32,768 pairs is far larger than one GPU can hold. Most interview questions about CLIP, and most bugs in real implementations, live in those two places.

Here is the design from §30.1, restated. An image tower and a text tower each end in a vector of the same width $d$. Both vectors are scaled to length one, and the score of an image and a caption is their dot product (their cosine) times a learned number called the logit scale. Training shows the model $B$ matching pairs at a time, and every image must pick its own caption out of the $B$ captions, while every caption must pick its own image. Each of those picks is the InfoNCE loss of §7.2, a softmax cross-entropy over the batch, read in both directions.

This section turns that outline into a working system, in five steps. We build the two towers and say how each turns a whole image or caption into one vector. We restate the loss and watch what its learned scale does. We measure what a batch of 32,768 buys. We show how 256 GPUs compute one loss in which every example meets every other, with exactly the right gradient. Last, we budget the run in FLOPs, check the budget against the paper, and read the scaling laws and the cheaper recipes that came after.

The model, piece by piece

Here is the whole model for one image and one caption, in the ViT-L/14 version that most later work uses. ViT-L/14 is a Vision Transformer, a Transformer that cuts an image into square patches and reads them as a sequence of tokens (§27.2). This one has 24 blocks of width 1,024 and patches of 14 by 14 pixels. Shapes are for one pair. In a batch, a leading dimension $B$ rides along in front.

image x (224 x 224 x 3)                              caption y ("a photo of a red fox in snow")
  │ cut into 16 x 16 = 256 patches of 14 x 14 px        │ lower-cased BPE, then [SOS] tokens [EOS], padded to 77
  │ + class token + positions, one extra LayerNorm      │ + positions
  ▼                                                     ▼
ViT-L/14, 24 blocks, width 1024        (257 x 1024)     Transformer, 12 blocks, width 768, causal mask   (77 x 768)
  │ take the class token, LayerNorm                      │ take the [EOS] position, LayerNorm
  ▼                                                     ▼
h^I  (1024)                                           h^T  (768)
  │ times W_I  (1024 x 768), no bias                     │ times W_T  (768 x 768), no bias
  ▼                                                     ▼
u (768), then u / |u|                                 v (768), then v / |v|
                         └──────── s = e^alpha · cos(u, v) ────────┘

Read it from top to bottom. Each side turns its input into a sequence of vectors and runs a Transformer over them. It then pools the result, reducing the sequence to one feature vector, $\mathbf h^{I}$ for the image and $\mathbf h^{T}$ for the text. A matrix projects that vector into the joint space of width $d = 768$, and the result is scaled to length one. The two unit vectors meet in one number, the score $s$, which is their cosine times the scale $e^{\alpha}$.

The image tower

CLIP trained two families of image towers, ResNets and Vision Transformers. Both end with the same question, which is how to turn many position vectors into one.

The Vision Transformer answers it with a class token. It embeds each patch as a token, puts one extra learned token in front, and runs Transformer blocks over the whole sequence, so the extra token can gather information from every patch. The pooled feature $\mathbf h^{I}$ is that token's final state after a LayerNorm. (A LayerNorm rescales each token vector to zero mean and unit variance, then applies a learned gain and bias, §5.2.) CLIP changes the standard ViT in only two small ways. It adds one LayerNorm before the first block, applied to the patch embeddings plus their position embeddings, and it uses a slightly different initialization.

The ResNet needed a new answer. A ResNet-50 is a 50-layer convolutional network with skip connections (§0.8). At 224 px its last stage is a 7 by 7 grid of feature vectors with 2,048 channels each, so 49 positions. The standard ResNet averages them, which gives every position the weight $1/49$, about 2 percent. If a small fox fills one position and snow fills the other 48, the snow makes up 98 percent of the pooled vector.

CLIP replaces the average with attention pooling, which lets the image choose. It forms one query from the average feature, scores every position against it, and averages the positions with softmax weights. Here are real numbers. Suppose the fox's position scores 4 higher than each of the 49 other keys (the 48 snow positions and the average itself). Its weight becomes $e^4/(e^4 + 49) = 54.6/103.6 \approx 0.53$, so the fox is half the pooled vector instead of 2 percent. With several heads, each head can weight a different region.

In plain words, attention pooling is one step of attention with a single query. Attention scores a query against keys, turns the scores into weights with a softmax, and averages the values with those weights (§4.3, with queries, keys, and values from §4.2 and heads from §4.5). With one query its cost is linear in the number of positions. The mean token is also one of the keys, so a head that finds nothing better to look at can fall back on the average feature.

An interviewer may ask why CLIP's ResNets do not just average their last feature map. The fox example is the answer. The same one-query attention is a common way to pool any set of vectors into one.

import numpy as np

def attention_pool(X, P, Wq, Wk, Wv, Wi, h):
    """CLIP's attention pooling. X (B, d_I, s, s), P (s*s+1, d_I). Returns (B, d)."""
    B, dI, s, _ = X.shape
    return X.mean(axis=(2, 3)) @ Wi          # global average pooling, the layer CLIP replaced

The text tower

The text tower is a Transformer that reads a caption left to right, the same design as the GPT-2 language model (§6.2). The base size has 12 blocks of width 512 with 8 heads, 63 million parameters in all. The caption is lower-cased and split by byte-pair encoding (BPE), which builds a vocabulary of frequent character sequences (§2.2). The released tokenizer has 49,408 ids, including a start marker [SOS] and an end marker [EOS]. Each caption is bracketed by the two markers and padded to 77 positions, so a caption keeps at most 75 tokens. For bigger image towers, CLIP widened the text tower in proportion but never deepened it, because performance was less sensitive to the text tower's size.

The text tower uses a causal mask, so each position can attend only to itself and the positions before it (§4.4). The text feature $\mathbf h^{T}$ is the final layer's state at the [EOS] position, after a LayerNorm. These two choices go together. Under a causal mask, position $t$ has read only tokens 1 to $t$, so [EOS], the last real position, is the only one that has read the whole caption. A ViT can read its answer at the class token in position 0 for the opposite reason, since without a mask position 0 sees everything.

Why use a mask at all in an encoder? The paper's reason is to keep the option of starting from a pretrained language model, or of adding a language-modeling loss. It left both for future work.

In a batch, [EOS] sits at a different position in every row, because captions have different lengths. Every position after it holds the padding id 0. The causal mask makes the padding harmless, since [EOS] cannot attend to anything after it. What remains is to find [EOS] in each row.

import numpy as np

def pack_captions(token_lists, sos_id, eos_id, ctx=77, pad_id=0):
    """Returns an int array (B, ctx): [SOS] tokens [EOS], then padding."""
    raise NotImplementedError

def pool_eos(ids, H, eos_id):
    """Hidden states at each row's [EOS], (B, w)."""

The projections, and what is left out

The two projections $\mathbf W_I \in \R^{d_I\times d}$ and $\mathbf W_T \in \R^{d_T\times d}$ are plain linear maps. The joint width $d$ is 512 for ViT-B and 768 for ViT-L/14. Self-supervised image models of the time put a small two-layer network, a projection head, in front of their contrastive loss, as in Chen et al. 2020 (§28.2). CLIP saw no difference in training efficiency between the two. The authors guessed that the nonlinear head suits details of image-only self-supervision.

Several things are absent. Neither tower starts from pretrained weights. The only data augmentation is a random square crop of the resized image. An earlier image-text model for chest X-rays sampled one sentence from each long report as text augmentation (Zhang et al. 2020). CLIP dropped that, since most of its captions are a single sentence. With 400 million pairs, overfitting is not a major concern, so the heavy image augmentations of self-supervised learning were dropped too.

Here is the whole model and its loss in one box. The loss is the symmetric InfoNCE loss of §30.1. Each image runs a softmax over its row of scores to guess its caption, each caption runs a softmax down its column to guess its image, and the loss averages the two cross-entropies.

The released models differ only in their towers and widths. Their exact shapes are in the paper's Tables 19 and 20.

The loss and its learned scale

To train, we need the gradient of this loss. It follows in two lines from the one-direction InfoNCE gradient of §7.2, and the result is short. Write $\mathbf A$ for the softmax of each row of $\mathbf S$, so $A_{ij}$ is how strongly image $i$ picks caption $j$. Write $\bar{\mathbf A}$ for the softmax of each column, so $\bar A_{ij}$ is how strongly caption $j$ picks image $i$.

In plain words, each logit is pushed down by the probability that its row and its column give it, and pushed up by 2 if it is a true pair, all divided by $2B$. The gradient for the log scale is the softmax-weighted average logit minus the positive logit, averaged over rows and columns. It is negative while the positives win, so gradient descent keeps raising the scale on an easy batch. Without a guard the scale would grow without limit. CLIP's guard is the clip at 100, a temperature of at least 0.01, which the paper found necessary to prevent training instability.

Every row of $\mathbf S$ has $B - 1$ negatives, the captions that are not its own, and so does every column. So one batch contains $B^2 - B$ mismatched pairs. Three practical facts about the scale go beyond §7.2.

• The scale is learned in log space. The parameter is $\alpha$, with $\alpha = \ln 14.29 = 2.659$ at the start. The exponential keeps the scale positive and makes $\partial s_{ij}/\partial\alpha = s_{ij}$.
• The clip is a clamp. OpenCLIP, the main open reimplementation of CLIP (Ilharco et al. 2021), clamps $\alpha$ to $\ln 100 = 4.6052$ after every optimizer step. On the forward pass this has the same effect as using $\min(e^{\alpha}, 100)$.
• The trained models sit at the clip. In the released CLIP checkpoints, $e^{\alpha} = 100$ exactly. By the end of training the model is confident enough that the temperature wants to go lower than it is allowed to. A contemporary model trained on noisier data, ALIGN, reports that its learned temperature converged to about $1/64$, a scale of about 64 (Jia et al. 2021).

Here is one batch of four, small enough to check by hand. The embeddings are unit vectors in $\R^8$. Image $i$ and caption $i$ are noisy copies of the same random concept, and the scale is 10.

The widget builds such batches at any size. Each cell shows a weight of its row's softmax $\mathbf A$, and the bars show every row's and column's loss. The curve below the matrix averages the loss of 64 batches at every size, and the last number is $\partial\mathcal L/\partial\alpha$.

The first experiment is the important one. At a fixed encoder the loss grows with $B$, because each image has more captions to confuse with its own. So a CLIP loss is a statement about one batch size. Comparing training losses across batch sizes compares different tasks, which is a common mistake when reading training curves, and a common interview trap.

The optimization recipe

CLIP's optimizer settings are short to state. Each choice is a standard one (§0.6), with a few adjustments for a large batch.

• Optimizer. AdamW, which is Adam with the weight decay applied directly to the weights rather than added to the gradient (§8.6). Weight decay is 0.2 on every weight that is not a gain or a bias. $\beta_1 = 0.9$, and $\beta_2 = 0.98$ for the ViTs (0.999 for the ResNets). $\epsilon = 10^{-6}$ for the ViTs ($10^{-8}$ for the ResNets). A lower $\beta_2$ makes Adam's running average of the squared gradient forget a spike faster, which matters at large batch. SigLIP, a successor that replaces the softmax with a sigmoid, later went further, to 0.95 (§30.4, and the loss spikes of §8.8).
• Learning rate. It warms up over 2,000 steps, then follows a cosine decay (§8.7). The peak is between $3.6\times10^{-4}$ and $5\times10^{-4}$, as in the table above.
• Batch and length. 32,768 pairs per step, for 32 epochs of 400 million pairs. That is 12.8 billion pairs seen and $12.8\times10^{9}/32{,}768 = 390{,}625$ steps.
• Memory savers. Mixed precision, which runs most arithmetic in 16-bit numbers (§8.5). Gradient checkpointing, which discards most activations in the forward pass and recomputes them in the backward pass (§10.1). Adam's statistics stored in half precision. And the text tower's weights stored in half precision with stochastic rounding.
• Tuning. The hyperparameters were tuned by grid search, random search, and hand tuning on a ResNet-50 trained for one epoch. They were then adapted by rule of thumb to larger models, because full-size sweeps were unaffordable.

Stochastic rounding is the unusual trick. Adjacent half-precision numbers near a weight sit a fixed gap apart. Rounding to the nearest one throws away every update smaller than half that gap, so small updates never land. Stochastic rounding rounds up or down at random, with probabilities that make the stored weight correct on average, as §8.5 shows.

A training step from scratch

Before running experiments we need a training step. The exercise below asks you to write one for linear towers, so that every gradient is visible. Two facts do all the work. The first is the logit gradient $\partial\mathcal L/\partial\mathbf S$ from the box above, which the exercise's setup gives you as a helper. The second is how a gradient passes back through the normalization $\hat{\mathbf u} = \mathbf u/\lVert\mathbf u\rVert$.

In plain words, remove the part of the gradient that would only change the length of $\mathbf u$, which normalization ignores, and divide by the length. So the gradient with respect to the raw embedding is always at right angles to it, because $\hat{\mathbf u}^\top(\mathbf I - \hat{\mathbf u}\hat{\mathbf u}^\top) = \mathbf 0^\top$. A quick test of any backward pass is that (dU * U).sum(1) is zero to rounding.

The whole backward pass is then four steps. Write $\mathrm d\mathbf S = \partial\mathcal L/\partial\mathbf S$.

1. From $\mathbf S = e^{\alpha}\hat{\mathbf U}\hat{\mathbf V}^\top$, with the clip applied, $\partial\mathcal L/\partial\hat{\mathbf U} = e^{\alpha}\,\mathrm d\mathbf S\,\hat{\mathbf V}$ and $\partial\mathcal L/\partial\hat{\mathbf V} = e^{\alpha}\,\mathrm d\mathbf S^\top\hat{\mathbf U}$.
2. The normalization Jacobian, row by row, gives $\partial\mathcal L/\partial\mathbf U$ and $\partial\mathcal L/\partial\mathbf V$.
3. From $\mathbf U = \mathbf H^{I}\mathbf W_I$, the rule for a matrix product (§0.2) gives $\partial\mathcal L/\partial\mathbf W_I = \mathbf H^{I\top}\,\partial\mathcal L/\partial\mathbf U$, and the same on the text side.
4. $\partial\mathcal L/\partial\alpha = \sum_{i,j}\mathrm dS_{ij}\,s_{ij}$ while the scale is below the clip, and 0 once it is clipped.

import numpy as np

def clip_step(Wi, Wt, alpha, Hi, Ht, lr, max_scale=100.0):
    """One SGD step of CLIP with linear towers. Returns (Wi_new, Wt_new, alpha_new, loss)."""
    raise NotImplementedError

The cell below is the same step with Adam in place of SGD ($\beta_2 = 0.98$, as for CLIP's ViTs), wrapped in a training loop on synthetic data. Each pair comes from a fresh random concept $\mathbf c \sim \mathcal N(\mathbf 0, \mathbf I_{24})$. The image side sees it as $\mathbf M_I^\top\mathbf c$ plus noise in $\R^{64}$, and the text side as $\mathbf M_T^\top\mathbf c$ plus noise in $\R^{48}$, with fixed random matrices $\mathbf M_I \in \R^{24\times64}$ and $\mathbf M_T \in \R^{24\times48}$. The towers project to a joint width of 16. Every step draws new pairs, so the model never sees a pair twice, there is no overfitting, and "pairs seen" is the only measure of data.

First, the temperature. The sign of $\partial\mathcal L/\partial\alpha$ predicts two behaviors. On data the model can separate, the learned scale keeps climbing until the clip stops it. On hard data it stops where a sharper softmax would cost more on the wrong answers than it saves on the right ones. To see the clip bind within a few hundred steps, this run lowers the clip from 100 to 30.

Real CLIP is the first case at the real clip, which is why its checkpoints sit at 100. In practice, a scale pinned at the clip means the softmax over class prompts is very sharp, and §30.3 shows how to recalibrate it. If an interviewer asks where the temperature ends up, the answer is at the clip, with this reason.

What the batch of 32,768 buys

A batch of $B$ gives every image $B - 1$ negatives in its row and every caption $B - 1$ in its column. That has two consequences, both from §7.2.

The first is a ceiling. Mutual information is the number of nats that seeing an image tells you, on average, about its caption (§0.5). §7.2 proves that $\ln B - \mathcal L_{\text{row}}$ is a lower bound on it. The row loss is never below 0, so the bound can never certify more than $\ln B$ nats. At $B = 32{,}768$ that is $\ln 32{,}768 = 10.40$ nats, or 15 bits. A small batch caps how much the loss can ask of the encoders.

The second is hardness. Among 32,767 random captions, a few will describe something close to the image, such as another red fox or another snowy field. These are hard negatives. The softmax puts most of the gradient on exactly those, because they have the highest logits. So a large batch finds hard negatives for free, with no separate search for them. A bigger batch is a harder test.

Against that, at a fixed budget of pairs seen, a bigger batch means fewer optimizer steps. The cell trains the same model at four batch sizes on the same 32,768 pairs. It evaluates every model the same way, retrieving the captions of 1,000 held-out images among 1,000 candidates.

What the literature found at real scale has the same shape.

• CLIP used 32,768 and did not report a batch-size ablation.
• ALIGN used 16,384, spread over 1,024 TPU cores with 16 pairs each. In its ablation, using only 50 or 25 percent of the in-batch negatives lowered both retrieval scores and nearest-neighbor accuracy on ImageNet (Jia et al. 2021).
• BASIC, a model that scaled data, model, and batch together, used 65,536. It gave a theoretical argument that a larger contrastive batch shrinks the gap between training and test performance (Pham et al. 2021).
• FLIP, which trains on images with most patches dropped, found zero-shot accuracy rising steadily from a batch of 16k (16,384) to 32k and 64k at a fixed masking ratio (Li et al. 2023).
• SigLIP swept batches from 512 to one million (Zhai et al. 2023). With a frozen pretrained image tower, both the softmax loss and its sigmoid loss saturated around 32k. With both towers trained from scratch, the sigmoid loss peaked at 32k, the softmax loss at 98k, and a batch of 307k hurt both. §30.4 tells the sigmoid side of that study.

Sharding the loss across devices

Now the systems problem. With $B = 32{,}768$ the logit matrix has $1.07\times10^9$ entries, 4.3 GB in float32, before its two softmaxes and their gradients. No single GPU should hold it. And under data parallelism no GPU holds all the images anyway. Data parallelism gives every GPU a full copy of the model and its own slice of the batch (§10.2). Three operations that move data between devices, called collectives, do the work here.

That last fact is what this whole part rests on. Say device $r$ holds $b = B/n_{\text{dp}}$ pairs, with index set $\mathcal I_r$. The paper describes its solution in one sentence. Each GPU computed only the pairwise similarities needed for its own local batch. Here is what that means.

1. Each device encodes its own $b$ images and $b$ captions.
2. An all-gather gives every device all $B$ image embeddings and all $B$ caption embeddings ($\hat{\mathbf U}$ and $\hat{\mathbf V}$, $B\times d$ each).
3. Device $r$ computes its row block $e^{\alpha}\hat{\mathbf U}_{\mathcal I_r}\hat{\mathbf V}^\top \in \R^{b\times B}$, its images against every caption. It also computes its column block $e^{\alpha}\hat{\mathbf U}\hat{\mathbf V}_{\mathcal I_r}^\top \in \R^{B\times b}$, every image against its captions. From these it gets the row losses $\ell^{\text{row}}_i$ of its images and the column losses $\ell^{\text{col}}_j$ of its captions.
4. Its partial loss is $\mathcal L^{(r)} = \frac{1}{2B}\big[\sum_{i\in\mathcal I_r}\ell^{\text{row}}_i + \sum_{j\in\mathcal I_r}\ell^{\text{col}}_j\big]$.

The cell splits the batch of four from before across two simulated devices.

The partial losses add up, but does the gradient come out right? Each device's embeddings also appear as negatives in the other devices' rows and columns. The theorem says that summing those contributions back, which is what a reduce-scatter does, gives exactly the full-batch gradient.

In plain words, device $r$'s embeddings appear in other devices' partial losses too, as negatives in their rows and columns, and the gradient from all those appearances has to come home. Each device differentiates its own partial loss and gets a full $B\times d$ gradient buffer for $\hat{\mathbf U}$ and for $\hat{\mathbf V}$. The reduce-scatter sums the buffers across devices, and each device keeps the rows that belong to it. Then each device backpropagates through its own encoder, and the usual all-reduce sums the parameter gradients. Normalization acts row by row, so it can sit on either side of the gather, and the same argument holds for the raw embeddings.

Now the same plan at CLIP's real scale, with every number worked out.

The widget draws the logit matrix of one global batch, cut into the devices' shares. Pick a device to see the stripes it computes, and switch the implementation to compare what each one holds and sends.

import numpy as np

def device_logits(U_shards, V_shards, r, scale):
    """Device r's row block (b_r, B) and column block (B, b_r) of the logit matrix."""
    raise NotImplementedError

def sharded_clip_loss(U_shards, V_shards, scale):
    """Returns (total, partials), one partial loss per device."""

import numpy as np

def sharded_clip_grads(U_shards, V_shards, scale):
    """Returns (dU_shards, dV_shards): device r's gradient of the full CLIP loss
    with respect to its own unit embeddings U_shards[r] and V_shards[r]."""
    raise NotImplementedError

Four ways to write it, and the bug interviewers love

Real code bases offer several variants, and they differ in subtle ways. OpenCLIP exposes two flags. --local-loss computes only the device's own rows and columns. --gather-with-grad uses an all-gather whose backward is the reduce-scatter. Two facts about data parallelism decide which combinations are right.

• Each device backpropagates its own loss through its own encoder. PyTorch's DistributedDataParallel (DDP), its standard data-parallel wrapper, then averages the parameter gradients over the $n_{\text{dp}}$ devices.
• OpenCLIP's local loss is the mean cross-entropy over the device's $b$ rows and $b$ columns. For equal shards that is $\frac{1}{2b}\big[\sum_{i\in\mathcal I_r}\ell^{\text{row}}_i + \sum_{j\in\mathcal I_r}\ell^{\text{col}}_j\big] = n_{\text{dp}}\,\mathcal L^{(r)}$, because $B = n_{\text{dp}}b$. So the gradients come out $n_{\text{dp}}$ times too large, and DDP's average divides that factor back out.

The first row is the one to use. It gives the exact gradient with $2bB$ logits per device, memory that grows only linearly in the batch, and the OpenCLIP documentation pairs the two flags for that reason.

The second row is the subtle one. Each device computes the whole loss from gathered embeddings, of which only its own shard carries gradient. So the gradient it gets for its own embeddings is exactly $\partial\mathcal L_{\text{CLIP}}/\partial\hat{\mathbf U}_r$, with no backward communication, at the price of $B^2$ memory and compute on every device. Then DDP averages, $\frac{1}{n_{\text{dp}}}\sum_r(\partial\mathcal L/\partial\hat{\mathbf U}_r)(\partial\hat{\mathbf U}_r/\partial\theta) = \frac{1}{n_{\text{dp}}}\,\partial\mathcal L/\partial\theta$, where $\theta$ stands for the encoders' parameters. Meanwhile every device computes the full $\partial\mathcal L/\partial\alpha$, whose average is itself.

Does a factor of $1/n_{\text{dp}}$ matter? Under Adam it cancels in the ratio of the first moment to the square root of the second (up to $\epsilon$), so this variant trains well in practice. Under SGD it would cut the encoders' learning rate by $n_{\text{dp}}$. And a gradient-clipping threshold would mean different things on different cluster sizes.

Gradient accumulation fails, and gradient caching works

Sometimes even $b$ pairs per device do not fit in memory. The usual remedy is gradient accumulation, which runs several small micro-batches one after another and sums their gradients. For CLIP that is wrong. The loss is not a sum over examples, since every row's normalizer involves every caption in the batch. So averaging the micro-batches' losses trains the loss with micro-batch negatives only, the last row of the table. In the batch of four, accumulating over two micro-batches of two pairs gives exactly the 0.0961 of the shortcut above, not 0.3501.

The fix is called gradient caching. First compute every embedding without keeping activations, then compute the loss and its gradient with respect to the embeddings, then re-encode one micro-batch at a time and push its slice of that gradient through the encoder. It was introduced for text retrieval (Gao et al. 2021), BASIC used the same idea under the name GradAccum (Pham et al. 2021), and OpenCLIP implements it as --accum-freq.

The price is one extra forward pass. A forward and backward pass cost about $6N$ FLOPs per token for a model with $N$ weights, and the extra forward adds $2N$, so a third more compute. The memory is that of a single micro-batch plus the $B\times d$ embeddings. It is the sharding theorem again, with micro-batches in sequence instead of devices in parallel. An interviewer who asks how to train CLIP at a batch of 32k on 8 GPUs wants this answer.

import numpy as np

def grad_cache_grads(Wi, Wt, alpha, Hi, Ht, micro, max_scale=100.0):
    """Exact full-batch CLIP gradients, re-encoding one micro-batch at a time.
    Returns (dWi, dWt, dalpha, loss)."""
    raise NotImplementedError

The compute budget

How much compute did CLIP ViT-L/14 use? The tool is a rule of thumb from §5.6. Training a Transformer costs about $6N$ FLOPs per token, $2N$ for the forward pass and $4N$ for the backward, where $N$ counts the weights used in matrix products. (Embedding lookups are not products, so the token table does not count.) CLIP has two Transformers that see different numbers of tokens, so the cost of one pair is

with $N_I$ and $T_I$ the image tower's matrix parameters and tokens, and $N_T$ and $T_T$ the text tower's. The image tower sees $T_I = 16^2 + 1 = 257$ tokens at 224 px, 256 patches and the class token. The text tower always runs over all $T_T = 77$ positions, padding included, because batches are padded to a fixed length.

The cell counts the parameters from the shapes above. A tower with $L$ layers of width $w$ has $N = 12Lw^2$ matrix parameters, $4w^2$ in attention and $8w^2$ in the MLP of each block. The other matrix products, the patch embedding and the two projections, add under 0.2 percent, so we leave them out. (§5.6 calls the width $d$, which in this chapter is the joint width.) The cell also adds the cost of attention itself, $12LTw$ FLOPs per token over $T$ tokens, from the same section.

Treat these as estimates. They leave out three things. The first is the extra epoch at 336 px. There 577 image tokens make a pair 2.25 times as expensive once attention's quadratic growth is counted, so that epoch adds $2.25/32$, about 7 percent, to the 32 epochs at 224 px, as Exercise 30.2d computes. The second is the similarity matrix, about 0.15 GFLOPs per pair over the forward and backward passes. The third is the masked half of text attention, about 0.33 GFLOPs per pair, which attention code that skips masked entries would save. Each is under 0.1 percent of the total. And "12 days" is rounded.

def clip_budget(img, txt, n_pairs, epochs, gpus, days, peak_flops, t_img_hi):
    """Returns (flops_per_pair, mfu, hfu, extra_hi)."""
    raise NotImplementedError

def plan_days(img, txt, n_pairs, epochs, gpus, peak_flops, mfu):
    raise NotImplementedError

On 512 A100s at an assumed 40 percent MFU, the same 32 epochs take 1.2 days, about ten times less wall-clock time than the original 12. The speedup splits into three factors. The A100's dense 16-bit peak of 312 TFLOP/s is 2.5 times the V100's 125. Twice the devices give another 2. And almost twice the utilization, 40 against 20.3 percent, gives the last 2. Splitting a speedup into hardware, scale, and efficiency this way is the quickest check on any quoted training time.

Scaling laws for contrastive training

A scaling law is a formula that predicts a model's error from the compute spent on it, fitted on small runs and used to plan large ones (§9.1). CLIP already found one. Its average zero-shot error across 39 evaluations on 36 datasets fell along a straight line in log-log coordinates against model compute, over a 44-fold range covering its five ResNet models.

Cherti et al. (2023) made this a proper scaling study with OpenCLIP. They trained on LAION, open datasets of web image-text pairs filtered with an earlier CLIP model (Schuhmann et al. 2021), at 80 million, 400 million, and 2 billion pairs. They trained for 3, 13, and 34 billion samples seen, with towers from ViT-B/32 to ViT-g/14.

Their method is worth knowing, because it is how you read any compute scaling plot. Among many runs, most waste compute, with a model too large for its samples seen or too small for its budget. So they fit only the best runs.

In plain words, the fit asks how fast the best achievable error falls as you spend more compute. A tenfold increase in compute multiplies the best error by $10^{-k}$. (The paper writes the same law as $\beta C^{\alpha}$ with a negative $\alpha$. We keep §9.1's letters, since $\alpha$ is the log of the logit scale in this chapter.)

Put numbers on that. Ten times more compute multiplies the ImageNet error by $10^{-0.11} = 0.78$ on LAION and by $10^{-0.16} = 0.69$ on WIT. Over a hundredfold increase the gap compounds, to 0.60 against 0.48. Two more findings matter for planning a run.

Samples seen can be the bottleneck. ViT-B/32 gained almost nothing, or lost, going from LAION-400M to LAION-2B at 3 or 13 billion samples seen (57.1 to 57.4 percent, and 63.2 to 62.5 percent). At 34 billion it gained 2.4 points. So a bigger dataset pays only once you train long enough to use it. ViT-L/14 hit the opposite wall. On LAION-400M it gained less than a point from 13 to 34 billion samples, while on LAION-2B it gained 2.1 points.

The fit predicts, but only within its data. From their law they predicted 79.1 percent zero-shot ImageNet accuracy for a ViT-g/14 trained on LAION-2B for 34 billion samples. Models trained partly on Google's private labeled JFT images report higher numbers. BASIC reached 85.7 percent. Locked-image tuning, which trains only a text tower against a frozen image tower pretrained on JFT, reached 85.2 (Zhai et al. 2022). CoCa, which adds a captioning loss to the contrastive one, reached 86.3 (Yu et al. 2022). A law fitted on one data distribution does not predict models trained on another.

import numpy as np

def fit_frontier_power_law(C, E):
    """Returns (k, a, frontier_idx) for E = a * C**(-k) fitted on the Pareto frontier."""
    raise NotImplementedError

def predict(k, a, C_new):
    raise NotImplementedError

Variants that changed the recipe

CLIP's recipe was a starting point. Four follow-ups each changed one part of it, and each is a standard interview reference.

Noisier data at larger scale (Jia et al. 2021, a model called ALIGN). It kept the loss and changed everything else. The towers are an EfficientNet-L2 convolutional network for images and BERT-Large for text (§7.1). It trained on 1.8 billion image and alt-text pairs that received only light filtering, by simple rules on image size, text length, and how often an alt-text repeats. It reached 76.4 percent zero-shot on ImageNet, against 76.2 for CLIP's best model. That showed scale can make up for curation, a story §30.5 continues.

Scale everything at once (Pham et al. 2021, called BASIC). It trained on 6.6 billion pairs, a 3-billion-parameter model, and a batch of 65,536, and reached 85.7 percent zero-shot on ImageNet. Its data was ALIGN's set (1.7 billion pairs by BASIC's count) plus 5 billion images from JFT, Google's internal labeled set, with each image's class names joined into a text. To fit that batch in memory it used gradient caching and model parallelism.

Encode fewer patches (Li et al. 2023, called FLIP). It randomly drops 50 or 75 percent of each training image's patches and runs the ViT only on the visible ones. That is the masking of masked autoencoders (§28.5) without the reconstruction. Masking half the patches halves the compute per image, and the memory saved buys a batch twice as large. For the same accuracy as an unmasked baseline, they report a 3.7 times speedup.

Shorten every sequence (Li, Wang and Xie 2023, called CLIPA). They found an inverse scaling law. The larger the encoders, the shorter the image and text sequences they can be trained on without losing much. A short phase at full length at the end recovers the rest. With 8 A100 GPUs, CLIPA reached 63.2 percent zero-shot on ImageNet in about 2 days.

In PyTorch

The block below checks three of this section's claims with autograd on a tiny model, with four simulated devices standing in for a cluster. The local loss with a differentiable gather reproduces the full-batch gradient. The same loss with a non-differentiable gather does not. And gradient caching, with only one micro-batch's activations alive at a time, reproduces it too. On a real cluster the gather is torch.distributed.nn.functional.all_gather, whose backward is the reduce-scatter, and DDP does the averaging.

Where it breaks

The recipe has known weak points, and the rest of the chapter takes them up.

• The temperature pins at the clip. Once $e^{\alpha} = 100$ the scale stops learning. Logits of 100 times a cosine make CLIP's softmax over class prompts very sharp, which is why zero-shot probabilities need recalibrating (§30.3).
• Large batches spike. Big-batch contrastive training is prone to sudden jumps in the loss, and lowering Adam's $\beta_2$ is the fix SigLIP adopted (§30.4).
• The softmax couples the batch. Every row's normalizer needs every caption, so the loss needs a global gather, and at thousands of devices even two all-gathers add up. SigLIP's sigmoid loss scores each pair on its own and computes its loss in a ring of exchanges between neighboring devices (§30.4).
• The text window is short. It has 77 positions, and a follow-up that extends it measured that CLIP uses fewer than 20 tokens effectively (Zhang et al. 2024). So long, detailed captions are cut off or ignored (§30.6).
• Compute is not the only budget. The scaling exponent depends on the data distribution, so what you train on matters as much as how long (§30.5).

Read more

• Radford et al. (2021). Learning Transferable Visual Models From Natural Language Supervision. Sections 2.3 to 2.5 and Appendix F give the architecture, the training details, and the hyperparameter tables this section quotes.
• Jia et al. (2021). Scaling Up Visual and Vision-Language Representation Learning With Noisy Text Supervision. ALIGN, with an ablation of embedding width, in-batch negatives, and temperature.
• Pham et al. (2021). Combined Scaling for Zero-shot Transfer Learning. BASIC, and the clearest explanation of why gradient accumulation fails for contrastive losses and how to fix it.
• Cherti et al. (2023). Reproducible scaling laws for contrastive language-image learning. Power laws fitted on the Pareto frontier of OpenCLIP runs, and the data dependence of the exponent.
• Li et al. (2023). Scaling Language-Image Pre-training via Masking. FLIP, patch masking for faster and larger-batch CLIP training.
• Li, Wang & Xie (2023). An Inverse Scaling Law for CLIP Training. CLIPA, and how to train a CLIP on 8 GPUs.
• Gao et al. (2021). Scaling Deep Contrastive Learning Batch Size under Memory Limited Setup. GradCache, the gradient-caching technique for any contrastive loss.
• Ilharco et al. OpenCLIP. The reference open implementation, with the local-loss and gather-with-gradient options.
• Zhai et al. (2023). Sigmoid Loss for Language Image Pre-Training. The batch-size sweep from 512 to one million, for both the softmax and the sigmoid loss.

Watch next

▶

▶