Self-Orthogonality Module: A Plug-and-Play Kernel Orthogonalization Module

Self-Orthogonality Module：一个即插即用的核正交化模块 苏剑林 · 2020-01-12 · Machine Learning model

A few days ago while browsing Arxiv I came across the paper Self-Orthogonality Module: A Network Architecture Plug-in for Learning Orthogonal Filters (hereafter "the paper"), which looked rather interesting, so I gave it a read. It turned out to be genuinely rewarding, and here I'll record and share my notes.

Adding an orthogonality-inducing regularization term to the kernels of fully-connected or convolutional models is something many models want to do — for instance, the famous BigGAN includes a regularization term of this kind. This paper introduces a new regularization term, and I found the whole analysis process quite interesting and worth reading.

Why do we want orthogonality?

Before we start, let's fix a convention: every one-dimensional vector appearing in this post represents a column vector. Now suppose we have an input sample $d$-dimensional $\boldsymbol{x}\in \mathbb{R}^d$, and when it passes through a fully-connected or convolutional layer, the core computation is:

\begin{equation}\boldsymbol{y}^{\top}=\boldsymbol{x}^{\top}\boldsymbol{W},\quad \boldsymbol{W}\triangleq (\boldsymbol{w}_1,\boldsymbol{w}_2,\dots,\boldsymbol{w}_k)\label{eq:k}\end{equation}

where $\boldsymbol{W}\in \mathbb{R}^{d\times k}$ is a matrix, called the "kernel" (fully-connected kernel / convolutional kernel), and $\boldsymbol{w}_1,\boldsymbol{w}_2,\dots,\boldsymbol{w}_k\in \mathbb{R}^{d}$ are its column vectors.

The above can also be written as

\begin{equation}\boldsymbol{y}=\begin{pmatrix}\boldsymbol{x}^{\top}\boldsymbol{w}_1 \\ \boldsymbol{x}^{\top}\boldsymbol{w}_2\\ \vdots \\ \boldsymbol{x}^{\top}\boldsymbol{w}_k\end{pmatrix}\end{equation}

Intuitively, we can think of $\boldsymbol{w}_1,\boldsymbol{w}_2,\dots,\boldsymbol{w}_k$ as representing $k$ different "viewpoints," and $\boldsymbol{y}$ as the observation of $\boldsymbol{x}$ from each of these $k$ viewpoints.

Given that there are $k$ viewpoints, in order to reduce redundancy among them (and make fuller use of the parameters across all viewpoints), naturally we'd like the different viewpoints to be mutually uncorrelated (as an extreme example, if two viewpoints are identical, we might as well keep only one of them). For vectors in a linear space, being uncorrelated essentially means being orthogonal, so we want

\begin{equation}\boldsymbol{w}_i^{\top}\boldsymbol{w}_j=0,\,\forall i\neq j\end{equation}

This is where the idea of orthogonalization comes from.

Common orthogonalization methods

Orthogonalizing a matrix is somewhat analogous to normalizing a vector, but the difficulty is quite different. For a nonzero vector $\boldsymbol{w}$, normalizing it is as simple as $\boldsymbol{w}/\Vert\boldsymbol{w}\Vert_2$, but there's no similarly simple trick for orthogonalizing a matrix. Readers might think of Gram-Schmidt orthogonalization, but this is fairly costly to compute, and its asymmetry is also an obvious drawback.

Of course, generally speaking we don't necessarily need strict orthogonality, so common approaches to matrix orthogonalization actually add an orthogonality-related regularization term. For instance, for an orthogonal matrix we have $\boldsymbol{W}^{\top}\boldsymbol{W}=\boldsymbol{I}$, so we can add the regularization term

\begin{equation}\left\Vert\boldsymbol{W}^{\top}\boldsymbol{W}-\boldsymbol{I}\right\Vert^2\label{eq:reg0}\end{equation}

Here the norm $\Vert\cdot\Vert$ can be the matrix 2-norm or the matrix $F$-norm (for the concept of matrix norms, see Lipschitz Constraints in Deep Learning: Generalization and Generative Models). Note that this regularization term doesn't just push toward orthogonality — it simultaneously wants normalization too (each vector having unit norm). If we only want orthogonality, we can mask out the diagonal part, giving

