Transformer Upgrade Path: 6. A Completeness Analysis of Rotary Position Embedding

Transformer升级之路：6、旋转位置编码的完备性分析 苏剑林 · 2022-12-28 · Machine Learning matrix attention positional encoding RoPE

In last year's article Transformer Upgrade Path: 2. Rotary Position Embedding, Drawing on the Best of Many Ideas, I introduced rotary position embedding (RoPE). Back then, my starting point was simply the feeling that implementing relative position via absolute position was "a rather fun thing to do" — I hadn't expected its practical performance to turn out so good, or that it would be so widely adopted. I have to say it was a genuinely pleasant surprise. Later, in Transformer Upgrade Path: 4. Two-Dimensional Rotary Position Embedding, I discussed the two-dimensional form of RoPE and studied the general solution of RoPE expressed via the matrix exponential.

Now that we have a general solution, a natural question arises: the RoPE we commonly use is just a block-diagonal matrix built from 2×2 rotation matrices as its basic unit. If we instead used the general solution, would the performance theoretically be even better? This post sets out to answer that question.

The General Exponential Solution

In Transformer Upgrade Path: 4. Two-Dimensional Rotary Position Embedding, we abstractly defined RoPE as any square matrix satisfying the equation below

\begin{equation}\boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n=\boldsymbol{\mathcal{R}}_{n-m}\label{eq:re}\end{equation}more

We then explored the following solution in matrix-exponential form

\begin{equation}\boldsymbol{\mathcal{R}}_n=\exp n\boldsymbol{B}\end{equation}

Here, the matrix exponential is not an element-wise activation-style operation like softmax, but the "Matrix Exponential" defined via a Taylor series. By the "Baker–Campbell–Hausdorff formula", we have

\begin{equation}\begin{aligned} \boldsymbol{\mathcal{R}}_m^{\top}\boldsymbol{\mathcal{R}}_n=&\,\big(\exp m\boldsymbol{B}\big)^{\top}\big(\exp n\boldsymbol{B}\big) = \big(\exp m\boldsymbol{B}^{\top}\big)\big(\exp n\boldsymbol{B}\big) \\ =&\,\exp\left(m\boldsymbol{B}^{\top} + n\boldsymbol{B} + \frac{1}{2}mn\left[\boldsymbol{B}^{\top}, \boldsymbol{B}\right]+\cdots\right) \end{aligned}\end{equation}

Here, both the terms omitted in $\left[\boldsymbol{A}, \boldsymbol{B}\right]=\boldsymbol{A}\boldsymbol{B}-\boldsymbol{B}\boldsymbol{A}$ and in $\cdots$ are of third order or higher in $m,n$. According to equation $\eqref{eq:re}$, the exponent on the right-hand side above should equal $(n-m)\boldsymbol{B}$, from which it follows that

\begin{equation}\boldsymbol{B}^{\top} = - \boldsymbol{B}\end{equation}

that is, $\boldsymbol{B}$ is required to be an antisymmetric matrix.

The General Orthogonal Solution

Going further, we have $(\exp \boldsymbol{B})^{\top}(\exp \boldsymbol{B})=\exp(\boldsymbol{B}-\boldsymbol{B}) = \boldsymbol{I}$ and $\exp n\boldsymbol{B} = \left(\exp\boldsymbol{B}\right)^n$. The former tells us that $\exp\boldsymbol{B}$ is an orthogonal matrix, while the latter suggests the question: can this be generalized to an arbitrary orthogonal matrix? It's not hard to verify that the answer is yes, and we arrive at the following conclusion:

For any orthogonal matrix $\boldsymbol{O}$, $\boldsymbol{\mathcal{R}}_n=\boldsymbol{O}^n$ is a solution satisfying equation $\eqref{eq:re}$.

