Contrastive loss and the InfoNCE gradient

Worked solutions

Problem 1

1. $-\log p_1 = -\log\dfrac{e^{z_1}}{\sum_j e^{z_j}} = -z_1 + \log\sum_j e^{z_j}$.$\log(a/b) = \log a - \log b$ and $\log e^{z_1} = z_1$.
2. $\partial\ell/\partial z_j = -\delta_{j1} + \dfrac{e^{z_j}}{\sum_k e^{z_k}} = p_j - y_j$.The first term contributes only for $j = 1$; the log-sum-exp has derivative $p_j$ by the chain rule, $\tfrac{d}{dz_j}\log S = \tfrac1S\,\partial S/\partial z_j$ with $S = \sum_k e^{z_k}$. This is the softmax page's $p - y$.
3. $z = c/\tau$, so $\partial z_j/\partial c_j = 1/\tau$ and $\nabla_c\ell = \tfrac1\tau\nabla_z\ell$.Each $c_j$ enters one logit, scaled by the constant $1/\tau$.
4. $\ell = -\log p_1$, $\nabla_z\ell = p - y$, $\nabla_c\ell = (p - y)/\tau$Every entry of $\nabla_c\ell$ is positive for a negative ($p_j > 0$) and negative for the positive ($p_1 - 1 < 0$): gradient descent raises the positive's similarity and lowers each negative's, in proportion to the softmax weight the loss currently gives it, divided by $\tau$. The rest of this page carries this vector through the similarities to the embeddings.

Problem 2

1. $c = \hat u^\top\hat v$, and $\hat v$ does not depend on $u$.Definition of the cosine; $v$ is held fixed.
2. $\nabla_uc = \Big(\dfrac{\partial\hat u}{\partial u}\Big)^\top\hat v = \dfrac1{\|u\|}(I - \hat u\hat u^\top)\hat v$.$c = a^\top\hat u$ with $a = \hat v$ constant, so $\nabla_{\hat u}c = \hat v$ and the chain rule through $\hat u(u)$ multiplies by the transposed Jacobian, which is symmetric (Before you start).
3. $(I - \hat u\hat u^\top)\hat v = \hat v - \hat u(\hat u^\top\hat v) = \hat v - c\,\hat u$.Distribute; $\hat u^\top\hat v$ is the scalar $c$.
4. By symmetry, $\nabla_vc = \dfrac1{\|v\|}(\hat u - c\,\hat v)$.$c(u, v) = c(v, u)$, so swap the roles.
5. $u^\top\nabla_uc = \dfrac{1}{\|u\|}\big(u^\top\hat v - c\,u^\top\hat u\big) = \dfrac1{\|u\|}\big(\|u\|c - c\|u\|\big) = 0$.$u^\top\hat v = \|u\|\hat u^\top\hat v = \|u\|c$ and $u^\top\hat u = \|u\|$.
6. $\|\nabla_uc\|^2 = \dfrac{1}{\|u\|^2}\big(\hat v^\top\hat v - 2c\,\hat u^\top\hat v + c^2\hat u^\top\hat u\big) = \dfrac{1 - 2c^2 + c^2}{\|u\|^2} = \dfrac{1 - c^2}{\|u\|^2}$.Expand $\|\hat v - c\hat u\|^2$ with unit vectors and $\hat u^\top\hat v = c$.
7. $\nabla_uc = \dfrac{\hat v - c\,\hat u}{\|u\|}$, $\nabla_vc = \dfrac{\hat u - c\,\hat v}{\|v\|}$; $u^\top\nabla_uc = 0$ and $\|\nabla_uc\| = \dfrac{\sqrt{1 - c^2}}{\|u\|}$The gradient is the part of $\hat v$ orthogonal to $\hat u$ (the direction that turns $u$ towards $v$), scaled by $1/\|u\|$: a long anchor turns more slowly. It vanishes when $c = \pm1$, where $u$ and $v$ are already aligned or opposed, and nothing in it can change the length of $u$.

Problem 3