\begin{equation}\left\Vert\left(\boldsymbol{W}^{\top}\boldsymbol{W}-\boldsymbol{I}\right)\otimes (1 - \boldsymbol{I})\right\Vert^2\label{eq:reg00}\end{equation}

This is the regularization term added in BigGAN.

The regularization term proposed by the paper

The paper also proposes a new orthogonality regularization term, and it involves some interesting discussion and derivation, along with experiments verifying its effectiveness.

Locality-sensitive hashing

The starting point of the paper is the following lemma:

Let $\boldsymbol{w}_i,\boldsymbol{w}_j\in\mathbb{R}^d$ be two given vectors, $\theta_{i,j}\in[0,\pi]$ their angle, $\mathcal{X}$ the $d$-dimensional unit hypersphere, and $\boldsymbol{x}\sim\mathcal{X}$ denote a vector chosen uniformly at random on $\mathcal{X}$. Then we have:
\begin{equation}\vartheta_{i,j}\triangleq \mathbb{E}_{\boldsymbol{x}\sim\mathcal{X}}\left[\text{sgn}\left(\boldsymbol{x}^{\top}\boldsymbol{w}_i\right)\text{sgn}\left(\boldsymbol{x}^{\top}\boldsymbol{w}_j\right)\right]=1-\frac{2\theta}{\pi}\label{eq:lsh}\end{equation}