It's worth pointing out that, over the reals, not every orthogonal matrix can be written in the form $\exp\boldsymbol{B}$, so $\boldsymbol{O}^n$ is in fact a broader class of solutions than the matrix-exponential form. From Appreciating the Identity det(exp(A)) = exp(Tr(A)) we know that $\det(\exp(\boldsymbol{A})) = \exp(\text{Tr}(\boldsymbol{A})) > 0$, so an orthogonal matrix that can be written in matrix-exponential form must necessarily have determinant greater than 0 (i.e., equal to 1). It turns out the converse also holds: any orthogonal matrix with determinant equal to 1 can necessarily be written in the form $\exp\boldsymbol{B}$, where $\boldsymbol{B}$ is an antisymmetric matrix. (See Why can any orthogonal matrix be written as O=e^A).

For an orthogonal matrix with $\det(\boldsymbol{O}) = -1$, we have $\boldsymbol{O}=\boldsymbol{O}_+ \boldsymbol{I}_-$, where $\boldsymbol{I}_-$ is a diagonal matrix with one diagonal entry equal to -1 and the rest equal to 1, and $\boldsymbol{O}_+$ is an orthogonal matrix with $\det(\boldsymbol{O}_+) = 1$, which can be written in the form $\exp\boldsymbol{B}$, giving us $\boldsymbol{O}^n = (\boldsymbol{O}_+ \boldsymbol{I}_-)^n = \boldsymbol{I}_-^n\exp n\boldsymbol{B}$ in this case. In other words, even for $\det(\boldsymbol{O}) = -1$ of type $\boldsymbol{O}^n$, it amounts to nothing more than a simple transformation of $\exp n\boldsymbol{B}$, so going forward we'll mainly focus on solutions of the form $\exp n\boldsymbol{B}$.

Completeness Analysis

As we all know, the RoPE position encoding commonly used in practice is a block-diagonal matrix of the following form:

\begin{equation}\scriptsize{\left(\begin{array}{cc:cc:cc:cc} \cos n\theta_0 & -\sin n\theta_0 & 0 & 0 & \cdots & \cdots & 0 & 0 \\ \sin n\theta_0 & \cos n\theta_0 & 0 & 0 & \cdots & \cdots & 0 & 0 \\ \hdashline 0 & 0 & \cos n\theta_1 & -\sin n\theta_1 & \cdots & \cdots & 0 & 0 \\ 0 & 0 & \sin n\theta_1 & \cos n\theta_1 & \cdots & \cdots & 0 & 0 \\ \hdashline \vdots & \vdots & \vdots & \vdots & \ddots & \ddots & \vdots & \vdots \\ \vdots & \vdots & \vdots & \vdots & \ddots & \ddots & \vdots & \vdots \\ \hdashline 0 & 0 & 0 & 0 & \cdots & \cdots & \cos n\theta_{d/2-1} & -\sin n\theta_{d/2-1} \\ 0 & 0 & 0 & 0 & \cdots & \cdots & \sin n\theta_{d/2-1} & \cos n\theta_{d/2-1} \\ \end{array}\right)}\end{equation}

which can be abbreviated as

\begin{equation}\begin{pmatrix} \boldsymbol{R}_{n\theta_0} & \boldsymbol{0} & \cdots & \boldsymbol{0} \\ \boldsymbol{0} & \boldsymbol{R}_{n\theta_1} & \cdots & \boldsymbol{0} \\ \vdots & \vdots & \ddots & \vdots \\ \boldsymbol{0} & \boldsymbol{0} & \cdots & \boldsymbol{R}_{n\theta_{d/2-1}} \\ \end{pmatrix} = \exp n\begin{pmatrix} \boldsymbol{J}_{\theta_0} & \boldsymbol{0} & \cdots & \boldsymbol{0} \\ \boldsymbol{0} & \boldsymbol{J}_{\theta_1} & \cdots & \boldsymbol{0} \\ \vdots & \vdots & \ddots & \vdots \\ \boldsymbol{0} & \boldsymbol{0} & \cdots & \boldsymbol{J}_{\theta_{d/2-1}} \\ \end{pmatrix}\end{equation}

where

