Source-linked AI summary

Sigmoid Loss for Language Image Pre-Training

Xiaohua Zhai, Basil Mustafa, Alexander Kolesnikov, Lucas Beyer

arXiv:2303.15343v4cs.CVcs.AI

TL;DR

Language-image pre-training commonly relies on batch-wide softmax normalization, creating costly distributed computation and leaving limited evidence about how performance scales with batch size. This paper introduces a pairwise sigmoid loss and finds it outperforms softmax at smaller batch sizes while performance largely saturates around a 32k batch size.

  • Problem

    Contrastive language-image training requires costly all-gathers and memory-intensive batch-wide similarity matrices, while evidence on batch-size scaling remains limited.

  • Method

    The paper replaces softmax normalization with a symmetric pairwise sigmoid loss that decouples batch size from task definition and simplifies distributed training.

  • Results

    Sigmoid loss performs better than softmax at smaller batch sizes, while scaling beyond 32 k yields only minor gains and performance saturates.

  • Takeaways & Limitations

    A 32 k batch size is sufficient for near-optimal language-image pre-training, and the sigmoid loss supports larger batches with improved memory efficiency.

Abstract

from arXiv · show

We propose a simple pairwise Sigmoid loss for Language-Image Pre-training (SigLIP). Unlike standard contrastive learning with softmax normalization, the sigmoid loss operates solely on image-text pairs and does not require a global view of the pairwise similarities for normalization. The sigmoid loss simultaneously allows further scaling up the batch size, while also performing better at smaller batch sizes. Combined with Locked-image Tuning, with only four TPUv4 chips, we train a SigLiT model that achieves 84.5% ImageNet zero-shot accuracy in two days. The disentanglement of the batch size from the loss further allows us to study the impact of examples vs pairs and negative to positive ratio. Finally, we push the batch size to the extreme, up to one million, and find that the benefits of growing batch size quickly diminish, with a more reasonable batch size of 32k being sufficient. We release our models at https://github.com/google-research/big_vision and hope our research motivates further explorations in improving the quality and efficiency of language-image pre-training.

1. Introduction

The paper introduces sigmoid loss for language-image pre-training as a simpler, more efficient alternative to batch-level softmax contrastive loss. SigLIP and SigLiT improve efficiency and achieve strong ImageNet zero-shot accuracy with modest hardware, while enabling analysis of batch-size effects.

  • Efficiency and results: 71.0% 0-shot accuracy on ImageNet is achieved by SigLIP using 16 TPU-v4 chips for three days on WebLI with a public B/16 image checkpoint.The result uses a public image checkpoint rather than a randomly initialized model.
  • Sigmoid loss: Sigmoid loss avoids full-batch operations, simplifying distributed training, reducing memory use, and decoupling batch size from task definition.It is symmetric and requires a single pass, unlike the batch-level softmax loss.
  • Batch-size effects: Below batch size 16 k, sigmoid loss performs significantly better than softmax loss; as batch size grows, their performance gap closes.This comparison is reported across multiple image-text learning setups.
  • Efficiency and results: 79.7% zero-shot accuracy on ImageNet is achieved by SigLiT in one day using four TPUv4 chips.The setup uses a frozen public B/8 image checkpoint and the LiT image-text dataset.
  • Fine-tuning: Disabling weight decay on the pre-trained vision backbone produces better SigLIP fine-tuning results.The paper refers to Figure 4 for details.

2. Related Work

Prior language-image pre-training largely uses softmax contrastive learning, while generative decoders and efficiency-focused methods provide alternative directions. Related work also finds sigmoid loss slightly more effective and robust than softmax loss in supervised classification.

  • Contrastive learning with the sigmoid loss: The sigmoid loss has prior use in unsupervised dimensionality reduction and is slightly more effective and robust than softmax loss in supervised classification.Most contrastive image-text learning works instead rely on the softmax-based InfoNCE loss.
  • Contrastive language-image pre-training: CLIP and ALIGN popularized softmax contrastive learning for large-scale image-text pre-training, achieving strong zero-shot classification and retrieval.Follow-up models also transfer effectively to fine-tuning, linear regression, object detection, semantic segmentation, and video tasks.
  • Generative language-image pre-training: Generative alternatives such as GIT, SimVLM, and LEMON pre-train language-image models with generative text decoders.CoCa combines a decoder with the discriminative CLIP/ALIGN setup, while BLIP uses a decoder for caption generation and pair filtering.
  • Efficient language-image pre-training: LiT and FLIP are notable efficiency efforts, while BASIC and LAION study larger batch sizes, reaching 16 k and 160 k respectively.LiT requires a pre-trained locked backbone, whereas FLIP randomly drops visual tokens and sacrifices quality; these scaling studies use many hundreds of chips.

