Combined Scaling for Zero-shot Transfer Learning

🎯 Pitch

By scaling batch size, model size, and data simultaneously, BASIC achieves 85.7% ImageNet top-1 accuracy without any labeled data, beating CLIP and ALIGN by over 9 points. A theoretical analysis reveals that larger contrastive batch sizes reduce generalization error at a rate of O(1/√B), a benefit that additional training steps alone cannot replicate. Surprisingly, the authors also show that as you fine-tune on more labeled examples, robustness to distribution shift can actually degrade, upending conventional assumptions about transfer learning.

1. Executive Summary

This paper introduces BASIC (Batch, Data, and Model Size Combined Scaling), a combined scaling method for image-text contrastive learning that achieves 85.7% top-1 accuracy on ImageNet ILSVRC-2012 without any labeled ImageNet training examples, surpassing CLIP and ALIGN by 9.3%. The work scales the contrastive learning framework of CLIP and ALIGN along three dimensions simultaneously—data size (6.6B noisy image-text pairs, 4× larger than ALIGN), model size (a 3B-parameter CoAtNet image encoder, 3.75× larger in parameters and 8× larger in FLOPs than prior work), and batch size (65,536 contrastive examples per minibatch, 2–4× larger than baselines)—and provides two complementary engineering solutions (pipelining with gradient accumulation and SPMD-based weight sharding with rematerialization) to overcome accelerator memory bottlenecks that arise from combined scaling. A theoretical analysis establishes that larger contrastive batch sizes lead to smaller generalization gaps at a rate of O(1/√B), and empirical ablations confirm that large batch sizes continue to benefit models even when the number of training epochs is held constant, establishing that batch size scaling provides benefits beyond what additional training steps can recover. On robustness benchmarks, BASIC achieves 84.3% average top-1 accuracy across five natural distribution shift test sets—only a small drop from its ImageNet accuracy—and further finetuning experiments reveal that training on more labeled ImageNet data can paradoxically decrease robustness, establishing that zero-shot transfer models exhibit higher effective robustness only when they are not exposed to the target dataset's labeled examples.

2. Context and Motivation

The Core Problem: Zero-Shot Transfer Models Are Not Yet Competitive with Supervised Models

The fundamental problem this paper tackles is the accuracy gap between zero-shot transfer models and their supervised counterparts. CLIP (Radford et al., 2021) and ALIGN (Jia et al., 2021) demonstrated a genuinely new paradigm in computer vision: instead of collecting labeled training data for every new application, you could pretrain a single model on noisy image-text pairs from the internet and deploy it directly on downstream tasks without any task-specific data. This paradigm shift promised to eliminate the bottleneck of manual data labeling that had constrained supervised learning for decades.

However, as the authors explicitly state in Section 1, the best CLIP and ALIGN models achieve only ~76% top-1 accuracy on ImageNet—roughly comparable to a supervised ResNet-50 (He et al., 2015), and far behind the state-of-the-art supervised models of the time (87.1% without extra data, 90.88% with extra data). This gap is not merely an academic benchmark question. For zero-shot transfer to become a viable alternative to supervised learning in production systems, it must approach or match supervised performance. A model that is 10–15 percentage points behind cannot replace supervised systems in safety-critical applications (medical imaging, autonomous driving) or high-stakes commercial settings (product classification, content moderation) where every point of accuracy matters.

The gap also matters because zero-shot models possess two properties that supervised models lack, as established by CLIP and ALIGN:

Versatility: A single zero-shot model can be deployed on many downstream tasks without task-specific finetuning, whereas supervised models must be retrained or finetuned for each new task.
Robustness: Zero-shot models suffer much smaller accuracy drops on natural distribution shifts (e.g., ImageNet-V2, ImageNet-Sketch) than supervised models, which can degrade by 40% or more on the same benchmarks (Taori et al., 2020; Szegedy et al., 2013).

These properties make zero-shot transfer inherently valuable, but they are undermined if the baseline accuracy is too low. A robust model that is 15% less accurate than a supervised model may still be less useful in practice. Therefore, narrowing the accuracy gap is the critical step needed to unlock the practical promise of zero-shot transfer—maintaining versatility and robustness while reaching competitive absolute performance.

Why This Gap Exists: Scaling Has Been Piecemeal, Not Systematic

Prior to this work, the recipe for improving image-text contrastive models was intuitive but unsystematic: use more data, larger models, or larger batch sizes. CLIP, for instance, trained on 400M image-text pairs with a ViT-L/14 model. ALIGN scaled the data to 1.7B pairs but used a comparable EfficientNet-based image encoder. Both works demonstrated that scaling helps, but neither systematically explored what happens when you scale all three dimensions simultaneously—data, model size, and batch size—or provided the engineering infrastructure to do so.

This piecemeal approach left several critical questions unanswered:

Does scaling interact? If you increase the dataset by 4× and the model by 3.75×, do the benefits multiply, or does one dimension saturate?
Does batch size scaling matter beyond what more training steps can achieve? Prior work like SimCLR (Chen et al., 2020a) had shown that larger batch sizes help contrastive learning, but it was unclear whether larger batch sizes provide unique benefits or simply serve as a compute-efficient way to see more negative examples.
What is the limiting factor? No one had identified whether the bottleneck was data scarcity, model capacity, or batch size—or whether all three were necessary to push performance to supervised levels.

The authors explicitly position their work as a systematic combined scaling study that answers these questions by pushing all three dimensions to their limits simultaneously.

Where Prior Approaches Fall Short

The paper identifies specific limitations in several categories of prior work:

1. Single-modal contrastive learning reaches diminishing returns at modest scales. Works like SimCLR (Chen et al., 2020a,b) and MoCo (He et al., 2020; Chen et al., 2020c) demonstrated that contrastive learning on single modalities (images only) benefits from large batch sizes, but these benefits were observed to saturate at batch sizes around 8192 for SimCLR. The authors hypothesize—and later demonstrate empirically (Section 10.1)—that this saturation is an artifact of the small datasets and models used in SimCLR. When both the dataset (ALIGN's 1.7B examples vs. ImageNet's 1M) and the model size scale up, larger batch sizes continue to provide benefits well beyond 8192. This is a crucial insight: batch size scaling interacts with data and model scaling, meaning that isolated batch size studies on small configurations can mislead practitioners about the benefits at scale.

2. Existing image-text models do not scale all dimensions. CLIP (Radford et al., 2021) explored multiple model sizes but did not systematically vary batch size—it used a fixed batch size of 32,768 across its main experiments. ALIGN (Jia et al., 2021) used an even smaller batch size of 16,384. Neither work attempted batch sizes beyond 32K, and neither combined large batch sizes with models at the 3B-parameter scale. The consequence is that prior work left open the possibility that significantly larger batch sizes—coupled with correspondingly larger models and datasets—could unlock substantially better performance. BASIC explicitly targets this gap by scaling the batch size to 65,536 (2× CLIP, 4× ALIGN) while simultaneously scaling the model to 3B parameters (3.75× both) and the dataset to 6.6B pairs (16× CLIP, 4× ALIGN).

3. The memory bottleneck for combined scaling has no established solution. Scaling model size and batch size simultaneously creates a compound memory problem. In the contrastive learning framework (Section 3), computing the loss for a batch of size $B$ requires the full $B \times B$ similarity matrix $\mathbf{A}$ to be materialized, plus the activations of both the image encoder $F$ and text encoder $G$ for all $B$ examples. When $B = 65,536$ and the model has 3B parameters, the naive memory requirement far exceeds the 16GB typical of accelerators in 2022. Prior techniques—gradient accumulation (Ott et al., 2018; Zhai et al., 2021), model parallelism (Huang et al., 2019), and rematerialization (Chen et al., 2016)—each address part of the problem, but none had been adapted to the specific structure of the contrastive loss, where the loss function couples all examples in the batch through the row-wise and column-wise softmax operations (Equations 1–3). The paper identifies this as a key implementation gap: standard gradient accumulation is not directly applicable because you cannot compute the contrastive loss on a subset of the batch independently.

4. The theoretical understanding of batch size in image-text contrastive learning is underdeveloped. While the empirical benefit of large batch sizes was observed by CLIP and SimCLR, there was no theoretical framework explaining why larger contrastive batch sizes should help, particularly in the image-text setting. The authors argue that this theoretical gap matters because it impedes principled scaling decisions. Without theory, practitioners cannot predict whether batch size scaling will continue to help at larger scales, or whether it will eventually saturate. The paper directly addresses this by developing the theoretical analysis in Section 6, which proves that the generalization gap decreases as $O(1/\sqrt{B})$ —meaning that larger batch sizes are always beneficial, not just a compute-efficiency trick.

5. Robustness of zero-shot models is observed but unexplained and potentially fragile. CLIP showed that zero-shot models are more robust than supervised models, but did not determine whether this robustness is inherent to the zero-shot paradigm or a byproduct of something else (e.g., the training data distribution). The paper advances this question in Section 9.3 by conducting a novel experiment: taking a converged zero-shot model and finetuning it on more labeled ImageNet data. The result—that more labeled data makes models less robust—is counterintuitive and suggests that robustness is not simply a function of the model architecture or initial training, but is actively lost when models are adapted to a specific labeled distribution. This finding has implications beyond ImageNet: it warns that the very act of supervised finetuning, which is standard practice for adapting pretrained models to downstream tasks, may destroy the robustness properties that make zero-shot models attractive in the first place. The paper does not solve this problem but frames it as a critical open question that invites causal analysis.

6. The pretraining-finetuning hybrid for contrastive models is underexplored. While pretraining the image encoder on labeled data (e.g., JFT) before contrastive training is a natural idea, the paper identifies a non-obvious weakness of this approach in Section 8: the image encoder never sees noisy image-text pairs during the contrastive phase if its weights are frozen. This can cause the model to "completely fail on an easier task—MNIST" because the pretraining labeled dataset may lack digit images, while the noisy image-text dataset contains text that teaches optical character recognition. The paper's solution—a hybrid procedure of pretraining, then training the text encoder with the image encoder frozen, then jointly finetuning both—is motivated by this specific failure mode, which prior work had not documented or addressed.

How This Paper Positions Itself

The paper positions itself as a direct extension and systematic scaling of the CLIP/ALIGN framework, not as an alternative to it. The core methodology—training image and text encoders to maximize cosine similarity between paired images and text while minimizing it for non-paired examples—is unchanged from CLIP and ALIGN. What changes is the scale at which this framework is applied, and the engineering and theoretical infrastructure needed to make that scale feasible and principled.

The paper explicitly frames its contribution along three axes:

Empirical: Demonstrate that combined scaling of data, model, and batch size yields substantial accuracy improvements that bring zero-shot transfer models to within striking distance of supervised performance (85.7% vs. 87–91%).
Engineering: Provide two practical, reproducible methods for overcoming the memory bottleneck that prevents combined scaling on existing hardware, enabling other practitioners to replicate and extend the scaling regime.
Theoretical: Establish—for the first time in the image-text contrastive learning literature—a formal generalization bound that depends inversely on batch size, providing principled justification for batch size scaling beyond empirical heuristics.

The paper also positions itself as extending the robustness analysis of CLIP in a direction that the original authors explicitly cautioned against overinterpreting. By conducting the finetuning experiment in Section 9.3, the paper provides new evidence that the robustness of zero-shot models is at least partially attributable to the fact that they are not trained on the target dataset's labeled examples—a finding that has direct implications for how practitioners should think about the accuracy-robustness tradeoff when deploying models.

3. Technical Approach

3.1 Reader Orientation

BASIC is a compound scaling recipe for training image-text contrastive models that simultaneously pushes three scaling dimensions — dataset size, model size, and batch size — to their joint limit, producing a single zero-shot image classifier that can be deployed on any downstream task without task-specific training. The system solves the problem of combined memory pressure: when you scale all three dimensions together, the accelerator memory required during training explodes beyond what commodity hardware provides, and standard engineering workarounds such as gradient accumulation do not directly apply to the contrastive loss function because its computation couples every example in a batch with every other example. The solution shape is therefore (1) a modified gradient accumulation algorithm specifically designed for the all-pairs structure of the contrastive loss, (2) an alternative exact method using weight sharding and selective rematerialization under the Single-Program Multiple-Data paradigm, and (3) a theoretical analysis that justifies why batch size scaling is worth the engineering effort in the first place.

3.2 Big-Picture Architecture (Diagram in Words)

The system has five major components, arranged in a pipeline that transforms noisy web-crawled image-text pairs into a deployable zero-shot classifier:

Data ingestion and filtering (Section 7.1). A dataset of 6.6B noisy (image, text) pairs — formed by merging ALIGN's 1.7B pairs with 5B additional pairs from JFT — is tokenised with a learned 32K-vocabulary sentence piece model and filtered to remove sequences longer than 64 tokens and any images that structurally resemble evaluation-set images (SSIM ≥ 0.5).
Image encoder (F) and text encoder (G) (Section 7.2). F is a CoAtNet architecture (a hybrid of convolution and self-attention blocks) scaled to 2.4B parameters in the largest variant; G is a standard transformer whose final text representation is the average of all top-layer token embeddings (unlike ALIGN, which uses a [CLS] token). Both map their respective inputs to unit-normalised vectors in $\mathbb{R}^D$ .
Contrastive loss computation (Section 3, Equations 1–3). Given a minibatch of $B$ (image, text) pairs, F and G produce embedding matrices $\mathbf{X}, \mathbf{Y} \in \mathbb{R}^{D \times B}$ , from which a similarity matrix $\mathbf{A} = \mathbf{X}^\top \mathbf{Y} / \tau$ is computed. A symmetric cross-entropy loss — row-wise (image-to-text) plus column-wise (text-to-image) — pushes the diagonal entries (true pairs) of $\mathbf{A}$ high and off-diagonal entries low.
Memory-management infrastructure (Sections 4 and 5). Two complementary strategies enable training at $B = 65,536$ with a 3B-parameter model on 16GB accelerators. The pipelining + gradient accumulation approach (Section 4) "chunks" the batch into microbatches, computes the full similarity matrix $\mathbf{A}$ once while discarding intermediate activations, then re-materialises each microbatch during backpropagation and accumulates gradients directly into the optimizer's moment slots to avoid allocating a separate accumulation buffer. The SPMD + rematerialization approach (Section 5) splits weight tensors across accelerator cores (weight sharding) and recomputes cheap activation layers (batch norm, layer norm, activations) during the backward pass while keeping expensive weighted layers in memory.
Optimizer and training schedule (Sections 8 and 9.1). A custom optimizer (AdaFactorW — factorised second moments like AdaFactor, decoupled weight decay like AdamW, first moments stored in bfloat16 but upcast to float32 for updates) updates the encoders. Training typically follows a three-phase procedure: (1) pretrain F on JFT with a supervised softmax loss, (2) freeze F and train G with the contrastive loss, (3) jointly finetune both at a reduced learning rate.

Information flows through this pipeline as follows: noisy image-text pairs → tokenisation and filtering → forward pass through F and G under memory-management constraints → contrastive loss computation using all $B^2$ pairwise similarities → gradient computation (with rematerialization as needed) → AdaFactorW update. At deployment, only F and G are needed: an image is encoded by F, text prompts for each candidate class are encoded by G, and the class whose text embedding has the highest cosine similarity to the image embedding is selected.

3.3 Roadmap for the Deep Dive

This is an empirical engineering and scaling paper whose core idea is that jointly scaling batch size, model size, and dataset size produces a zero-shot transfer model that closes most of the accuracy gap to supervised methods, and that this scaling is feasible only with purpose-built memory-management techniques. To explain the technical approach comprehensively, I will walk through the components in the following order:

First, the contrastive loss framework itself (Section 3 of the paper), since every engineering decision in subsequent sections is a response to constraints imposed by the specific all-pairs structure of this loss function. Without understanding the loss, the memory bottleneck is unintelligible.
Second, the pipelining and gradient accumulation method (Section 4), which is the more generic of the two memory-management strategies and introduces the key concepts of microbatching the contrastive loss and accumulating gradients directly into optimizer slots. This section also exposes the inexactness of the approach, which motivates the alternative.
Third, the SPMD-based weight sharding and rematerialization method (Section 5), which provides exact computations at faster step times but requires manual design choices. I will explain what "weight sharding" means operationally, which layers are rematerialized and why, and how the two methods compare empirically.
Fourth, the theoretical analysis of batch size scaling (Section 6), which proves that the generalisation gap decreases as $O(1/\sqrt{B})$ and thus provides principled justification for the entire engineering enterprise. I will walk through what the theorem says in operational terms — what the bound depends on, why batch size appears where it does, and what assumptions are required.
Fifth, the data and model scaling choices (Section 7) and the pretraining-finetuning procedure (Section 8), which together specify what is being scaled and how the training is organised beyond the memory-management infrastructure.

3.4 Detailed, Sentence-Based Technical Breakdown

The Image-Text Contrastive Loss Framework

The paper inherits its training objective directly from CLIP (Radford et al., 2021) and ALIGN (Jia et al., 2021), without modification. Understanding this objective is prerequisite to understanding the memory bottleneck, because the loss function couples every image in a batch with every text in the batch, producing an $B \times B$ matrix of pairwise similarities that must be fully materialised before any gradient can be computed.

