nt-xent-loss/NT-Xent Loss.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "db411784",
   "metadata": {},
   "source": [
    "# Normalized Temperature-scaled Cross Entropy Loss (NT-Xent)\n",
    "\n",
    "Links:\n",
    "\n",
    "* https://paperswithcode.com/method/nt-xent\n",
    "* https://github.com/KevinMusgrave/pytorch-metric-learning/issues/6\n",
    "\n",
    "### The setup\n",
    "\n",
    "This loss function is applied to N feature vectors (one feature vector computed for each input). N here being the batch size.\n",
    "\n",
    "The $1^{st}$ and $2^{nd}$ feature vectors in the set are considered to be similar to each other (as used in SimCLR for self-supervised learning). The rest are considered to be dissimilar to the $1^{st}$ one. Similarly, the $3^{rd}$ and $4^{th}$ feature vectors are considered to be similar to each other and the rest are considered to be dissimilar to the $3^{rd}$ one.\n",
    "\n",
    "### Simple interpretation\n",
    "\n",
    "NT-Xent is just a very fancy way to say the following:\n",
    "\n",
    "1. The all-pairs Cosine Similarity score is computed for each of the $2N$ vectors produced by the SimCLR model. This results in $(2N)^2$ similarity scores represented as a $2N x 2N$ matrix\n",
    "2. Comparison results between the same value are discarded (since a feature vector is perfectly similar to itself and can't possibly allow the model to learn anything useful)\n",
    "3. Each value (cosine similarity) is scaled by a temperature parameter $\\tau$ (which is a hyper-parameter)\n",
    "4. Cross Entropy Loss is applied to each row of the resulting matrix above\n",
    "5. Typically, the mean of these losses (one loss per element in a batch) is used for backpropagation\n",
    "\n",
    "### What this notebook will show\n",
    "\n",
    "The first part will show an intuitive and efficient implementation of the NT-Xent loss in PyTorch for single-label (only 1 pair can be positive) contrastive learning tasks.\n",
    "\n",
    "The second part will show an intuitive and efficient implementation of the NT-BXent (binary cross entropy) loss (defined below) in PyTorch for multi-label (multiple pairs can be positive) contrastive learning tasks."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 111,
   "id": "5777a3b7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import torch\n",
    "from torch import nn\n",
    "from torch.nn import functional as F\n",
    "_ = torch.manual_seed(21)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "66b2fa12",
   "metadata": {},
   "source": [
    "In this example, we assume that the input (and output) batch size is 8."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 112,
   "id": "1b24002e",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[1., 0., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 1., 0., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 1., 0., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 1., 0., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 1., 0., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 1., 0., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 1., 0.],\n",
       "         [0., 0., 0., 0., 0., 0., 0., 1.]]),\n",
       " tensor([[False,  True,  True,  True,  True,  True,  True,  True],\n",
       "         [ True, False,  True,  True,  True,  True,  True,  True],\n",
       "         [ True,  True, False,  True,  True,  True,  True,  True],\n",
       "         [ True,  True,  True, False,  True,  True,  True,  True],\n",
       "         [ True,  True,  True,  True, False,  True,  True,  True],\n",
       "         [ True,  True,  True,  True,  True, False,  True,  True],\n",
       "         [ True,  True,  True,  True,  True,  True, False,  True],\n",
       "         [ True,  True,  True,  True,  True,  True,  True, False]]))"
      ]
     },
     "execution_count": 112,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "eye = torch.eye(8)\n",
    "eye, ~eye.bool()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "57df1f1d",
   "metadata": {},
   "source": [
    "Let's assume that the model's output is an `(8, 2)` tensor, with each element being a 2d vector in space."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 113,
   "id": "25254709",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[-0.2386, -1.0934],\n",
       "        [ 0.1558,  0.1750],\n",
       "        [-0.9526, -0.5442],\n",
       "        [ 1.1985,  0.9604],\n",
       "        [-1.1074, -0.8403],\n",
       "        [-0.0020,  0.2240],\n",
       "        [ 0.8766, -0.5379],\n",
       "        [-0.2994,  0.9785]])"
      ]
     },
     "execution_count": 113,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "x = torch.randn(8, 2)\n",
    "x"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a471abf",
   "metadata": {},
   "source": [
    "Let's compute the cosine similarity between every pair of vectors. Since we have 8 vectors, we expect `8x8=64` similarity scores."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 114,
   "id": "71918128",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[ 1.0000, -0.8714,  0.6698, -0.7773,  0.7604, -0.9751,  0.3293, -0.8719],\n",
       "        [-0.8714,  1.0000, -0.9479,  0.9860, -0.9812,  0.7408,  0.1763,  0.5195],\n",
       "        [ 0.6698, -0.9479,  1.0000, -0.9878,  0.9916, -0.4883, -0.4806, -0.2203],\n",
       "        [-0.7773,  0.9860, -0.9878,  1.0000, -0.9997,  0.6183,  0.3381,  0.3696],\n",
       "        [ 0.7604, -0.9812,  0.9916, -0.9997,  1.0000, -0.5973, -0.3628, -0.3449],\n",
       "        [-0.9751,  0.7408, -0.4883,  0.6183, -0.5973,  1.0000, -0.5306,  0.9588],\n",
       "        [ 0.3293,  0.1763, -0.4806,  0.3381, -0.3628, -0.5306,  1.0000, -0.7495],\n",
       "        [-0.8719,  0.5195, -0.2203,  0.3696, -0.3449,  0.9588, -0.7495,  1.0000]])"
      ]
     },
     "execution_count": 114,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "xcs = F.cosine_similarity(x[None,:,:], x[:,None,:], dim=-1)\n",
    "xcs"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b84af806",
   "metadata": {},
   "source": [
    "Let's replace all the diagonal elements (cosine similarity of an element with itself) with `-inf` so that when we compute the softmax later, it will show up as `0.0`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 115,
   "id": "f05f5dac",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[   -inf, -0.8714,  0.6698, -0.7773,  0.7604, -0.9751,  0.3293, -0.8719],\n",
       "        [-0.8714,    -inf, -0.9479,  0.9860, -0.9812,  0.7408,  0.1763,  0.5195],\n",
       "        [ 0.6698, -0.9479,    -inf, -0.9878,  0.9916, -0.4883, -0.4806, -0.2203],\n",
       "        [-0.7773,  0.9860, -0.9878,    -inf, -0.9997,  0.6183,  0.3381,  0.3696],\n",
       "        [ 0.7604, -0.9812,  0.9916, -0.9997,    -inf, -0.5973, -0.3628, -0.3449],\n",
       "        [-0.9751,  0.7408, -0.4883,  0.6183, -0.5973,    -inf, -0.5306,  0.9588],\n",
       "        [ 0.3293,  0.1763, -0.4806,  0.3381, -0.3628, -0.5306,    -inf, -0.7495],\n",
       "        [-0.8719,  0.5195, -0.2203,  0.3696, -0.3449,  0.9588, -0.7495,    -inf]])"
      ]
     },
     "execution_count": 115,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = xcs.clone()\n",
    "y[eye.bool()] = float(\"-inf\")\n",
    "y"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7fc604c2",
   "metadata": {},
   "source": [
    "Before we jump into using Cross Entropy loss to compute the loss, let's visualize the ground truth labels we wish to assign to every element in the cosine similarity matrix above.\n",
    "\n",
    "The matrix `target` below will contain the ground-truth labels we will feed in to the `cross_entropy` function to compute the cross entropy loss. However, when visualizing it, we wish to view it as a 2d matrix, and not a 1d matrix of labels per element.\n",
    "\n",
    "When interpreting the matrix `y` above, we need to keep in mind that the first row represents the Cosine Similarity between the element at index 0 and at every index between 0 and 7. The second row represents the Cosine Similarity between the element at index 1 and at every index between 0 and 7, and so on.\n",
    "\n",
    "The ground truth labels can be visualized as shown below:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 116,
   "id": "92998c74",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([1, 0, 3, 2, 5, 4, 7, 6])"
      ]
     },
     "execution_count": 116,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "target = torch.arange(8)\n",
    "target[0::2] += 1\n",
    "target[1::2] -= 1\n",
    "target"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2a35ce3c",
   "metadata": {},
   "source": [
    "The tensor `index` is going to be a column tensor (1 row, 8 columns) with the same elements as `target` that will be used to scatter the ones into the zeros 2d tensor `ground_truth_labels`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 117,
   "id": "a663273a",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[1],\n",
       "        [0],\n",
       "        [3],\n",
       "        [2],\n",
       "        [5],\n",
       "        [4],\n",
       "        [7],\n",
       "        [6]])"
      ]
     },
     "execution_count": 117,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "index = target.reshape(8, 1).long()\n",
    "index"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 118,
   "id": "7627c1ad",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor([[0, 1, 0, 0, 0, 0, 0, 0],\n",
       "        [1, 0, 0, 0, 0, 0, 0, 0],\n",
       "        [0, 0, 0, 1, 0, 0, 0, 0],\n",
       "        [0, 0, 1, 0, 0, 0, 0, 0],\n",
       "        [0, 0, 0, 0, 0, 1, 0, 0],\n",
       "        [0, 0, 0, 0, 1, 0, 0, 0],\n",
       "        [0, 0, 0, 0, 0, 0, 0, 1],\n",
       "        [0, 0, 0, 0, 0, 0, 1, 0]])"
      ]
     },
     "execution_count": 118,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "ground_truth_labels = torch.zeros(8, 8).long()\n",
    "src = torch.ones(8, 8).long()\n",
    "ground_truth_labels = torch.scatter(ground_truth_labels, 1, index, src)\n",
    "ground_truth_labels"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b3ca47b4",
   "metadata": {},
   "source": [
    "What follows is a more detailed explanation of why we create `target` the way we do above.\n",
    "\n",
    "If we wish to compute the row (or column) wise cross-entropy-loss, we observe that element at index 1 in the first row (index 0) is expected to have label 1 (for similar) and the rest of the values are dis-similar from the 0th indexed result. This is because the elements at index 0 and 1 are expected to be augmentations of the same image when using SimCLR.\n",
    "\n",
    "For the 2nd row (index 1), the element at index 0 should have the label 1 (mirror-image case of the one above), and the rest are dissimilar to the element at index 1.\n",
    "\n",
    "Similarly, we observe that the elements 3, 2, 5, 4, 7, 6 will have a label `1`.\n",
    "\n",
    "The tensor `y` contains the ground-truth labels for feeding into the cross entropy loss function.\n",
    "\n",
    "The way that the cross entropy loss is used here is slightly different from how it is used in classification tasks. In classification tasks, the network's classification head is trained to learn one of N classes for a given input image. In the SimCLR algorithm, we aren't dealing with classes. We compute the all-pairs cosine similarity of each image with all the other images in the batch, and then compare if the resulting distribution of probabilities (of how probable it is for an image in the batch to be similar to the i'th image) is compare with the ground-truth distribution using cross entropy loss. We could have use KL divergence too. However, since Cross Entropy is the special case when the target probabilities are either 0 or 1, we just use cross entropy loss.\n",
    "\n",
    "Here's what the documentation says about the [cross_entropy](https://pytorch.org/docs/stable/generated/torch.nn.functional.cosine_similarity.html) function:\n",
    "\n",
    "> This criterion computes the cross entropy loss between input logits and target."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 119,
   "id": "17a5ed8d",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "tensor(2.8555)"
      ]
     },
     "execution_count": 119,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "F.cross_entropy(y, target, reduction=\"mean\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5cf90eb8",
   "metadata": {},
   "source": [
    "## Final implementation of the NT-Xent Loss\n",
    "\n",
    "Now that we have seen how to build up the NT-Xent loss, we can condense the complete implementation (including the computation of the all-pairs Cosine Similarity) into this one simple function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 120,
   "id": "092837c1",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Temperature:  0.01, Loss: 167.33396911621094\n",
      "Temperature:  0.10, Loss: 16.916988372802734\n",
      "Temperature:  1.00, Loss: 2.8555006980895996\n",
      "Temperature: 10.00, Loss: 2.0152008533477783\n",
      "Temperature: 20.00, Loss: 1.979940414428711\n"
     ]
    }
   ],
   "source": [
    "def nt_xent_loss(x, temperature):\n",
    "    assert len(x.size()) == 2\n",
    "    \n",
    "    # Cosine similarity\n",
    "    xcs = F.cosine_similarity(x[None,:,:], x[:,None,:], dim=-1)\n",
    "    xcs[torch.eye(x.size(0)).bool()] = float(\"-inf\")\n",
    "\n",
    "    # Ground truth labels\n",
    "    target = torch.arange(8)\n",
    "    target[0::2] += 1\n",
    "    target[1::2] -= 1\n",
    "    \n",
    "    # Standard cross entropy loss\n",
    "    return F.cross_entropy(xcs / temperature, target, reduction=\"mean\")\n",
    "\n",
    "for t in (0.01, 0.1, 1.0, 10.0, 20.0):\n",
    "    print(f\"Temperature: {t:5.2f}, Loss: {nt_xent_loss(x, temperature=t)}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "afeef020",
   "metadata": {},
   "source": [
    "# Multi-label loss for contrastive learning (NT-BXent)\n",
    "\n",
    "In the previous single-label setting, there's exactly 1 positive example in each row, with the rest being negative examples. In a multi-label setting, it's possible for more than one positive example in a given row, with the rest being negative examples. This can occur when we set up SimCLR so that each image is augmented N times (N > 2), and hence, it results in multiple (> 1) pairs of positive examples.\n",
    "\n",
    "In the single-label case, we could use cross-entropy loss. However, in the multi-label setting, we can't use cross entropy loss since cross entropy is effectively KL divergence, which measures the difference in 2 distributions. We can't have a distribution with the total probability summing up to a value > 1.0 and we don't want the positive examples to have a probability significantly less than one, which would have to happen if want all the probabilities in a row to sum up to 1.0. Hence, we turn our attention to Binary Cross Entropy Loss, and treat each sample as independently being 0.0 (negative) or 1.0 (positive). This leads to a slightly different formulation in code.\n",
    "\n",
    "For each pair $(i,j)$, we compute $s_{ij}$ as the cosine similarity between that pair of elements in the all-pairs cosine similarity matrix. The value $l_{ij}$ computes the loss for a specific element.\n",
    "\n",
    "$$ l_{ij} = - y_{ij}.\\log{\\sigma{ ( {s_{ij}} / \\tau ) }} - (1 - y_{ij}).\\log{\\sigma((1 - s_{ij}) / \\tau )}$$\n",
    "\n",
    "We can then compute the mean or weighted mean across all the element-wise losses. The unweighted mean is.\n",
    "\n",
    "$$ L_i^{unweighted} = \\frac{1}{N} \\Sigma_{j=1}^N {l_{ij}} $$\n",
    "\n",
    "However, there's a class imbalance problem with the formulation above. Suppose that in a set of N=100 samples, there are only 2 positive pairs and 98 negative pairs, then the loss from the negative pairs will overpower the loss from the positive pairs. Hence, we need to weight the losses from the positive and negative pairs accordingly. In the formula below, $1_{ij}^{pos}$ is $1$ if $ij$ is a positive pair, and $0$ otherwise. Similarly, $1_{ij}^{neg}$ is $1$ if $ij$ is a negative pair, and $0$ otherwise. $N_{pos}$ and $N_{neg}$ are the number of positive and negative pairs respectively.\n",
    "\n",
    "$$ L_i^{weighted} = \\frac{1}{N_{pos}} \\Sigma_{j=1}^N {1_{ij}^{pos}l_{ij}} + \\frac{1}{N_{neg}} \\Sigma_{j=1}^N {1_{ij}^{neg}l_{ij}}$$\n",
    "\n",
    "We can call this formulation the **NT-BXent** for **Normalized Temperature-scaled Binary Cross Entropy Loss**."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 121,
   "id": "094c11b9",
   "metadata": {
    "scrolled": false
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Positive indexes list\n",
      "tensor([[0, 0],\n",
      "        [0, 2],\n",
      "        [0, 4],\n",
      "        [1, 4],\n",
      "        [1, 6],\n",
      "        [1, 1],\n",
      "        [2, 3],\n",
      "        [3, 7],\n",
      "        [4, 3],\n",
      "        [7, 6],\n",
      "        [0, 0],\n",
      "        [1, 1],\n",
      "        [2, 2],\n",
      "        [3, 3],\n",
      "        [4, 4],\n",
      "        [5, 5],\n",
      "        [6, 6],\n",
      "        [7, 7]])\n",
      "\n",
      "Ground Truth labels for positive and negative pairs for BCE Loss\n",
      "tensor([[1., 0., 1., 0., 1., 0., 0., 0.],\n",
      "        [0., 1., 0., 0., 1., 0., 1., 0.],\n",
      "        [0., 0., 1., 1., 0., 0., 0., 0.],\n",
      "        [0., 0., 0., 1., 0., 0., 0., 1.],\n",
      "        [0., 0., 0., 1., 1., 0., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 1., 0., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 1., 0.],\n",
      "        [0., 0., 0., 0., 0., 0., 1., 1.]])\n"
     ]
    }
   ],
   "source": [
    "# The input to loss function will be 3 values:\n",
    "#\n",
    "# [1] target: Tensor with shape (Batch,NumFeatures). The feature vector for\n",
    "# each input is computed and returned by the model in its forward() pass.\n",
    "#\n",
    "# [2] pos_indices: Tensor with shape (N,2), where N is the number of positive\n",
    "# pairs. Each positive pair is a 0-indexed (Row,Col) position within the matrix\n",
    "# x above.\n",
    "#\n",
    "# [3] temperature: float. The temperature to scale the raw logits by before\n",
    "# applying the sigmoid activation function.\n",
    "#\n",
    "target = torch.zeros(8, 8)\n",
    "pos_indices = torch.tensor([\n",
    "    (0, 0), (0, 2), (0, 4),\n",
    "    (1, 4), (1, 6), (1, 1),\n",
    "    (2, 3),\n",
    "    (3, 7),\n",
    "    (4, 3),\n",
    "    (7, 6),\n",
    "])\n",
    "# Add indexes of the principal diagonal as positive indexes. This will be useful\n",
    "# since we will use the BCELoss in PyTorch, which will expect a value for the\n",
    "# elements on the principal diagonal as well.\n",
    "pos_indices = torch.cat([pos_indices, torch.arange(8).reshape(8, 1).expand(-1, 2)], dim=0)\n",
    "print(\"\\nPositive indexes list\")\n",
    "print(pos_indices)\n",
    "\n",
    "\n",
    "# Set the values in the target vector to 1.\n",
    "target[pos_indices[:,0], pos_indices[:,1]] = 1\n",
    "print(f\"\\nGround Truth labels for positive and negative pairs for BCE Loss\")\n",
    "print(target)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "4f42c7d5",
   "metadata": {},
   "source": [
    "## Documentation for binary cross entropy loss computation\n",
    "\n",
    "* [torch.nn.functional.binary_cross_entropy_with_logits](https://pytorch.org/docs/stable/generated/torch.nn.functional.binary_cross_entropy_with_logits.html)\n",
    "* [torch.nn.functional.binary_cross_entropy](https://pytorch.org/docs/stable/generated/torch.nn.functional.binary_cross_entropy.html)\n",
    "\n",
    "The `_with_logits` version accepts raw logits, but we use the one which accepts values in the range 0.0 - 1.0 (i.e. we manually apply the Sigmoid activation function) since the one that accepts raw-logits uses the [log-sum-exp-trick](https://gregorygundersen.com/blog/2020/02/09/log-sum-exp/), and causes a `NaN` result when fed with $\\infty$ or $-\\infty$ as input. [This comment](https://github.com/pytorch/pytorch/issues/49844#issuecomment-1574693686) on a related github issue provides details of why the log-sum-exp-trick causes a problem with such values whereas running Sigmoid first avoids these issues."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 122,
   "id": "2825d2cd",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "(tensor([[    inf, -0.8714,  0.6698, -0.7773,  0.7604, -0.9751,  0.3293, -0.8719],\n",
       "         [-0.8714,     inf, -0.9479,  0.9860, -0.9812,  0.7408,  0.1763,  0.5195],\n",
       "         [ 0.6698, -0.9479,     inf, -0.9878,  0.9916, -0.4883, -0.4806, -0.2203],\n",
       "         [-0.7773,  0.9860, -0.9878,     inf, -0.9997,  0.6183,  0.3381,  0.3696],\n",
       "         [ 0.7604, -0.9812,  0.9916, -0.9997,     inf, -0.5973, -0.3628, -0.3449],\n",
       "         [-0.9751,  0.7408, -0.4883,  0.6183, -0.5973,     inf, -0.5306,  0.9588],\n",
       "         [ 0.3293,  0.1763, -0.4806,  0.3381, -0.3628, -0.5306,     inf, -0.7495],\n",
       "         [-0.8719,  0.5195, -0.2203,  0.3696, -0.3449,  0.9588, -0.7495,     inf]]),\n",
       " tensor([[1.0000, 0.2950, 0.6615, 0.3149, 0.6814, 0.2739, 0.5816, 0.2949],\n",
       "         [0.2950, 1.0000, 0.2793, 0.7283, 0.2726, 0.6772, 0.5440, 0.6270],\n",
       "         [0.6615, 0.2793, 1.0000, 0.2713, 0.7294, 0.3803, 0.3821, 0.4451],\n",
       "         [0.3149, 0.7283, 0.2713, 1.0000, 0.2690, 0.6498, 0.5837, 0.5914],\n",
       "         [0.6814, 0.2726, 0.7294, 0.2690, 1.0000, 0.3550, 0.4103, 0.4146],\n",
       "         [0.2739, 0.6772, 0.3803, 0.6498, 0.3550, 1.0000, 0.3704, 0.7229],\n",
       "         [0.5816, 0.5440, 0.3821, 0.5837, 0.4103, 0.3704, 1.0000, 0.3209],\n",
       "         [0.2949, 0.6270, 0.4451, 0.5914, 0.4146, 0.7229, 0.3209, 1.0000]]))"
      ]
     },
     "execution_count": 122,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "y = xcs.clone()\n",
    "# In our predicted all-pairs cosine similarity, we set the diagonal elements\n",
    "# to infinity so that the sigmoid activation when applied to these elements\n",
    "# will set them to 1.0. We're going to set the diagonal elements to infinity\n",
    "# so that each element is considered to be similar to itself, and won't generate\n",
    "# gradients due to this specific pair when the loss is propagated into the\n",
    "# weights that contributed that specific loss value.\n",
    "y[eye.bool()] = float(\"inf\")\n",
    "y, y.sigmoid()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 123,
   "id": "131f0404",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "Positive pairs mask\n",
      "tensor([[ True, False,  True, False,  True, False, False, False],\n",
      "        [False,  True, False, False,  True, False,  True, False],\n",
      "        [False, False,  True,  True, False, False, False, False],\n",
      "        [False, False, False,  True, False, False, False,  True],\n",
      "        [False, False, False,  True,  True, False, False, False],\n",
      "        [False, False, False, False, False,  True, False, False],\n",
      "        [False, False, False, False, False, False,  True, False],\n",
      "        [False, False, False, False, False, False,  True,  True]])\n",
      "\n",
      "Negative pairs mask\n",
      "tensor([[False,  True, False,  True, False,  True,  True,  True],\n",
      "        [ True, False,  True,  True, False,  True, False,  True],\n",
      "        [ True,  True, False, False,  True,  True,  True,  True],\n",
      "        [ True,  True,  True, False,  True,  True,  True, False],\n",
      "        [ True,  True,  True, False, False,  True,  True,  True],\n",
      "        [ True,  True,  True,  True,  True, False,  True,  True],\n",
      "        [ True,  True,  True,  True,  True,  True, False,  True],\n",
      "        [ True,  True,  True,  True,  True,  True, False, False]])\n",
      "\n",
      "Positive pairs only (loss)\n",
      "tensor([[-0.0000e+00, 0.0000e+00, 1.2330e-03, 0.0000e+00, 4.9830e-04, 0.0000e+00,\n",
      "         0.0000e+00, 0.0000e+00],\n",
      "        [0.0000e+00, -0.0000e+00, 0.0000e+00, 0.0000e+00, 9.8123e+00, 0.0000e+00,\n",
      "         1.5834e-01, 0.0000e+00],\n",
      "        [0.0000e+00, 0.0000e+00, -0.0000e+00, 9.8779e+00, 0.0000e+00, 0.0000e+00,\n",
      "         0.0000e+00, 0.0000e+00],\n",
      "        [0.0000e+00, 0.0000e+00, 0.0000e+00, -0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "         0.0000e+00, 2.4520e-02],\n",
      "        [0.0000e+00, 0.0000e+00, 0.0000e+00, 9.9966e+00, -0.0000e+00, 0.0000e+00,\n",
      "         0.0000e+00, 0.0000e+00],\n",
      "        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, -0.0000e+00,\n",
      "         0.0000e+00, 0.0000e+00],\n",
      "        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "         -0.0000e+00, 0.0000e+00],\n",
      "        [0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00, 0.0000e+00,\n",
      "         7.4957e+00, -0.0000e+00]])\n",
      "\n",
      "Negative pairs only (loss)\n",
      "tensor([[0.0000e+00, 1.6428e-04, 0.0000e+00, 4.2090e-04, 0.0000e+00, 5.8235e-05,\n",
      "         3.3295e+00, 1.6345e-04],\n",
      "        [1.6428e-04, 0.0000e+00, 7.6416e-05, 9.8602e+00, 0.0000e+00, 7.4085e+00,\n",
      "         0.0000e+00, 5.2004e+00],\n",
      "        [6.6989e+00, 7.6416e-05, 0.0000e+00, 0.0000e+00, 9.9165e+00, 7.5463e-03,\n",
      "         8.1472e-03, 1.0479e-01],\n",
      "        [4.2090e-04, 9.8602e+00, 5.1321e-05, 0.0000e+00, 4.5539e-05, 6.1852e+00,\n",
      "         3.4143e+00, 0.0000e+00],\n",
      "        [7.6046e+00, 5.4778e-05, 9.9165e+00, 0.0000e+00, 0.0000e+00, 2.5423e-03,\n",
      "         2.6218e-02, 3.1272e-02],\n",
      "        [5.8235e-05, 7.4085e+00, 7.5463e-03, 6.1852e+00, 2.5423e-03, 0.0000e+00,\n",
      "         4.9488e-03, 9.5880e+00],\n",
      "        [3.3295e+00, 1.9212e+00, 8.1472e-03, 3.4143e+00, 2.6218e-02, 4.9488e-03,\n",
      "         0.0000e+00, 5.5561e-04],\n",
      "        [1.6345e-04, 5.2004e+00, 1.0479e-01, 3.7205e+00, 3.1272e-02, 9.5880e+00,\n",
      "         0.0000e+00, 0.0000e+00]])\n",
      "\n",
      "Element-wise positive pairs loss\n",
      "tensor([1.7313e-03, 9.9706e+00, 9.8779e+00, 2.4520e-02, 9.9966e+00, 0.0000e+00,\n",
      "        0.0000e+00, 7.4957e+00])\n",
      "\n",
      "Element-wise negative pairs loss\n",
      "tensor([ 3.3303, 22.4694, 16.7360, 19.4602, 17.5812, 23.1968,  8.7048, 18.6452])\n",
      "\n",
      "Number of elements similar to the i'th element\n",
      "tensor([3., 3., 2., 2., 2., 1., 1., 2.])\n",
      "\n",
      "Number of elements dissimilar to the i'th element\n",
      "tensor([5., 5., 6., 6., 6., 7., 7., 6.])\n",
      "\n",
      "Overall loss: 4.851151943206787\n"
     ]
    }
   ],
   "source": [
    "temperature = 0.1\n",
    "loss = F.binary_cross_entropy((y / temperature).sigmoid(), target, reduction=\"none\")\n",
    "\n",
    "target_pos = target.bool()\n",
    "target_neg = ~target_pos\n",
    "\n",
    "print(\"\\nPositive pairs mask\")\n",
    "print(target_pos)\n",
    "\n",
    "print(\"\\nNegative pairs mask\")\n",
    "print(target_neg)\n",
    "\n",
    "# loss_pos and loss_neg below contain non-zero values only for those elements\n",
    "# that are positive pairs and negative pairs respectively.\n",
    "loss_pos = torch.zeros(x.size(0), x.size(0)).masked_scatter(target_pos, loss[target_pos])\n",
    "loss_neg = torch.zeros(x.size(0), x.size(0)).masked_scatter(target_neg, loss[target_neg])\n",
    "\n",
    "print(\"\\nPositive pairs only (loss)\")\n",
    "print(loss_pos)\n",
    "\n",
    "print(\"\\nNegative pairs only (loss)\")\n",
    "print(loss_neg)\n",
    "\n",
    "# loss_pos and loss_neg now contain the sum of positive and negative pair losses\n",
    "# as computed relative to the i'th input.\n",
    "loss_pos = loss_pos.sum(dim=1)\n",
    "loss_neg = loss_neg.sum(dim=1)\n",
    "\n",
    "print(\"\\nElement-wise positive pairs loss\")\n",
    "print(loss_pos)\n",
    "\n",
    "print(\"\\nElement-wise negative pairs loss\")\n",
    "print(loss_neg)\n",
    "\n",
    "# num_pos and num_neg below contain the number of positive and negative pairs\n",
    "# computed relative to the i'th input. In an actual setting, this number should\n",
    "# be the same for every input element, but we let it vary here for maximum\n",
    "# flexibility.\n",
    "num_pos = target.sum(dim=1)\n",
    "num_neg = target.size(0) - num_pos\n",
    "\n",
    "print(\"\\nNumber of elements similar to the i'th element\")\n",
    "print(num_pos)\n",
    "\n",
    "print(\"\\nNumber of elements dissimilar to the i'th element\")\n",
    "print(num_neg)\n",
    "\n",
    "# Compute the weighted overall loss as seen in the formula above.\n",
    "overall_loss = ((loss_pos / num_pos) + (loss_neg / num_neg)).mean()\n",
    "print(f\"\\nOverall loss: {overall_loss.item()}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "696675c7",
   "metadata": {},
   "source": [
    "## Final implementation of NT-BXent Loss\n",
    "\n",
    "Now that we have seen how to build up the NT-BXent loss, we can condense the complete implementation (including the computation of the all-pairs Cosine Similarity) into this one simple function below."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 124,
   "id": "e9b32c56",
   "metadata": {
    "scrolled": true
   },
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Temperature:  0.01, Loss: 62.898780822753906\n",
      "Temperature:  0.10, Loss: 4.851151943206787\n",
      "Temperature:  1.00, Loss: 1.0727109909057617\n",
      "Temperature: 10.00, Loss: 0.9827173948287964\n",
      "Temperature: 20.00, Loss: 0.982099175453186\n"
     ]
    }
   ],
   "source": [
    "def nt_bxent_loss(x, pos_indices, temperature):\n",
    "    assert len(x.size()) == 2\n",
    "\n",
    "    # Add indexes of the principal diagonal elements to pos_indices\n",
    "    pos_indices = torch.cat([\n",
    "        pos_indices,\n",
    "        torch.arange(x.size(0)).reshape(x.size(0), 1).expand(-1, 2),\n",
    "    ], dim=0)\n",
    "    \n",
    "    # Ground truth labels\n",
    "    target = torch.zeros(x.size(0), x.size(0))\n",
    "    target[pos_indices[:,0], pos_indices[:,1]] = 1.0\n",
    "\n",
    "    # Cosine similarity\n",
    "    xcs = F.cosine_similarity(x[None,:,:], x[:,None,:], dim=-1)\n",
    "    # Set logit of diagonal element to \"inf\" signifying complete\n",
    "    # correlation. sigmoid(inf) = 1.0 so this will work out nicely\n",
    "    # when computing the Binary Cross Entropy Loss.\n",
    "    xcs[torch.eye(x.size(0)).bool()] = float(\"inf\")\n",
    "\n",
    "    # Standard binary cross entropy loss. We use binary_cross_entropy() here and not\n",
    "    # binary_cross_entropy_with_logits() because of https://github.com/pytorch/pytorch/issues/102894\n",
    "    # The method *_with_logits() uses the log-sum-exp-trick, which causes inf and -inf values\n",
    "    # to result in a NaN result.\n",
    "    loss = F.binary_cross_entropy((xcs / temperature).sigmoid(), target, reduction=\"none\")\n",
    "    \n",
    "    target_pos = target.bool()\n",
    "    target_neg = ~target_pos\n",
    "    \n",
    "    loss_pos = torch.zeros(x.size(0), x.size(0)).masked_scatter(target_pos, loss[target_pos])\n",
    "    loss_neg = torch.zeros(x.size(0), x.size(0)).masked_scatter(target_neg, loss[target_neg])\n",
    "    loss_pos = loss_pos.sum(dim=1)\n",
    "    loss_neg = loss_neg.sum(dim=1)\n",
    "    num_pos = target.sum(dim=1)\n",
    "    num_neg = x.size(0) - num_pos\n",
    "\n",
    "    return ((loss_pos / num_pos) + (loss_neg / num_neg)).mean()\n",
    "\n",
    "pos_indices = torch.tensor([\n",
    "    (0, 0), (0, 2), (0, 4),\n",
    "    (1, 4), (1, 6), (1, 1),\n",
    "    (2, 3),\n",
    "    (3, 7),\n",
    "    (4, 3),\n",
    "    (7, 6),\n",
    "])\n",
    "for t in (0.01, 0.1, 1.0, 10.0, 20.0):\n",
    "    print(f\"Temperature: {t:5.2f}, Loss: {nt_bxent_loss(x, pos_indices, temperature=t)}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}