1. $\nabla_u\ell = \sum_j\dfrac{\partial\ell}{\partial c_j}\nabla_uc_j = \dfrac1\tau\sum_j(p_j - y_j)\,\nabla_uc_j$.$\ell$ depends on $u$ only through the $K$ similarities; chain rule over all of them with Problem 1.
2. $= \dfrac{1}{\tau\|u\|}\sum_j(p_j - y_j)\big(\hat v_j - c_j\hat u\big)$.Problem 2 for each $c_j$; the common factor $1/\|u\|$ comes out.
3. $\sum_j(p_j - y_j)\hat v_j = \sum_jp_j\hat v_j - \hat v_1$.$y$ is one-hot on the positive.
4. $\sum_j(p_j - y_j)c_j\hat u = \Big(\sum_jp_jc_j - c_1\Big)\hat u$.The same split; $\hat u$ does not depend on $j$.
5. $\nabla_u\ell = \dfrac{1}{\tau\|u\|}\Big[\sum_jp_j\hat v_j - \hat v_1 - \Big(\sum_jp_jc_j - c_1\Big)\hat u\Big]$The step $-\eta\nabla_u\ell$ pulls $u$ towards the positive direction $\hat v_1$ and pushes it away from the softmax-weighted average of all candidates, $\sum_jp_j\hat v_j$, which is dominated by the negatives it currently confuses with the positive; the $\hat u$ term removes whatever component of that lies along $u$, so that $u^\top\nabla_u\ell = 0$ (Problem 7). The $\sum_j(p_j - y_j) = 0$ identity does not kill the $\hat u$ term, because each term carries its own $c_j$.

Problem 4

1. $\ell$ depends on $v_j$ only through $c_j$.$v_j$ appears in no other similarity.
2. $\nabla_{v_j}\ell = \dfrac{\partial\ell}{\partial c_j}\nabla_{v_j}c_j = \dfrac{p_j - y_j}{\tau}\cdot\dfrac{\hat u - c_j\hat v_j}{\|v_j\|}$.Problem 1 for the first factor and Problem 2 (the $\nabla_v$ form) for the second.
3. Positive: $\nabla_{v_1}\ell = \dfrac{p_1 - 1}{\tau\|v_1\|}(\hat u - c_1\hat v_1)$, with $p_1 - 1 < 0$.$y_1 = 1$ and $0 < p_1 < 1$.
4. Negative: $\nabla_{v_j}\ell = \dfrac{p_j}{\tau\|v_j\|}(\hat u - c_j\hat v_j)$, with $p_j > 0$.$y_j = 0$.
5. $\nabla_{v_j}\ell = \dfrac{p_j - y_j}{\tau\|v_j\|}(\hat u - c_j\hat v_j)$: the positive moves along $+(\hat u - c_1\hat v_1)$, towards the anchor's direction; each negative moves along $-(\hat u - c_j\hat v_j)$, away from it, with weight $p_j$$\hat u - c_j\hat v_j$ is the component of $\hat u$ orthogonal to $\hat v_j$, the direction that turns $v_j$ towards $u$ (Problem 2). A negative with tiny $p_j$, one the model already separates, is barely touched; a negative with $p_j$ close to $p_1$ is pushed as hard as the positive is pulled. In a batch where every example is also a negative for the others, each $v_j$ collects one such term from every anchor (Problem 9).

Problem 5