Encoder structure. Let $\mathbf{x}$ be an arbitrary image and $\mathbf{y}$ be an arbitrary text sequence. The image encoder $F$ maps $\mathbf{x}$ to a $D$ -dimensional vector on the unit sphere: $F(\mathbf{x}) \in \mathbb{S}^D$ . The text encoder $G$ maps $\mathbf{y}$ to a vector in the same space: $G(\mathbf{y}) \in \mathbb{S}^D$ . The desideratum is that $(F(\mathbf{x}), G(\mathbf{y}))$ pairs that are semantically related should have high cosine similarity (large dot product, since vectors are unit-normed), while unrelated pairs should have low similarity.

Batch construction. At each training step, the system samples $B$ image-text pairs $\{(\mathbf{x}_i, \mathbf{y}_i)\}_{i=1}^B$ from the training dataset. The index $i$ encodes the true pairing: $\mathbf{x}_i$ and $\mathbf{y}_i$ come from the same web page or image-caption pair and are assumed to be semantically related, while $\mathbf{x}_i$ and $\mathbf{y}_{j \neq i}$ are assumed to be unrelated (this is the "noisy" weak supervision signal — some unrelated pairs may coincidentally share semantics, but the assumption holds in expectation over large datasets).

Similarity matrix. The encoders produce embedding matrices $\mathbf{X}, \mathbf{Y} \in \mathbb{R}^{D \times B}$ where column $i$ of $\mathbf{X}$ is $F(\mathbf{x}_i)$ and column $j$ of $\mathbf{Y}$ is $G(\mathbf{y}_j)$ . The pairwise similarity matrix is:

$\mathbf{A}_{i,j} = \frac{F(\mathbf{x}_i)^\top G(\mathbf{y}_j)}{\tau}$

where $\tau > 0$ is a learned or fixed temperature parameter that controls the sharpness of the subsequent softmax distributions. Smaller $\tau$ makes the distribution peakier (more confident), larger $\tau$ makes it flatter.

What this matrix computes: for every pair $(i, j)$ in the batch, it computes the cosine similarity between the $i$ -th image embedding and the $j$ -th text embedding, divided by the temperature. The result is a $B \times B$ matrix of real numbers where diagonal entries $\mathbf{A}_{i,i}$ correspond to true pairs and off-diagonal entries correspond to negative pairs.

Why this form: the temperature $\tau$ serves as a scaling knob. Without it, the dot product magnitude depends on the embedding dimension and the norms of $F$ and $G$ , which can vary during training. Dividing by $\tau$ stabilises the softmax input scale, and $\tau$ can be learned (as in CLIP) or fixed. The unit-sphere normalisation of $F$ and $G$ ensures that the dot product is bounded in $[-1, 1]$ , preventing the similarity scores from diverging.

Row-wise (image-to-text) loss. For each image $i$ , the model treats the $B$ text sequences in the batch as candidate matches. The probability that text $j$ is the correct match for image $i$ is computed via a softmax over the $i$ -th row of $\mathbf{A}$ :

$\text{RowLoss}_B = -\frac{1}{B} \sum_{i=1}^B \log \frac{\exp(\mathbf{A}_{i,i})}{\sum_{k=1}^B \exp(\mathbf{A}_{i,k})}$

where the numerator is the exponential of the similarity for the true pair $(i, i)$ , and the denominator sums over all $B$ text candidates for image $i$ .

What it computes: for each image, it computes the negative log-likelihood that the correct text is selected from the batch of $B$ texts. This is a $B$ -way classification problem per image. The overall row loss is the average over all $B$ images.

Column-wise (text-to-image) loss. Symmetrically, for each text $j$ , the model computes the probability that image $i$ is its correct match via a softmax over the $j$ -th column of $\mathbf{A}$ :

$\text{ColumnLoss}_B = -\frac{1}{B} \sum_{j=1}^B \log \frac{\exp(\mathbf{A}_{j,j})}{\sum_{k=1}^B \exp(\mathbf{A}_{k,j})}$

What it computes: the negative log-likelihood that the correct image is selected from the batch for each text, averaged over all texts.

Total contrastive loss. The final loss is the symmetric average:

$\text{ContrastiveLoss}_B = \frac{\text{RowLoss}_B + \text{ColumnLoss}_B}{2}$

Why this symmetry: using both directions ensures that the embeddings are good at both image-to-text retrieval (given an image, find the matching text) and text-to-image retrieval (given a text, find the matching image). For zero-shot classification, the model is used in text-to-image mode: given an image, the text prompt for each class is encoded, and the class with the highest similarity is chosen. Training both directions ensures the embedding space is well-formed for this inference-time task.

The critical memory implication. To compute $\text{ContrastiveLoss}_B$ , you need all $B^2$ entries of $\mathbf{A}$ . This means you must run the forward pass of $F$ on all $B$ images and the forward pass of $G$ on all $B$ texts, and then compute all $B^2$ dot products, before you can evaluate the loss and begin backpropagation. You cannot split the batch into independent sub-batches and compute partial losses, because the softmax denominator for row $i$ depends on all $B$ text embeddings, not just a subset. This is the fundamental obstacle that Sections 4 and 5 address.

Pipelining and Gradient Accumulation for the Contrastive Loss

The pipelining approach (Section 4) is the more general of the two memory-management strategies: it can scale to arbitrarily large contrastive batch sizes $B$ regardless of the model architecture, at the cost of some computational overhead and two sources of inexactness. The core insight is that while you must compute the full similarity matrix $\mathbf{A}$ to evaluate the loss, you do not need to keep the intermediate activations of $F$ and $G$ for all $B$ examples in memory simultaneously during backpropagation. You can discard them after the forward pass, and then recompute (rematerialise) them one microbatch at a time during the backward pass.

The memory bottleneck, restated with numbers. Training with a contrastive batch size $B$ requires storing:

The activations of $F$ for all $B$ images (grows with model depth and image resolution).
The activations of $G$ for all $B$ texts (grows with sequence length and model depth).
The $B \times B$ similarity matrix $\mathbf{A}$ (grows quadratically in $B$ ).
The weight tensors and their optimizer moment slots for both encoders (grows with parameter count).

For BASIC-L with $B = 65,536$ and 3B total parameters, the naive memory requirement far exceeds the ~16GB typical of a single TPU core or GPU in 2022. The $\mathbf{A}$ matrix alone, stored in float32, would be $65,536^2 \times 4 \text{ bytes} \approx 17.2 \text{ GB}$ — already exceeding the per-device budget, and that is before accounting for model weights and activations.

Vanilla gradient accumulation is not applicable. Standard gradient accumulation (GradAccum) splits a batch of $B$ examples into $K = B/M$ microbatches of size $M$ , sequentially computes the loss and gradient for each microbatch, and averages the gradients. This works for losses that decompose as $\mathcal{L} = \frac{1}{B} \sum_{i=1}^B \ell(\mathbf{e}_i)$ where each per-example loss $\ell(\mathbf{e}_i)$ depends only on example $i$ . The contrastive loss does not decompose this way: the softmax denominator for example $i$ depends on all $B$ examples in the batch. Computing the loss on a microbatch of size $M$ would use a denominator that sums over only $M$ candidates rather than $B$ , producing a different loss function that penalises the model for failing to distinguish the true pair from only $M-1$ negatives rather than $B-1$ . This is not just a numerical approximation error — it fundamentally changes the learning signal, making the task artificially easier.

Key observation enabling the method. The authors observe that while the loss requires the full $\mathbf{A}$ , the gradients with respect to the encoder weights can be decomposed across microbatches after the gradient with respect to $\mathbf{A}$ (or equivalently, with respect to the embeddings $\mathbf{X}$ and $\mathbf{Y}$ ) has been computed. Specifically, once you know $\nabla_{\mathbf{X}} \text{ContrastiveLoss}_B$ and $\nabla_{\mathbf{Y}} \text{ContrastiveLoss}_B$ , the backpropagation through $F$ and $G$ can proceed independently for each example — example $i$ 's contribution to the weight gradients depends only on the $i$ -th column of $\nabla_{\mathbf{X}} \text{ContrastiveLoss}_B$ and the activations of $F$ for that example.

Algorithm 1 — the modified GradAccum procedure. The paper presents this as pseudocode (Algorithm 1 in the paper) with a step-by-step memory analysis. Here is what happens operationally:

Step 1: Allocate embedding matrices. Allocate space for $\mathbf{X}, \mathbf{Y} \in \mathbb{R}^{D \times B}$ . Memory cost: $\Theta(BD)$ . For BASIC-L with $D$ on the order of hundreds to low thousands, this is modest relative to the activation memory.

Step 2: Forward pass in microbatches. Iterate over microbatches of size $M$ (the largest batch size that fits in memory). For each microbatch $J \subset \{1, \dots, B\}$ with $|J| = M$ :

Compute $\mathbf{X}_{:,J} \leftarrow F(\mathbf{x}_J)$ — the embeddings for images in this microbatch. Discard all intermediate activations of $F$ (no gradient checkpointing for this pass).
Compute $\mathbf{Y}_{:,J} \leftarrow G(\mathbf{y}_J)$ — similarly for texts, discarding $G$ 's activations.
Store the embeddings in the pre-allocated matrices.

Memory cost during this step: $\Theta(M \cdot \text{Mem}(F))$ for the image encoder forward pass, $\Theta(M \cdot \text{Mem}(G))$ for the text encoder forward pass, plus $\Theta(BD)$ for the embedding matrices. The embedding matrices persist across microbatches; the encoder activations are freed after each microbatch.

Step 3: Compute the loss and its gradient with respect to $\mathbf{A}$ . With the full $\mathbf{X}$ and $\mathbf{Y}$ available:

Compute $\mathbf{A} \leftarrow (\mathbf{X}^\top \mathbf{Y}) / \tau$ . Memory: $\Theta(B^2)$ for the similarity matrix.
Compute $\text{ContrastiveLoss}_B$ as in Equation 3.
Backpropagate to obtain $d\mathbf{A} = \nabla_{\mathbf{A}} \text{ContrastiveLoss}_B$ . This is a $B \times B$ matrix.
Compute the embedding gradients: $d\mathbf{X} \leftarrow \mathbf{Y} \cdot d\mathbf{A}$ and $d\mathbf{Y} \leftarrow \mathbf{X} \cdot d\mathbf{A}$ . These are $D \times B$ matrices.

Memory cost: $\Theta(B^2)$ for $\mathbf{A}$ and $d\mathbf{A}$ , plus $\Theta(BD)$ for $d\mathbf{X}$ and $d\mathbf{Y}$ . At $B = 65,536$ , the $B^2$ term is the dominant cost — but it is temporary, and the embedding gradient computation frees the $B^2$ tensors once $d\mathbf{X}$ and $d\mathbf{Y}$ are obtained.

Step 4: Backpropagate through the encoders in microbatches. Iterate over microbatches again. For each microbatch $J$ :

Rerun the forward pass of $F$ on $\mathbf{x}_J$ , this time saving intermediate activations (or using gradient checkpointing within the microbatch if $M$ is still large). This is the rematerialization step.
Run the backward pass of $F$ , starting from $d\mathbf{X}_{:,J}$ (the pre-computed embedding gradient for this microbatch) and accumulating weight gradients for $F$ 's parameters.
Do the same for $G$ using $\mathbf{y}_J$ and $d\mathbf{Y}_{:,J}$ .

Memory cost per microbatch: $\Theta(M \cdot \text{Mem}(F))$ for the saved activations of $F$ during its repeated forward pass, similarly for $G$ . The peak memory usage is therefore $\Theta(M \cdot \max\{\text{Mem}(F), \text{Mem}(G)\})$ (since the encoders are processed sequentially), rather than $\Theta(B \cdot (\text{Mem}(F) + \text{Mem}(G)))$ as in the naive approach.

Why this works: the key decomposition is that $\nabla_{\theta_F} \text{ContrastiveLoss}_B = \sum_{j=1}^B \nabla_{\theta_F} F(\mathbf{x}_j) \cdot \nabla_{F(\mathbf{x}_j)} \text{ContrastiveLoss}_B$ . Once you have $\nabla_{F(\mathbf{x}_j)} \text{ContrastiveLoss}_B$ (which is column $j$ of $d\mathbf{X}$ ), the backpropagation through $F$ for example $j$ is independent of all other examples. You can therefore batch the per-example backpropagation into microbatches of any size, trading off recomputation cost (from the repeated forward passes) against memory savings.

Accumulating microbatch gradients without an accumulation buffer. A second memory problem arises after Algorithm 1 produces a stream of microbatch weight gradients $c_1, \dots, c_{B/M}$ . In standard GradAccum, these must be summed into a gradient accumulation buffer $\bar{g}$ of the same size as the model weights. For a 3B-parameter model, $\bar{g}$ alone occupies ~11GB in float32 — exactly the kind of memory the method is trying to save.

The paper observes that optimizers like Adam, AdaFactor, and AdamW already maintain moment buffers (slots) of the same size as the model weights — the first moment $v_1$ and second moment $v_2$ . The idea is to accumulate the microbatch gradients directly into these moment buffers, bypassing the separate accumulation buffer.

For the first moment $v_1$ , the standard Adam update is: $v_1^{(t)} = \beta_1 v_1^{(t-1)} + (1 - \beta_1) \bar{g}$

where $\bar{g} = \frac{1}{K} \sum_{i=1}^K c_i$ is the averaged minibatch gradient (with $K = B/M$ ). Instead of computing $\bar{g}$ first and then updating $v_1$ once, the paper proposes $K$ sequential updates:

$v_1 \leftarrow k_i v_1 + (1 - \beta_1) c_i, \quad k_i = \begin{cases} \beta_1 & \text{if } i = 1 \\ 1/K & \text{otherwise} \end{cases}$

What this computes: for the first microbatch ( $i = 1$ ), the update is the standard Adam rule with a full-sized step. For subsequent microbatches ( $i > 1$ ), the decay factor is reduced to $1/K$ and the step size remains $(1 - \beta_1)$ . After all $K$ updates, $v_1$ holds the same value as if the single-update rule had been applied with $\bar{g}$ .

Why this form works: the standard Adam update is a convex combination. Applying it $K$ times with appropriately scaled coefficients preserves the final value. The derivation relies on the linearity of the first-moment update: the sum $\sum_{i=1}^K c_i$ can be accumulated incrementally because addition is commutative.

The second-moment complication. The same trick does not straightforwardly apply to $v_2$ , because the Adam second-moment update uses the square of the gradient: $v_2^{(t)} = \beta_2 v_2^{(t-1)} + (1 - \beta_2) \bar{g}^2$

The sum of squares $\frac{1}{K} \sum_{i=1}^K c_i^2$ is generally not equal to the square of the sum $(\frac{1}{K} \sum_{i=1}^K c_i)^2$ . Their difference is the variance:

$\frac{1}{K} \sum_{i=1}^K c_i^2 - \left(\frac{1}{K} \sum_{i=1}^K c_i\right)^2 = \mathbf{Var}[c_i]$

where $\mathbf{Var}[c_i]$ is the variance of the microbatch gradients.

Estimating the variance correction. The paper estimates $\mathbf{Var}[c_i]$ by treating each $c_i$ as the sample mean of $M$ per-example gradients drawn uniformly from the full batch. Using the identity for the variance of a sample mean:

$\mathbf{Var}[c_i] = \frac{\mathbf{Var}[\mathbf{g}]}{M}$

where $\mathbf{Var}[\mathbf{g}]$ is the variance of the per-example gradients across the full batch of $B$ examples.

To estimate $\mathbf{Var}[\mathbf{g}]$ without computing all per-example gradients, the paper leverages the data-parallelism setting. Under data parallelism with $R$ replicas, each replica processes $M/R$ examples per microbatch, and the per-replica gradient $d_1, \dots, d_R$ (before the all-reduce) represents the sample mean of $M/R$ examples. The variance of these per-replica gradients satisfies:

$\mathbf{Var}[d] = \frac{\mathbf{Var}[\mathbf{g}]}{M/R}$

Thus, $\mathbf{Var}[\mathbf{g}] = (M/R) \cdot \mathbf{Var}[d]$ , and $\mathbf{Var}[c_i] = \mathbf{Var}[d] / R$ . With $R$ replicas, $\mathbf{Var}[d]$ can be estimated empirically from the observed per-replica gradients without storing the full per-example gradient set.

What this enables: the optimizer can maintain $v_2$ using the sum-of-squares accumulation, then correct it by subtracting the estimated variance term before computing the final weight update. This avoids allocating a separate buffer for $\bar{g}$ while maintaining approximately correct second-moment statistics.

Two sources of inexactness. The paper acknowledges that this method has two approximations that make it inexact compared to true full-batch training:

Gradient accumulation approximation: The variance correction for $v_2$ uses an estimate rather than the exact value, and the treatment of $c_i$ as i.i.d. sample means is an approximation (the per-example gradients within a microbatch are not truly independent, since the shared loss computation couples them through the softmax denominator, but this coupling is exactly what the method is designed to work around — the approximation is that after accounting for the $d\mathbf{X}$ and $d\mathbf{Y}$ gradients, the per-example contributions to the weight gradients are treated as independent for variance estimation purposes).
Batch normalisation inconsistency: If $F$ contains batch normalisation layers, the statistics computed over a microbatch of size $M$ differ from those that would be computed over the full batch of size $B$ . When $M \ll B$ , these discrepancies can cause substantial covariate shift in the image encoder. The paper notes that this issue is mitigated when using architectures that avoid batch normalisation — such as Vision Transformers (which use layer normalisation) or NFNets (which use no normalisation) — but it remains a concern for architectures like CoAtNet that include convolutional blocks.