3. Method

The method replaces softmax-based contrastive learning with an independently computed pairwise sigmoid loss, then distributes its terms across devices without global normalization or all-gathers. A learnable bias addresses the severe initial imbalance between matching and nonmatching pairs.

  • Softmax contrastive objective: Softmax contrastive learning aligns matched embeddings and separates unmatched embeddings, with normalization performed independently across images and texts.The image and text embeddings are produced by vision-transformer and transformer encoders, respectively.
  • Pairwise sigmoid loss: The sigmoid loss treats matching image-text pairs as positive labels and every nonmatching combination as a negative label in independent binary-classification terms.This removes the need to compute global normalization factors over pairwise similarities.
  • Pairwise sigmoid loss: The learnable bias b is initialized to −10, while t′ is initialized to log 10, preventing large early updates caused by the many negative pairs.The bias plays a role analogous to the temperature parameter t.
  • Efficient implementation: The distributed implementation avoids expensive all-gathers and materializing the memory-intensive |B| × |B| similarity matrix required by conventional contrastive training.Each device computes positive and local negative terms, swaps neighboring-device representations for additional negatives, and sums losses across devices.

4. Results

Results show that SigLIP performs well across batch sizes, multilingual settings, and training interventions, with benefits generally saturating around 32k examples. The experiments also identify effective stabilization, negative-pair selection, bias initialization, and robustness to data noise.

  • Batch size: 32k is optimal for SigLIP, while softmax requires 98k and still does not outperform sigmoid; 307k hurts both losses.As batch size increases, the performance gap between sigmoid and softmax losses diminishes.
  • Multilingual pre-training: A bottleneck with K = 96 and W = 768 causes only about a half-percent ImageNet zero-shot quality drop versus the full 250k vocabulary.The bottleneck replaces the N × W lookup table with N × K embeddings and a K × W projection.
  • Training stability: Removing weight decay from pre-trained weights significantly improves SigLIP, while applying it only to randomly initialized text-model weights stabilizes the recipe.Large-batch transformer pre-training becomes increasingly unstable because of large spikes in training dynamics.
  • Multilingual pre-training: 42.6% image-retrieval recall@1 and 54.1% text-retrieval recall@1 establish state-of-the-art results for the scaled-up multilingual mSigLIP ViT-B on XM3600.The Large model in slightly outperforms image retrieval with 42.96% recall@1.
  • Pair composition: Hard negative pairs preserve nearly all performance, whereas random negative removal deteriorates it and retaining easy negatives fails entirely.Longer training on hard negatives slightly increases performance, while the positive-negative imbalance itself is not a major concern.
  • Bias initialization: A −10 bias initialization consistently improves performance by starting training near the prior and preventing early over-correction.Random initialization, including 0, produces significantly worse results, especially with a small temperature t′.
  • Data noise robustness: Sigmoid-trained models become increasingly robust to all kinds of added noise as corruption likelihood increases.This finding comes from experiments varying the likelihood of data corruption.

5. Conclusion

The study evaluates sigmoid loss in SigLiT and SigLIP, finding it outperforms the softmax baseline, particularly with small train batch sizes. Its memory efficiency also enables larger train batch sizes without additional resources.

  • Sigmoid loss was studied in two language-image pretraining instances: SigLiT and SigLIP.
  • Sigmoid loss performs better than the softmax baseline, particularly for small train batch sizes.
  • Memory efficiency enables larger train batch sizes without requiring additional resources.

A. More results for SigLiT

SigLiT training used precomputed ViT-g image embeddings with fixed-resolution resizing and a standard base text tower. An under-a-day, four-chip setup instead used LION with a staged learning-rate schedule.

  • Training setup: Precomputed ViT-g image embeddings were trained with resize-only augmentation to a fixed 288 × 288 resolution.The image embeddings came from a ViT-g vision model.
  • Training setup: The standard base text tower used ScalingViT-Adafactor with β1 = 0.9, β2 = 0.95, learning rate 0.001, and weight decay 0.0001.Learning rate warmup covered the first 200 M examples, followed by cosine decay to zero.
  • Training setup: Under-a-day SigLiT training on 4 chips used LION with peak learning rate 1×10−4 and weight decay 1×10−7.The learning rate warmed linearly for 6.5 k steps before cosine decay over the remaining 58.5 k steps.