1. $\dfrac{p_j}{p_k} = \dfrac{e^{c_j/\tau}}{e^{c_k/\tau}} = e^{(c_j - c_k)/\tau}$.The common denominator of the softmax cancels.
2. With $\tau = 0.1$, a negative whose similarity is $0.1$ higher than another's gets $e^{1} \approx 2.7$ times the weight; with $\tau = 0.01$, $e^{10} \approx 22{,}000$ times.Step 1 with $c_j - c_k = 0.1$.
3. As $\tau \to 0$, $p \to e_{j^*}$ where $j^*$ is the candidate with the largest similarity.The softmax page, Problem 9: $\operatorname{softmax}(c/\tau)$ concentrates on the largest entry as the temperature goes to $0$ (assuming a unique maximum).
4. Then $\nabla_c\ell = (p - y)/\tau \to (e_{j^*} - e_1)/\tau$: zero if the positive is the top candidate, and otherwise a push on the single hardest negative and a pull on the positive, each of size $1/\tau$.Problem 1 with step 3. The loss itself tends to $\max(0, \max_{j\neq1}c_j - c_1)/\tau$, a hinge on the hardest negative.
5. As $\tau \to \infty$, $p \to \tfrac1K\mathbf{1}$ and $\nabla_c\ell \to (\tfrac1K\mathbf{1} - y)/\tau \to 0$.All logits go to $0$, so the softmax is uniform; the $1/\tau$ then sends the gradient to zero.
6. $p_j/p_k = e^{(c_j - c_k)/\tau}$; $\tau \to 0$ gives a hinge on the hardest negative with gradient $(e_{j^*} - e_1)/\tau$; $\tau \to \infty$ gives uniform weights and a vanishing gradientThe temperature is a hardness dial: small $\tau$ makes the loss attend to the few negatives nearest the anchor, large $\tau$ treats all negatives alike. Typical values ($0.05$ to $0.1$ for SimCLR, about $0.01$ for a trained CLIP) sit near the hard end, which is why these losses need many negatives: with few, the hardest negative is often not hard.

Problem 6

1. $\partial\ell^{\text{row}}_i/\partial Z_{ij} = (P_r)_{ij} - \delta_{ij}$, and $\ell^{\text{row}}_i$ does not depend on any other row of $Z$.Problem 1 applied to row $i$ of $Z$ with label $i$: its softmax is row $i$ of $P_r$.
2. $\nabla_Z\sum_i\ell^{\text{row}}_i = P_r - I$.Stack step 1 over the rows: entry $(i, j)$ gets its contribution from $\ell^{\text{row}}_i$ alone.
3. $\partial\ell^{\text{col}}_i/\partial Z_{ji} = (P_c)_{ji} - \delta_{ji}$, and $\ell^{\text{col}}_i$ depends only on column $i$.The same with column $i$ of $Z$ and label $i$; its softmax is column $i$ of $P_c$.
4. $\nabla_Z\sum_i\ell^{\text{col}}_i = P_c - I$.Stack over the columns: entry $(j, i)$ comes from $\ell^{\text{col}}_i$.
5. $\nabla_ZL = \dfrac{1}{2N}\big(P_r + P_c - 2I\big)$$N\times N$, the shape of $Z$. Each row of $P_r - I$ and each column of $P_c - I$ sums to $0$, so all the entries of $\nabla_ZL$ sum to $0$. The diagonal entries $(P_r)_{ii} + (P_c)_{ii} - 2$ are negative (pull the pairs together) and the off-diagonal entries are positive (push every non-pair apart). $P_c \neq P_r^\top$ in general: the two directions normalise over different sets, so neither half can be dropped and recovered "by symmetry".

Problem 7

1. $c(\alpha u, v) = \dfrac{\alpha u^\top v}{\alpha\|u\|\|v\|} = c(u, v)$.$\|\alpha u\| = \alpha\|u\|$ for $\alpha > 0$; the $\alpha$ cancels. Every $c_j$, hence $z$, $p$ and $\ell$, is unchanged.
2. $\tfrac{d}{d\alpha}\ell(\alpha u) = u^\top\nabla_u\ell(\alpha u)$, and the left side is $0$ for all $\alpha$.Chain rule along the ray $\alpha \mapsto \alpha u$; step 1 says $\ell$ is constant on it. At $\alpha = 1$, $u^\top\nabla_u\ell = 0$.
3. $\nabla_u\ell = \dfrac1\tau\sum_j(p_j - y_j)\nabla_uc_j$, and each $\nabla_uc_j$ at $\alpha u$ is $\dfrac{\hat v_j - c_j\hat u}{\alpha\|u\|}$, with $p_j$, $c_j$, $\hat u$ unchanged.Problem 3, step 1, and Problem 2 with $\|\alpha u\| = \alpha\|u\|$ and $\widehat{\alpha u} = \hat u$.
4. So $\nabla_u\ell(\alpha u) = \tfrac1\alpha\nabla_u\ell(u)$.Every term is divided by $\alpha$.
5. $\nabla_W\ell = (\nabla_u\ell)\,x^\top$.$u = Wx$ is linear in $W$ with $x$ fixed: $\partial\ell/\partial W_{ab} = \sum_i(\partial\ell/\partial u_i)(\partial u_i/\partial W_{ab}) = (\nabla_u\ell)_ax_b$, the one-hidden-layer page's weight-gradient pattern.
6. $\ell(\alpha u) = \ell(u)$; $u^\top\nabla_u\ell = 0$; $\nabla_u\ell(\alpha u) = \tfrac1\alpha\nabla_u\ell(u)$; $\nabla_W\ell = (\nabla_u\ell)x^\top$The loss sees only the direction of $u$, so its gradient can only turn $u$, and it turns a long $u$ less: the effective learning rate on the direction falls as $1/\|u\|^2$ (one factor from the gradient, one from the angle a fixed step subtends). Nothing in the loss shrinks $\|u\|$, and $\|Wx\|$ tends to grow under noisy gradient steps, so a scale-invariant loss relies on weight decay to keep its effective learning rate from decaying (the weight-decay page).