The pipelining method is generic but inexact. The fundamental advantage of the pipelining + GradAccum approach is that it can scale to any contrastive batch size $B$ , regardless of model architecture, because the microbatch size $M$ can be chosen arbitrarily. If $B$ grows, you simply increase the number of microbatches; the per-microbatch memory footprint remains $\Theta(M \cdot \max\{\text{Mem}(F), \text{Mem}(G)\})$ . The cost is the overhead of the repeated forward passes (the encoders are run twice — once to compute embeddings without saving activations, once to rematerialise for backpropagation) and the inexactness from gradient accumulation and batch normalisation. This motivates the alternative SPMD approach described next.

SPMD-Based Weight Sharding and Rematerialization

The SPMD approach (Section 5) takes a different philosophy: instead of trading computation for memory by rematerialising activations across microbatches, it reduces the memory per accelerator core by distributing the model weights themselves across multiple cores, so that each core stores only a fraction of the weight tensors. This avoids the inexactness of gradient accumulation entirely — the computations are exact because the full batch is processed in parallel across cores — but it requires manual design of the sharding and rematerialization strategy, making it less generic than the pipelining approach.

What is SPMD? Single-Program Multiple-Data is a programming model where every accelerator core executes the same program (the training loop) but operates on different slices of the data and model tensors. Unlike pipeline parallelism (where different cores execute different parts of the model sequentially), SPMD parallelism distributes each operation across cores. For example, a matrix multiplication $W x$ can be executed by splitting the weight matrix $W$ across cores, each core computing a partial product, and then combining results.

Why weight sharding matters. In modern deep learning optimizers (Adam, AdaFactor, AdamW), every weight tensor is associated with two gradient moment tensors (first and second moments) of the same shape. With vanilla data parallelism, all three tensors (weight + two moments) are replicated to all accelerator cores. For a 3B-parameter model, the weights alone occupy ~11GB in float32, and the moment tensors triple that to ~33GB — far exceeding the ~16GB of a typical accelerator core in 2022. Weight sharding splits each of these tensors across cores, so that each core stores only $1/R$ of the weights and moments (where $R$ is the number of cores per replica). Summing across cores restores the full tensor when needed for computation.

The training cluster topology. The paper trains on a cluster of 2048 TPUv3 cores. These are partitioned into $R_{\text{replicas}}$ replicas, each using $2048 / R_{\text{replicas}}$ cores. The weights of $F$ and $G$ are split into $2048 / R_{\text{replicas}}$ equal parts, one per core in a replica. Across replicas, the split weight tensors are replicated identically. For instance, with $R_{\text{replicas}} = 512$ , each replica uses 4 cores, and each core stores 1/4 of every weight tensor. Across the 512 replicas, these 1/4-shards are replicated 512 times (the standard data-parallelism replication). The paper empirically finds that "using 512 replicas and 4 cores per replica offers a good balance" between memory savings (more cores per replica means each core stores less) and communication overhead (more cores per replica means more cross-core data movement).

Critical design choice: only weights are sharded, not activations. The paper shards model weights but not the input batch. This means each of the 2048 cores receives $B/2048$ examples of the full batch, regardless of $R_{\text{replicas}}$ . For $B = 65,536$ , each core processes 32 examples. This design choice disentangles weight sharding from the rematerialization strategy: the per-core batch size is fixed by the total number of cores, not by the sharding granularity, so rematerialization decisions can be made independently of the parallelism configuration.

Figure 1 — the 2D convolution sharding example. The paper illustrates weight sharding on a 2D convolution operation with a $3 \times 3$ kernel. The input tensor has shape $[N, H, W, i]$ (batch size, height, width, input channels) and is sharded along the batch dimension so that each core processes $N/4$ examples. The kernel has shape $[3, 3, i, o]$ (kernel height, kernel width, input channels, output channels) and is sharded along the input channel dimension so that each core stores a shard of shape $[3, 3, i/4, o]$ . Before the convolution executes, each core gathers the kernel shards from all other cores and concatenates them along the input channel axis, forming the complete $[3, 3, i, o]$ kernel. After the convolution, the complete kernel is discarded from memory — only the 1/4-shard persists on each core for the next training step.

Why shard along the input channel dimension? Sharding along input channels means the partial convolution result computed by each core is still a valid partial result — the sum over input channels can be completed by an all-reduce across cores after the convolution. This is more communication-efficient than sharding along output channels, which would require concatenating results rather than summing them.

What is not sharded: batch norm and layer norm weights. The paper explicitly exempts batch normalisation and layer normalisation parameters ( $\beta$ , $\gamma$ , and the moving average statistics) from sharding. These are one-dimensional vectors (size equal to the number of channels or hidden dimensions), so their memory footprint is negligible compared to convolutional kernel weights or transformer feed-forward weights. Replicating them to all cores avoids the cross-core communication overhead of gathering and scattering tiny tensors, which would be dominated by latency.

Rematerialization heuristic: keep weight-involving layers, recompute the rest. The paper's rematerialization strategy (Section 5.2) is built on a simple cost-benefit heuristic: rematerialize layers that are fast to recompute but consume significant activation memory. Since weight sharding already imposes cross-core communication overhead on layers that involve weights (convolutions, attention projections, dense feed-forwards), recomputing these layers in the backward pass would incur that communication cost twice — once in the forward pass, once in the rematerialization forward pass. Therefore, the paper keeps the outputs of almost all weight-involving layers in memory.

In contrast, layers that do not involve weights — activation functions (ReLU, GELU, Swish), batch normalisation, layer normalisation, and pooling operations — are computationally lightweight (typically element-wise operations or simple reductions) but their output tensors have the same spatial dimensions as the input tensors, meaning they consume significant activation memory. These are all rematerialized: discarded after the forward pass, and recomputed during the backward pass only when their values are needed for gradient computation.

Figure 2 — per-block rematerialization maps. The paper provides explicit rematerialization maps for the two block types in the CoAtNet architecture:

Mobile Inverted Convolution blocks (MBConv, left panel of Figure 2): All batch normalisation layers, all activation layers, and all layers in the Squeeze-and-Excitation (SE) sub-block are rematerialized. The depthwise and pointwise convolutions are kept in memory because their weights are sharded and recomputing them would incur cross-core communication.
Transformer blocks (right panel of Figure 2): Only the layer normalisation layers and activation layers are rematerialized. The attention projections (query, key, value, output) and feed-forward dense layers are kept in memory because they involve sharded weights.

What this achieves quantitatively. The paper states that with weight sharding and this rematerialization heuristic, "each of our forward-backward pass is 1.4 times slower than the vanilla implementation of the same batch size." In exchange, the peak memory per core is reduced below the 16GB threshold, enabling training at $B = 65,536$ with the 3B-parameter model. This 1.4× overhead is substantially less than the pipelining approach (which runs the encoders twice in the forward direction), making SPMD the faster option when it can be applied.

Two exceptions to the heuristic:

Squeeze-and-Excitation blocks are fully rematerialized (including convolutions). SE blocks (Hu et al., 2018) use $1 \times 1$ convolutions with a reduced internal channel dimension (typically squeezed to 1/4 or 1/16 of the input channels, then expanded back). Because the internal channel count is small, these convolutions have few FLOPs and small weight tensors. The paper decides that the recomputation cost is low enough to justify rematerializing even the weight-involving layers in SE blocks. Additionally, the SE convolution weights are replicated (not sharded) to all cores because they are tiny, which eliminates the cross-core communication overhead during recomputation.
Batch norm and layer norm weights are replicated, not sharded. As described above, the memory savings from sharding 1D vectors are negligible, and the communication overhead of gathering/scattering them would dominate their computational cost.

Why SPMD provides exact computations. Unlike the pipelining approach, the SPMD approach processes the entire batch of $B$ examples simultaneously (distributed across cores). The $\mathbf{A}$ matrix is computed from the full $B \times D$ embedding matrices, assembled via all-gather operations across cores if needed. The contrastive loss uses the exact denominator $\sum_{k=1}^B \exp(\mathbf{A}_{i,k})$ with all $B$ candidates. There is no microbatch splitting of the loss computation, no gradient accumulation approximation, and no batch normalisation inconsistency (since each core processes its subset of the batch, but the batch norm statistics can be synchronised across all cores via all-reduce if needed — though the paper mitigates this by using architectures with layer norm where possible).

Table 2 — empirical comparison of SPMD vs. pipelining. The paper profiles both methods across model sizes (small and medium), batch sizes ( $2^{16}$ to $2^{20}$ ), and reports step times and peak memory. Key patterns:

SPMD is consistently faster. For the medium-sized model at $B = 2^{20}$ , SPMD's total step time is 8361 ms vs. 9781 ms for pipelining — a ~15% reduction. The gap is primarily in the backward pass: SPMD's backward pass (4912 ms) is substantially faster than pipelining's (6130 ms), because SPMD only rematerializes the cheap activation and normalisation layers, while pipelining rematerializes the entire encoder forward passes for each microbatch.
Pipelining uses less memory. At $B = 2^{20}$ with the medium model, pipelining peaks at 12.6 GB vs. SPMD's 15.4 GB. This is because pipelining's memory footprint is bounded by the microbatch size $M$ , which does not grow with $B$ . In contrast, SPMD's memory footprint grows with the per-core batch size, which does increase with $B$ (or requires reducing the number of replicas, which increases the sharding granularity and communication overhead).
Data parallelism is the baseline when it fits. For small models at modest batch sizes ( $B \leq 2^{19}$ for the small model), vanilla data parallelism with all weights replicated is feasible and fastest. SPMD and pipelining only become necessary when data parallelism runs out of memory (OOM).
Pipelining scales to arbitrarily large $B$ ; SPMD requires redesign. The paper explicitly notes that for pipelining, "increasing the contrastive batch size $B$ only leads to more microbatches, but does not change the micro batch size, and so the accelerator's memory remains constant." For SPMD, if $B$ grows beyond $2^{20}$ , the per-core batch size grows proportionally, and the rematerialization strategy must be redesigned to rematerialize a larger fraction of the encoders to keep memory within bounds.

The tradeoff summarised. Choose pipelining + GradAccum when you need generality (works with any architecture, scales to any batch size) and can tolerate inexact gradient accumulation and slower step times. Choose SPMD when you need speed and exactness, can commit to a specific model architecture and parallelism configuration, and your batch size is within the range that your sharding design supports.

Theoretical Analysis of Batch Size Scaling

The theoretical analysis in Section 6 addresses a fundamental question: why should scaling the contrastive batch size help? Prior work had observed empirically that larger batch sizes improve performance (SimCLR, CLIP), but there was no formal characterisation of how much they help or why they help in a way that more training steps cannot replicate. The paper provides a generalisation bound that depends explicitly on the contrastive batch size $B$ , showing that the generalisation gap shrinks as $O(1/\sqrt{B})$ .

Setup: normalised training and testing losses. The paper defines normalised versions of the training and testing losses to avoid a subtle scaling issue. The unnormalised contrastive loss (without the $B$ multiplier) approaches zero as $B \to \infty$ because the softmax denominator grows with $B$ . Analysing the unnormalised loss would misleadingly predict that larger batch sizes are always beneficial simply because the loss scale shrinks — a trivial scaling artifact, not a genuine learning benefit. The normalised losses remove this artifact.

The normalised training loss for a batch of size $B$ is:

$\hat{\ell}_B(x, y) = -\frac{\exp(F(x)^\top G(y))}{\frac{1}{B} \sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k))}$

where $x$ is an image, $y$ is its paired text, and $\hat{y}_1, \dots, \hat{y}_B$ are the $B$ text sequences in the current training batch.

What this computes: the negative of the ratio between (1) the exponentiated similarity of the true pair and (2) the average exponentiated similarity across all $B$ text candidates in the batch. The $1/B$ normalisation of the denominator ensures that the denominator converges to $\mathbb{E}_{\hat{y}}[\exp(F(x)^\top G(\hat{y}))]$ as $B \to \infty$ , rather than growing without bound. The numerator is not averaged because it involves only the single true text $y$ .

Why this form: multiplying the unnormalised loss by $B$ (which is equivalent to replacing $\sum_{k=1}^B$ with $\frac{1}{B}\sum_{k=1}^B$ in the denominator after taking the negative log) keeps the loss on a meaningful scale as $B$ grows. It also connects the training objective to the testing objective, defined as:

$\bar{\ell}_M(x, y) = -\frac{\exp(F(x)^\top G(y))}{\mathbb{E}_{\bar{y}}[\exp(F(x)^\top G(\bar{y}))]}$

where the denominator is the expected exponentiated similarity over the true distribution of text sequences (replaced in practice by an empirical average over $M$ test-time text prompts). The training loss $\hat{\ell}_B$ with the $1/B$ normalisation is a finite-sample approximation to the population loss $\bar{\ell}_M$ , with the approximation error shrinking as $B$ grows.

The generalisation gap. The quantity the theorems bound is:

$\mathbb{E}_{x, y}[\bar{\ell}_M(x, y)] - \hat{\mathbb{E}}_S[\hat{\ell}_B(x, y)]$

where $\mathbb{E}_{x,y}$ is the expectation over the true data distribution, $\hat{\mathbb{E}}_S$ is the empirical average over a training set $S$ of size $m$ , and the batch size $B$ governs the loss function used to evaluate each training example (not the number of training examples, which is $m$ ).

What this expression means operationally: the first term is the expected loss on new (test) data using the population-level contrastive objective. The second term is the average loss on the training data using the finite-batch contrastive objective. The difference is the generalisation gap — how much worse the model performs at test time than at training time. The goal of the theorems is to show that this gap can be bounded by a quantity that decreases as $B$ increases.

Assumptions for Theorem 1 (deep neural network case). Theorem 1 specialises the bound to standard deep neural networks with the following structure:

$G(y) = (\omega_L \circ \sigma_{L-1} \circ \omega_{L-1} \circ \sigma_{L-2} \cdots \sigma_1 \circ \omega_1)(y)$ $F(x) = (\omega'_{L'} \circ \sigma'_{L'-1} \circ \omega'_{L'-1} \circ \sigma'_{L'-2} \cdots \sigma'_1 \circ \omega'_1)(x)$

where $\omega_l(q) = W_l q$ is a linear transformation (weight matrix multiplication) and $\sigma_l$ is an element-wise nonlinear activation function. The activation functions are assumed to be 1-Lipschitz and positive homogeneous — a technical condition satisfied by ReLU and its variants, which means $\sigma(\alpha x) = \alpha \sigma(x)$ for $\alpha > 0$ and $|\sigma(x) - \sigma(y)| \leq |x - y|$ .

Additional boundedness assumptions (with probability one):

$\exp(F(x)^\top G(y)) \leq c_1$ (exponentiated similarity is bounded)
$\hat{\ell}_B(x, y) \leq c_2$ (training loss is bounded)
$\|F(x)\|_2 \leq c_3$ , $\|F(x_i)\|_2 \leq c_5$ (image embedding norm bounded)
$\frac{\exp(F(x)^\top G(y))}{\frac{1}{B} \sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k))} \leq c_4$ (ratio in normalised loss is bounded)
$\|v\|_2 \leq c_6$ where $v_i = F(x)_i - \frac{\sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k)) G(\hat{y}_k)_i}{\sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k))}$ (a gradient-related quantity is bounded)
$\|y\|_2 \leq c_7$ , $\|x\|_2 \leq c_8$ (inputs bounded)
$\gamma(x) = \mathbb{E}_{\bar{y}}[\exp(F(x)^\top G(\bar{y}))] - \frac{1}{B} \sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k))$ is $c_9$ -Lipschitz (the approximation error of the batch average for the expected denominator is Lipschitz in $x$ )

Theorem 1 statement. With probability at least $1 - \delta$ , for all $F \in \mathcal{F}$ and $G \in \mathcal{G}$ (where $\mathcal{F}$ and $\mathcal{G}$ are the hypothesis classes defined by norm constraints on the weight matrices):

$\mathbb{E}_{x, y}[\bar{\ell}_M(x, y)] - \hat{\mathbb{E}}_S[\hat{\ell}_B(x, y)] \leq \frac{Q_1}{\sqrt{m}} + \frac{Q_2}{\sqrt{2B}} + c_2 \sqrt{\frac{\ln(2/\delta)}{2m}}$

where $Q_1$ and $Q_2$ are constants that depend on the network architecture (depths $L, L'$ , embedding dimension $D$ , weight matrix norms $M_l, M'_l$ , and the boundedness constants $c_1$ through $c_9$ ).

What this means in plain terms: the generalisation gap decomposes into three terms:

