Margin-Enforcing Projection

In this post we introduce a mathematical operation called "Margin-Enforcing Projection (MEP)". It splits a vector into two parts according to a specified rule, and then requires that the margin between these two parts be at least $m$ (Margin).

Background

As we know, classification tasks are about picking out the correct category, so the training objective is usually just "the positive class score should be larger than the negative class score." But in certain scenarios we want more than that: we not only want the positive score to exceed the negative score, but to exceed it by at least a specified margin $m > 0$.

There are mainly two such scenarios. The first is when we want the classification result to be more robust and not easily disturbed by random noise—especially in low-precision inference settings—so we'd like the margin of the predicted result to be noticeably larger. The second is when we're not really using a classification model at all, but rather using classification as a proxy to learn features, with the ultimate goal being retrieval using those features. In this case, if no margin is enforced, retrieval results near the decision boundary can easily go wrong. more

Actually, both of these scenarios were already quite mature many years ago—especially the second one, whose typical representative is training feature models for face recognition—so this post is really something of an "archaeological dig." The standard approach to this problem is to design various margin losses, such as hinge loss, margin softmax, AM-Softmax, and so on. We touched on this briefly before in A Sentence Similarity Model Based on GRU and AM-Softmax and From the Triangle Inequality to Margin Softmax.

The idea in this post is: if we had some operation that could project the predicted scores $\boldsymbol{x}$ onto scores $\boldsymbol{y}$ satisfying the margin requirement, then we could simply train the model directly toward this projection as the target, e.g., by minimizing $\Vert\boldsymbol{y} - \boldsymbol{x}\Vert_2^2$. This projection operation is exactly what we're going to study next.

Mathematical Definition

Let $\boldsymbol{x}=(x_1,x_2,\cdots,x_n), \boldsymbol{z}=(z_1,z_2,\cdots,z_n)$, where $\boldsymbol{x}$ denotes the model's predicted scores. For simplicity, suppose the first $k$ scores represent the positive class and the last $n-k$ represent the negative class. We define

\begin{equation}\newcommand{argmin}{\mathop{\text{argmin}}}\mathcal{P}_m(\boldsymbol{x})\triangleq \argmin_{\boldsymbol{z}\in\mathbb{R}^n} d(\boldsymbol{z}, \boldsymbol{x})\quad\text{s.t.}\quad \min(\boldsymbol{z}_{\leq k}) - \max(\boldsymbol{z}_{> k}) \geq m\end{equation}

where $\boldsymbol{z}_{\leq k} = (z_1,\cdots,z_k), \boldsymbol{z}_{> k} = (z_{k+1},\cdots,z_n)$, $k$ are positive integers less than $n$, $m > 0$ is the given margin, and $d(\boldsymbol{z}, \boldsymbol{x})$ is the distance function to be minimized—we'll consider both L1 distance and L2 distance below.

This definition is quite intuitive: it's about finding the vector closest to the predicted scores while satisfying the margin requirement, which matches the general idea of a projection operation—hence the name "Margin-Enforcing Projection (MEP)." Using such a projected result as the learning target should, in theory, let the model take the easiest and quickest path, while also "braking" in time once the target is reached, avoiding overtraining.

Common Results

If $\boldsymbol{x}$ already satisfies $\min(\boldsymbol{x}_{\leq k}) - \max(\boldsymbol{x}_{> k}) \geq m$, then clearly $\boldsymbol{z}^* = \boldsymbol{x}$—this case is trivial. Without loss of generality, we assume $\min(\boldsymbol{x}_{\leq k}) - \max(\boldsymbol{x}_{> k}) < m$ from here on. It's not hard to see that, regardless of whether $d$ chooses the L1 or L2 distance, the optimal solution $\boldsymbol{z}^*$ must be attained at

\begin{equation}\min(\boldsymbol{z}^*_{\leq k}) - \max(\boldsymbol{z}^*_{> k}) = m\end{equation}