Problem 8

1. $\ell = -\dfrac{c_1}{\tau} + \log\sum_je^{c_j/\tau}$ with the $c_j$ fixed.Problem 1's loss with $z_j = c_j/\tau$ written out.
2. $\dfrac{\partial}{\partial\tau}\Big(-\dfrac{c_1}{\tau}\Big) = \dfrac{c_1}{\tau^2}$.$\tfrac{d}{d\tau}\tau^{-1} = -\tau^{-2}$.
3. $\dfrac{\partial}{\partial\tau}\log\sum_je^{c_j/\tau} = \sum_jp_j\cdot\dfrac{\partial}{\partial\tau}\dfrac{c_j}{\tau} = -\dfrac{1}{\tau^2}\sum_jp_jc_j$.Chain rule through each logit: $\partial\ell/\partial z_j = p_j$ for the log-sum-exp part (Problem 1, step 2), and $\partial z_j/\partial\tau = -c_j/\tau^2$.
4. $\dfrac{\partial\ell}{\partial\tau} = \dfrac{c_1 - \sum_jp_jc_j}{\tau^2}$.Add steps 2 and 3.
5. $\tau = e^{-t}$, so $\dfrac{d\tau}{dt} = -e^{-t} = -\tau$ and $\dfrac{\partial\ell}{\partial t} = \dfrac{\partial\ell}{\partial\tau}\cdot(-\tau) = -\dfrac{c_1 - \sum_jp_jc_j}{\tau}$.$s = e^t = 1/\tau$ inverts to $\tau = e^{-t}$; chain rule.
6. $\dfrac{\partial\ell}{\partial\tau} = \dfrac{c_1 - \sum_jp_jc_j}{\tau^2}$ and $\dfrac{\partial\ell}{\partial t} = \dfrac{\sum_jp_jc_j - c_1}{\tau}$; when $c_1 > \sum_jp_jc_j$, $\partial\ell/\partial t < 0$, so gradient descent increases $t$ and decreases $\tau$When the positive beats the weighted average, sharpening the softmax lowers the loss, so the model learns a smaller temperature; when the positive is losing, it learns a larger one to soften the penalty. Left unbounded, a model that is winning sharpens without limit, which is why CLIP clips the logit scale at $100$ ($\tau \ge 0.01$). The $\sum_jp_jc_j$ is the loss's own expectation of the similarity under $p$.

Problem 9

Write $\Gamma = \nabla_ZL = \tfrac1{2N}(P_r + P_c - 2I)$ from Problem 6.

