Sigmoid Loss for Language Image Pre-Training

In plain words

Models like CLIP learn to put pictures and their captions close together in a shared space, so that you can later classify an image by comparing it to text such as "a photo of a dog". The standard training signal is a softmax over a whole batch: for each image, the model must pick its own caption out of every caption in the batch, like a multiple-choice question with thousands of options. That needs every device to see every other device's embeddings and to build one huge table of all image-caption scores.

SigLIP asks a simpler question for each image-caption pair on its own: "do these two belong together, yes or no?" Each pair becomes a tiny binary classification problem, scored with a sigmoid. Because no pair needs to know about the others, the work can be split into small pieces that devices pass around, and memory stays small. Since almost every pair in a batch is a mismatch, the authors add a learned offset that starts the model off predicting "no", so early training is not dominated by fixing that imbalance.

The result is a loss that does as well as or better than the softmax version, clearly better at small batch sizes, and lets you train with far fewer accelerators. The authors also find that pushing the batch size towards a million brings little: around 32 thousand pairs per batch is enough.

TL;DR

• Replace the two-direction softmax (InfoNCE) loss of CLIP/LiT with a per-pair sigmoid binary cross-entropy over all |B|² image-text pairs, with logit t·<x_i, y_j> + b, learnable t = exp(t') and b, initialised t' = log 10, b = -10 §3.2 alg. 1 §4.9.
• The loss is a sum of independent terms, so it can be computed device-by-device with text embeddings rotated between devices instead of an all-gather and a full |B|×|B| matrix; per-device memory goes from |B|² to b² §3.3 fig. 1.
• Sigmoid beats softmax clearly below 16k batch; both saturate around 32k, and 1M batch gives almost nothing over 256k §4.1 §4.2 fig. 2.
• Headline (author-reported): SigLiT g/14 84.5% ImageNet 0-shot with 4 TPUv4 chips in 2 days; SigLIP B/16 from scratch 73.4% with 32 chips in 5 days; SigLIP SO-400M at 384px 83.2% ImageNet 0-shot tab. 1 tab. 3.
• Also: set Adam/Adafactor β2 = 0.95 to survive large-batch gradient spikes §4.7.

Problem

CLIP-style contrastive pre-training normalises pairwise similarities with a batch softmax, applied twice (image→text and text→image) because the softmax is asymmetric §3.1. A stable softmax needs an extra pass to subtract the maximum §1. In data-parallel training, every device must all-gather all embeddings and materialise the |B|×|B| similarity matrix, which is memory heavy and scales quadratically with batch size §3.3. This makes large batches expensive and couples the batch size to the definition of the task §1. Like any contrastive setup, the method assumes that the text belonging to image j is unrelated to image i for i ≠ j, which is noisy in practice §3.

Method

Both towers produce embeddings that are L2-normalised: x_i = f(I_i)/‖f(I_i)‖₂, y_i = g(T_i)/‖g(T_i)‖₂ §3.1. The image tower is a ViT, the text tower a transformer §3.1.

Softmax baseline (CLIP/LiT) §3.1:

-\frac{1}{2|\mathcal{B}|}\sum_{i=1}^{|\mathcal{B}|}\left(\log\frac{e^{t\,x_i\cdot y_i}}{\sum_{j=1}^{|\mathcal{B}|} e^{t\,x_i\cdot y_j}} + \log\frac{e^{t\,x_i\cdot y_i}}{\sum_{j=1}^{|\mathcal{B}|} e^{t\,x_j\cdot y_i}}\right) §3.1

with t = exp(t') and t' a global learnable scalar §3.1.

Sigmoid loss §3.2. Every pair (i, j) in the batch is an independent binary classification; z_ij = 1 if image i and text j are paired, -1 otherwise. As printed in the paper:

-\frac{1}{|\mathcal{B}|}\sum_{i=1}^{|\mathcal{B}|}\sum_{j=1}^{|\mathcal{B}|}\underbrace{\log\frac{1}{1+e^{z_{ij}(-t\,x_i\cdot y_j + b)}}}_{\mathcal{L}_{ij}} §3.2

Note the normaliser is 1/|B| (number of positives), not 1/|B|² §3.2. The bias b is a learnable scalar like t; it counters the heavy imbalance of |B| positives against |B|² − |B| negatives, which would otherwise dominate the loss and cause large corrective steps early on §3.2. See Pitfalls for the sign of b in this equation versus Algorithm 1.

Algorithm 1 from the paper alg. 1:

# img_emb    : image model embedding [n, dim]
# txt_emb    : text model embedding [n, dim]
# t_prime, b : learnable temperature and bias
# n          : mini-batch size

t = exp(t_prime)
zimg = l2_normalize(img_emb)
ztxt = l2_normalize(txt_emb)
logits = dot(zimg, ztxt.T) * t + b
labels = 2 * eye(n) - ones(n) # -1 with diagonal 1
l = -sum(log_sigmoid(labels * logits)) / n

Chunked multi-device implementation §3.3 fig. 1. With D devices and per-device batch b = |B|/D (here b is a batch size, not the bias), the loss is rewritten as:

-\frac{1}{|\mathcal{B}|}\underbrace{\sum_{d_i=1}^{D}}_{\text{A: all devices}}\;\underbrace{\sum_{d_j=1}^{D}}_{\text{B: swap negs}}\;\underbrace{\sum_{i=b d_i}^{b(d_i+1)}\;\sum_{j=b d_j}^{b(d_j+1)}\mathcal{L}_{ij}}_{\text{C: per-device loss}} §3.3

Algorithm: chunked sigmoid loss  [§3.3] [fig. 1]
Input: on each device d, local image embeddings X_d and text embeddings Y_d (b each)
1. Each device computes the b×b block X_d · Y_dᵀ: b positives (diagonal) and
   b−1 negatives per row; accumulate its loss terms (sum C).
2. Permute text embeddings across devices, so device d now holds the texts of
   its neighbour; compute the b×b all-negative block and accumulate (sum B).
3. Repeat step 2 until every device has seen every other device's texts
   (D−1 permutes in total).
4. Sum the per-device losses across devices (sum A) and divide by |B|.
Output: the full sigmoid loss; only a b×b block is ever materialised.

The authors argue D collective permutes are faster than two all-gathers and that per-device memory drops from |B|² to b² §3.3. This is what enabled batch sizes above one million on relatively few devices §3.3.

Recipes studied.

• SigLiT: a frozen, pre-trained image tower (image embeddings precomputed) plus a text tower trained from scratch with the sigmoid loss §1 §4.1.
• SigLIP: both towers trained from scratch, or the image tower initialised from a pre-trained, unlocked checkpoint §4.2 §4.5.
• mSigLIP: multilingual SigLIP on all ~100 WebLI languages with a 250k vocabulary and a bottlenecked token embedding: an N×K embedding matrix followed by a K×W projection, with K ≪ W §4.3.
• When fine-tuning a pre-trained image tower inside SigLIP, use a 0.1 learning-rate multiplier on it and apply no weight decay to it; weight decay only on the randomly initialised text tower §4.5.

Hyperparameters

Results

All numbers are author-reported. ImageNet accuracy is zero-shot top-1 (%).

Small-hardware setups tab. 1. SigLiT uses frozen public checkpoints; SigLIP B/16 in the first row starts from a public unlocked checkpoint, the last two rows are from scratch.

For comparison the authors cite CLIP at about 2500 TPUv3-days for 72.6% §4.5. Note that the text says the B/8 SigLiT reached 79.7% §4.4 while Table 1 lists 79.8.

SigLiT batch-size sweep (ImageNet 0-shot, sigmoid / softmax) tab. 8:

The best SigLiT with a B-sized text model reaches 84.7%, vs 85.2% reported by LiT with a 10× larger g-sized text model §4.1. With short schedules, very large batches lose because they take few optimizer steps §4.1 fig. 3.

SigLIP batch-size sweep (ImageNet 0-shot) tab. 5:

Sigmoid peaks at 32k, softmax needs 98k and still does not exceed the sigmoid peak; 307k hurts both §4.2. On four TPUv4 chips, Base SigLIP fits batch 4096 vs 2048 for the equivalent CLIP §4.2.

mSigLIP batch-size sweep (Base, 30B examples; XM = Crossmodal-3600 text→image recall@1) tab. 2:

The 32k result (34.9%) is compared against 28.5% reported for a LiT model with a four-billion-parameter ViT-e §4.3. The scaled-up mSigLIP ViT-B (§4.6 recipe, "*32k") averages 54.1 image→text and 42.6 text→image recall@1 over 36 languages; per-language numbers are in tab. 9 and fig. 8. The authors note a concurrent Large model reaching 42.96% image retrieval recall@1 §4.6.

Comparison with released models tab. 3 (ImageNet-1k zero-shot; COCO zero-shot recall@1):

Bias and temperature initialisation (Base, 8k batch, 900M examples, SigLIP setup; 0-shot accuracy) tab. 4:

Hyperparameter robustness (SigLIP from scratch, 3B examples; ImageNet 0-shot) tab. 6. The sweep covered LR {0.0003, 0.001, 0.003} × WD {0.00003, 0.0001, 0.0003} × batch {8k, 16k, 32k} app. C.

Run-to-run variance (32k batch, 3B examples, five runs; ImageNet 0-shot) tab. 7:

A linear LR schedule instead of cosine gave 69.9% app. C.

Other findings.

• Fine-tuning a pre-trained image tower with weight decay degrades its representation (ImageNet 10-shot falls); disabling weight decay on it fixes this §4.5 fig. 4.
• Batch composition (SigLiT, 16k batch, 900M examples): masking negatives at random hurts more as the ratio approaches balance; keeping only the easiest negatives fails; keeping the hardest nearly preserves quality, and matching total pairs seen with hard negatives slightly improves it §4.8 fig. 6. The learned bias and mean logits rise as fewer negatives remain §4.8.
• Under image, text, batch-alignment or combined corruption, sigmoid-trained M-scale models stay ahead of softmax across corruption probabilities §4.10 fig. 7.

Paper vs code

• Sign of the bias. The printed sigmoid-loss equation puts b inside z_ij(−t·x·y + b), i.e. an effective logit t·x·y − b §3.2. Algorithm 1 and the code use logit t·x·y + b alg. 1 siglip.py. With b = −10 meant to push initial predictions towards "negative" §3.2, only the Algorithm 1 / code form is consistent.
• Chunked loss. The paper describes D collective permutes passing text embeddings to the neighbouring device §3.3 fig. 1. The current trainer has no chunked version ("TODO: implement chunked version with shard_map") and computes the full logit matrix under jit on sharded arrays siglip.py siglip.py. A chunked variant exists in the deprecated contrastive trainer; it loops over devices and fetches each device's text embeddings with a masked psum (an all-reduce), with a comment that a ppermute version was memory intensive _deprecated_contrastive.py.
• Example config differs from paper defaults. The released minimal SigLIP config (siglip_lit_coco.py, COCO captions) uses bias_init = -2.71, temperature_init = 10, lr = 0.001, wd = 0.01, Adam for a BERT-base text tower, gradient clipping at norm 1.0 and warmup of max(3% of steps, 100) siglip_lit_coco.py. The paper's bias init is −10 §4.9; the paper does not mention gradient clipping. The configs for the paper's own WebLI runs were not found at this commit.
• β2 in Adafactor. big_vision's Adafactor uses a step-dependent second-moment decay min(beta2_cap, 1 − (step+1)^(−0.8)) with beta2_cap = 0.999 by default optax.py. The paper's "β2 = 0.95" §4.7 presumably corresponds to setting this cap to 0.95; the paper does not say so.
• Normalisation epsilon. The code divides by ‖z‖ + 1e-8 two_towers.py; the paper writes a plain L2 normalisation §3.1.
• Loss normaliser matches: the code takes the row-sum of per-pair negative log-likelihoods and averages over images, i.e. divides by the number of positives siglip.py, as in alg. 1.

Pitfalls

• Bias sign. Implement logits = t·<x, y> + b with b initialised at −10, as in Algorithm 1 and the code, not the printed equation's −t·x·y + b form (see Paper vs code). Earlier arXiv versions had a pseudocode mistake and t vs t' typos, per the acknowledgements p. 10; check which version you are reading.
• t vs t'. The learnable parameter is t'; the multiplier is t = exp(t'). Initialise t' = log 10, not t = log 10 §4.9. Initialising with t' = log 1 and b = 0 costs 9.3 ImageNet points in the ablation tab. 4.
• Do not divide by |B|². The loss is summed over all |B|² pairs but divided only by |B| §3.2 alg. 1. Dividing by |B|² changes the effective learning rate as batch size changes.
• Use log_sigmoid(labels * logits), not log(sigmoid(...)), for numerical stability alg. 1 siglip.py.
• Distributed negatives. Positives are only on the local diagonal; every cross-device block is all negatives §3.3. If you shard the batch, make sure cross-device blocks get label −1 everywhere and that each pair is counted exactly once.
• The chunked loss needs D−1 communication rounds per step and a loop over devices; "faster than all-gather" is the authors' claim, and the released chunked code actually uses all-reduces §3.3 _deprecated_contrastive.py.
• Large-batch instability. At large batch, gradient-norm spikes cause large parameter jumps and loss spikes; β2 = 0.95 fixed it for all their runs §4.7 fig. 5. Occasional spikes still occur fig. 5.
• Bigger batch is not better. Gains saturate near 32k; 307k hurt both SigLIP losses §4.2 and >32k hurt mSigLIP retrieval §4.3. Large batches also need long schedules to pay off fig. 3.
• Fine-tuning a pre-trained image tower: without disabling weight decay on it, results are barely above from-scratch §4.5 fig. 4. Use the 0.1 LR multiplier too §4.5.
• Unstated: learning-rate warmup and schedule for the from-scratch SigLIP runs, weight-decay form (coupled vs decoupled) for Adafactor runs, text-tower architecture details beyond size, and the exact WebLI filtering are not stated.
• SigLiT numbers depend on precomputed frozen embeddings of a strong image tower; the 4-chip results do not include the cost of training that tower or computing the embeddings §4.4.

Lineage

• Builds on: CLIP and ALIGN, whose batch-softmax loss it replaces; LiT, whose locked-image recipe gives SigLiT §1 §4.1.
• Related efficiency work: FLIP (arXiv:2212.00794) randomly drops visual tokens; BASIC and LAION scaled batch size only up to 16k and 160k respectively §2.
• Compared against: CLIP, OpenCLIP, EVA-CLIP (arXiv:2303.15389), CLIPA-v2 (arXiv:2306.15658) tab. 3.