Otherwise we could always move some $z_i$ closer to $x_i$ to reduce the objective value further. Based on this observation, let $\min(\boldsymbol{z}^*_{\leq k})=\ell$, so that $\max(\boldsymbol{z}^*_{> k}) = \ell - m$, and we obtain

\begin{equation}\boldsymbol{z}^*_{\leq k} = \max(\boldsymbol{x}_{\leq k}, \ell),\qquad \boldsymbol{z}^*_{> k} = \min(\boldsymbol{x}_{> k}, \ell - m)\end{equation}

At this point

\begin{equation}\begin{aligned} \boldsymbol{z}^* - \boldsymbol{x} =&\, [\max(\boldsymbol{x}_{\leq k}, \ell) - \boldsymbol{x}_{\leq k}, \min(\boldsymbol{x}_{> k}, \ell - m) - \boldsymbol{x}_{> k}] \\[4pt] =&\, [\max(\ell - \boldsymbol{x}_{\leq k}, 0), \min(\ell - m - \boldsymbol{x}_{> k}, 0)] \\[4pt] =&\, [\max(\ell - \boldsymbol{x}_{\leq k}, 0), -\max(\boldsymbol{x}_{> k} + m - \ell, 0)] \end{aligned}\end{equation}

Next we need to solve for $\ell$ based on the choice of $d$.

L2 Distance

Let's first consider the L2 distance. Here we have

\begin{equation}\Vert\boldsymbol{z}^* - \boldsymbol{x}\Vert_2^2 = \sum_{i=1}^k\max(\ell - x_i, 0)^2 + \sum_{j=k+1}^n \max(x_j + m - \ell, 0)^2 \triangleq f(\ell)\end{equation}

Taking the derivative gives

\begin{gather}f'(\ell) = 2\sum_{i=1}^k\max(\ell - x_i, 0) - 2\sum_{j=k+1}^n \max(x_j + m - \ell, 0) \\[5pt] f''(\ell) = 2\#\{\ell > x_i\} + 2\#\{\ell < x_j + m\} \end{gather}

where $\#$ is the counting function, with the convention $1\leq i\leq k < j\leq n$. Clearly $f''(\ell) \geq 0$, but we can actually strengthen this to $f''(\ell) > 0$.

This is because $f''(\ell) = 0$ would imply $\#\{\ell > x_i\}=0$ and $\#\{\ell < x_j + m\}=0$, i.e., both $\min(\boldsymbol{x}_{\leq k}) \geq \ell$ and $\max(\boldsymbol{x}_{> k})\leq \ell - m$ hold simultaneously, which contradicts the assumption $\min(\boldsymbol{x}_{\leq k}) - \max(\boldsymbol{x}_{> k}) < m$. So $f''(\ell) > 0$, meaning $f(\ell)$ is strictly convex. Combined with the continuity of $f(\ell)$ and $f'(\ell)$, along with $f'(-\infty)=-\infty$ and $f'(\infty)=\infty$, we can conclude that $f(\ell)$ has a unique minimum point, and it must be attained at $f'(\ell)=0$.

Since $f'(\ell)$ is a composition of the piecewise linear function $\max(x, 0)$, $f'(\ell)$ is also a piecewise linear function of $\ell$, with breakpoints given by all the $x_i$ and $x_j + m$. To solve for $f'(\ell)=0$, we first sort the breakpoints $\{x_1,\cdots,x_k,x_{k+1}+m,\cdots,x_n+m\}$ in increasing order, giving $n-1$ intervals. Within a single interval, $f'(\ell)$ is a straight line; we go through all the intervals $[a, b]$, find the one where $f'(a) \leq 0$ and $f'(b) \geq 0$, and solve for the zero of the line within that interval.

L1 Distance

Now let's consider the L1 distance. Here we have

\begin{equation}\Vert\boldsymbol{z}^* - \boldsymbol{x}\Vert_1 = \sum_{i=1}^k\max(\ell - x_i, 0) + \sum_{j=k+1}^n \max(x_j + m - \ell, 0) \triangleq g(\ell)\end{equation}