1. $Z = \tfrac1\tau\hat U\hat V^\top$ has $\hat U$ as its left factor and $\hat V^\top$ as its right factor.A matrix product scaled by a constant.
2. $\nabla_{\hat U}L = \tfrac1\tau\Gamma\hat V$, $N\times d$.The attention page's left-factor pattern: for $C = PR$, $\nabla_PL = (\nabla_CL)R^\top$, with $R = \hat V^\top$ and the constant $1/\tau$ carried along.
3. $\nabla_{\hat V}L = \tfrac1\tau\Gamma^\top\hat U$, $N\times d$.The right-factor pattern $\nabla_RL = P^\top\nabla_CL$ gives $\nabla_{\hat V^\top}L = \tfrac1\tau\hat U^\top\Gamma$; transpose to get the gradient with respect to $\hat V$.
4. Row $i$ of $\hat U$ depends on $u_i$ alone, so $\nabla_{u_i}L = \Big(\dfrac{\partial\hat u_i}{\partial u_i}\Big)^\top(\nabla_{\hat U}L)_{i,:}^\top = \dfrac{1}{\|u_i\|}(I - \hat u_i\hat u_i^\top)(\nabla_{\hat U}L)_{i,:}^\top$.Chain rule through the normalisation with the Jacobian from Before you start, which is symmetric; the upstream gradient at $\hat u_i$ is row $i$ of step 2, as a column.
5. $\nabla_{\hat U}L = \tfrac1\tau\Gamma\hat V$, $\nabla_{\hat V}L = \tfrac1\tau\Gamma^\top\hat U$, $\Gamma = \tfrac1{2N}(P_r + P_c - 2I)$; then $\nabla_{u_i}L = \dfrac{1}{\|u_i\|}(I - \hat u_i\hat u_i^\top)\big(\tfrac1\tau\Gamma\hat V\big)_{i,:}^\top$, and likewise for $v_i$ with $\Gamma^\top\hat U$Row $i$ of $\Gamma\hat V$ is $\sum_j\Gamma_{ij}\hat v_j$: the positive $\hat v_i$ with negative weight, every other $\hat v_j$ with positive weight, which is Problem 3 for anchor $u_i$ plus the terms from $u_i$ serving as a negative for the other anchors through $P_c$. The projection then drops the part along $\hat u_i$ and the $1/\|u_i\|$ scales it, exactly as in Problem 2. The whole backward pass is two matrix products and a row-wise projection; the $N\times N$ matrix $\Gamma$ is the only thing that scales with the batch.

Problem 10

1. $\nabla_u\ell = \sum_j(p_j - y_j)\nabla_uz_j = \dfrac1\tau\sum_j(p_j - y_j)v_j$.Problem 1's $\nabla_z\ell = p - y$ and $\nabla_u(u^\top v_j/\tau) = v_j/\tau$.
2. $u^\top\nabla_u\ell = \dfrac1\tau\sum_j(p_j - y_j)u^\top v_j = \sum_j(p_j - y_j)z_j = \sum_jp_jz_j - z_1$.$u^\top v_j/\tau = z_j$, and $y$ picks out $z_1$.
3. $\tfrac{d}{d\alpha}\ell(\alpha u)\big|_{\alpha=1} = u^\top\nabla_u\ell = \sum_jp_jz_j - z_1$.Chain rule along the ray, as in Problem 7, step 2. Here it is not zero.
4. If $z_1 > z_j$ for every $j \neq 1$ then $\sum_jp_jz_j < z_1$, so the derivative is negative.A weighted average of the $z_j$ with weights $p_j > 0$ summing to $1$ is strictly less than the largest of them when the others are smaller.
5. The same holds at every $\alpha > 0$, since $z_j(\alpha u) = \alpha z_j(u)$ keeps the ordering.Scaling $u$ scales every logit by the same positive factor; the positive stays on top and step 4 applies again.
6. $\nabla_u\ell = \tfrac1\tau\sum_j(p_j - y_j)v_j$, $u^\top\nabla_u\ell = \sum_jp_jz_j - z_1$; when the positive leads, $\ell(\alpha u)$ decreases in $\alpha$, so the loss rewards growing $\|u\|$With dot-product logits, the encoder can drive $\ell \to 0$ on every example it already gets right by inflating embedding norms, which sharpens the softmax without learning anything new about directions; the gradient has a component along $u$ that does exactly this. Normalising removes that component (Problem 7) and the temperature puts the sharpness back under control (Problem 5): cosine plus $\tau$ is dot product with the norm fixed at $1/\tau$.