From the Triangle Inequality to Margin Softmax

从三角不等式到Margin Softmax 苏剑林 · 2021-09-01 · Machine Learning loss function similarity

In Sentence Similarity Models Based on GRU and AM-Softmax] we introduced AM-Softmax, a kind of softmax with a margin, which is commonly used in scenarios where classification is used to perform retrieval. Back then, we explained—somewhat informally, with a diagram—that the margin is introduced because of a "non-equivalence between classification and ranking," but we did not give a very quantitative account of where this non-equivalence comes from.

In this post, let's revisit this topic and derive and understand the necessity of the margin from the perspective of the triangle inequality for distances.

The Triangle Inequality

In everyday usage, "distance" usually refers to the intuitive "Euclidean distance," but mathematically, distance—also called a "metric"—has an axiomatic definition: it is a binary function $d(x,y)$ defined on some set, satisfying:

1. Non-negativity: $d(x,y)\geq 0$;
2. Identity of indiscernibles: $d(x,y)=0\Leftrightarrow x = y$;
3. Symmetry: $d(x,y)=d(y,x)$;
4. Triangle inequality: $d(x,y)\leq d(x,z) + d(z,y)$.

As the name suggests, a distance is used to measure the degree of difference between $x,y$. In theory, as long as the first two conditions are satisfied, a function can be used to measure difference—for example, the KL divergence commonly used in probability satisfies only the first two conditions. The reason conditions 3 and 4 are added is essentially to make such a defined distance closer to the ordinary Euclidean distance we're familiar with. For instance, symmetry embodies the idea that "distance has no direction," while the triangle inequality embodies the idea that "a straight line is the shortest path between two points." These properties help us reason about more general distances by analogy with Euclidean distance.

From this definition, it turns out that deep learning rarely encounters distances satisfying all four of the above conditions. For example, ordinary classification directly uses an inner product plus softmax, and the inner product only satisfies condition 3; cosine distance $1-\cos(x,y)$ satisfies only conditions 1 and 3, failing conditions 2 and 4—although if we regard all vectors pointing in the same direction as equal, then it also satisfies condition 2.

However, for certain functions we can tweak the definition slightly so that they become a proper distance. For instance, since we know Euclidean distance satisfies the triangle inequality,

\begin{equation}\left\Vert \frac{x}{\Vert x\Vert} - \frac{y}{\Vert y\Vert}\right\Vert = \sqrt{2 - 2\cos(x,y)}\end{equation}

must also satisfy the triangle inequality. So, while the cosine distance $1-\cos(x,y)$ does not satisfy the triangle inequality, the modified form $\sqrt{1-\cos(x,y)}$ does.

Classification vs. Ranking

In scenarios like face recognition or sentence similarity, at prediction time we use the learned features to perform ranking: we naturally want that, given any sample, we can retrieve all other samples of the same class. This requires that "the intra-class gap be smaller than the inter-class gap." However, if we train the model as a classification task, this goal may not necessarily be achieved, because the objective of a classification task is simply "being closest to the center of the class it belongs to." A concrete example can be seen in the figure below:

A possible classification outcome, where red points represent class centers and other points represent samples

In this figure, $z_1,z_3$ belongs to class $c_1$, and $z_2$ belongs to class $c_2$. From a classification standpoint we have $d(z_1, c_1) < d(z_1, c_2)$ and $d(z_2, c_2) < d(z_2, c_1)$, so both are correctly classified. But $d(z_1, z_2) < d(z_1, z_3)$, so if we use $z_1$ for retrieval, what we find is $z_2$ from a different class, rather than $z_3$ from the same class.

We can describe this discrepancy more quantitatively using the triangle inequality: we want to achieve $d(z_1, z_3) < d(z_1, z_2)$. By the triangle inequality we have $d(z_1,z_3)\leq d(z_1, c_1) + d(z_3, c_1)$, so a sufficient condition is

\begin{equation}d(z_1, c_1) + d(z_3, c_1) < d(z_1, z_2) \end{equation}

