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

Locality Sensitive Hashing

The starting point of the original 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] be the angle between them, \mathcal{X} be the d-dimensional unit hypersphere, and \boldsymbol{x} \sim \mathcal{X} represent a vector randomly selected from \mathcal{X}. We then have the following result: \vartheta_{i,j} \triangleq \mathbb{E}_{\boldsymbol{x}\sim\mathcal{X}}\left[\operatorname{sgn}\left(\boldsymbol{x}^{\top}\boldsymbol{w}_i\right)\operatorname{sgn}\left(\boldsymbol{x}^{\top}\boldsymbol{w}_j\right)\right]=1-\frac{2\theta}{\pi} \label{eq:lsh}

Where \operatorname{sgn} is the sign function, i.e., \operatorname{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" (LSH) for cosine similarity. LSH originates from the paper "Similarity Estimation Techniques from Rounding Algorithms". If one wishes to trace the proof, they can follow that path.

At first glance, Eq. [eq:lsh] is just an ordinary mathematical formula. However, it actually contains richer meaning. It allows us to (approximately) transform the similarity of two continuous real-valued vectors into the similarity of two binary vectors (-1 and 1). Once transformed into binary vectors, it is equivalent to transforming into a "term-document" matrix, which allows us to build indices to accelerate retrieval. In other words, this can effectively improve the retrieval speed of continuous real-valued vectors!

Form of the Optimization Objective

Looking directly at the definition in Eq. [eq:lsh], its derivative is always 0. However, we can obtain a smooth approximation of it. Assuming we have obtained a smooth approximation of \vartheta, we can use it to construct an orthogonal regularization term. The regularization term constructed in the original paper is: \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} Clearly, this regularization term encourages \vartheta_{i,j}=0. Since \vartheta_{i,j}=0 implies \theta_{i,j}=\pi/2, this means the vectors in \boldsymbol{W} are pairwise perpendicular. Relatively speaking, the term controlled by \lambda_1 is softer, as it only requires the mean of \vartheta_{i,j} to be 0, while \lambda_2 is more rigid, requiring all \theta_{i,j} to be 0.

Considering that practical problems may be complex, we should not impose overly rigid constraints on the model. Therefore, the original paper sets \lambda_1 > \lambda_2, specifically \lambda_1 = 100, \lambda_2 = 1.

Integration into the Model

Now let us consider the practical estimation of \vartheta_{i,j}.

First, let’s understand Eq. [eq:lsh] from another perspective. If we sample b samples \boldsymbol{x}_1, \boldsymbol{x}_2, \dots, \boldsymbol{x}_b to estimate \vartheta_{i,j}, we have: \begin{aligned} \vartheta_{i,j}\approx&\frac{1}{b}\sum_{\alpha=1}^b\left[\operatorname{sgn}\left(\boldsymbol{x}_{\alpha}^{\top}\boldsymbol{w}_i\right)\operatorname{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} Here: \boldsymbol{y}=\begin{pmatrix} \operatorname{sgn}\left(\boldsymbol{x}_{1} ^{\top}\boldsymbol{w}\right)\\ \operatorname{sgn}\left(\boldsymbol{x}_{2}^{\top}\boldsymbol{w}\right)\\ \vdots\\ \operatorname{sgn}\left(\boldsymbol{x}_{b}^{\top}\boldsymbol{w}\right) \end{pmatrix}=\operatorname{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} The cleverest part of this transformation is that since the elements of \boldsymbol{y} are either 1 or -1, the length of \boldsymbol{y} is exactly \sqrt{b}. Therefore, the factor 1/b is exactly equivalent to normalizing both \boldsymbol{y}_i and \boldsymbol{y}_j!

Furthermore, it is worth noting that neither Eq. [eq:lsh] nor Eq. [eq:lsh-2] is related to the magnitude of each \boldsymbol{x}_{\alpha}, because \operatorname{sgn}(x)=\operatorname{sgn}(|\lambda|x). The reason the previous lemma required sampling on the "unit hypersphere" was merely to emphasize the uniformity of the sampling direction (rather than the magnitude).

With this understanding, we can clarify the estimation process for \vartheta_{i,j}:

Estimation Process for \vartheta_{i,j}

1. Randomly initialize a d \times b matrix \boldsymbol{X} (when viewed as b d-dimensional vectors, the magnitude is unrestricted, but the directions should be as uniform as possible);

2. Calculate \boldsymbol{X}^{\top}\boldsymbol{w}_i and \boldsymbol{X}^{\top}\boldsymbol{w}_j to obtain two b-dimensional vectors, activate them with the \operatorname{sgn} function, perform l_2 normalization on each, and finally calculate their inner product;

3. If a smooth approximation is required, use \operatorname{sgn}(x) \approx \tanh(\gamma x). The original paper used \gamma=10.

How should \boldsymbol{X} be chosen? The original paper directly chooses it as the input of the current batch. Returning to Eq. [eq:k], generally, the input to a neural network is a b \times d matrix, which we can treat as \boldsymbol{X}^{\top}. In this case, b is the batch size. The neural network then multiplies this by \boldsymbol{W} \in \mathbb{R}^{d \times k} to obtain the output \boldsymbol{Y} \in \mathbb{R}^{b \times k}. This exactly corresponds to the k b-dimensional vectors \boldsymbol{X}^{\top}\boldsymbol{w}_1, \boldsymbol{X}^{\top}\boldsymbol{w}_2, \dots, \boldsymbol{X}^{\top}\boldsymbol{w}_k calculated for the k kernel vectors in the "Estimation Process." In this way, we even save most of the computational work in the estimation process, as we can estimate directly based on the current layer’s output.

Note: If readers look at the original paper, they will find that the description in that part differs from the description in this blog (mainly in the two paragraphs above Section 3 Experiments in the original paper). Based on my understanding of the overall logic of the article, I believe the description in those paragraphs of the original paper is incorrect (mainly confusing the meanings of D and d), while the version in this blog is correct.

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

1. The input of the current layer is \boldsymbol{X}^{\top} \in \mathbb{R}^{b \times d}, and the kernel matrix is \boldsymbol{W} \in \mathbb{R}^{d \times k}. After matrix multiplication, the output is \boldsymbol{Y} \in \mathbb{R}^{b \times k};

2. Activate \boldsymbol{Y} \in \mathbb{R}^{b \times k} using \tanh(\gamma x), and then perform l_2 normalization along the b dimension (i.e., the batch size dimension);

3. Calculate \boldsymbol{Y}^{\top}\boldsymbol{Y} to obtain a k \times k matrix, which contains all \vartheta_{i,j};

4. Once \vartheta_{i,j} are obtained, substitute them into Eq. [eq:reg] to calculate the regularization term. Since the regularization term is constructed using the model’s own output, it is called the "Self-Orthogonalization Regularization Term."

Connection to BN

Additionally, the authors of the original paper speculate that the operation of "performing l_2 normalization along the b dimension (batch size dimension)" is somewhat similar to Batch Normalization (BN). Therefore, after adding the self-orthogonalization regularization term, the model might not need BN. Personally, I find this speculation a bit forced because this operation is only used when calculating the regularization term and does not affect the model’s normal forward propagation process; thus, the necessity of BN cannot be ruled out. Furthermore, in the original "estimation process," we require the directions of the vectors in \boldsymbol{X} to be as uniform as possible. However, when we directly select the (transpose of) the current layer’s input as \boldsymbol{X}, we cannot effectively guarantee directional uniformity. Adding BN theoretically helps make the directions of the input vectors more uniform, making BN even more indispensable. In fact, the experiments in the original paper do not fully support the authors’ speculation.