Clearly, $g(\ell)$ is itself a piecewise linear function, and a function of this kind can only attain its minimum at a breakpoint. So the most straightforward solution is to go through all the breakpoints $\{x_1,\cdots,x_k,x_{k+1}+m,\cdots,x_n+m\}$ and pick the one that minimizes $g(\ell)$, with a complexity of $\mathcal{O}(n^2)$. But we can simplify this further. First, taking the derivative gives

\begin{equation}\begin{aligned} g'(\ell) =&\, \#\{\ell > x_i\} - \#\{\ell < x_j + m\} \\[4pt] =&\, \#\{\ell > x_i\} + \#\{\ell \geq x_j + m\} - (n - k) \end{aligned}\end{equation}

The second equality uses the identity $\#\{\ell < x_j + m\} + \#\{\ell \geq x_j + m\} = n - k$. Now it's easy to see that $g'(\ell)$ is monotonically increasing from negative to positive. Of course, $g'(\ell)$ is not continuous, so we can't guarantee finding a point where $g'(\ell)=0$.

However, if we redefine $\#\{\ell > x_i\}$ as $\#\{\ell \geq x_i\}$ (since the derivative at a discontinuous breakpoint can be treated as arbitrary, this adjustment is allowed), then the meaning of $g'(\ell)=0$ becomes exactly "the number of breakpoints less than or equal to $\ell$ is exactly $n-k$." If we further assume all the breakpoints are pairwise distinct, then $l$ is exactly the $n-k$-th breakpoint when all breakpoints are sorted in increasing order! So in the L1 case, a single sort suffices to obtain $\ell$—rather elegant.

Reference Implementation

Reference implementations for both versions of MEP are given below:

import jax
import jax.numpy as jnp

@jax.jit
def l2mep(inputs, mask, margin):
    x, m = inputs, margin
    u = jnp.where(mask, x, x + m).sort()[:, None]
    v = jnp.where(mask, jnp.fmax(u - x, 0), -jnp.fmax(x + m - u, 0)).sum(axis=1)
    i = ((v[:-1] < 0) & (v[1:] >= 0)).argmax()
    l = (u[i + 1] * v[i] - u[i] * v[i + 1]) / (v[i] - v[i + 1])
    return jnp.where(mask, jnp.fmax(x, l), jnp.fmin(x, l - m))

@jax.jit
def l1mep(inputs, mask, margin):
    x, m = inputs, margin
    u = jnp.where(mask, x, x + m).sort(axis=-1)
    l = jnp.take_along_axis(u, (~mask).sum(axis=-1, keepdims=True) - 1)
    return jnp.where(mask, jnp.fmax(x, l), jnp.fmin(x, l - m))

A brief note: in practice, the positive class usually isn't simply arranged as the first $k$ entries—its count and positions may both be unfixed. So in the implementation above, we use a mask vector to mark the positive and negative classes.

l2mep The current implementation only supports 1D input. If you need to support a batch dimension, just wrap it with jax.vmap. This algorithm finds the zero-crossing interval by iterating over all intervals; each step has complexity $\mathcal{O}(n)$, giving an overall complexity of $\mathcal{O}(n^2)$. To improve efficiency, one could use binary search to locate the zero-crossing interval instead, which would reduce this to $\mathcal{O}(n\log n)$.

l1mep Thanks to its simplicity, this algorithm already supports an arbitrary batch dimension as written. Its core operation is just sorting, with complexity $\mathcal{O}(n \log n)$, which is already quite satisfying in terms of both simplicity and speed. Absent other considerations, the L1 version is recommended in practice.

Summary

This post introduced the "Margin-Enforcing Projection (MEP)" operation: given a score vector, it projects it onto the nearest vector satisfying the requirement that "the minimum score of the positive class is at least $m$ larger than the maximum score of the negative class." This may offer some new ideas for traditional margin learning.

English translation of a post from 科学空间 | Scientific Spaces by 苏剑林. Original: https://kexue.fm/archives/11784
Translated automatically with claude-sonnet-5; all equations are reproduced verbatim from the source. Copyright remains with the original author.