add sigmoid contrastive loss

Author: lucidrains

README.md

@@ -220,6 +220,14 @@ music = musiclm('the crystalline sounds of the piano in a ballroom', num_samples
 }
 ```
 
+```bibtex
+@inproceedings{Zhai2023SigmoidLF,
+    title   = {Sigmoid Loss for Language Image Pre-Training},
+    author  = {Xiaohua Zhai and Basil Mustafa and Alexander Kolesnikov and Lucas Beyer},
+    year    = {2023}
+}
+```
+
 *The only truth is music.* - Jack Kerouac
 
 *Music is the universal language of mankind.* - Henry Wadsworth Longfellow

musiclm_pytorch/musiclm_pytorch.py

@@ -1,3 +1,4 @@
+import math
 from functools import wraps
 
 import torch
@@ -248,6 +249,76 @@ def forward(
 
         return x, torch.stack(layers[:-1])
 
+# contrastive losses
+
+class SoftmaxContrastiveLearning(nn.Module):
+    def __init__(
+        self,
+        *,
+        layers = 1,
+        decoupled_contrastive_learning = False,
+        init_temp = 10
+    ):
+        super().__init__()
+        self.temperatures = nn.Parameter(torch.ones(layers, 1, 1) * math.log(init_temp))
+        self.decoupled_contrastive_learning = decoupled_contrastive_learning
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, sims):
+        batch = sims.shape[-1]
+
+        if sims.ndim == 2:
+            sims = rearrange(sims, 'i j -> 1 i j')
+
+        sims = sims * self.temperatures.exp()
+
+        cosine_sims_exp = sims.exp()
+
+        numerator = matrix_diag(cosine_sims_exp)
+
+        if self.decoupled_contrastive_learning:
+            eye = torch.eye(batch, device = self.device, dtype = torch.bool)
+            cosine_sims_exp = cosine_sims_exp.masked_fill(eye, 0.)
+
+        denominator_i = reduce(cosine_sims_exp, 'l i j -> l i', 'sum')
+        denominator_j = reduce(cosine_sims_exp, 'l i j -> l j', 'sum')
+
+        contrastive_loss = -log(numerator) + 0.5 * (log(denominator_i) + log(denominator_j))
+
+        contrastive_loss = reduce(contrastive_loss, 'l n -> l', 'mean')
+        return contrastive_loss.sum()
+
+class SigmoidContrastiveLearning(nn.Module):
+    """ https://arxiv.org/abs/2303.15343 """
+
+    def __init__(
+        self,
+        *,
+        layers = 1,
+        init_temp = 10,
+        init_bias = -10
+    ):
+        super().__init__()
+        self.temperatures = nn.Parameter(torch.ones(layers, 1, 1) * math.log(init_temp))
+        self.bias = nn.Parameter(torch.ones(layers, 1, 1) * init_bias)
+
+    @property
+    def device(self):
+        return next(self.parameters()).device
+
+    def forward(self, sims):
+        if sims.ndim == 2:
+            sims = rearrange(sims, 'i j -> 1 i j')
+
+        n = sims.shape[-1]
+        sims = sims * self.temperatures.exp() + self.bias
+        labels = 2 * rearrange(torch.eye(n), 'i j -> 1 i j') - torch.ones_like(sims)
+
+        return -F.logsigmoid(labels * sims).sum() / n
+
 # Audio Spectrogram Transformer - https://arxiv.org/abs/2104.01778
 
 def pair(t):
@@ -539,7 +610,8 @@ def __init__(
         text_dim,
         dim_latent,
         layers,
-        decoupled_contrastive_learning = False
+        decoupled_contrastive_learning = False,
+        sigmoid_contrastive_loss = False
     ):
         super().__init__()
         self.layers = layers
@@ -554,9 +626,8 @@ def __init__(
         self.text_latent_weight = nn.Parameter(torch.randn(layers, text_dim, dim_latent))
         self.text_latent_bias = nn.Parameter(torch.randn(layers, 1, dim_latent))
 
-        self.temperatures = nn.Parameter(torch.ones(layers, 1, 1))
-
-        self.decoupled_contrastive_learning = decoupled_contrastive_learning
+        klass = SigmoidContrastiveLearning if sigmoid_contrastive_loss else partial(SoftmaxContrastiveLearning, decoupled_contrastive_learning = decoupled_contrastive_learning)
+        self.contrast = klass(layers = layers)
 
     def forward(self, *, audio_layers, text_layers):
         device, batch = audio_layers.device, audio_layers.shape[1]
@@ -571,23 +642,9 @@ def forward(self, *, audio_layers, text_layers):
         text_latents = einsum('l b d, l d e -> l b e', text_embeds, self.text_latent_weight) + self.text_latent_bias
         text_latents = l2norm(text_latents)
 
-        cosine_sims = einsum('l i d, l j d -> l i j', audio_latents, text_latents) * self.temperatures.exp()
-
-        cosine_sims_exp = cosine_sims.exp()
+        cosine_sims = einsum('l i d, l j d -> l i j', audio_latents, text_latents)
 
-        numerator = matrix_diag(cosine_sims_exp)
-
-        if self.decoupled_contrastive_learning:
-            eye = torch.eye(batch, device = device, dtype = torch.bool)
-            cosine_sims_exp = cosine_sims_exp.masked_fill(eye, 0.)
-
-        denominator_i = reduce(cosine_sims_exp, 'l i j -> l i', 'sum')
-        denominator_j = reduce(cosine_sims_exp, 'l i j -> l j', 'sum')
-
-        contrastive_loss = -log(numerator) + 0.5 * (log(denominator_i) + log(denominator_j))
-
-        contrastive_loss = reduce(contrastive_loss, 'l n -> l', 'mean')
-        return contrastive_loss.sum()
+        return self.contrast(cosine_sims)
 
 # main classes
 
@@ -600,20 +657,21 @@ def __init__(
         dim_latent = 128,                       # they use 128
         decoupled_contrastive_learning = True,  # think this was used, make it optional
         hierarchical_contrastive_loss = False,
-        hierarchical_contrastive_loss_layers = None
+        hierarchical_contrastive_loss_layers = None,
+        sigmoid_contrastive_loss = False
     ):
         super().__init__()
         self.dim_latent = dim_latent
 
         self.audio = audio_transformer
         self.text = text_transformer
 
-        self.temperature = nn.Parameter(torch.tensor(1.))
 
         self.text_to_latents = nn.Linear(self.text.dim, dim_latent)
         self.audio_to_latents = nn.Linear(self.audio.dim, dim_latent)
 
-        self.decoupled_contrastive_learning = decoupled_contrastive_learning
+        klass = SigmoidContrastiveLearning if sigmoid_contrastive_loss else partial(SoftmaxContrastiveLearning, decoupled_contrastive_learning = decoupled_contrastive_learning)
+        self.contrast = klass()
 
         self.multi_layer_contrastive_learning = None
 
@@ -629,7 +687,8 @@ def __init__(
                 text_dim = self.text.dim,
                 dim_latent = dim_latent,
                 layers = num_layers,
-                decoupled_contrastive_learning = decoupled_contrastive_learning
+                decoupled_contrastive_learning = decoupled_contrastive_learning,
+                sigmoid_contrastive_loss = sigmoid_contrastive_loss
             )
 
     def get_audio_latents(
@@ -688,21 +747,7 @@ def forward(
         if return_pairwise_similarities:
             return cosine_sim
 
-        cosine_sim = cosine_sim * self.temperature.exp()
-
-        cosine_sim_exp = cosine_sim.exp()
-
-        numerator = cosine_sim_exp.diag()
-
-        if self.decoupled_contrastive_learning:
-            eye = torch.eye(batch, device = device, dtype = torch.bool)
-            cosine_sim_exp = cosine_sim_exp.masked_fill(eye, 0.)
-
-        denominator_i = reduce(cosine_sim_exp, 'i j -> i', 'sum')
-        denominator_j = reduce(cosine_sim_exp, 'i j -> j', 'sum')
-
-        contrastive_loss = -log(numerator) + 0.5 * (log(denominator_i) + log(denominator_j))
-        cl_loss = contrastive_loss.mean()
+        cl_loss = self.contrast(cosine_sim)
 
         if not exists(self.multi_layer_contrastive_learning):
             return cl_loss

setup.py

@@ -3,7 +3,7 @@
 setup(
   name = 'musiclm-pytorch',
   packages = find_packages(exclude=[]),
-  version = '0.1.2',
+  version = '0.2.0',
   license='MIT',
   description = 'MusicLM - AudioLM + Audio CLIP to text to music synthesis',
   author = 'Phil Wang',