\begin{equation}\boldsymbol{R}_{\theta} = \begin{pmatrix} \cos\theta & -\sin\theta \\ \sin\theta & \cos\theta \end{pmatrix},\quad \boldsymbol{J}_{\theta} = \begin{pmatrix} 0 & -\theta \\ \theta & 0\end{pmatrix}\end{equation}

This choice is arguably the simplest possible one, and the fundamental reason for it is to reduce computational cost. So the so-called completeness question is really asking: compared to the fully general $\exp n\boldsymbol{B}$, does this special case of a block-diagonal matrix lack any expressive capability? In other words, if we set aside computational cost and replace $\boldsymbol{B}$ with a general antisymmetric matrix, might performance actually improve?

This question isn't hard to answer. In fact, any antisymmetric matrix of even order can be diagonalized into the block-diagonal form

\begin{equation}\boldsymbol{\Lambda} = \begin{pmatrix} \boldsymbol{J}_{\theta_0} & \boldsymbol{0} & \cdots & \boldsymbol{0} \\ \boldsymbol{0} & \boldsymbol{J}_{\theta_1} & \cdots & \boldsymbol{0} \\ \vdots & \vdots & \ddots & \vdots \\ \boldsymbol{0} & \boldsymbol{0} & \cdots & \boldsymbol{J}_{\theta_{d/2-1}} \\ \end{pmatrix}\end{equation}

This result can be found in reference Skew-symmetric matrix. That is, there exists an invertible matrix $\boldsymbol{P}$ such that $\boldsymbol{B}=\boldsymbol{P}\boldsymbol{\Lambda}\boldsymbol{P}^{-1}$, and thus

\begin{equation}\exp n\boldsymbol{B} = \exp \left(n\boldsymbol{P}\boldsymbol{\Lambda}\boldsymbol{P}^{-1}\right) = \boldsymbol{P}(\exp n\boldsymbol{\Lambda})\boldsymbol{P}^{-1}\end{equation}

In other words, any general $\exp n\boldsymbol{B}$ differs from the block-diagonal $\exp n\boldsymbol{\Lambda}$ only by a similarity transformation. And when we apply RoPE in self-attention, we have

\begin{equation}\boldsymbol{q}^{\top}\big(\exp (n-m)\boldsymbol{B}\big)\boldsymbol{k} = \big(\boldsymbol{P}^{\top}\boldsymbol{q}\big)^{\top}\big(\exp (n-m)\boldsymbol{\Lambda}\big)\big(\boldsymbol{P}^{-1}\boldsymbol{k}\big)\end{equation}

Since $\boldsymbol{q},\boldsymbol{k}$ generally comes from the input $\boldsymbol{x}$ via some learnable linear transformation, $\boldsymbol{P}^{\top},\boldsymbol{P}^{-1}$ can in principle be absorbed into the trainable parameters of that linear transformation. So setting it directly to $\boldsymbol{q}^{\top}\big(\exp (n-m)\boldsymbol{\Lambda}\big)\boldsymbol{k}$ theoretically loses no generality.

So for self-attention, the answer to the question is no. However, if we consider linear attention, the answer is slightly different, because in linear attention an activation function is applied to $\boldsymbol{q},\boldsymbol{k}$:

\begin{equation}\phi(\boldsymbol{q})^{\top}\big(\exp (n-m)\boldsymbol{B}\big)\varphi(\boldsymbol{k}) = \big(\boldsymbol{P}^{\top}\phi(\boldsymbol{q})\big)^{\top}\big(\exp (n-m)\boldsymbol{\Lambda}\big)\big(\boldsymbol{P}^{-1}\varphi(\boldsymbol{k})\big)\end{equation}

This means that $\boldsymbol{P}^{\top},\boldsymbol{P}^{-1}$ cannot necessarily be absorbed into the trainable parameters of the linear transformation, so adding two extra parameter matrices to linear attention could plausibly bring an improvement.

Summary

This post has given a brief analysis of the completeness of RoPE, showing that for self-attention, the currently used block-diagonal form of RoPE loses no generality.