where $\text{sgn}$ is the sign function, i.e. $\text{sgn}(x)=\left\{\begin{aligned}1,&\,x > 0\\ -1,&\, x\leq 0\end{aligned}\right.$. This lemma is a direct corollary of "locality-sensitive hashing" for cosine similarity, and locality-sensitive hashing (LSH) originates from the paper Similarity Estimation Techniques from Rounding Algorithms; if you want to trace through the proof, that's the route to follow.

At first glance $\eqref{eq:lsh}$ looks like just an ordinary mathematical result, but in fact it carries much richer implications: it lets us (approximately) convert the similarity between two continuous real-valued vectors into the similarity between two binary vectors (−1 and 1). Once converted into binary vectors, this is essentially equivalent to converting into a "word-document" matrix, which allows us to build an index to speed up retrieval. In other words, this can effectively speed up retrieval of continuous real-valued vectors!

The form of the optimization objective

Looking directly at the definition of expression $\eqref{eq:lsh}$, its derivative is identically zero, but we can obtain some smooth approximation to it. Suppose we've already obtained a smooth approximation to $\vartheta$; then we can use it to construct an orthogonality regularization term. The regularization term the paper constructs is:

\begin{equation}\mathcal{R}_{\vartheta}\triangleq \lambda_1\left(\sum_{i\neq j}\vartheta_{i,j}\right)^2 + \lambda_2\sum_{i\neq j}\vartheta_{i,j}^2\label{eq:reg}\end{equation}

Clearly, this regularization term wants $\vartheta_{i,j}=0$, and $\vartheta_{i,j}=0$ means $\theta_{i,j}=\pi/2$, i.e. that the vectors of $\boldsymbol{W}$ are pairwise orthogonal. By comparison, the regularization controlled by $\lambda_1$ is gentler — it only asks that the mean of $\vartheta_{i,j}$ be zero — while $\lambda_2$ is stricter, requiring all the $\theta_{i,j}$ to equal zero.

Considering that real-world problems can be fairly complex, we shouldn't impose overly strict constraints on the model, so the paper sets $\lambda_1 > \lambda_2$, with a specific value of $\lambda_1 = 100, \lambda_2 = 1$.

Plugging it into the model

Now let's consider the practical problem of estimating $\vartheta_{i,j}$.

First, let's look at expression $\eqref{eq:lsh}$ from a different angle. Suppose we sample $b$ samples $\boldsymbol{x}_1,\boldsymbol{x}_2,\dots,\boldsymbol{x}_b$ to estimate $\vartheta_{i,j}$; then we have

\begin{equation}\begin{aligned}\vartheta_{i,j}\approx&\frac{1}{b}\sum_{\alpha=1}^b\left[\text{sgn}\left(\boldsymbol{x}_{\alpha}^{\top}\boldsymbol{w}_i\right)\text{sgn}\left(\boldsymbol{x}_{\alpha}^{\top}\boldsymbol{w}_j\right)\right]\\ =&\left(\frac{\boldsymbol{y}_i}{\Vert\boldsymbol{y}_i\Vert_2}\right)^{\top}\left(\frac{\boldsymbol{y}_j}{\Vert\boldsymbol{y}_j\Vert_2}\right) \end{aligned}\label{eq:lsh-2}\end{equation}

where

\begin{equation}\boldsymbol{y}=\begin{pmatrix} \text{sgn}\left(\boldsymbol{x}_{1} ^{\top}\boldsymbol{w}\right)\\ \text{sgn}\left(\boldsymbol{x}_{2}^{\top}\boldsymbol{w}\right)\\ \vdots\\ \text{sgn}\left(\boldsymbol{x}_{b}^{\top}\boldsymbol{w}\right) \end{pmatrix}=\text{sgn}\left(\boldsymbol{X}^{\top}\boldsymbol{w}\right),\,\,\boldsymbol{X}=(\boldsymbol{x}_1,\boldsymbol{x}_2,\dots,\boldsymbol{x}_b)\in\mathbb{R}^{d\times b}\end{equation}

The cleverest part of this reformulation is that, since the entries of $\boldsymbol{y}$ are either 1 or −1, the norm of $\boldsymbol{y}$ is exactly $\sqrt{b}$, so the factor $1/b$ is exactly equivalent to normalizing $\boldsymbol{y}_i,\boldsymbol{y}_j$!

It's also worth noting that neither $\eqref{eq:lsh}$ nor $\eqref{eq:lsh-2}$ actually depends on the norms of the individual $\boldsymbol{x}_{\alpha}$, since $\text{sgn}(x)=\text{sgn}(|\lambda|x)$. The reason the earlier lemma requires sampling on the "unit hypersphere" is only to emphasize the uniformity of the sampling direction (rather than the norm).

With this understanding, we can now lay out the estimation procedure for $\vartheta_{i,j}$ clearly:

Estimation procedure for $\vartheta_{i,j}$
1. Randomly initialize a $d\times b$ matrix $\boldsymbol{X}$ (viewed as $b$ vectors of dimension $d$, with unrestricted norm but as uniform a direction as possible);
2. Compute $\boldsymbol{X}^{\top}\boldsymbol{w}_i, \boldsymbol{X}^{\top}\boldsymbol{w}_j$ to obtain two $b$-dimensional vectors, then apply the $\text{sgn}$ activation, then perform $l_2$ normalization on each separately, and finally compute the inner product;
3. If a smooth approximation is needed, one can use $\text{sgn}(x)\approx \tanh(\gamma x)$; the paper uses $\gamma=10$.

How should $\boldsymbol{X}$ be chosen? The paper directly chooses it to be the input of the current batch. Going back to $\eqref{eq:k}$, generally speaking the input to a neural network is just a matrix of shape $b\times d$, and we can treat it as $\boldsymbol{X}^{\top}$, in which case $b$ is the batch size; the network then multiplies it by $\boldsymbol{W}\in \mathbb{R}^{d\times k}$ to obtain the output $\boldsymbol{Y}\in\mathbb{R}^{b\times k}$, which corresponds exactly to the $k$ vectors of dimension $b$, $\boldsymbol{X}^{\top}\boldsymbol{w}_1,\boldsymbol{X}^{\top}\boldsymbol{w}_2,\dots,\boldsymbol{X}^{\top}\boldsymbol{w}_k$, computed from the $k$ kernel vectors $\boldsymbol{w}_1,\boldsymbol{w}_2,\dots,\boldsymbol{w}_n$ in the "estimation procedure for $\vartheta_{i,j}$." This way, we save most of the computation in the estimation procedure for $\vartheta_{i,j}$ — we can estimate directly from the current layer's output.

Note: if readers look at the original paper, they'll find that its description in this part doesn't quite match the description in this post (mainly the two paragraphs just above Section 3, Experiments). Based on my understanding of the overall logic of the paper, I believe the paper's description in that passage is mistaken (mainly the meanings of $D$ and $d$ get mixed up), and the version given in this post is the correct one.

In summary, the final scheme for estimating $\vartheta_{i,j}$ is:

1. The current layer's input $\boldsymbol{X}^{\top}\in \mathbb{R}^{b\times d}$ and kernel matrix $\boldsymbol{W}\in \mathbb{R}^{d\times k}$ are multiplied to give output $\boldsymbol{Y}\in\mathbb{R}^{b\times k}$;
2. Apply $\tanh(\gamma x)$ activation to $\boldsymbol{Y}\in\mathbb{R}^{b\times k}$, then normalize along the $b$ dimension (i.e. the batch-size dimension) with $l_2$;
3. Compute $\boldsymbol{Y}^{\top}\boldsymbol{Y}$, obtaining a matrix of shape $k\times k$ — this is the set of all $\vartheta_{i,j}$.
4. Once we have $\vartheta_{i,j}$, we can plug it into expression $\eqref{eq:reg}$ to compute the regularization term. Since this regularization term is built using the model's own outputs, it's called the "self-orthogonality regularization term."

Connection to BN

Additionally, the authors of the paper speculate that "normalizing along the $b$ dimension (i.e., the batch-size dimension) with $l_2$" is somewhat similar to BN, so that after adding the self-orthogonality regularization term, the model might no longer need BN. Personally, I find this conjecture a bit of a stretch, since this operation is only used when computing the regularization term and doesn't affect the model's normal forward pass, so BN's necessity can't be ruled out. Furthermore, in the "estimation procedure for $\vartheta_{i,j}$" itself, we require the directions of the individual vectors of $\boldsymbol{X}$ to be as uniform as possible, but later we directly take the (transposed) input of the current layer as $\boldsymbol{X}$, which can't effectively guarantee direction uniformity — and once BN is added, it should in theory help make the directions of the input vectors more uniform, which makes BN even less dispensable. In fact, the paper's own experiments don't fully support the authors' conjecture either.

Experiments and my own analysis

After all this length and all these formulas, we've finally derived the regularization term from the paper. The authors did run a fair number of experiments to verify its effectiveness, and the overall conclusion is that it does bring the distribution of pairwise angles between kernel vectors closer to pairwise orthogonality, and it also brings a certain (modest) improvement — unlike existing orthogonality regularization terms, which enforce orthogonality but usually come at the cost of accuracy.

I'll leave the specific experimental results for readers to check in the original paper — reproducing them here wouldn't add much. That said, although the authors ran quite a few experiments, I still feel they're not thorough enough, since most of the experiments are on point clouds, and the only conventional classification experiment is on CIFAR-10, which feels too limited.

Finally, why might this orthogonality regularization term be (seemingly) more effective? My personal guess is that it's because the new regularization term is relatively gentler. Whether it's $\eqref{eq:reg0}$ or $\eqref{eq:reg00}$, they both penalize individual inner products (angles) one at a time, whereas the paper's $\eqref{eq:reg}$ tends instead to achieve the orthogonality penalty from a more holistic view of the overall angle distribution. Moreover, the new regularization term involves $\tanh$, which has a saturation region — meaning that, much like hinge loss, it truncates the penalty, which further makes the penalty gentler.

A brief summary

This post has mainly given a simple introduction to a recent Arxiv paper. The paper points out that existing orthogonality regularization terms fail to improve model accuracy, so the authors introduce a new orthogonality regularization term and evaluate it accordingly, concluding that their regularization term not only promotes orthogonality but also brings a certain improvement in results.

Finally, since I hadn't previously looked into this area (especially the "locality-sensitive hashing" part earlier on), and only stumbled upon this paper on Arxiv by chance and found it quite interesting, I decided to share my notes here. If there are any mistakes or omissions, I ask for readers' understanding and welcome corrections.