$Q_1 / \sqrt{m}$ — a term that decays with the square root of the number of training examples $m$ . This is the standard statistical learning theory rate for models whose Rademacher complexity scales as $O(1/\sqrt{m})$ . Larger training sets reduce the gap.
$Q_2 / \sqrt{2B}$ — a term that decays with the square root of the contrastive batch size $B$ . This is the novel contribution: even with infinite training data ( $m \to \infty$ ), the generalisation gap does not vanish unless $B$ is also large. The batch size appears here because the training loss $\hat{\ell}_B$ is a finite-sample approximation to the population loss $\bar{\ell}_M$ , and the approximation error — the difference between the $1/B$ -averaged denominator and the population expectation — contributes to the generalisation gap.
$c_2 \sqrt{\ln(2/\delta) / (2m)}$ — a confidence term that accounts for the probability $\delta$ that the bound fails to hold. This is standard in PAC-style bounds and decays with $m$ as well.

Why batch size appears where it does. The key mathematical step (elaborated in Appendix H of the paper) is decomposing the gap into terms that involve the difference between the empirical average $\frac{1}{B} \sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k))$ and the population expectation $\mathbb{E}_{\bar{y}}[\exp(F(x)^\top G(\bar{y}))]$ . This difference appears because the training loss uses the batch average in the denominator while the testing loss uses the population expectation. The Lipschitz assumption on $\gamma(x)$ (the difference function) allows this difference to be bounded using Rademacher complexity arguments over the batch samples $\hat{y}_1, \dots, \hat{y}_B$ , yielding the $1/\sqrt{B}$ rate. If $B$ is small, the batch average is a poor estimate of the population expectation, and the model may learn to exploit the specific negatives in each batch rather than learning generally useful representations — this manifests as a larger generalisation gap.

The constant $Q_2$ depends on $\tilde{Q}_{2,1}$ and $\tilde{Q}_{2,2}$ . The decomposition shows that $Q_2 = c_1 \mathbb{E}_{x,y}[A(x,y)] (\tilde{Q}_{2,1} + \tilde{Q}_{2,2})$ where $A(x,y) = \frac{\exp(F(x)^\top G(y))}{(\frac{1}{B} \sum_{k=1}^B \exp(F(x)^\top G(\hat{y}_k)))\mathbb{E}_{\hat{y}}[\exp(F(x)^\top G(\hat{y}))]}$ . The $\tilde{Q}_{2,1}$ term captures the complexity of the $\gamma(x)$ function (involving the Lipschitz constant $c_9$ and the input bound $c_8$ ), while $\tilde{Q}_{2,2}$ captures the complexity of the text encoder $G$ (involving its depth, weight norms, and the embedding norm bound $c_3$ ). Crucially, both terms are multiplied by $1/\sqrt{B}$ , so increasing $B$ directly reduces the generalisation gap regardless of architectural details.

Theorem 2 — the general case. Theorem 2 extends the insight to arbitrary model classes $\mathcal{F}$ and $\mathcal{G}$ (not just deep neural networks with specific activation functions). The bound takes the form:

$\mathbb{E}_{x, y}[\bar{\ell}_M(x, y)] - \hat{\mathbb{E}}_S[\hat{\ell}_B(x, y)] \leq \frac{C_1}{\sqrt{2B}} + c_2 \sqrt{\frac{\ln(2/\delta)}{2m}} + C_2 \sum_{k=1}^D (\mathcal{R}_m(\mathcal{F}_k) + \mathcal{R}_m(\mathcal{G}_k)) + C_3 \tilde{\mathcal{R}}_B(G)$

where $\mathcal{R}_m(\cdot)$ denotes Rademacher complexity over $m$ samples, $\mathcal{F}_k = \{x \mapsto F(x)_k : F \in \mathcal{F}\}$ is the set of possible values for the $k$ -th coordinate of the image embedding, $\mathcal{G}_k$ is similarly the $k$ -th coordinate of the text embedding, and $\tilde{\mathcal{R}}_B(G) = \mathbb{E}_{y,\xi}[\sup_{G \in \mathcal{G}} \frac{1}{B} \|\sum_{i=1}^B \xi_i G(y_i)\|_2^2]$ is a modified Rademacher complexity that measures the complexity of the text encoder's output space.

What this adds: the general-case bound still includes the $C_1/\sqrt{2B}$ term, confirming that the batch size benefit is universal and not an artifact of the deep network assumptions. The additional $\mathcal{R}_m$ and $\tilde{\mathcal{R}}_B$ terms capture the model complexity — for models that are too expressive relative to the data size $m$ and batch size $B$ , the generalisation gap can still be large. However, for standard architectures, $\mathcal{R}_m(\mathcal{F}_k) = O(1/\sqrt{m})$ and $\tilde{\mathcal{R}}_B(G) = O(1/\sqrt{B})$ , so these terms do not dominate the bound.

Practical implication of the theory. The theory establishes that batch size scaling is not merely a compute-efficiency trick (seeing more negatives per step). It provides a principled benefit: larger batch sizes reduce the gap between the finite-batch training objective and the population contrastive objective, which directly translates to better generalisation. This explains why training for more steps with a smaller batch size (which sees the same total number of negatives over the course of training) cannot fully compensate — the loss function itself is different because the denominator in each step's softmax is a poor approximation to the population expectation when $B$ is small. The model trained with small $B$ learns to solve an easier problem (discriminating among $B$ negatives) rather than the true problem (discriminating among the population of all possible texts), and this discrepancy manifests as a generalisation gap that no amount of additional small-batch training can close.

Data and Model Scaling Design Choices

With the memory-management infrastructure established, the paper specifies what is being scaled. The data scaling (Section 7.1) and model scaling (Section 7.2) dimensions involve concrete design choices that are explained here.

Data scaling: expanding ALIGN with JFT. The ALIGN dataset (Jia et al., 2021) contains 1.7B noisy image-text pairs collected from the web. The paper expands this by adding 5B pairs derived from the JFT dataset (a large-scale image classification dataset with hierarchical labels). For each JFT image, which is associated with one or more class labels, the text sequence is constructed as: " $\{class\_1\}$ and $\{class\_2\}$ and ... and $\{class\_k\}$ ". This is a simple template-based conversion that treats multi-label images as having conjunctive text descriptions. The combined dataset, ALIGN+JFT, contains 6.6B pairs — roughly 4× the size of ALIGN and 16× the size of CLIP's 400M-pair dataset.

Why this matters: the JFT text sequences are clean and structured (class names joined by "and"), unlike the noisy alt-text from ALIGN. This provides a complementary signal: ALIGN teaches the model to handle natural, noisy language, while JFT teaches it to associate images with precise, decontextualised class names — which is exactly what zero-shot classification with prompt templates requires at test time. The paper speculates (Section 10.2) that this complementarity explains why training on ALIGN+JFT outperforms training on ALIGN alone by 5.3–5.8 percentage points.

Tokenisation. The authors train a sentence piece model (Kudo and Richardson, 2018) with a 32K-token vocabulary on a random 200M-sentence subset of ALIGN+JFT. Using a tokeniser learned directly from the training data (rather than borrowing a pre-existing tokeniser from a language modelling corpus) is described as boosting "our top-1 accuracy on ImageNet ILSVRC-2012 by more than 1%." This is because the tokeniser's subword splits are adapted to the distribution of text in web image-text pairs, which differs from standard text corpora in vocabulary (more nouns, object names, descriptive adjectives) and style (fragments, alt-text, hashtag-like strings). Sequences longer than 64 tokens are filtered out — a pragmatic choice to bound the memory and computation cost of the text encoder.

Contamination filtering. All training instances that have a structural similarity index (SSIM; Wang et al., 2004) of at least 0.5 with any image from the evaluation benchmarks are removed. This is a conservative filter designed to prevent the model from memorising evaluation-set images during training, which would inflate accuracy numbers without genuine generalisation. An SSIM threshold of 0.5 catches near-duplicate and heavily visually similar images.

Model scaling: CoAtNet for the image encoder. The paper selects CoAtNet (Dai et al., 2021), a hybrid architecture that interleaves convolutional blocks (MBConv, from EfficientNet) with transformer blocks (self-attention + feed-forward). The motivation is that CoAtNet has been shown to have higher learning capacity than pure convolutional networks (EfficientNet) or pure transformers (ViT) at comparable FLOP budgets, making it a better vehicle for absorbing the 6.6B-example dataset.

Three model sizes are defined, as shown in Table 5 of the paper:

Model	Image Encoder	Image Params	Image FLOPs	Text Layers	Text Hidden Dim	Text Params	Text FLOPs
BASIC-S	CoAtNet-0	25M	4.2B	6	1024	108M	10.7B
BASIC-M	CoAtNet-3	168M	34.7B	12	1024	184M	49.4B
BASIC-L	CoAtNet-7	2.4B	495.8B	12	2048	670M	212.6B

Two key architectural choices:

Asymmetric scaling favours the image encoder. The authors "find that for the same computational budget, it is more beneficial to invest in scaling up the image encoder, rather than the text encoder." This is reflected in the numbers: BASIC-L's image encoder has 2.4B parameters and 496B FLOPs, while its text encoder has 670M parameters and 213B FLOPs — roughly a 3.6× parameter ratio and 2.3× FLOP ratio in favour of the image encoder. This asymmetry is motivated by the observation that image understanding is the harder sub-problem for zero-shot classification: text prompts are short and structured, while images contain rich, high-dimensional information that requires deeper processing.
Text representation via averaging, not [CLS] token. Unlike ALIGN, which uses a special [CLS] token (following BERT's convention) whose final-layer embedding serves as the text representation, BASIC averages the representations across all token positions at the top layer of the text transformer. This means each token contributes equally to the final text embedding, rather than the model learning to route information through a dedicated summary token. The paper does not ablate this choice, but it is consistent with findings in the sentence embedding literature that mean pooling often outperforms [CLS] pooling for semantic similarity tasks.

The pretraining-finetuning training procedure (Section 8). The training proceeds in three phases, motivated by both memory efficiency and representational quality:

Phase 1 — Supervised pretraining of the image encoder. The image encoder $F$ is pretrained on the JFT labelled dataset (5B images, 29K classes) using a standard softmax classification loss. This phase does not involve the text encoder at all. The motivation is computational: supervised training on JFT does not require the $B \times B$ contrastive matrix, so it can run with standard data parallelism without the memory-management infrastructure of Sections 4 and 5. After pretraining, $F$ has already learned strong visual representations from a cleanly labelled dataset.

Phase 2 — Contrastive training of the text encoder (image encoder frozen). The weights of $F$ are frozen. The text encoder $G$ is trained from scratch using the contrastive loss on ALIGN+JFT. Because $F$ 's weights are fixed, the backpropagation only flows through $G$ , which reduces peak memory (the optimizer does not need moment buffers for $F$ , and the compiler can free $F$ 's intermediate activations as soon as the image embeddings are computed). This phase teaches $G$ to align text representations with the frozen visual representations learned in Phase 1.

Phase 3 — Joint finetuning. Both $F$ and $G$ are unfrozen and jointly trained for additional steps at a reduced learning rate. This phase requires the full memory-management infrastructure (GradAccum or SPMD) because gradients flow through both encoders. The motivation is that Phase 2 may leave some misalignment between $F$ and $G$ because $F$ was never exposed to the text-modality training signal. Joint finetuning allows both encoders to adjust to each other, closing the remaining gap. The paper reports that this phase adds 1.4% ImageNet accuracy for BASIC-S, 0.6% for BASIC-M, and 0.4% for BASIC-L.

A critical failure mode of the pretraining-finetuning approach. The paper documents a non-obvious weakness (Section 8): "while some pretrained-and-finetuned models achieve similar accuracy to their contrastive counterparts on ImageNet or CIFAR, they completely fail on an easier task — MNIST." The root cause is that the JFT pretraining dataset consists primarily of natural images (photographs of objects, scenes, animals) and contains very few images of handwritten digits. The frozen image encoder $F$ therefore never learns to extract digit-relevant features during Phase 1, and Phase 2 only trains $G$ to align with $F$ 's frozen representations — it cannot teach $F$ new visual features. Meanwhile, the noisy image-text dataset ALIGN contains many instances of digits with accompanying text (e.g., an image of a check with the text "pay to the order of"), which teach optical character recognition when both encoders are trained jointly from scratch. The three-phase procedure partially mitigates this by including Phase 3 (joint finetuning), which does expose $F$ to the image-text data, but if the joint finetuning phase is too short or the learning rate too low, $F$ may not fully recover the capabilities it would have learned from scratch.

AdaFactorW optimizer. The paper designs a custom optimizer combining two existing methods:

Factorised second moments from AdaFactor (Shazeer and Stern, 2018): Instead of storing a full $D \times D$ second-moment matrix for each weight tensor (as Adam does), AdaFactor factorises the second moment into a row vector and column vector, reducing memory from $O(n^2)$ to $O(n)$ for an $n \times n weight matrix. For very large weight tensors (e.g., the feed-forward layers in the transformer with hidden dimension 2048), this is a substantial memory saving.
Decoupled weight decay from AdamW (Loshchilov and Hutter, 2019): The weight decay regularisation is applied directly to the weights ( $\theta \leftarrow \theta - \eta \lambda \theta$ ) rather than being incorporated into the gradient via L2 regularisation. This decoupling ensures that the adaptive learning rate scaling (which depends on the second moment) does not interfere with the regularisation strength.
First moments stored in bfloat16: Following Zhai et al. (2021), the first gradient moments are stored in the bfloat16 half-precision format to save memory. However, the paper notes a critical detail: "while we can store these moments in bfloat16, we need to convert them into float32 prior to computing our weight updates to avoid numerical instability." This is because the weight update involves multiplying the first moment by a learning rate and dividing by the square root of the second moment — operations that can underflow or overflow in bfloat16 due to its limited exponent range.

Hyperparameters (Table 6). Key numbers across model sizes:

Hyperparameter	BASIC-S (Pretraining)	BASIC-S (Contrastive)	BASIC-M,L (Pretraining)	BASIC-M,L (Contrastive)
Optimizer	AdaFactorW	AdaFactorW	AdaFactorW	AdaFactorW
Batch size	16,384	65,536	16,384	65,536
Training steps	500K	500K	1.2M	500K
Warm-up steps	25K	25K	25K	25K
Max learning rate	1e-3	1e-3	4e-4	2.5e-4
Min learning rate	1e-5	1e-5	2e-5	1e-5
Learning decay	Cosine	Cosine	Linear	Cosine
Weight decay	0.005	0.0025	0.01	0.0025

Notable patterns: Batch size for contrastive training is uniformly 65,536. Maximum learning rates are lower for larger models (2.5e-4 for BASIC-M,L vs. 1e-3 for BASIC-S during contrastive training), reflecting the increased sensitivity of larger models to optimisation instability. Weight decay is higher during pretraining (0.005–0.01) than during contrastive training (0.0025), possibly because the contrastive loss itself provides implicit regularisation through the large number of negative examples.

No regularisation beyond weight decay. The paper explicitly states that "other than the decoupled weight decay in AdaFactorW, we do not use any other regularization technique." Adding stochastic depth (Huang et al., 2017) or dropout (Srivastava et al., 2014) causes "our ImageNet top-1 accuracy [to drop] substantially." The authors hypothesise that at the scale of ALIGN+JFT (6.6B examples), the dataset itself provides sufficient regularisation — overfitting is not a concern, and regularisation techniques only add optimisation noise that slows convergence without improving generalisation.

A subtlety about rematerialization and stochastic regularisation. The paper notes a second reason to avoid dropout-like regularisation: it would make the rematerialization steps in Algorithm 1 inconsistent. In the pipelining approach, the forward pass is run twice — once to compute embeddings (Lines 2–5 in Algorithm 1) and once to rematerialize activations for backpropagation (Lines 13–16). If dropout randomly zeros out different units in these two passes, the activation values would differ, making the rematerialized backward pass compute gradients with respect to a different computational graph than the one that produced the loss. While this could be treated "as a form of regularization noise," the paper observes empirically that with such noise, "our training loss stays relatively large throughout the course of training," suggesting optimisation difficulties. Hence, the pragmatic decision is to forgo stochastic regularisation entirely.

Image resolution. All training and evaluation use 224 × 224 resolution. The paper acknowledges that increasing resolution is a known method to gain performance (Tan and Le, 2019, 2021; Touvron et al., 2019), but chooses to "reserve our computational resources for scaling up our model and our batch size" instead. This is a deliberate allocation tradeoff: given a fixed compute budget, the paper hypothesises that combined scaling of data, model, and batch size provides larger returns than resolution scaling. The results (85.7% at 224px) suggest this tradeoff was reasonable, though the paper does not ablate it.

Summary of Design Choices and Their Justifications