Adding $d(z_2, c_2)$ to both sides, and using the triangle inequality $d(z_1, z_2) + d(z_2, c_2)\geq d(z_1, c_2)$, we find that a sufficient condition for the above is

\begin{equation}d(z_1, c_1) + d(z_3, c_1) + d(z_2, c_2) < d(z_1, c_2) \end{equation}

Note that a plain classification task only requires $d(z_1, c_1) < d(z_1, c_2)$ for $z_1$, whereas the above inequality has an extra term $d(z_3, c_1) + d(z_2, c_2)$—and it is precisely this extra term that is the margin.

Notice that $d(z_3, c_1),d(z_2, c_2)$ are, respectively, the distances from samples $z_3,z_2$ to their own class centers, so we can regard $d(z_3, c_1) + d(z_2, c_2)$ as the "average class diameter." It should be close to some constant $m$, which we can treat as a hyperparameter to tune. If we want to adjust it adaptively, we could first train $m=0$ for a while, then estimate the "average class diameter" to use as $m$ and continue training, then re-estimate $m$ and train again, and so on.

AM-Softmax

From the derivation above, we know that in order to ensure the features learned by a classification model are also usable for ranking, each sample must not only be closest to its own class center, but must remain closest to its class center even after adding a margin $m$ to the distance. That is, if $z_1$ belongs to class $c_1$, we require:

\begin{equation}\begin{aligned} d(z_1, c_1) +&\, m < d(z_1, c_2) \\ d(z_1, c_1) +&\, m < d(z_1, c_3) \\ &\vdots \\ d(z_1, c_1) +&\, m < d(z_1, c_k) \end{aligned}\end{equation}

Following the idea from Generalizing "Softmax + Cross-Entropy" to Multi-Label Classification]: whenever we want $s_i < s_j$, we can construct a loss by inserting $e^{s_i - s_j}$ into $\log$. So we can construct the following loss:

\begin{equation}\log\left(1+\sum_{i=2}^k e^{s\cdot[d(z_1, c_1) + m - d(z_1, c_i)]}\right)\end{equation}

This is precisely the cross-entropy with an additive margin, where $s$ is a scaling factor, playing the role of a temperature parameter for the softmax.

However, don't forget that the derivation above relies on $d$ satisfying the triangle inequality, whereas the scoring functions we normally use do not. For training retrieval models, we usually use cosine distance as the score, and as mentioned earlier, cosine distance can be made to satisfy the triangle inequality by taking a square root. So the corresponding requirement becomes (taking $i=2$ as an example):

\begin{equation}\begin{array}{c} \sqrt{1-\cos(z_1, c_1)} + m < \sqrt{1-\cos(z_1, c_2)}\\ \Downarrow\\ \sqrt{1-\cos(z_1, c_2)} - \sqrt{1-\cos(z_1, c_1)} > m \\ \end{array}\end{equation}

Multiplying both sides by $\sqrt{1-\cos(z_1, c_2)} + \sqrt{1-\cos(z_1, c_1)}$ gives

\begin{equation}\cos(z_1, c_1) - \cos(z_1, c_2) > m\left(\sqrt{1-\cos(z_1, c_2)} + \sqrt{1-\cos(z_1, c_1)}\right)\end{equation}

Clearly, the right-hand side has an upper bound, so by suitably adjusting $m$, we can make

\begin{equation}\cos(z_1, c_1) - \cos(z_1, c_2) > m\end{equation}

a sufficient condition, and the corresponding margin cross-entropy is then

\begin{equation}\log\left(1+\sum_{i=2}^k e^{s\cdot[\cos(z_1, c_i) + m - \cos(z_1, c_1)]}\right)\end{equation}

This is exactly AM-Softmax].

Review and Summary

This post derived, from the perspective of the triangle inequality, the necessity of the margin when using classification models for ranking tasks. Under the assumption that the scoring function used satisfies the triangle inequality, we were able to derive the relevant results in a fairly natural way.