B. More results for SigLIP

SigLIP Base is evaluated with 3 billion and 9 billion training examples across batch sizes from 512 to 307k. At 9 billion examples, sigmoid loss peaks at a smaller batch size than softmax loss, alongside a memory-efficiency advantage.

  • Training examples and batch sizes: SigLIP Base experiments use 3 billion and 9 billion training examples across batch sizes from 512 to 307k.ImageNet zero-shot accuracy is reported for these settings.
  • Loss comparison: At 9 billion examples, sigmoid loss peaks at 32k batch size, earlier than softmax loss at 98k.The comparison concerns ImageNet zero-shot accuracy.
  • Loss comparison: The earlier sigmoid-loss peak, together with its memory-efficient advantage, enables training the best language-image model with much larger batches.The supplied table passage states this conclusion but truncates the final quantitative detail.

C. Robustness of SigLIP results

SigLIP results are robust across batch sizes, repeated runs, and optimizer choices. Default hyperparameters remain effective from batch sizes 512 to 1024k, while repeated training produces very small variation and AdamW gives similar results.

  • Hyperparameter robustness: Default 0.001 learning rate and 0.0001 weight decay work across batch sizes from 512 to 1024k.This applies to both SigLiP and SigLiT setups.
  • Run-to-run stability: Five runs with the recommended 32k batch size and 3B seen examples show very small standard deviations for sigmoid and softmax.Table 7 reports the average and standard deviation across the five runs.
  • Alternative optimizers: AdamW produces very similar results and standard deviations, while a linear learning-rate scheduler achieves 69.9% accuracy.The comparison repeats the same experiment and contrasts linear scheduling with the default cosine scheduler.

D. More results for mSigLIP

This section reports mSigLIP Base crossmodal retrieval results on Crossmodal-3600 across all 36 languages, alongside SigLiT ImageNet zero-shot experiments comparing sigmoid and softmax losses across training scales and batch sizes.

  • Crossmodal retrieval: mSigLIP Base retrieval is evaluated across all 36 Crossmodal-3600 languages using image-to-text and text-to-image recall@1.Figure 8 reports both retrieval directions for every language.
  • ImageNet zero-shot accuracy: SigLiT ImageNet zero-shot accuracy is compared between sigmoid loss and a softmax baseline across 450 M, 900 M, 3 B, and 18 B train examples.The experiments also vary train batch sizes from 512 to 1 M.
  • Crossmodal retrieval: Figure 8 marks 32 k as the scaled-up batch-size result for the Crossmodal-3600 retrieval experiments.Colors indicate batch sizes, and the asterisk identifies the scaled-up 32 k results described in Section 4.6.

E. Label noise experiments

The experiments used M/16 image and M text towers trained from random initialization for 3.6B examples with batch size 16384. Training employed cosine learning-rate decay with 10% linear warmup to a peak learning rate of 0.001.

  • Experimental setup: 3.6B examples were seen during training with a batch size of 16384.All models used random initialization, an M/16 image tower, and an M text tower.
  • Experimental setup: 0.001 was the peak learning rate under a cosine schedule with 10% linear warmup.The warmup preceded the peak learning rate.
  • Crossmodal-3600 evaluation: 30 B total examples were used to train mSigLIP models at different batch sizes for Crossmodal-3600 retrieval evaluation across all 36 languages.The evaluation reported image-to-text and text-to-image zero-shot recall@1, with 32 k denoting scaled-up results.

F. Model Card

The model card describes SigLIP as a two-encoder contrastive model using sigmoid loss, with matched vision and language encoder sizes. It specifies the models’ inputs, outputs, intended uses, training and evaluation data, and implementation environment.

  • Model Architecture: SigLIP uses vision-transformer and language-transformer encoders with matched sizes selected from ViT-B, ViT-L, and SoViT-400M.The model is trained with contrastive pre-training using sigmoid loss.
  • Inputs and Outputs: Vision inputs include 224 × 224×3, 256×256×3, 384×384×3, and 512×512×3 images, while text is cropped to 64 tokens.Both encoders output d-dimensional feature vectors, with d equal to 768, 1024, and 1152 for ViT-B, ViT-L, and SoViT-400M, respectively.
  • Use and Training Data: The models support multimodal research, including zero-shot image classification and image-text retrieval, with both English-only and multilingual-trained variants.SigLIP uses mostly English-filtered WebLI data, whereas mSigLIP uses WebLI without language filters.
Loading 2303.15343v4…