From bd90295b1dc2dbe7f42f1fd2d9a45bc12895298f Mon Sep 17 00:00:00 2001 From: 24GUNV Date: Sun, 10 May 2026 04:19:03 -0400 Subject: [PATCH] Add SigLIP loss and forward template --- benchmarks/diffusiondb2m/siglip-vitb.py | 220 ++++++++++++++++++ docs/source/api/forward.rst | 28 +++ docs/source/api/losses.rst | 2 + docs/source/references.bib | 7 + stable_pretraining/forward.py | 63 +++++ stable_pretraining/losses/__init__.py | 3 +- stable_pretraining/losses/multimodal.py | 74 ++++++ .../tests/unit/test_forward_functions.py | 64 +++++ .../tests/unit/test_siglip_loss.py | 91 ++++++++ 9 files changed, 551 insertions(+), 1 deletion(-) create mode 100644 benchmarks/diffusiondb2m/siglip-vitb.py create mode 100644 stable_pretraining/tests/unit/test_siglip_loss.py diff --git a/benchmarks/diffusiondb2m/siglip-vitb.py b/benchmarks/diffusiondb2m/siglip-vitb.py new file mode 100644 index 000000000..63dd1f5cc --- /dev/null +++ b/benchmarks/diffusiondb2m/siglip-vitb.py @@ -0,0 +1,220 @@ +import argparse +from functools import partial + +import lightning as pl +import torch +import torch.nn.functional as F +from lightning.pytorch.callbacks import LearningRateMonitor, ModelCheckpoint +from lightning.pytorch.loggers import WandbLogger +from transformers import ( + AutoImageProcessor, + AutoTokenizer, + SiglipTextModel, + SiglipVisionModel, +) + +import stable_pretraining as spt +from stable_pretraining import forward + + +parser = argparse.ArgumentParser() +parser.add_argument("--lr", type=float, default=0.001) +parser.add_argument("--num_devices", type=int, default=8) +parser.add_argument("--global_batch", type=int, default=4096) +parser.add_argument("--num_epochs", type=int, default=8) +parser.add_argument("--val_percent", type=float, default=0.10) +parser.add_argument("--resume_ckpt_path", type=str, default=None) +args = parser.parse_args() + +lr = args.lr +num_devices = args.num_devices +global_batch = args.global_batch +batch_size = global_batch // num_devices +num_epochs = args.num_epochs +val_percent = args.val_percent +resume_ckpt_path = args.resume_ckpt_path + +model_name = "google/siglip-base-patch16-224" +tokenizer = AutoTokenizer.from_pretrained(model_name) +image_processor = AutoImageProcessor.from_pretrained(model_name) +vision_model = SiglipVisionModel.from_pretrained(model_name) +text_model = SiglipTextModel.from_pretrained(model_name) + + +def tokenize(text: str, tokenizer: AutoTokenizer): + data = tokenizer( + text, + return_tensors="pt", + padding="max_length", + max_length=tokenizer.model_max_length, + truncation=True, + ) + return data["input_ids"].squeeze(0), data["attention_mask"].squeeze(0) + + +image_transform = spt.data.transforms.Compose( + spt.data.transforms.Resize((224, 224)), + spt.data.transforms.ToImage( + mean=image_processor.image_mean, + std=image_processor.image_std, + ), + spt.data.transforms.LambdaTransform( + fn=partial(tokenize, tokenizer=tokenizer), + source="prompt", + targets=("tokenized_prompt", "attention_mask"), + ), +) + + +train_base = spt.data.HFDataset( + "poloclub/diffusiondb", + "2m_all", + split="train", + transform=image_transform, + remove_columns=[ + "timestamp", + "user_name", + "prompt_nsfw", + "image_nsfw", + "sampler", + ], +) + +size = len(train_base) +val_n = int(size * val_percent) +val_dataset = spt.data.Subset(train_base, range(0, val_n)) +train_dataset = spt.data.Subset(train_base, range(val_n, size)) + + +train_dataloader = torch.utils.data.DataLoader( + dataset=train_dataset, + batch_size=batch_size, + num_workers=16, + shuffle=True, + drop_last=True, + pin_memory=True, + persistent_workers=True, + prefetch_factor=4, +) +val_dataloader = torch.utils.data.DataLoader( + dataset=val_dataset, + batch_size=batch_size, + num_workers=8, + shuffle=False, + pin_memory=True, + persistent_workers=True, + prefetch_factor=4, +) + +data = spt.data.DataModule(train=train_dataloader, val=val_dataloader) + + +class SigLIPMonitor(pl.Callback): + """Log retrieval and pairwise sigmoid statistics for SigLIP training.""" + + def __init__(self, log_every_n_steps: int = 10): + super().__init__() + self.every = log_every_n_steps + + @torch.no_grad() + def _log(self, trainer: pl.Trainer, pl_module, outputs: dict, stage: str): + img = F.normalize(outputs["image_embeds"], dim=-1) + txt = F.normalize(outputs["text_embeds"], dim=-1) + + loss_fn = pl_module.siglip_loss + logits = loss_fn.logit_scale.exp() * (img @ txt.T) + loss_fn.logit_bias + batch_size = logits.size(0) + diag = torch.arange(batch_size, device=logits.device) + + r1_i2t = (logits.argmax(dim=1) == diag).float().mean() + r1_t2i = (logits.argmax(dim=0) == diag).float().mean() + pos_prob = torch.sigmoid(logits[diag, diag]).mean() + cos_pos = F.cosine_similarity(img, txt, dim=-1).mean() + + metrics = { + f"{stage}/retrieval/R@1_i2t": float(r1_i2t.cpu()), + f"{stage}/retrieval/R@1_t2i": float(r1_t2i.cpu()), + f"{stage}/contrast/pos_prob": float(pos_prob.cpu()), + f"{stage}/align/cos_pos": float(cos_pos.cpu()), + f"{stage}/config/logit_scale": float(loss_fn.logit_scale.exp().cpu()), + f"{stage}/config/logit_bias": float(loss_fn.logit_bias.cpu()), + } + if batch_size > 1: + neg = logits.masked_fill( + torch.eye(batch_size, dtype=torch.bool, device=logits.device), + float("-inf"), + ) + top_neg = neg.max(dim=1).values + margin = logits[diag, diag] - top_neg + metrics[f"{stage}/contrast/margin"] = float(margin.mean().cpu()) + + trainer.logger.log_metrics(metrics, step=trainer.global_step) + + def on_train_batch_end(self, trainer, pl_module, outputs, batch, batch_idx): + if trainer.global_step % self.every == 0: + self._log(trainer, pl_module, outputs, "train") + + def on_validation_batch_end( + self, trainer, pl_module, outputs, batch, batch_idx, dataloader_idx=0 + ): + self._log(trainer, pl_module, outputs, "val") + + +module = spt.Module( + vision_model=vision_model, + text_model=text_model, + forward=forward.siglip_forward, + siglip_loss=spt.losses.SigLIPLoss(), + optim={ + "optimizer": { + "type": "AdamW", + "lr": lr, + "weight_decay": 1.0e-6, + "betas": (0.9, 0.98), + }, + "scheduler": { + "type": "LinearWarmupCosineAnnealing", + "total_steps": (len(train_dataloader) // num_devices) * num_epochs, + "peak_step": 0.1, + }, + "interval": "step", + }, +) + +wandb_logger = WandbLogger( + entity="stable-pretraining", + project="diffusiondb2m-siglip", + name="siglip-vit-b16-diffusiondb2m-32k", + log_model=False, +) + +trainer = pl.Trainer( + max_epochs=num_epochs, + num_sanity_val_steps=0, + callbacks=[ + ModelCheckpoint( + monitor="fit/loss_step", + mode="min", + every_n_epochs=1, + save_top_k=-1, + dirpath="/your/path/to/checkpoints", + ), + LearningRateMonitor(logging_interval="step"), + SigLIPMonitor(log_every_n_steps=10), + ], + precision="bf16-mixed", + logger=wandb_logger, + enable_checkpointing=True, + devices=num_devices, + accelerator="gpu", + strategy="ddp", +) + +manager = spt.Manager( + trainer=trainer, + module=module, + data=data, + ckpt_path=resume_ckpt_path, +) + +manager() diff --git a/docs/source/api/forward.rst b/docs/source/api/forward.rst index 14d6cfe1e..0f02e0484 100644 --- a/docs/source/api/forward.rst +++ b/docs/source/api/forward.rst @@ -247,6 +247,34 @@ Barlow Twins _target_: stable_pretraining.losses.BarlowTwinsLoss lambda_: 0.005 +SigLIP +~~~~~~ + +.. autofunction:: siglip_forward + +**Required Module Attributes:** + +- ``vision_model``: Image encoder returning ``image_embeds`` or ``pooler_output`` +- ``text_model``: Text encoder returning ``text_embeds`` or ``pooler_output`` +- ``siglip_loss``: SigLIP loss function + +**Expected Batch Keys:** + +- ``image``: Image tensor passed to the vision model +- ``tokenized_prompt``: Token ids passed to the text model +- ``attention_mask``: Optional text attention mask + +**Example Config:** + +.. code-block:: yaml + + module: + forward: stable_pretraining.forward.siglip_forward + vision_model: ... + text_model: ... + siglip_loss: + _target_: stable_pretraining.losses.SigLIPLoss + Supervised ~~~~~~~~~~ diff --git a/docs/source/api/losses.rst b/docs/source/api/losses.rst index 0b487a840..675f51444 100644 --- a/docs/source/api/losses.rst +++ b/docs/source/api/losses.rst @@ -12,3 +12,5 @@ stable_pretraining.losses NegativeCosineSimilarity VICRegLoss BarlowTwinsLoss + CLIPLoss + SigLIPLoss diff --git a/docs/source/references.bib b/docs/source/references.bib index 60a02f87a..343e58003 100644 --- a/docs/source/references.bib +++ b/docs/source/references.bib @@ -96,6 +96,13 @@ @inproceedings{radford2021learning organization={PmLR} } +@article{zhai2023sigmoid, + title={Sigmoid loss for language image pre-training}, + author={Zhai, Xiaohua and Mustafa, Basil and Kolesnikov, Alexander and Beyer, Lucas}, + journal={arXiv preprint arXiv:2303.15343}, + year={2023} +} + @article{caron2020unsupervised, title={Unsupervised learning of visual features by contrasting cluster assignments}, author={Caron, Mathilde and Misra, Ishan and Mairal, Julien and Goyal, Priya and Bojanowski, Piotr and Joulin, Armand}, diff --git a/stable_pretraining/forward.py b/stable_pretraining/forward.py index 5716efa5b..59be3ef6e 100644 --- a/stable_pretraining/forward.py +++ b/stable_pretraining/forward.py @@ -1074,3 +1074,66 @@ def dinov2_forward(self, batch, stage): ) return out + + +def _get_embedding_output(outputs, primary_key: str, fallback_key: str): + embedding = getattr(outputs, primary_key, None) + if embedding is None: + embedding = getattr(outputs, fallback_key, None) + if embedding is None: + raise ValueError( + f"Expected model output to expose '{primary_key}' or '{fallback_key}'." + ) + return embedding + + +def siglip_forward(self, batch, stage): + """Forward function for SigLIP image-text pretraining. + + SigLIP learns aligned image-text representations with an independent + sigmoid loss over every image-text pair in the batch. Matching pairs are + expected to be aligned along the batch diagonal. + + Args: + self: Module instance with ``vision_model``, ``text_model``, and + ``siglip_loss`` attributes. ``vision_model`` must return + ``image_embeds`` or ``pooler_output``; ``text_model`` must return + ``text_embeds`` or ``pooler_output``. + batch: Paired image-text batch dictionary. Must contain ``image`` and + ``tokenized_prompt``. May contain ``attention_mask``. + stage: Training stage ('train', 'val', or 'test') + + Returns: + Dictionary containing normalized ``image_embeds`` and ``text_embeds``. + During training, also contains ``loss``. + """ + out = {} + + vision_outputs = self.vision_model(pixel_values=batch["image"]) + image_embeds = _get_embedding_output( + vision_outputs, "image_embeds", "pooler_output" + ) + image_embeds = torch.nn.functional.normalize(image_embeds, dim=-1) + + text_outputs = self.text_model( + input_ids=batch["tokenized_prompt"], + attention_mask=batch.get("attention_mask"), + ) + text_embeds = _get_embedding_output(text_outputs, "text_embeds", "pooler_output") + text_embeds = torch.nn.functional.normalize(text_embeds, dim=-1) + + out["image_embeds"] = image_embeds + out["text_embeds"] = text_embeds + + if self.training: + out["loss"] = self.siglip_loss(image_embeds, text_embeds) + self.log( + f"{stage}/loss", + out["loss"], + on_step=True, + on_epoch=True, + sync_dist=True, + prog_bar=True, + ) + + return out diff --git a/stable_pretraining/losses/__init__.py b/stable_pretraining/losses/__init__.py index effd83660..513f86610 100644 --- a/stable_pretraining/losses/__init__.py +++ b/stable_pretraining/losses/__init__.py @@ -19,7 +19,7 @@ ) # Multimodal losses -from .multimodal import CLIPLoss +from .multimodal import CLIPLoss, SigLIPLoss # Reconstruction losses from .reconstruction import mae @@ -43,6 +43,7 @@ "BarlowTwinsLoss", "NTXEntLoss", "CLIPLoss", + "SigLIPLoss", # Reconstruction "mae", # Utils diff --git a/stable_pretraining/losses/multimodal.py b/stable_pretraining/losses/multimodal.py index a1caae8d4..82cf7a1e5 100644 --- a/stable_pretraining/losses/multimodal.py +++ b/stable_pretraining/losses/multimodal.py @@ -4,7 +4,11 @@ particularly for image-text contrastive learning like CLIP. """ +import math + import torch +import torch.distributed as dist +import torch.distributed.nn.functional as dist_nn from typing import Optional from .joint_embedding import InfoNCELoss @@ -49,3 +53,73 @@ def forward( ) return 0.5 * (loss_i + loss_j) + + +class SigLIPLoss(torch.nn.Module): + """Sigmoid Loss for Language Image Pre-Training. + + Computes the pairwise SigLIP objective from + :cite:`zhai2023sigmoid`. Positive image-text pairs are expected on + the batch diagonal; all off-diagonal pairs are treated as negatives. + + Args: + init_logit_scale: Initial value for the learnable log-space scale + parameter. The effective multiplier is ``exp(logit_scale)``, so + it remains positive during training. Defaults to ``log(10)``. + init_logit_bias: Initial value for the learnable additive bias applied + to every image-text logit. Defaults to ``-10``. + normalize: Whether to L2-normalize image and text features before + computing logits. Defaults to ``True``. + gather_distributed: Whether to gather image and text features across + distributed workers before computing the loss. Defaults to ``True``. + """ + + def __init__( + self, + init_logit_scale: float = math.log(10.0), + init_logit_bias: float = -10.0, + normalize: bool = True, + gather_distributed: bool = True, + ): + super().__init__() + self.logit_scale = torch.nn.Parameter(torch.tensor(init_logit_scale)) + self.logit_bias = torch.nn.Parameter(torch.tensor(init_logit_bias)) + self.normalize = normalize + self.gather_distributed = gather_distributed + + def _gather_features( + self, + image_features: torch.Tensor, + text_features: torch.Tensor, + ) -> tuple[torch.Tensor, torch.Tensor]: + if self.gather_distributed and dist.is_available() and dist.is_initialized(): + image_features = torch.cat(dist_nn.all_gather(image_features), dim=0) + text_features = torch.cat(dist_nn.all_gather(text_features), dim=0) + return image_features, text_features + + def forward( + self, + image_features: torch.Tensor, + text_features: torch.Tensor, + ) -> torch.Tensor: + if image_features.shape[0] != text_features.shape[0]: + raise ValueError( + "SigLIPLoss expects paired image/text batches with the same batch size." + ) + + if self.normalize: + image_features = torch.nn.functional.normalize(image_features, dim=-1) + text_features = torch.nn.functional.normalize(text_features, dim=-1) + + image_features, text_features = self._gather_features( + image_features, text_features + ) + + scale = self.logit_scale.exp() + logits = scale * (image_features @ text_features.T) + self.logit_bias + + labels = -torch.ones_like(logits) + labels.fill_diagonal_(1) + + loss = -torch.nn.functional.logsigmoid(labels * logits).sum(dim=-1).mean() + return loss diff --git a/stable_pretraining/tests/unit/test_forward_functions.py b/stable_pretraining/tests/unit/test_forward_functions.py index f735a6af0..ddc0c6227 100644 --- a/stable_pretraining/tests/unit/test_forward_functions.py +++ b/stable_pretraining/tests/unit/test_forward_functions.py @@ -1,5 +1,7 @@ """Unit tests for all forward functions using benchmark transforms.""" +import types + import pytest import torch import torch.nn as nn @@ -439,3 +441,65 @@ def test_supervised_forward(self): assert "embedding" in result assert "logits" in result assert "loss" in result + + def test_siglip_forward(self): + """Test siglip_forward with paired image-text batch.""" + torch.manual_seed(0) + batch = { + "image": torch.randn(2, 3, 224, 224), + "tokenized_prompt": torch.ones(2, 8, dtype=torch.long), + "attention_mask": torch.ones(2, 8, dtype=torch.long), + } + image_embeds = torch.randn(2, 16) + text_embeds = torch.randn(2, 16) + + module = Mock() + module.vision_model = Mock( + return_value=types.SimpleNamespace(image_embeds=image_embeds) + ) + module.text_model = Mock( + return_value=types.SimpleNamespace(text_embeds=text_embeds) + ) + module.siglip_loss = Mock(return_value=torch.tensor(0.5)) + module.training = True + module.log = Mock() + + result = forward_module.siglip_forward(module, batch, "train") + + module.vision_model.assert_called_once_with(pixel_values=batch["image"]) + module.text_model.assert_called_once_with( + input_ids=batch["tokenized_prompt"], + attention_mask=batch["attention_mask"], + ) + module.siglip_loss.assert_called_once() + assert "image_embeds" in result + assert "text_embeds" in result + assert "loss" in result + assert result["image_embeds"].shape == image_embeds.shape + assert result["text_embeds"].shape == text_embeds.shape + + def test_siglip_forward_accepts_pooler_outputs(self): + """Test siglip_forward with Hugging Face SigLIP-style pooler outputs.""" + batch = { + "image": torch.randn(2, 3, 224, 224), + "tokenized_prompt": torch.ones(2, 8, dtype=torch.long), + } + image_embeds = torch.randn(2, 16) + text_embeds = torch.randn(2, 16) + + module = Mock() + module.vision_model = Mock( + return_value=types.SimpleNamespace(pooler_output=image_embeds) + ) + module.text_model = Mock( + return_value=types.SimpleNamespace(pooler_output=text_embeds) + ) + module.siglip_loss = Mock(return_value=torch.tensor(0.5)) + module.training = True + module.log = Mock() + + result = forward_module.siglip_forward(module, batch, "train") + + assert "image_embeds" in result + assert "text_embeds" in result + assert "loss" in result diff --git a/stable_pretraining/tests/unit/test_siglip_loss.py b/stable_pretraining/tests/unit/test_siglip_loss.py new file mode 100644 index 000000000..cd2a2f83e --- /dev/null +++ b/stable_pretraining/tests/unit/test_siglip_loss.py @@ -0,0 +1,91 @@ +import math + +import pytest +import torch +import torch.nn.functional as F + +from stable_pretraining.losses import SigLIPLoss + + +@pytest.mark.unit +class TestSigLIPLoss: + """Unit tests for the SigLIPLoss function.""" + + def test_loss_is_lower_for_matched_pairs_than_mismatched_pairs(self): + """Loss should be lower when positive pairs are aligned on the diagonal.""" + torch.manual_seed(0) + batch_size, dim = 4, 8 + feats_i = F.normalize(torch.eye(dim)[:batch_size], dim=-1) + feats_j = feats_i.clone() + loss_fn = SigLIPLoss() + + matched_loss = loss_fn(image_features=feats_i, text_features=feats_j) + mismatched_loss = loss_fn( + image_features=feats_i, + text_features=torch.flip(feats_j, dims=[0]), + ) + + assert matched_loss.ndim == 0 + assert matched_loss < mismatched_loss + + def test_matches_reference_softplus_formula(self): + """Loss should match the positive/negative softplus decomposition.""" + feats_i = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + feats_j = torch.tensor([[1.0, 0.0], [0.0, 1.0]]) + loss_fn = SigLIPLoss( + init_logit_scale=math.log(2.0), + init_logit_bias=-0.5, + normalize=False, + ) + + loss = loss_fn(image_features=feats_i, text_features=feats_j) + + logits = 2.0 * (feats_i @ feats_j.T) - 0.5 + positive_logits = logits.diagonal() + negative_mask = ~torch.eye(logits.size(0), dtype=torch.bool) + negative_logits = logits[negative_mask] + expected = ( + F.softplus(-positive_logits).sum() + F.softplus(negative_logits).sum() + ) / feats_i.size(0) + assert torch.allclose(loss, expected, atol=1e-7, rtol=0) + + def test_invariance_to_feature_magnitude(self): + """Loss should be identical regardless of input vector magnitude.""" + torch.manual_seed(123) + batch_size, dim = 8, 256 + feats_i = torch.randn(batch_size, dim) + feats_j = torch.randn(batch_size, dim) + loss_fn = SigLIPLoss() + + loss1 = loss_fn(image_features=feats_i, text_features=feats_j) + loss2 = loss_fn( + image_features=feats_i * 100.0, + text_features=feats_j * 0.01, + ) + + assert torch.allclose(loss1, loss2, atol=1e-7, rtol=1e-6) + + def test_logit_scale_and_bias_receive_gradients(self): + """Learnable logit scale and bias should receive gradients.""" + torch.manual_seed(42) + batch_size, dim = 4, 128 + feats_i = torch.randn(batch_size, dim, requires_grad=True) + feats_j = torch.randn(batch_size, dim, requires_grad=True) + loss_fn = SigLIPLoss() + + loss = loss_fn(image_features=feats_i, text_features=feats_j) + loss.backward() + + assert loss_fn.logit_scale.grad is not None + assert loss_fn.logit_bias.grad is not None + assert loss_fn.logit_scale.grad.abs().item() > 0 + assert loss_fn.logit_bias.grad.abs().item() > 0 + + def test_raises_for_unpaired_batches(self): + """Loss should reject unequal image/text batch sizes.""" + loss_fn = SigLIPLoss() + feats_i = torch.randn(2, 8) + feats_j = torch.randn(3, 8) + + with pytest.raises(ValueError, match="same batch size"): + loss_fn(image_features=feats_i, text_features=feats_j)