Pipelining + GradAccum for generality, SPMD for speed and exactness. The two memory-management methods serve different use cases: pipelining is the fallback that always works, SPMD is the optimised path for architectures where manual sharding design is feasible. The paper uses both and provides the first systematic comparison of their tradeoffs.
Gradient accumulation directly into optimizer slots rather than into a separate buffer: eliminates the need for an 11GB gradient buffer for the 3B-parameter model, at the cost of a variance estimation approximation for the second moment. This is a practical engineering trick motivated by the specific memory constraints of large-scale training.
Rematerialization of activation and normalisation layers, not weighted layers: exploits the fact that element-wise and reduction operations are computationally cheap but memory-intensive, while sharded weight operations are expensive to recompute due to cross-core communication. The 1.4× step-time overhead is the price paid for fitting the model in memory.
Asymmetric model scaling (image encoder larger than text encoder): based on the empirical observation that image understanding is the harder sub-problem and benefits more from additional capacity at equivalent FLOP budgets.
Mean pooling over token embeddings for text representation, not [CLS]: follows the sentence embedding literature's finding that mean pooling produces better semantic similarity representations. Not ablated but consistent with external evidence.
Three-phase training (supervised pretrain → frozen-image contrastive → joint finetune): motivated by memory efficiency (Phase 2 avoids training $F$ ) and by the failure mode where frozen image encoders trained only on natural images fail on specialised domains like MNIST. Phase 3 partially recovers these capabilities.
No stochastic regularisation: justified by the scale of the dataset (6.6B examples are their own regulariser) and the incompatibility of dropout-like randomness with the double-forward-pass structure of rematerialization.
Learned tokeniser on training data: the 1%+ accuracy boost from adapting the tokeniser to the web image-text distribution (rather than using a standard text corpus tokeniser) likely reflects better handling of noisy text features like hashtags, filenames, and alt-text fragments that are common in the training data but rare in standard text corpora.
SSIM-based contamination filtering with threshold 0.5: a conservative filter that removes near-duplicate images between training and evaluation sets, preventing inflated accuracy numbers from memorisation. The threshold of 0.5 is standard in the literature for detecting visually similar images.

4. Key Insights and Innovations

Innovation 1: Batch Size as a Generalisation Dimension, Not a Compute-Efficiency Knob

The field’s default assumption about batch size in contrastive learning — inherited from SimCLR (Chen et al., 2020a) and carried into CLIP and ALIGN — was that larger batches help by providing more negative examples per gradient step, which is a compute-efficiency argument: you get more learning signal per step, so training converges faster in wall-clock time. The implication of this view is that if you simply train for more steps with a smaller batch, you should eventually match the performance of a larger-batch model, because the total number of negative examples seen over training would be equalised.

BASIC’s most consequential intellectual move is to reject that framing and establish batch size as an independent generalisation dimension with a unique, theoretically grounded benefit that cannot be recovered by additional training steps. The theoretical analysis (Section 6, Theorem 1) proves that the generalisation gap between the finite-batch training objective and the population contrastive objective shrinks as O(1/√B), and that this term remains even as the number of training examples m → ∞. In operational terms: a model trained with B = 4,096 on infinite data would still have a larger generalisation gap than one trained with B = 65,536, because the loss function itself is different — the small-batch model learns to solve an easier problem (discriminating among 4K negatives) while the large-batch model learns to approximate the true population-level discrimination task.

The empirical confirmation is Figure 5 and Table 4 (Section 10.1): models trained with larger batch sizes for fewer steps (but the same number of total examples seen) achieve higher ImageNet accuracy than models trained with smaller batch sizes for more steps. The gap cannot be closed by extending training. This finding is important because it converts batch size from a training-speed hyperparameter (tune it for throughput) into an accuracy-critical scaling dimension (tune it for final performance, alongside data and model size). It also explains why SimCLR’s observation that batch size benefits "saturate at 8192" was an artifact of its small dataset and model — when data and model scale up, the optimal batch size shifts upward, creating a coupling among the three scaling dimensions that the paper’s combined scaling approach is designed to exploit.

This is not an incremental refinement. It is a fundamental reframing of what batch size means in contrastive learning, backed by both theory and controlled experiments. It transforms batch size from an implementation detail into a first-class design axis, directly motivating the entire engineering enterprise of Sections 4 and 5: if large batch sizes only helped training speed, the memory-management infrastructure would be a convenience; because they improve final accuracy, it is a necessity.

Innovation 2: The Discovery That Supervised Finetuning Destroys Zero-Shot Robustness

CLIP (Radford et al., 2021) established that zero-shot models are more robust to natural distribution shifts than supervised ImageNet-trained models, and coined the term "effective robustness" to describe the phenomenon. However, CLIP left the mechanism deliberately unresolved, cautioning against "generalizing too far from initial findings." The dominant interpretation in the community was that something about the contrastive pretraining objective or the diverse web-scale data distribution confers robustness — perhaps the training data itself is more varied, or the language supervision provides a richer supervisory signal than one-hot class labels.

BASIC upends this interpretation with a clean, diagnostic experiment (Section 9.3, Figure 4): take a converged zero-shot model that already exhibits high effective robustness, then finetune it on more labeled ImageNet data. If robustness were a property of the model architecture or the pretraining data, finetuning on a subset of ImageNet labels should either preserve or slightly improve robustness (since the model sees more relevant data). Instead, the paper finds the opposite: as the model is exposed to 1%, then 10%, then 20%, then 50% of ImageNet’s labeled examples, its ImageNet accuracy rises but its robustness benchmark accuracy declines or stagnates. In the extreme case, a 3% accuracy gain on ImageNet is accompanied by an 8.3% accuracy drop on ImageNet-R.

This is a negative result with deep implications, not a metric gain. It suggests that the robustness of zero-shot models is not an inherent property of how they are trained but rather a fragile state that is actively lost when models are adapted to a specific labeled distribution. The causal mechanism is not identified — the paper explicitly frames this as an open question inviting "further causal analysis on the effects of ImageNet’s labeled data" — but the observation itself is intellectually significant because it changes how practitioners should think about the accuracy-robustness tradeoff. The standard deployment pipeline (pretrain a large model, then finetune on the target task’s labeled data) may systematically destroy the very robustness that made the pretrained model attractive. This is not a call to abandon finetuning but a warning that the field does not yet understand what finetuning costs in terms of distribution-shift resilience.

The finding is strengthened by a methodological choice: the finetuning uses the contrastive loss with class-name prompts (not a linear classifier), which controls for the confounding possibility that linear probing behaves differently from zero-shot transfer. The robustness degradation is therefore attributable to the data distribution of the finetuning set, not to the classifier head architecture. This clean experimental design makes the finding harder to dismiss as a measurement artifact.

Innovation 3: The Pretraining-Finetuning Hybrid as Both an Enabler and a Trap

Large-scale contrastive training of both image and text encoders from scratch is expensive, so it is natural to ask whether pretraining the image encoder on a labeled dataset — where supervised training is cheaper and well-understood — can accelerate the process. This idea is not itself novel; what BASIC contributes is a careful characterisation of when this shortcut works and when it catastrophically fails, revealing a tradeoff that prior work had not documented.

The paper’s three-phase procedure (Section 8) — supervised pretrain on JFT, contrastive train the text encoder with the image encoder frozen, then jointly finetune both — is pragmatic and yields the strongest final results. But the paper also identifies a specific failure mode: a model whose image encoder was pretrained on JFT (which consists almost entirely of natural images) and then frozen during contrastive training "completely fails on an easier task — MNIST" because JFT contains few digit images while the noisy image-text data (ALIGN) contains digit-related text that teaches optical character recognition. The frozen image encoder never learns to extract digit features, and the text encoder — trained only to align with the frozen visual representations — cannot compensate.

This is a diagnostic insight about representation learning under modality asymmetry. It reveals that the image-text data is not merely providing a weak supervision signal for categories that could be learned from labels; it is teaching fundamentally different visual features than supervised ImageNet-style training. The handwritten digit domain is absent from JFT but present in ALIGN’s noisy web text, meaning the contrastive phase does not just align modalities — it can teach the image encoder to see things that supervised pretraining missed entirely. The joint finetuning phase (Phase 3) partially recovers this capability, but the finding implies that pretraining on a labeled dataset and then switching to contrastive training is not a free lunch — it biases the visual representations toward the pretraining distribution in ways that may be invisible on standard benchmarks (where MNIST is typically excluded) but matter for real-world deployment on diverse input types.

This is an incremental but practically important refinement of the pretraining-finetuning paradigm. It changes the default recommendation from "always pretrain on labeled data first" to "pretrain if the labeled dataset covers your deployment domain; otherwise, the shortcut may cost you capabilities that contrastive training from scratch would have provided."

Innovation 4: The Two-Method Memory Infrastructure as a Generality-vs-Optimality Design Spectrum

Scaling up model size and batch size simultaneously creates a compound memory problem that no single prior technique fully solved for contrastive learning. Gradient accumulation (Ott et al., 2018; Zhai et al., 2021) does not apply directly because the contrastive loss couples all examples in the batch through the softmax denominator. Model parallelism (Huang et al., 2019; Shazeer et al., 2018) distributes computation but does not address the B × B similarity matrix. Rematerialization (Chen et al., 2016) trades compute for memory but does not specify what to rematerialize for contrastive models.

BASIC’s contribution here is not any single technique but the articulation of two complementary approaches that span a spectrum from generality to optimality, each with a clear characterisation of its exactness, speed, and scalability limits. This is an engineering design contribution — it provides practitioners with a decision framework, not just methods.

The pipelining + GradAccum approach (Section 4) is generic: it works for any model architecture, scales to arbitrarily large B by adjusting the microbatch size, and requires no manual sharding design. Its costs are (1) inexact gradient accumulation (the variance correction for the second moment is an approximation), (2) batch normalisation inconsistency when M ≪ B, and (3) slower step times due to double forward passes. The SPMD approach (Section 5) is optimal: computations are exact, step times are faster (Table 2 shows ~15% improvement over pipelining at B = 2^20), and the rematerialization heuristic (recompute cheap activation/normalisation layers, keep expensive weighted layers) is computationally efficient. Its cost is limited generality — it requires manual design of the sharding strategy and rematerialization schedule, and its memory footprint grows with the per-core batch size, meaning it must be redesigned if B increases beyond the supported range.

Prior work on large-scale training infrastructure typically advocates for a single method. BASIC’s contribution is showing that both methods have a role, and providing the first systematic comparison (Table 2) of their step times, memory footprints, and scaling behaviour at identical model sizes and batch sizes. This comparison — showing that SPMD is faster but uses more memory, while pipelining is slower but has constant memory regardless of B — is what enables practitioners to make an informed choice based on their hardware (available memory per core, inter-core communication bandwidth) and scaling ambitions (will B increase further in future experiments?).

This is an incremental but practically significant contribution because it converts what could have been a one-off engineering solution for the paper’s specific training run into a reusable design pattern for the broader contrastive learning community.

5. Experimental Analysis

Evaluation Methodology

Dataset. All experiments use the MATH benchmark (Hendrycks et al., 2021), a collection of high-school competition-level mathematics problems with ground-truth answers that can be graded via exact string matching. The paper uses the specific split from Lightman et al. (2022): 12,000 training questions and 500 test questions. The choice of MATH is motivated by the observation (Section 4) that test-time compute is expected to help most when the model already possesses the necessary knowledge and the challenge lies in drawing complex, multi-step inferences — mathematical reasoning fits this profile because it requires logical deduction rather than novel factual recall, making it an ideal testbed for studying how variations in inference strategy affect problem-solving success.

Base model. All experiments use PaLM 2-S* (Codey) (Anil et al., 2023), a model that the authors argue is "representative of the capabilities of many contemporary LLMs" and operates in a useful performance regime for the study: non-trivial but far-from-saturated performance on MATH (roughly 10–19% pass@1 depending on prompt and sampling configuration), leaving substantial room for test-time compute strategies to demonstrate differential effects. For the FLOPs-matched comparison in Section 7, a second model with approximately 14× more parameters is used as the pretraining-scaled baseline, with its pretraining FLOPs and inference FLOPs accounted for using standard scaling-law approximations.

Metrics. The primary metric throughout the paper is MATH test accuracy (%), defined as the fraction of the 500 test questions for which the model's selected final answer matches the ground truth as determined by the grading function released by Lightman et al. (2022). When analyzing difficulty-dependent behavior, accuracy is reported within each of the five difficulty quintiles separately. For the FLOPs-matched comparison, relative improvement or degradation versus the larger model is reported as a percentage change. The revision model is additionally evaluated using pass@1 at each sequential revision step to track the trajectory of improvement throughout the revision chain.

Baselines. The paper employs several baselines, each designed to isolate specific effects:

Majority voting: select the most common final answer among N independently sampled solutions, without any learned verifier. This serves as a lower bound on what can be achieved by consensus alone, establishing the value added by learned verification.
ORM best-of-N weighted: score N complete solutions using an Outcome Reward Model (a single correctness score per solution) and apply best-of-N weighted selection (Li et al., 2023), where solutions arriving at the same final answer have their scores summed and the answer with the greatest total sum is selected. This baseline isolates the contribution of the process-level versus outcome-level verification signal.
PRM best-of-N weighted: score N complete solutions using the Process Reward Model and apply best-of-N weighted selection. This is the primary competitive baseline — it represents the standard "scale up parallel sampling with a good verifier" approach that was the dominant inference-time strategy prior to this work. It uses the same PRM as the search methods, ensuring that differences reflect search strategy rather than verifier quality.
Parallel sampling (for revisions): generate N independent solutions from the revision model and select the best via verifier or majority voting. This baseline isolates the benefit of sequential revision from the benefit of simply using the revision model (which may produce better individual answers than the base model).

Generation budget / compute accounting. The paper measures test-time compute in units of "generations," where one generation equals one complete sampled answer from the base LLM. For best-of-N and beam search, the budget is N (the number of beams or samples). For lookahead search with k lookahead steps, the cost is N × (k+1) to account for the additional rollouts. Budgets are swept across powers of 2, typically from 2^0 to 2^9 (1 to 512 generations). This accounting ensures fair comparison across methods — a method that does more computation per candidate (like lookahead search) is charged proportionally. For the FLOPs-matched comparison in Section 7, the total FLOP budget is computed using standard approximations from the scaling laws literature: pretraining FLOPs = 6ND_pretrain and inference FLOPs = 2ND_inference, where N is the number of model parameters, D_pretrain is pretraining tokens, and D_inference is total inference tokens generated.

Cross-validation / statistical protocol. To avoid the circularity of selecting the best strategy and evaluating it on the same data, the paper uses two-fold cross-validation within each difficulty bin on the 500-question test set. The best-performing strategy (e.g., which search algorithm, which sequential-to-parallel ratio) is selected on one fold and evaluated on the other, with results averaged. Difficulty bins are constructed by sampling 2,048 solutions per question from the base model, computing the pass@1 rate (oracle) or averaging the PRM's predicted final-answer correctness (predicted), and binning into five quintiles. The key consequence of this protocol: the compute-optimal policy is selected based on roughly 50 questions per fold per bin, a relatively small sample that could introduce variance in the selected strategies. The paper does not report confidence intervals on the compute-optimal scaling curves.

Main Quantitative Results

Search Against PRM Verifiers

The headline result for search (Section 5) is that beam search significantly outperforms best-of-N at low generation budgets, but its advantage diminishes or reverses at high budgets, and that the optimal search strategy depends critically on problem difficulty, enabling a compute-optimal allocation policy that achieves ~4× efficiency gains over best-of-N. All search results are presented in Figures 3 and 4, with the PRM training details in Appendices D and E, and the comparison to ORM in Appendix F.

Aggregate search comparison (Figure 3, left). Across all 500 test questions with a maximum budget of 256 generations, beam search with M = 4 (fixed beam width) dominates at low generation counts: at 4 generations, beam search achieves roughly 27% accuracy versus roughly 16% for PRM best-of-N weighted — a gap of approximately 11 percentage points. However, as the budget increases, this advantage erodes. At 512 generations, PRM best-of-N weighted reaches approximately 38% while beam search (M = 4) plateaus around 34%. This crossover is a central empirical finding: it demonstrates that more aggressive search optimization is not always better, and that the relationship between search intensity and accuracy is non-monotonic.

Beam search with M = √N (growing beam width) performs similarly to M = 4 at low budgets but does not reach the same peak, suggesting that the fixed budget of beam expansions matters more than the specific beam width parameterization. Lookahead search — the most powerful optimizer because it uses additional rollout computation to improve step-level scoring — paradoxically performs worst across nearly all budget levels. At 256 generations, 3-step lookahead search with M = 4 achieves accuracy below both beam search and best-of-N, despite consuming the same generation budget. The paper attributes this to verifier over-optimization: lookahead search is better at finding solutions that score highly under the PRM, but these solutions are more likely to be adversarial examples that exploit the verifier's blind spots rather than genuinely correct answers.

Majority voting trails all verifier-based methods substantially, reaching only about 29% at 512 generations — approximately 9 percentage points below PRM best-of-N weighted. This gap quantifies the value added by the learned verifier over simple consensus.

Difficulty-dependent behavior of search (Figure 3, right). The per-difficulty breakdown — shown for beam search (M = 4) versus PRM best-of-N weighted at four budget levels (4, 16, 64, 256 generations) — reveals qualitatively different scaling behaviors across difficulty bins that explain and reconcile the aggregate results:

Bin 1 (easiest problems): Beam search accuracy actually decreases with increasing budget — from roughly 78% at 4 generations to roughly 77% at 256 generations — while PRM best-of-N weighted increases from roughly 68% to roughly 88%. This is the clearest evidence of PRM over-optimization in the paper: on problems where the model already produces many correct answers, aggressive search finds solutions that exploit the verifier signal rather than genuinely correct ones, degrading performance. The verifier's guidance is counterproductive when the model's base capability is already high.
Bin 2: Both methods improve with budget, but PRM best-of-N weighted improves faster — from roughly 14% to roughly 60% at 256 generations — while beam search improves more modestly — from roughly 14% to roughly 32%. The verifier's signal is not reliable enough to justify aggressive optimization, and the diversity of independent sampling outweighs the focused search.
Bin 3: Beam search consistently outperforms PRM best-of-N weighted across all budget levels, reaching roughly 34% versus roughly 23% at 256 generations. This is the regime where the verifier signal is most valuable: the model has non-trivial capability (correct solutions exist in the sampling distribution) but cannot find them reliably through random sampling alone. Guided search navigates toward higher-quality regions of the solution space.
Bin 4: Beam search shows its strongest relative advantage, reaching roughly 17% versus roughly 10% for PRM best-of-N at 256 generations. Even here, absolute performance remains low, suggesting that while search helps, the base model's capability on these problems is fundamentally limited.
Bin 5 (hardest problems): Both methods hover at 1–3% accuracy regardless of budget. No method makes meaningful progress. The base model simply cannot produce correct solutions for these problems at any non-trivial rate, and no search strategy can find what is not in the sampling distribution.

This breakdown is the paper's most important empirical result for search: it demonstrates that no single strategy is universally optimal. Beam search is best on medium problems, best-of-N is best on easy problems, and nothing works on hard problems. A uniform strategy applied to all problems will underperform on some subset of the distribution.

Compute-optimal search (Figure 4). By selecting the best search strategy per difficulty bin at each budget level, the compute-optimal policy achieves substantial efficiency gains:

At 16 generations, compute-optimal oracle achieves approximately 27% accuracy, roughly matching PRM best-of-N weighted at 64 generations. This is the ~4× compute reduction: matching the performance of 64 generations of the best baseline with only 16 generations by adaptively choosing between beam search and best-of-N per difficulty bin.
At 256 generations, compute-optimal oracle reaches approximately 39.5%, surpassing PRM best-of-N weighted at the same budget (roughly 37%) and majoritatively voting (roughly 29%).
The compute-optimal policy with predicted difficulty bins (using the PRM's own scores to estimate difficulty without ground-truth labels) tracks the oracle version closely, particularly at lower budgets. The two curves "largely overlap" per the authors, with the predicted version reaching approximately 37% at 256 generations versus approximately 39.5% for oracle — a small gap that grows at higher budgets, likely because the PRM-based difficulty estimates become noisier as the estimation targets diverge.
Both compute-optimal variants consistently outperform ORM best-of-N weighted (which peaks around 34% at 512 generations), confirming that the PRM provides a stronger verification signal.

PRM versus ORM comparison (Figure 14, Appendix F). At 2,048 samples, PRM best-of-N weighted achieves approximately 40% accuracy versus roughly 35% for ORM best-of-N weighted and roughly 30% for majority voting. The gap between PRM and ORM widens with the number of samples, confirming that process-level supervision provides better scaling properties than outcome-level supervision. The PRM's advantage is attributed to its step-level training acting as beneficial representation learning, even though the "last" aggregation method effectively reduces it to ORM-like behavior at aggregation time (Appendix E, Figure 13).

Step-wise aggregation strategy (Figure 13, Appendix E). Comparing "min," "prod," and "last" aggregation at 256 samples: "last" achieves roughly 37%, "min" achieves roughly 35%, and "prod" achieves roughly 27%. The superiority of "last" aggregation — which effectively uses only the PRM's final-step prediction, making it functionally similar to an ORM at selection time — is noteworthy because it contradicts prior work (Lightman et al., 2023; Wang et al., 2023) that found "min" to be superior. The authors hypothesize that this reversal stems from their use of soft Monte Carlo rollout labels (rather than binary correctness labels) in PRM training, which changes how per-step scores distribute. This is an important calibration detail: PRM training methodology interacts with aggregation strategy in non-obvious ways, and findings from one labeling scheme do not transfer directly to another.

Revision Model Results

The headline result for revisions (Section 6) is that sequential revisions (conditioning the model on its own previous incorrect answers) marginally outperform parallel sampling when aggregated across all questions, and that the optimal ratio of sequential to parallel sampling depends on difficulty, with purely sequential being best on easy questions and a balanced ratio being best on hard questions. A compute-optimal allocation of the sequential-to-parallel ratio achieves ~4× efficiency gains over the parallel-only baseline. All revision results are presented in Figures 5–8, with training details in Section 6.1 and Appendices H–I.

Revision model pass@1 trajectory (Figure 6, left). Starting from approximately 18.2% pass@1 at the first step, the revision model's per-step accuracy improves to roughly 24–25% by steps 15–20 and remains in the 23–25% range out to 64 steps. This gradual improvement and subsequent plateau demonstrate two things: (1) the model has learned a generalizable revision skill that extends well beyond its training horizon of 4 previous answers in context, and (2) the benefits of revision eventually saturate — additional revision steps beyond roughly 15–20 provide marginal gains. The fact that the model does not degrade significantly even after 64 steps (despite only being trained on sequences with up to 4 previous answers) is evidence that the revision skill transfers to longer contexts without catastrophic forgetting or mode collapse.

Sequential versus parallel comparison (Figure 6, right). At 64 generations, the head-to-head comparison shows:

Sequential + best-of-N weighted: approximately 41.5%
Parallel + best-of-N weighted: approximately 39%
Sequential + majority: approximately 38%
Parallel + majority: approximately 35%

Sequential revisions outperform parallel sampling under both verifier-based and majority-based selection. The gap under verifier-based selection (roughly 2.5 percentage points) is narrower than the gap under majority-based selection (roughly 3 percentage points), suggesting that the verifier somewhat equalizes the quality of the two approaches — but sequential revisions still hold an edge. The advantage of sequential over parallel is consistent but modest in aggregate, which is why the difficulty-dependent analysis that follows is essential for understanding when each approach should be deployed.

Sequential-to-parallel ratio sweep (Figure 7, left). For a fixed total generation budget, the ratio of sequential depth (chain length) to parallel breadth (number of independent chains) is varied. At 256 generations:

The optimal ratio is around 2^1 to 2^3 (2:1 to 8:1 sequential-to-parallel), achieving approximately 43–44% accuracy.
Fully parallel (leftmost point of the curve) yields approximately 40%.
Fully sequential (rightmost point) yields approximately 42%.

The optimal intermediate ratio suggests that neither extreme is ideal in aggregate: some parallel diversity is valuable for exploring different high-level solution strategies, while some sequential depth is valuable for refining promising candidates within each strategy. At lower budgets (8–32 generations), the curves are monotonically increasing with the sequential-to-parallel ratio — fully sequential is optimal, likely because at low budgets there are insufficient parallel samples to provide meaningful diversity, and depth is the better investment.

Difficulty-dependent optimal ratio (Figure 7, right). At a fixed budget of 128 generations, broken out by difficulty quintile:

Bin 1 (easiest): Performance is essentially flat across all ratios, around 90–92%. Easy questions are insensitive to the allocation strategy — the model produces correct answers regardless of whether resources are spent sequentially or in parallel.
Bin 2: Slight advantage for higher sequential ratios — approximately 63% at fully sequential versus roughly 58% at fully parallel. Sequential refinement of already-strong candidates is beneficial.
Bin 3: A clear optimal ratio emerges at moderate sequential-to-parallel values (around 2^1 to 2^3), reaching approximately 42% versus roughly 35% at the extremes. This is the regime where the complementary benefits of exploration (parallel) and exploitation (sequential) are most pronounced.
Bin 4: A similar pattern, with the peak at a moderate ratio achieving roughly 18% versus roughly 14% at fully parallel. However, the absolute performance is low, indicating limited base model capability.
Bin 5 (hardest): All ratios produce roughly 2–3% accuracy. No allocation strategy helps on problems where the base model cannot produce correct solutions.

This mirroring of the search results — easy problems benefit from exploitation (sequential refinement), hard problems benefit from a balance of exploration and exploitation (mixed sequential-parallel), and the hardest problems benefit from nothing — reinforces the paper's central thesis that difficulty is the key conditioning variable for test-time compute allocation.

Compute-optimal revisions (Figure 8). By selecting the optimal sequential-to-parallel ratio per difficulty bin at each budget level:

At 64 generations, compute-optimal oracle achieves approximately 40%, roughly matching parallel best-of-N weighted at 256 generations — a ~4× efficiency gain.
At 256 generations, compute-optimal oracle reaches approximately 44%, compared to roughly 41% for best-of-N weighted and approximately 37% for the parallel-only baseline.
Compute-optimal with predicted difficulty bins performs slightly below oracle bins at high budgets (approximately 41% at 256 generations versus roughly 44% for oracle), but still substantially outperforms the parallel baseline.
Notably, the parallel baseline appears to plateau around 36–37% at high budgets, while compute-optimal scaling continues to improve. The paper interprets this as evidence that adaptive allocation not only improves efficiency at fixed budgets but also extends the scaling ceiling — the gains from switching between strategies compound as budgets increase.

Majority voting for revisions (Figure 10, Appendix B). The sequential-to-parallel ratio trends observed with verifier-based selection replicate qualitatively with majority voting: easy questions are insensitive to ratio, hard questions show an optimal intermediate ratio, and fully sequential marginally outperforms fully parallel in aggregate. This replication with a non-learned selection mechanism confirms that the benefit of sequential sampling is not solely attributable to the verifier having access to more revision context — the revision model itself produces better answers when conditioned on its own history.

Verifier choice for revision outputs (Figure 15a, Appendix J). The base-LM PRM — trained on base model outputs — underperforms the revision-specific ORM when scoring revision model outputs. At 64 generations, sequential revisions scored with the base-LM PRM achieve roughly 40%, while the same chain scored with the revision-specific ORM achieves roughly 42%. This ~2 percentage-point gap confirms that distribution shift between the base model's outputs and the revision model's outputs is a genuine practical concern — verifiers should be trained on the output distribution they will evaluate at test time.

Revision history in verifier context (Figure 15b, Appendix J). Including previous revisions in the ORM's input context provides a small improvement of approximately 1–2 percentage points at 64 generations over not including revision history. Both variants of the ORM (with and without revision history) outperform the parallel baseline, confirming that the benefit of sequential sampling is not solely due to the verifier seeing more context. The revision model genuinely produces better candidates in sequence.

FLOPs-Matched Comparison: Test-Time Compute versus Pretraining Compute

The headline result for the FLOPs-matched comparison (Section 7) is that a smaller model (PaLM 2-S*) with compute-optimal test-time strategies can outperform a ~14× larger model on easy-to-medium difficulty problems when the inference-to-pretraining token ratio is low (R ≪ 1), but pretraining is substantially more effective on hard problems and at high R values. This comparison establishes boundary conditions on the substitutability of test-time and pretraining compute. All results are presented in Figure 9 and the bar charts in Figure 1.

Experimental setup. The comparison pits PaLM 2-S* with compute-optimal test-time scaling against a model with approximately 14× more parameters (greedy decoding, no extra test-time compute). Three values of R = D_inference / D_pretrain are tested: 0.16 (R ≪ 1, representing scenarios where inference tokens are a small fraction of pretraining tokens, such as one-time evaluation or self-improvement pipelines), 0.79 (R ≈ 1, roughly equal inference and pretraining token counts), and 22 (R ≫ 1, representing high-volume production deployments where inference dominates total compute). At each R value, the total FLOPs budget of the larger model is matched by the smaller model, with the FLOP savings from reduced pretraining reallocated to additional inference-time generations.

Revisions versus the larger model (Figure 9, left; Figure 1, top-right bar chart). Across difficulty levels and R values, with compute-optimal revisions on the smaller model versus greedy decoding on the 14× larger model:

Difficulty	R ≪ 1 (0.16)	R ≈ 1 (0.79)	R ≫ 1 (22)
Easy (bin 1)	+11.8%	+3.5%	−11.9%
Medium (bins 2–3)	+27.8%	+16.7%	+5.4%
Hard (bins 4–5)	+21.6%	−(implied negative)	−37.2%

(Numbers are relative improvements of test-time compute over pretraining, from the bar chart in Figure 1, top-right. Positive values indicate the smaller model with test-time compute outperforms the larger model; negative values indicate the larger model wins.)

The pattern is clear and asymmetric. At R ≪ 1, test-time compute outperforms pretraining across all difficulty levels, with the largest advantage on medium questions (+27.8%). At R ≈ 1, test-time compute is still preferable on easy and medium questions but the margins shrink, and hard questions already favor pretraining. At R ≫ 1, test-time compute is preferable only on easy questions (+5.4% on medium is marginal), and pretraining dominates on hard questions by a substantial margin (−37.2%). This R-dependence is intuitive: when R ≪ 1, the pretraining cost savings from using a smaller model are enormous relative to the inference budget, giving the smaller model a large generation budget to work with. When R ≫ 1, the larger model's per-query inference cost is the dominant factor, and the pretraining savings from going smaller are relatively modest — the smaller model cannot afford enough extra inference compute to close the gap.

PRM search versus the larger model (Figure 9, right; Figure 1, bottom-right bar chart):

Difficulty	R ≪ 1 (0.16)	R ≈ 1 (0.79)	R ≫ 1 (22)
Easy	+19.1%	+2.2%	+2.0%
Medium	0.0%	−35.3%	−30.8%
Hard	−3.6%	−35.3%	−52.9%

PRM search shows substantially weaker benefits than revisions for the FLOPs-matched comparison. On easy questions, test-time compute remains preferable across all R regimes (+19.1% at R ≪ 1, narrowing to +2.0% at R ≫ 1). On medium questions, pretraining dominates at moderate and high R values (−35.3% at R ≈ 1). On hard questions, pretraining dominates across all R values, with the disadvantage reaching −52.9% at R ≫ 1. The paper does not explicitly discuss why search underperforms revisions in this comparison, but the aggregate search results in Figure 3 provide a likely explanation: beam search degrades at high budgets on easy problems (due to verifier over-optimization) and provides minimal gains on hard problems (due to base model capability limits), meaning the extra inference budget allocated in the FLOPs-matched comparison is spent less efficiently than with revisions, which avoid over-optimization and provide more consistent scaling.

Figure 9 detail. The line plots in Figure 9 show accuracy per difficulty bin as test-time compute scales (x-axis: generation budget on log scale). The 14× larger model's greedy performance is indicated by stars placed at three x-axis positions corresponding to the three R values. Where the compute-optimal scaling line (solid curve) is above the star, test-time compute wins; where it is below, pretraining wins. On Bin 1 (easiest, purple curve): the scaling line is above or near all three stars for revisions, and above all three stars for search. On Bin 5 (hardest, blue curve): the scaling line is below all three stars and essentially flat near 0–5% accuracy for both revisions and search, confirming that no amount of test-time compute helps on problems where the base model's capability is near zero.

Ablation Studies and Robustness Checks

PRM step-wise aggregation strategy (Figure 13, Appendix E): "Last" (using only the PRM's prediction at the final step) achieves approximately 37% at 256 samples, outperforming "min" (~35%) and "prod" (~27%). The superiority of "last" contradicts prior work that found "min" to be best (Lightman et al., 2023; Wang et al., 2023). The paper attributes this reversal to their use of soft Monte Carlo rollout labels rather than binary correctness labels in PRM training. This is an important finding for practitioners: PRM training methodology and aggregation strategy are coupled, and transferring aggregation heuristics across labeling schemes can be counterproductive.

PRM versus ORM scaling (Figure 14, Appendix F): The PRM consistently outperforms the ORM, with the gap widening at higher sample counts. At 2,048 samples, PRM best-of-N weighted reaches approximately 40% versus ORM's approximately 35%, confirming that process-level training provides better scaling properties even when the "last" aggregation effectively collapses it to ORM-like behavior at selection time. This implies that the benefit of PRM training is primarily in the learned representations, not in the per-step scoring at aggregation time.

Revision model verifier choice (Figure 15a, Appendix J): The base-LM PRM — trained on base model outputs — underperforms a revision-specific ORM when scoring revision model outputs. Sequential + base-LM PRM achieves approximately 40% at 64 generations versus sequential + revision ORM at approximately 42%. The ~2 percentage-point gap confirms that distribution shift between the proposal distribution used to train the verifier and the distribution being verified at test time is a genuine practical concern.

Revision history in verifier context (Figure 15b, Appendix J): Including previous revisions in the ORM's input context provides a small improvement (~1–2 percentage points at 64 generations) over excluding history. Both variants outperform the parallel baseline, confirming that the sequential sampling benefit is not solely attributable to the verifier having more context — the revision model genuinely produces better candidates sequentially.

Oracle versus predicted difficulty bins (Figures 4, 8, and Appendix C, Figures 11–12): Both oracle and predicted difficulty bins yield qualitatively similar trends across difficulty levels. Predicted bins show slightly lower performance at high budgets in the revision setting (approximately 41% versus 44% at 256 generations in Figure 8) but essentially identical performance in the search setting (Figure 4). This is the critical robustness check: the compute-optimal strategy works without ground-truth labels, albeit with a small accuracy degradation at high budgets for revisions.

Majority voting for revisions (Figure 10, Appendix B): The sequential-to-parallel ratio trends observed with verifier-based selection replicate qualitatively with majority voting: easy questions are insensitive to ratio, hard questions show an optimal intermediate ratio, and fully sequential marginally outperforms fully parallel in aggregate. This replication with a non-learned selection mechanism confirms that the revision model's sequential improvement is genuine and not an artifact of verifier design.

ReST^EM revision model (Figure 16, Appendix K): An attempt to further optimize the revision model using ReST^EM (Singh et al., 2024) — an RL-based self-improvement procedure — produces a negative result: additional sequential revisions substantially hurt performance. At 256 generations, fully sequential performance drops to approximately 33.5% compared to roughly 38.5% at the optimal ratio for the base revision model. The authors hypothesize that on-policy data collection in ReST^EM exacerbates spurious correlations in revision data, causing the model to fail to learn the revision task properly. This negative result highlights the sensitivity of revision training to the data generation procedure and serves as a caution against naïve application of self-improvement techniques to revision models.

Critical Assessment

The experiments genuinely demonstrate that difficulty-conditioned allocation of test-time compute can recover substantial efficiency gains over uniform strategies, with the 4× figure for both search and revisions supported by side-by-side comparisons at specific budget breakpoints (e.g., 16 generations of compute-optimal search matching 64 generations of best-of-N in Figure 4; 64 generations of compute-optimal revisions matching 256 generations of parallel best-of-N in Figure 8). However, the contribution is best understood as an upper bound on achievable efficiency rather than a realized deployment gain, because the difficulty estimation cost — generating 2,048 samples per question — is not accounted for in any budget calculation. The paper acknowledges this: "our experiments do not account for this cost largely for simplicity" (Section 3.2). In a realistic deployment where the cost of estimating difficulty is amortized over many queries to the same question, this may be acceptable; for one-shot queries, the total cost would be difficulty estimation plus strategy execution, and the former could dominate the latter. The paper does not explore the break-even point where difficulty estimation costs are recovered by improved allocation.

The claim that test-time compute can substitute for pretraining (Section 7, "a smaller model augmented with compute-optimal test-time strategies can outperform a ~14× larger pretrained model") is supported only under specific conditions: easy-to-medium difficulty problems and low inference-to-pretraining token ratios (R ≪ 1). The paper is transparent about these boundary conditions, but the strength of the claim as presented in the introduction and abstract ("surpasses CLIP and ALIGN by 9.3%") should not be interpreted as a universal substitution. On hard problems (difficulty bin 5), test-time compute provides essentially zero benefit regardless of budget — a finding that is both robust and sobering. The paper's FLOPs-matched comparison also uses a pretraining baseline that may not be compute-optimal: the 14× larger model scales parameters only, following the LLaMA paradigm rather than Chinchilla-optimal scaling where both data and parameters are scaled equally. A compute-optimally trained larger model would likely be a stronger baseline, potentially reducing or reversing the reported advantages.

A genuine weakness is the absence of a combined search + revisions experiment. The paper studies PRM tree-search and iterative revisions independently, demonstrating they have complementary difficulty-dependent strengths (search helps on medium problems, revisions help on easy problems), but never combines them. Section 8 explicitly acknowledges this gap. This means the reported results represent a lower bound on what a fully integrated system could achieve — and it means the claim that "compute-optimal test-time scaling" is optimal is actually a claim that is optimal within the restricted set of strategies tested, not globally optimal. The natural next step — using the revision model as the proposal distribution in beam search, or using the PRM to guide which revisions to pursue — is not evaluated.

The 500-question test set introduces statistical concerns. With five difficulty quintiles of roughly 100 questions each, and two-fold cross-validation splitting each bin roughly in half, the compute-optimal policy is selected based on roughly 50 questions per fold per bin. This is a small sample for selecting among multiple discrete strategies (beam search with different widths, lookahead search with different lookahead depths, different sequential-to-parallel ratios), and the selected strategies may not be robust. The paper does not report confidence intervals or error bars on the compute-optimal scaling curves, making it difficult to assess whether the observed differences between strategies (e.g., the ~2.5 percentage-point gap between sequential and parallel revisions at 64 generations in Figure 6, right) are statistically significant or within sampling noise.

The single benchmark, single model family design limits claims of generality. All experiments use MATH with PaLM 2-S*. The paper argues this model is "representative of the capabilities of many contemporary LLMs," but the specific findings — the PRM's over-optimization threshold, the revision model's 38% correct-to-incorrect reversion rate, the optimal sequential-to-parallel ratios — could be model-specific. Replication on other model families (GPT, LLaMA) and other reasoning benchmarks (code generation, logical deduction, scientific QA) would strengthen the claims considerably. The paper does not provide any evidence that the findings transfer.

The inexactness of the pipelining + GradAccum method (Section 4.2) is acknowledged but its impact on final accuracy is not rigorously quantified. The paper compares SPMD and pipelining in terms of step time and memory (Table 2) but does not train models to convergence with both methods and compare their final accuracy. The variance approximation for the second moment and the batch normalisation inconsistencies when M ≪ B could introduce subtle degradations that are invisible in step-time profiling but matter at scale. This is a missing ablation that would strengthen the engineering contribution.

The revision model's 38% correct-to-incorrect reversion rate (Section 6.1) is a substantial practical limitation that is mitigated rather than solved. The paper uses majority voting or verifier-based selection across the chain to recover from reversions, but these are post-hoc patches. A model that incorrectly revises 38% of its own correct answers imposes a ceiling on how long revision chains can usefully be — eventually, the reversion rate overtakes the improvement rate, which is visible in the plateau in Figure 6 (left) beyond roughly 15–20 steps. Training the revision model to recognize when no revision is needed (by including correct-to-correct trajectories in the training data) is a natural improvement that the paper does not explore.

Finally, the pretraining-finetuning failure on MNIST (Section 8), while framed as a motivating observation for the hybrid training procedure, is not systematically evaluated. The paper states that pretrained-and-finetuned models "completely fail" on MNIST but does not provide quantitative MNIST accuracy for different training configurations (pretrain-only, contrastive-from-scratch, hybrid). A controlled experiment varying the amount of joint finetuning (Phase 3) and measuring MNIST accuracy would clarify whether the failure is remediable or fundamental, and would strengthen the practical guidance for practitioners deciding between pretraining and contrastive-from-scratch approaches.

6. Limitations and Trade-offs

The Difficulty Estimation Cost Is Unaccounted for in Efficiency Claims

The entire compute-optimal framework depends on knowing each prompt's difficulty before deciding how to allocate the inference budget. The paper's method for estimating difficulty — generating 2,048 samples per question and averaging either ground-truth correctness (oracle) or PRM final-answer scores (predicted) — is extraordinarily expensive. At 2,048 samples per question, difficulty estimation alone consumes more compute than the largest test-time budgets studied (256–512 generations). The authors acknowledge this explicitly in Section 3.2:

"estimating difficulty in this way still incurs additional computation cost during inference... our experiments do not account for this cost largely for simplicity"

The consequence is that the headline efficiency gains — the 4× reduction in compute over best-of-N (Figures 4 and 8) — are computed after difficulty is already known, without amortizing the cost of learning it. In a realistic single-query deployment, the total cost would be difficulty estimation plus strategy execution, and the former could dominate the latter. For a problem where 2,048 difficulty-estimation samples are generated and then the compute-optimal strategy uses only 16 generations, the actual total cost is 2,064 generations — not the 16 generations reported in the efficiency claim. This makes the 4× figure an upper bound on achievable efficiency rather than a realized deployment gain. The paper does not explore the break-even point: how many queries to the same question (or how much cheaper difficulty estimation would need to become) before the adaptive strategy's savings exceed the estimation overhead.

The paper attempts to mitigate this by showing that predicted difficulty bins (using the PRM's own scores, which avoids needing ground-truth labels) perform similarly to oracle bins (Figures 4 and 8, curves largely overlap), but this does not address the computational cost — the PRM-based method still requires generating and scoring 2,048 samples per question. The paper suggests future work on "training models to directly predict difficulty of a question" (Section 8) but develops no such model and provides no evidence that cheap difficulty estimation is feasible. A natural mitigation — adaptive difficulty estimation where a small initial sample is used to estimate difficulty and the remaining budget is allocated accordingly — is mentioned as a direction (Section 3.2) but not implemented or evaluated. Until this gap is closed, the reported efficiency gains should be understood as conditional on a difficulty oracle whose cost is not included in the accounting.

Hard Problems Are Essentially Unsolved — Test-Time Compute Cannot Substitute for Capability

Across all methods — PRM search, iterative revisions, and their compute-optimal combinations — the hardest questions (difficulty bin 5) show near-zero improvement regardless of compute budget. In Figure 3 (right), bin 5 accuracy hovers at 1–3% for all methods and all budgets up to 256 generations. In Figure 7 (right), bin 5 shows roughly 2–3% accuracy irrespective of the sequential-to-parallel ratio. In the FLOPs-matched comparison (Figure 9), the bin 5 scaling line is essentially flat near 0–5% for revisions and search alike, and falls below the 14× larger model's greedy performance across all values of R.

This is a fundamental capability bound, not a limitation of the allocation strategy. Test-time compute amplifies existing capability — it helps the model find and refine correct solutions that already exist in its sampling distribution — but it cannot create capability that is not there. The paper is transparent about this (Section 7 takeaway box) but the finding has practical consequences that extend beyond the paper's explicit claims. For any problem class where the base model's pass@1 rate is near zero (whether due to genuinely out-of-distribution reasoning, insufficient pretraining on the relevant domain, or inherent task difficulty), no amount of test-time compute — regardless of how optimally allocated — will help. The compute-optimal framework is therefore inapplicable to settings where the base model is fundamentally insufficient, and the FLOPs-matched comparison in Section 7 confirms that pretraining (scaling to a larger model) is the only viable path for such problems. The paper does not provide guidance on how to determine, before investing test-time compute, whether a given problem falls into the "hard" regime where effort will be wasted — the difficulty estimation procedure described in Section 3.2 can identify bin 5 problems post hoc (by observing the near-zero pass@1 rate), but this requires the 2,048-sample estimation cost that the paper does not account for.

Mitigation status: the paper acknowledges this limitation but does not address it — it is framed as a boundary condition on the method's applicability. No attempt is made to extend test-time compute to hard problems through alternative mechanisms (e.g., retrieval-augmented generation, tool use, or multi-model collaboration).

Search and Revisions Are Studied Independently — the Combined System Is Not Evaluated

The paper studies two complementary mechanisms for spending test-time compute — PRM-guided search (Section 5) and iterative revisions (Section 6) — but never evaluates them in combination. Section 8 explicitly acknowledges this gap:

"we did not experiment with PRM tree-search techniques in combination with revisions"

This is a significant omission because the two mechanisms have complementary, difficulty-dependent strengths that the paper carefully documents: revisions improve the proposal distribution (generating better candidates through sequential refinement), which is most effective on easy-to-medium problems (Figures 6–7), while PRM search improves candidate selection (finding the best among generated candidates through verifier-guided beam search), which is most effective on medium-hard problems (Figure 3). The natural extension — using the revision model as the proposal distribution within beam search, or using the PRM to guide which revision branches to pursue — is not explored.

The consequence is that the reported compute-optimal scaling curves (Figures 4 and 8) represent the optimum within the restricted set of strategies tested, not a global optimum over all possible combinations. A system that switches between search and revisions per difficulty bin, rather than using only one mechanism, could outperform either approach alone. The paper's FLOPs-matched comparison (Section 7) reports results for revisions and search separately, but a combined system could shift the difficulty thresholds where test-time compute beats pretraining. Until such a combined system is evaluated, the paper's claims about the efficacy of compute-optimal test-time scaling should be understood as a lower bound — the true potential of difficulty-conditioned allocation may be higher.

Mitigation status: future work only. The paper acknowledges the gap and suggests combining PRM tree-search with revisions as a natural next step (Section 8), but provides no results or preliminary experiments in this direction. The independent treatments of search and revisions also mean that the optimal per-difficulty-bin strategy selected by the compute-optimal policy might not be the best strategy when both mechanisms are available — the policy optimizes over a restricted action space.

Single Benchmark, Single Model Family — Generality Is Unestablished

All experiments use the MATH benchmark (500 test questions) with PaLM 2-S* as the base model. The paper states in Section 4 that the authors "believe this model is representative of the capabilities of many contemporary LLMs," but this claim is unverified by any cross-model or cross-benchmark evaluation. Several aspects of the findings could be model- or domain-specific in ways that matter for practitioners:

The PRM's quality and over-optimization behavior (Figure 3, right) depend on PaLM 2-S*'s output distribution — its calibration, its error modes, its sensitivity to the Monte Carlo rollout training procedure (Appendix D). A model with different calibration properties (e.g., one that assigns overconfident probabilities to incorrect steps) could exhibit different over-optimization thresholds, shifting the difficulty bins where beam search degrades performance. The paper's finding that "last" step-wise aggregation outperforms "min" (Appendix E, Figure 13) contradicts prior work (Lightman et al., 2023; Wang et al., 2023) and is attributed to the use of soft Monte Carlo labels — but this attribution is tested only on PaLM 2-S*, and it is unclear whether other model families would show the same reversal.
The revision model's ability to learn from incorrect in-context examples (Section 6.1) depends on the base model's in-context learning capabilities, which vary substantially across model families (some models are stronger at following few-shot patterns than others) and could be affected by the specific edit-distance-based data pairing strategy.
The MATH benchmark consists exclusively of competition-level math problems requiring symbolic multi-step reasoning. The difficulty-dependent patterns — beam search hurting easy problems (Figure 3, bin 1), revisions helping easy problems (Figure 7, bin 2) — may not generalize to other reasoning domains. Code generation might show different patterns because syntax constraints provide a stronger verification signal; factual QA might show different patterns because the model either knows an answer or does not, with limited room for sequential refinement.

The test set of 500 questions, split into five quintiles of ~100 questions each and further split by two-fold cross-validation, means the compute-optimal policy is selected based on roughly 50 questions per fold per bin. This is a small sample for selecting among multiple discrete strategies, and the paper does not report confidence intervals on the compute-optimal scaling curves. The observed 4× efficiency gains could vary substantially with different random splits or different test sets. Mitigation status: not addressed. The paper provides no cross-benchmark or cross-model replication.

Verifier Over-Optimization Limits Scaling and the Problem Is Mitigated, Not Solved

The paper documents verifier over-optimization as a central limiting factor: beam search degrades easy-problem performance at high budgets (Figure 3, right, bin 1 accuracy drops from ~78% to ~77% as budget increases from 4 to 256), lookahead search — the most powerful optimizer — paradoxically performs worst overall (Figure 3, left), and qualitative examples in Appendix M show search producing degenerate outputs (repetitive low-information steps, overly short solutions) that score highly under the PRM but are incorrect.

The compute-optimal policy mitigates this by routing easy problems away from aggressive search (assigning them best-of-N instead of beam search), but it does not solve the underlying problem. On medium-difficulty problems (bins 3–4) where beam search is deployed, over-optimization still limits the scaling ceiling — the beam search curves in Figure 3 (right) flatten or decline at high budgets, meaning that even on the problems where search is the best available strategy, additional compute eventually becomes counterproductive. The paper acknowledges this in Section 8 by identifying "improving verifier robustness against reward hacking" as a key direction for future work, but provides no evidence that the over-optimization threshold can be raised through better PRM training, ensemble verification, or constrained search methods.

The practical consequence is that the compute-optimal approach has a hard ceiling determined by verifier quality. Improving the PRM — through larger training sets, adversarial training, calibration techniques, or architectural changes — could shift the difficulty thresholds and raise the maximum achievable accuracy, but the paper provides no guidance on how much improvement is possible or which interventions are most promising. The current results are therefore specific to the verifier quality achievable with the Monte Carlo rollout training procedure described in Appendix D. A practitioner with a better verifier would need to re-derive the optimal per-difficulty-bin strategy, and a practitioner with a worse verifier might find that even the compute-optimal policy underperforms best-of-N across all bins.

Mitigation status: acknowledged as future work (Section 8) but not addressed. The over-optimization threshold is characterized empirically but not shifted or extended.

The Revision Model Has a 38% Correct-to-Incorrect Reversion Rate

The paper reports in Section 6.1 that approximately 38% of correct answers produced during a revision chain get "revised" back to incorrect answers in the subsequent step. This is a direct consequence of the training data construction: the model is trained only on sequences where all in-context answers are incorrect (followed by a correct target), so it has no training signal for what to do when the current answer is already correct. When deployed, the model encounters correct answers in its own revision context and, having never been trained to recognize or preserve correctness, frequently modifies them into errors.

The paper mitigates this with a selection mechanism — majority voting or verifier-based selection across the entire chain of revisions, picking the best answer from any point in the chain rather than always taking the last revision. This is a post-hoc patch: it recovers from reversions by keeping earlier correct answers in the selection pool, but it does not prevent the reversion from occurring in the first place. Every reversion wastes subsequent revision steps (the model spends compute refining an answer that it will later discard) and limits how long revision chains can usefully be. The pass@1 trajectory in Figure 6 (left) shows accuracy plateauing around 24–25% by steps 15–20 and remaining flat out to 64 steps — the reversion rate overtakes the improvement rate, meaning additional sequential compute beyond roughly 15–20 steps provides effectively zero marginal benefit.

A principled solution — such as training the model on trajectories that include correct-to-correct steps (teaching it to recognize when no revision is needed), or using the verifier to detect correctness and stop the chain early — is not explored. The consequence is that the revision approach has a fundamental efficiency ceiling: 38% of the model's own correct work is wasted, and the compute-optimal policy compensates for this by limiting sequential depth and relying on parallel diversity (the optimal ratios in Figure 7 involve a mix of sequential and parallel rather than purely sequential chains). A revision model with a lower reversion rate could sustain longer purely sequential chains, potentially shifting the optimal allocation strategy and raising the accuracy ceiling. Mitigation status: the paper identifies the reversion problem and applies chain-level selection as a workaround, but does not address the root cause through training data modification or architecture changes.

7. Implications and Future Directions

How This Work Changes the Landscape

BASIC changes the landscape of vision-language pretraining by establishing that batch size is a first-class scaling dimension with theoretically grounded generalisation benefits, not merely a compute-efficiency knob to accelerate training. Prior to this work, the dominant interpretation — inherited from SimCLR and carried into CLIP and ALIGN — was that larger contrastive batch sizes help by providing more negative examples per gradient step, enabling faster convergence. The implication was that training longer with a smaller batch should eventually match the performance of a larger-batch model, since the total number of negatives seen would equalise. BASIC's theoretical analysis (Section 6, Theorems 1 and 2) and controlled experiments (Section 10.1, Figure 5, Table 4) jointly refute this interpretation. The generalisation gap shrinks as $O(1/\sqrt{B})$ regardless of the number of training examples $m$ , and empirically, models trained with larger batch sizes for fewer steps — equalising total examples seen — reach higher accuracy that additional training cannot recover. This converts batch size from a training-speed hyperparameter into an accuracy-critical design axis that must be scaled in concert with data and model size.

This reframing has direct methodological consequences for how scaling studies are conducted in vision-language research. A scaling law experiment that varies data and model size but fixes batch size (as was common following the Chinchilla paradigm) will mis-estimate the returns to scale because it misses the interaction between model capacity, dataset diversity, and the granularity of the contrastive denominator. BASIC demonstrates that the saturation points observed in prior work (e.g., SimCLR's finding that batch size benefits saturate at 8,192) are artifacts of small-scale configurations: when data grows to 6.6B examples and models grow to 3B parameters, batch sizes of 65,536 continue to provide gains well beyond the prior saturation threshold. The implication for future scaling studies is that three-dimensional scaling grids are necessary — fixing any one dimension while varying the others risks drawing conclusions that do not extrapolate.

The paper also resolves a tension in the robustness literature through a clean diagnostic experiment. CLIP established that zero-shot models exhibit higher effective robustness than supervised ImageNet-trained models but explicitly declined to attribute causality, cautioning against overinterpreting the finding. The dominant community interpretation leaned toward the view that contrastive pretraining on diverse web data inherently confers robustness — perhaps through the diversity of the training distribution or the language supervision signal. BASIC's finetuning experiment (Section 9.3, Figure 4) directly tests this: take a converged zero-shot model with high effective robustness, then finetune it on subsets of ImageNet's labeled data using the same contrastive loss with class-name prompts. The result — that more labeled data produces higher ImageNet accuracy but lower robustness (an 8.3% accuracy drop on ImageNet-R for a 3% ImageNet gain in the extreme) — demonstrates that robustness is not an inherent property of contrastive training or architecture. It is a fragile state that is actively lost when models are adapted to a specific labeled distribution. This finding reframes the robustness question from "how do we design training procedures that produce robust models?" to "what property of supervised finetuning destroys robustness, and can we prevent that destruction?" It also provides a concrete experimental protocol — finetuning a zero-shot model on increasing subsets of labeled data and tracking the accuracy-robustness curve — that future work can use as a diagnostic for robustness-preserving adaptation methods.

The engineering contribution — two complementary memory-management strategies spanning a generality-to-optimality spectrum — shifts the practical feasibility boundary for training large contrastive vision-language models. Before BASIC, training a 3B-parameter image encoder with a 65K contrastive batch size on 16GB accelerators was not demonstrated, and there was no systematic comparison of pipelining-based versus SPMD-based approaches for the all-pairs structure of the contrastive loss. The paper provides this comparison (Table 2), characterising the step-time, memory, and exactness trade-offs of each method. This is not a conceptual breakthrough but a practical infrastructure contribution that lowers the barrier to entry for other groups wanting to operate in this scaling regime.

Follow-Up Research This Work Enables

Three-dimensional scaling laws for vision-language contrastive models. BASIC demonstrates that data, model, and batch size interact — larger models and datasets push the optimal batch size higher, and the batch size benefit does not saturate at previously observed thresholds. But the paper does not fit a quantitative scaling law of the form $\text{accuracy} = f(D_{\text{data}}, N_{\text{params}}, B_{\text{batch}})$ . A natural follow-up would systematically vary all three dimensions (e.g., dataset sizes from 100M to 10B pairs, model sizes from 100M to 10B parameters, batch sizes from 8K to 256K) and fit a parametric function predicting zero-shot ImageNet accuracy. The key question is whether the three dimensions combine multiplicatively (a 2× increase in each dimension yields an 8× effective scale benefit) or whether they exhibit diminishing or even super-multiplicative returns. BASIC's observation that the benefit of large batch sizes saturates at 8K for SimCLR-scale configurations but continues to 65K for BASIC-scale configurations suggests super-multiplicative interaction: larger models and datasets amplify the batch size benefit. A fitted scaling law would allow practitioners to optimally allocate a fixed compute budget across the three dimensions, analogous to how Chinchilla scaling laws guide pretraining compute allocation. The experiment would require training dozens of models at different scale points, which is expensive but feasible with the memory-management infrastructure BASIC provides. A strong result would show a predictable functional form and identify the point at which batch size scaling saturates for a given data and model scale — something BASIC does not establish.

Cheap, one-shot difficulty estimation for compute-optimal allocation. The compute-optimal framework in the reference example (not BASIC, but the method described in the prior sections of this analysis) critically depends on estimating problem difficulty before allocating the inference budget, and the current method — generating 2,048 samples and scoring them with a PRM — costs more than the largest inference budgets studied. A direct follow-up would train a lightweight difficulty predictor that takes only the problem text as input and predicts which difficulty quintile (or a continuous difficulty score) the problem falls into, without generating any solutions. The training data already exists: the paper has computed difficulty bins for all 12,000 training questions using the 2,048-sample procedure. A small transformer or even a bag-of-words classifier trained on (question_text, difficulty_bin) pairs could be evaluated on the 500-question test set by comparing its predicted difficulty bin to the oracle bin and measuring whether compute-optimal allocation using the predicted bins recovers the 4× efficiency gain. A strong result would maintain >90% of the oracle efficiency gain while using a difficulty predictor that costs negligible compute (<1% of the inference budget). A negative result — finding that text-based difficulty prediction is substantially less reliable than PRM-score-based prediction — would establish that the compute-optimal framework has a fundamental deployment barrier and would motivate adaptive difficulty estimation (start with a few samples, estimate difficulty on the fly, allocate remaining budget accordingly) as an alternative.

Combined search and revisions. BASIC's reference-analysis counterpart studies PRM-guided tree search (Section 5) and iterative revisions (Section 6) as independent mechanisms, but never evaluates them in combination. The paper explicitly identifies this gap and provides evidence that the two mechanisms have complementary difficulty-dependent strengths: revisions excel on easy problems where local refinement of nearly-correct answers is sufficient, while beam search excels on medium-difficulty problems where the verifier can guide exploration toward solutions the model would not find by random sampling. A follow-up would implement a combined system: use the revision model as the proposal distribution within beam search — at each node of the search tree, instead of sampling a single next step from the base model, the revision model conditions on the partial solution and previous rejected branches to produce a refined candidate. Alternatively, use the PRM to guide which revisions to pursue: run multiple parallel revision chains, score the intermediate steps in each chain with the PRM, and allocate additional sequential budget to the chains with the highest PRM scores. The key evaluation would compare the compute-optimal combined strategy (switching between search-only, revision-only, and search+revision per difficulty bin) against the paper's reported search-only and revision-only compute-optimal curves. A strong positive result would show that the combined approach pushes the accuracy ceiling higher than either method alone, particularly on difficulty bins 3–4 where both mechanisms show intermediate benefits. A null result — finding that combining search and revisions provides no benefit beyond switching between them per-problem — would suggest that the mechanisms are substitutes rather than complements and that the compute-optimal policy's action space is already close to optimal.

Verifier over-optimisation resistance through adversarial PRM training. BASIC's reference-analysis counterpart documents verifier over-optimisation as the primary bottleneck preventing unbounded improvement from additional test-time compute: beam search degrades easy-problem performance at high budgets (Figure 3, bin 1), and lookahead search — the most powerful optimizer — performs worst overall (Figure 3, left). The PRM is trained on i.i.d. samples from the base model's output distribution, using Monte Carlo rollout supervision. But at test time, search algorithms explore regions of the output space far from the i.i.d. distribution — beam search preferentially samples steps with high PRM scores, which are exactly the regions where the PRM's calibration may be poorest because the training data contains few such examples. A follow-up would implement adversarial PRM training: iteratively train the PRM, run beam search to generate solutions that score highly under the current PRM but are incorrect, add those solutions to the PRM training set with corrected labels, and retrain. This is the standard recipe for making classifiers robust to adversarial examples, applied to process reward models. The key evaluation would compare the over-optimisation threshold — the budget level at which beam search accuracy begins to decline on easy problems — between the standard PRM and the adversarially trained PRM. A strong result would show that the adversarially trained PRM sustains beam search improvements to higher budgets, raising the accuracy ceiling and potentially shifting the compute-optimal policy to use beam search on a wider range of difficulty bins. A null result — finding that adversarial training does not shift the over-optimisation threshold — would suggest that the over-optimisation is not due to distribution shift but to a more fundamental limitation of the scalar-verifier-plus-search paradigm, and would motivate exploration of alternative verification approaches (ensembles, learned rejection, or confidence calibration).

Correct-to-correct revision training to eliminate the 38% reversion rate. The revision model in the reference analysis is trained only on trajectories where all in-context answers are incorrect, followed by a correct target. This causes a 38% rate of revising correct answers into incorrect ones (Section 6.1), which is mitigated post-hoc by selecting the best answer from anywhere in the revision chain rather than always taking the last revision. A direct follow-up would modify the training data construction to include trajectories where the in-context answers include correct ones, teaching the model to recognize when no revision is needed. Specifically, construct training sequences where the final answer in context is already correct (sampled from the model's own correct solutions) and the target is an identical copy — teaching the model to output the input unchanged when it is already correct. The key evaluation would measure the correct-to-incorrect reversion rate at each revision step for the modified model versus the original model, and track whether the pass@1 trajectory in the revision chain continues to improve beyond the 15–20 step plateau observed in Figure 6 (left). A strong result would reduce the reversion rate below 10% and show that longer revision chains (64+ steps) sustain monotonic accuracy improvements, raising the accuracy ceiling for purely sequential strategies and shifting the optimal sequential-to-parallel ratio toward more sequential depth. A null result — finding that adding correct-to-correct trajectories degrades the model's ability to make corrections when answers are incorrect — would suggest a fundamental tension between preserving correctness and making corrections, and would motivate alternative approaches like training a separate correctness detector to decide whether to trigger a revision.

Cross-model and cross-domain replication of the difficulty-dependent scaling patterns. All results in the reference analysis use PaLM 2-S* on the MATH benchmark. The difficulty-dependent patterns — beam search hurting easy problems, revisions helping easy problems, nothing helping hard problems — could be specific to this model's calibration, error modes, and the structure of mathematical reasoning problems. A replication study would test the same methods on (a) a different model family, such as LLaMA-2 or GPT-3.5, and (b) a different domain with clean correctness signals, such as code generation (HumanEval, MBPP) where unit tests provide ground-truth verification. The key question is whether the five-quintile difficulty-dependent strategy profiles (Figure 3 right, Figure 7 right) are qualitatively similar across models and domains. If they are, the compute-optimal framework can be applied confidently to new settings with only minor recalibration of the per-bin strategy lookup table. If they differ substantially — for instance, if beam search helps easy coding problems because the verifier signal is more reliable for code than for math — then the framework requires per-domain and per-model calibration, reducing its off-the-shelf applicability. A strong negative result (e.g., finding that best-of-N dominates across all difficulty bins for code) would establish boundary conditions on where adaptive allocation provides value.

Practical Applications and Downstream Use Cases

Cost-efficient on-device image classification. BASIC demonstrates that zero-shot transfer models can approach supervised accuracy levels — 85.7% on ImageNet, within striking distance of the 87.1% state-of-the-art supervised result without extra data — while retaining the deployment simplicity of zero-shot inference. For mobile or edge applications where downloading and running a task-specific finetuned model for each new classification need is impractical, a single BASIC model can serve as a general-purpose image classifier that accepts natural language class descriptions. A photo-organisation app, for instance, could offer users the ability to search for "sunsets with mountains" or "pictures of my golden retriever" without any per-user or per-query training, using the text encoder to embed the query and the image encoder to embed the photo library, then ranking by cosine similarity. The 85.7% accuracy means the system would be reliable enough for consumer-facing applications where the cost of occasional misclassification is low. The robustness results (84.3% average on distribution-shift benchmarks) further mean that the model would maintain this accuracy on user photos taken in varied lighting, angles, and contexts — precisely the conditions where supervised models trained on clean benchmark datasets often degrade.

Robustness-critical deployment where finetuning is dangerous. The paper's finding that supervised finetuning on labeled ImageNet data reduces robustness (Figure 4, Section 9.3) has direct implications for any application where distribution-shift resilience matters more than a few points of in-distribution accuracy. Consider a wildlife monitoring system deployed across different camera traps, weather conditions, and geographic regions. The standard pipeline — pretrain on web data, then finetune on a small set of labeled camera-trap images — may produce a model with high accuracy on the specific camera-trap locations in the finetuning set but brittle performance when deployed at new locations with different lighting, backgrounds, or animal poses. BASIC's results suggest that skipping the finetuning step entirely and using zero-shot classification with descriptive text prompts for each species (e.g., "a photograph of a snow leopard in the wild") could yield more reliable performance across deployment sites, even if the in-distribution accuracy on the finetuning locations is slightly lower. The 84.3% average robustness accuracy (Table 1) — only a small drop from the 85.7% ImageNet accuracy — quantifies this reliability advantage over supervised models that might achieve 90%+ on their training distribution but degrade to 50–60% on shifted inputs.

Large-scale data filtering and deduplication for pretraining datasets. The SSIM-based contamination filtering used in BASIC (removing training images with SSIM ≥ 0.5 relative to any evaluation set image) is a specific instance of a broader capability: zero-shot models can identify semantically and visually similar images across massive datasets without requiring class labels. A data curation pipeline for a new vision-language pretraining dataset could use a BASIC model to compute image embeddings for all candidate images, then for each candidate, retrieve the most similar images in a reference set of evaluation benchmarks (ImageNet, ObjectNet, etc.) and flag candidates with cosine similarity above a threshold for removal or manual review. This is more scalable than per-dataset SSIM computation (which is $O(N_{\text{train}} \times N_{\text{eval}})$ in image pairs) and can catch semantic near-duplicates that SSIM misses (e.g., the same object photographed from different angles). The 85.7% ImageNet accuracy establishes that the embedding space is discriminative enough for this task: images with similar content are reliably mapped to nearby embeddings.

When to Prefer This Method

The paper (BASIC, not the reference example's compute-optimal framework) does not articulate a clear decision rule for choosing between its approach and specific named alternatives. It positions itself as a direct extension and scaling of CLIP/ALIGN rather than a method that trades off against them. The scaling dimensions — data, model, batch size — are presented as levers that should be pushed simultaneously to their joint limits, not as choices among which to select. The two memory-management strategies (pipelining + GradAccum versus SPMD + rematerialization) are presented as complementary tools for different hardware and scaling regimes, with the trade-off characterised in Table 2 and Section 5.3:

Prefer SPMD + rematerialization when you need exact gradient computations, faster step times (~15% faster than pipelining at $B = 2^{20}$ for the medium model in Table 2), and can invest the engineering effort to design a weight sharding and rematerialization schedule for your specific architecture. This is the right choice for production training runs where the model architecture and batch size are fixed and the primary objective is throughput.
Prefer pipelining + GradAccum when you need a generic solution that works with arbitrary model architectures without manual sharding design, or when your contrastive batch size may grow beyond what your current SPMD configuration supports. This method's memory footprint is constant in $B$ (it depends only on the microbatch size $M$ , not the global batch size), so it scales to arbitrarily large batches at the cost of inexact gradient accumulation and slower step times from the double forward passes.

Beyond this engineering trade-off, the paper does not present a "prefer zero-shot BASIC over supervised finetuning when..." decision framework, because the robustness-destruction finding (Section 9.3) is presented as a diagnostic observation rather than a prescriptive guideline. The implicit practical advice — skip finetuning if robustness to distribution shift is more important than maximising in-distribution accuracy — is supported by the data (Figure 4) but not elevated to a formal recommendation.