What Should "KL Divergence" Look Like Under GlobalPointer?

GlobalPointer下的“KL散度”应该是怎样的？ 苏剑林 · 2022-04-15 · Mathematics loss function adversarial training NER regularization

A reader recently mentioned wanting to test the effect of combining GlobalPointer with R-Drop, but wasn't sure how to compute the KL divergence for GlobalPointer. Regularization techniques like R-Drop or virtual adversarial training require computing the KL divergence between probability distributions, but GlobalPointer's output is not a probability distribution, so it can't be computed directly.

After some exploration, I arrived at a usable formulation and verified its feasibility through a simple experiment. Let me walk through the analysis.

Symmetric Divergence

KL divergence is a function of two probability distributions, and it is asymmetric—that is, $KL(p\Vert q)$ generally does not equal $KL(q\Vert p)$. In practice, we usually use a symmetrized version of KL divergence:

\begin{equation}D(p,q) = KL(p\Vert q) + KL(q\Vert p)\end{equation}more

Substituting in the definition of KL divergence, $KL(p\Vert q)=\sum\limits_i p_i\log\frac{p_i}{q_i}$, we can simplify this to

\begin{equation}D(p,q) = \sum_i (p_i - q_i)(\log p_i - \log q_i)\end{equation}

Given that $p,q$ is typically obtained via softmax, we define

\begin{equation}p_i = \frac{e^{s_i}}{\sum\limits_j e^{s_j}},\quad q_i = \frac{e^{t_i}}{\sum\limits_j e^{t_j}}\end{equation}

Substituting this in gives

\begin{equation}\begin{aligned} D(p,q) =&\, \sum_i (p_i - q_i)(s_i - t_i) + \sum_i (p_i - q_i)\left(\log\sum_j e^{t_j} - \log\sum_j e^{s_j}\right) \\ =&\, \sum_i (p_i - q_i)(s_i - t_i) + \left(\sum_i p_i - \sum_i q_i\right)\left(\log\sum_j e^{t_j} - \log\sum_j e^{s_j}\right) \\ =&\, \sum_i (p_i - q_i)(s_i - t_i) \end{aligned}\label{eq:kl-0}\end{equation}

Drawing an Analogy

As we can see, at the logits level, the symmetric KL divergence takes the form

\begin{equation}D(s, t) = \sum_i (f(s_i) - f(t_i))(s_i - t_i) = \langle f(s) - f(t), s -t \rangle\label{eq:kl}\end{equation}

where $f$ is the softmax operation and $\langle\cdot,\cdot\rangle$ denotes the inner product of vectors. Structurally, this is the inner product of two vectors: one is the difference of the logits, and the other is the difference of the logits after applying the transformation $f$. What is special about the transformation $f$? We know that softmax is essentially a smooth approximation of $\text{onehot}(\text{argmax}(\cdot))$ (see Musings on Function Smoothing: Differentiable Approximations of Non-Differentiable Functions). For classification, the maximum corresponds to the target class we want to predict, so in plain terms, softmax is really a smooth approximation of "setting the target class to 1 and all non-target classes to 0."

With this abstract viewpoint in hand, we can construct a "KL divergence" for GlobalPointer by analogy. GlobalPointer's output can likewise be understood as logits, but the loss function it uses is the multi-label cross-entropy proposed in Extending "Softmax + Cross-Entropy" to Multi-Label Classification. So this is essentially a question of how to compute KL divergence within the framework of multi-label cross-entropy. Moreover, the target classes output by GlobalPointer are not the ones with the largest logits, but rather all classes whose logits exceed 0.

So, for GlobalPointer, its symmetric divergence can retain the form of equation $\eqref{eq:kl}$, but $f$ should be replaced with a smooth approximation of "setting values greater than 0 to 1, and values less than 0 to 0." The sigmoid function $\sigma(x)=1/(1+e^{-x})$ happens to satisfy exactly this property, so we can design GlobalPointer's symmetric KL divergence as

\begin{equation}D(s, t) = \sum_i (\sigma(s_i) - \sigma(t_i))(s_i - t_i) = \langle \sigma(s) - \sigma(t), s -t \rangle\label{eq:gp-kl}\end{equation}

A Twist

Interestingly, I later discovered that equation $\eqref{eq:gp-kl}$ is in fact equivalent to applying $\sigma$ to each logit separately, computing the KL divergence for each resulting binary probability independently, and then summing them up.

This is easy to prove. Notice that the binary distribution $[\sigma(s),1 - \sigma(s)]$ constructed via the function $\sigma$ is equivalent to the binary distribution built by treating $[s, 0]$ as logits and applying softmax—that is, $[\sigma(s),1 - \sigma(s)]=softmax([s, 0])$. So, by formula $\eqref{eq:kl-0}$, we directly obtain

\begin{equation}\begin{aligned} &\,D\big([\sigma(s_i),1 - \sigma(s_i)],[\sigma(t_i),1 - \sigma(t_i)]\big) \\ =&\,(\sigma(s_i)-\sigma(t_i))(s_i - t_i) + \big((1-\sigma(s_i))-(1-\sigma(t_i))\big)(0 - 0)\\ =&\,(\sigma(s_i)-\sigma(t_i))(s_i - t_i) \end{aligned}\end{equation}

Summing over all components gives us equation $\eqref{eq:gp-kl}$.

This equivalence tells us something interesting: although treating multi-label classification as multiple independent binary classification problems introduces class-imbalance issues, if our goal is merely to evaluate the continuity of the results (rather than to classify), there's no such class-imbalance problem to speak of, since this isn't classification at all. So in this context, it's still valid to treat it as multiple binary classification problems and compute their standard KL divergence.

Experimental Results

Both I and some readers ran simple comparative experiments. The results show that using equation $\eqref{eq:gp-kl}$ as the KL divergence and applying R-Drop to GlobalPointer does yield a modest improvement, whereas directly applying softmax to GlobalPointer's logits and computing the standard KL divergence actually hurts performance. This confirms the reasonableness of equation $\eqref{eq:gp-kl}$.

That said, it should be noted that equation $\eqref{eq:gp-kl}$ merely provides a workable scheme for applying R-Drop or virtual adversarial training within GlobalPointer—whether it actually improves results in a given setting is not guaranteed, just as applying R-Drop to ordinary classification problems doesn't always yield gains either. This requires experimentation, and in particular, careful tuning of the regularization weight coefficient.

Summary

This post mainly discussed how to compute "KL divergence" under GlobalPointer, providing a usable form of KL divergence for applying R-Drop or virtual adversarial training to GlobalPointer.