R ELATIVE P OSITIONAL E NCODING F AMILY VIA U NITARY T RANSFORMATION
Anonymous authors
Paper under double-blind review
A BSTRACT
Relative positional encoding is widely used in vanilla and linear transformers to represent positional information.
However, existing encoding methods of a vanilla transformer are not always directly applicable to a linear transformer, because the latter requires a decomposition of the query and key representations into separate kernel functions.
Nevertheless, principles to design encoding methods suitable for linear transformers remain under-studied.
In this work, we put together a variety of existing encoding approaches under a canonical form and further propose a family of relative positional encoding algorithms via unitary transformation .
Our formulation leads to a principled framework that can be used to develop new relative positional encoding methods that preserve the linear space-time complexity.
Equipped with different parameters, the proposed linearized relative positional encoding (LRPE) family derives effective encoding for various applications.
Experiments show that compared with existing methods, LRPE achieves competitive performance on language modeling and various challenging downstream tasks, e .
g ., machine translation and text classification.
In the meantime, it highlights a general paradigm to design broadly more relative positional encoding methods, applicable inclusively to linear and vanilla transformers.
I NTRODUCTION
Transformers have achieved remarkable progress in natural language processing (Devlin et al., 2019;
Radford et al., 2019; Brown et al., 2020), computer vision (Dosovitskiy et al., 2020;
Liu et al., 2021; Arnab et al., 2021) and audio processing (Gulati et al., 2020).
As an important ingredient in transformers, positional encoding assigns a unique representation for each position of a token in a sequence so that the transformers can sense the position of input tokens.
Among these encoding methods, absolute positional encoding (Vaswani et al., 2017; Sukhbaatar et al., 2015; Devlin et al., 2019; Liu et al., 2020) maps each individual position index into a continuous encoding.
Whereas relative positional encoding (Shaw et al., 2018; Su et al., 2021; Horn et al., 2021; Liutkus et al., 2021;
Huang et al., 2020; Raffel et al., 2019) generates encoding for each query-key pair, representing their relative positional offset.
We focus on relative positional encoding as they are not constrained by input lengths (Chen, 2021) while showing superior performance (Shaw et al., 2018).
Linear transformers Chen (2021); Qin et al.
; Su et al.
(2021) attract more attention recently as they can achieve linear space-time complexity with respect to input sequence length, while maintaining comparable performance with vanilla transformers.
Most existing linear transformers use absolute positional encoding methods to encode positional information, since most existing relative positional encoding methods are designed for vanilla transformers and are not directly applicable to linear transformers.
The main cause behind this limitation is that linear transformers decompose key and value representations in the self-attention modules into separate kernel functions to achieve linear space-time complexity.
Such an additional requirement on the decomposibility is not always satisfied by existing relative positional encoding methods.
On the other hand, despite some individual works (Qin et al., 2022; Chen, 2021), general principles to design relative positional encoding for linear transformers remain largely under-studied.
A recent work, RoPE Su et al. (2021) proposes a new set of multiplicative encoding solutions based on rotate positional encoding and can be applied to linear transformers. In Section C.7, we show that RoPE can be seen as a special form of LRPE.
[Figure]
Q
Tensors in the same dashed line are associated for computation.
In the vanilla relative positional encoding, query key attention has to be calculated first, leading to a quadratic complexity.
W t − s refers to relative positional encoding, where t, s are two positional indices on the query and key, respectively.
Our LRPE achieves a decomposable encoding, i .
e ., W t and W s are only dependent on positions of the query and key, making it fully compatible with linear transformers.
When dealing with long sequences, d ≪ n , the computation complexity is dominated by n , rendering d negligible.
In this work, we aim to bridge this gap and study principal framework to develop relative positional encoding applicable for both linear and vanilla transformers. To this end, we start by presenting a canonical form of relative positional encoding, which reveals that differences in existing encoding methods boil down to choices of a set of query, key and relative positional matrix primitives .
By properly selecting and composing these primitives, we could derive various existing encoding methods for vanilla (Vaswani et al., 2017; Huang et al., 2020; Shaw et al., 2018) and linear (Qin et al., 2022) transformers.
Taking advantage of the canonical form, we introduce the main contribution of our work, i .
e ., a special family of relative positional encoding methods called linearized relative positional encoding (LRPE).
Specifically, we supply a sufficient condition to design compatible encoding methods specially for linear transformers and prove that the linearized relative positional encoding is unitary transformation.
Benefits of using unitary transformation are two-fold.
On one side, since it is derived from the decomposable positional matrix, it can maintain the linear space-time complexity as shown in Fig.
Second, the property of the unitary transformation allows us to effectively derive the family of closed-form solutions.
In particular, we show that a number of encoding methods pertain to the LRPE family, including those used in RoPE (Su et al., 2021) and PermuteFormer (Chen, 2021).
Furthermore, LRPE sheds light on a simple yet flexible theoretical paradigm to develop new effective relative positional encodings.
To demonstrate this, we derive non-exhaustively three additional
LRPE encoding methods by parameterizing the generic solution differently, including solutions living in either real or complex domains.
Since unitary transformations are special cases of relative positional matrix, LRPE are applicable in both linear and vanilla transformers, and exclusively suitable within encoder and/or decoder layers.
We experimentally demonstrate the effectiveness of the LRPE family on autoregressive and bidirectional language modelling, and on challenging downstream tasks, including machine translation and text classification.
Results show that LRPE achieves competitive capability in representing relative positional information, commonly resulting in superior performance than previous encoding methods.
In summary, our main contributions are three-fold:
• We present a canonical form of relative positional encoding, which derives most existing relative positional encoding methods as its special case, including those used in linear and vanilla transformers.
• Based on the canonical form, we propose linearized relative position encoding (LRPE), a simple yet principal formulation to derive an encoding family that respect the linear spacetime complexity in linear transformers, while being also applicable to vanilla transformers.
We show several existing relative positional encoding methods in linear transformers are in LRPE family.
We also provide additional particular solutions from this generic form.
• Experiments on various downstream tasks, including language modeling, machine translation and text classification show that the LRPE family show more robust and commonly superior results across tasks than previous relative encoding methods, are flexible in be- ing plugged into linear/vanilla models, in encoder and/or decoder layers.
In addition, it is generic to derive existing and potentially new encoding methods.
B ACKGROUND AND P RELIMINARY
In this section, we provide preliminary knowledge and describe related work to facilitate the rest discussions.
In the following, we denote the k -th row of matrix M as m T k , the d -dimensional identity matrix as I d .
We omit the subscript d when it is unambiguous from the context.
The complete list of notations can be found in Appendix A.
2. T RANSFORMER AND ITS LINEARIZATION
We first briefly review vanilla transformer (Vaswani et al., 2017) and its linearization (Katharopoulos et al., 2020).
The key component of transformer models is the self-attention block, which involves three matrices Q ( Query ), K ( Key ) and V ( Value ); each of them is a linear projection taking X ∈ R n × d as input:
[Equation]
The output O ∈ R n × d is computed using the Softmax weighted sum:
[Equation]
The computation overhead of the vanilla transformer grows quadratically with respect to the sequence length n , which becomes the bottleneck for transformers to handle long input sequences.
Linearization of self-attention aims to reduce the computation complexity to linear (Katharopoulos et al., 2020; Ke et al., 2021;
Qin et al., 2022; Vyas et al., 2020; Peng et al., 2021; Xiong et al., 2021), typically achieved via a decomposable kernel function ϕ : R d → R ¯ d .
Specifically, the output of linear attention is computed as:
[Equation]
The key property of linear attention is the decomposability of the kernel function.
This enables to compute ϕ ( K ) T V ∈ R
d
× d first, which leads to the O ( nd 2 ) complexity, further reducing to O ( n ) with longer inputs ( d ≪ n ).
See Appendix B for a detailed discussion.
2. P OSITIONAL ENCODING
Self-attention is capable of parallel sequence processing but cannot capture positional information of each token.
To address this issue, positional encoding methods are proposed, which can be generally categorized into two groups: absolute positional encoding and relative positional encoding.
Absolute positional encoding employs handcraft functions (Vaswani et al., 2017; Sukhbaatar et al., 2015) or learnable encoding lookup tables P ∈ R n × d (Devlin et al., 2019; Liu et al., 2020) to represent position indices as encodings.
These encodings are then combined with the context vector additively:
[Equation]
where the encoding formulation only depends on the absolute position index s , and the positional encoding size is restricted by the input sequence length.
Relative positional encoding considers relative position offsets between two input tokens (Shaw et al., 2018), i .
e .,
[Equation]
where s, t are the two positional indexes, e st denotes the attention score before softmax.
Compared to absolute positional encoding, relative positional encoding generally achieves better performance as it can handle variable input length (Chen, 2021).
However, extra cost on computation and memory makes it not so efficient than absolute positional encoding (Likhomanenko et al., 2021).
Most existing relative positional encoding methods (Raffel et al., 2019; Shaw et al., 2018; Huang et al., 2020) require computing query-key attention QK T and combine with relative positional information, which incurs quadratic complexity.
In contrast, linear attention avoids such a query-key product to achieve the linear complexity.
Therefore, common relative positional encoding methods are usually not applicable in linear transformers.
O UR M ETHOD
In this section, we present our main technical contribution on linearized relative positional encoding, which is an encoding family that preserve the linear space-time complexity.
Specifically, we start by presenting a canonical form of relative positional encoding, and show that existing encoding methods can be derived by instantiating the canonical form with different choices of so-called primitive queries, keys and positional matrices in Section 3.1.
When imposing the decomposability constraint on this canonical form, we obtain a sufficient condition for linearized relative positional encoding (LRPE) and derive a family of concrete solutions in real and complex domains in Section 3.2.
We provide an implementation sketch in Section 3.3.
3. C ANONICAL FORM OF RELATIVE POSITIONAL ENCODING
In order to better establish connections between existing relative positional encoding methods and understand their design principles, in this section, we first present a canonical form of relative positional encoding.
In particular, given a query q s and key k s pair, their relative positional encoding f rel : C d × C d → C can be represented as:
[Equation]
where H represents conjugate transposition and m represents number of primitives.
We refer ˆq ( l ) s ∈ C d ( l ) 1 , ˆk ( l ) t ∈ C d [Equation] ) 2 , W
[Equation]
l ) t − s ∈ C d ( l ) 1 × d [Equation] l ) 2 as query, key and relative positional matrix primitives , respectively, used as constituent components to construct the relative positional encoding.
Note that query primitives do not always indicate a reliance on query embeddings, similarly for other primitives.
For example, an identify matrix can also serve as primitives, as we will show shortly in Section 3.1.1.
To demonstrate Eq.
6 is a generic formulation, we show that it flexibly induces a wide range of existing relative encoding methods (Shaw et al., 2018;
Su et al., 2021; Horn et al., 2021; Liutkus et al., 2021;
Huang et al., 2020; Raffel et al., 2019) by selecting and compositing different choices of primitives.
Among them, we highlight two examples in the following section, and leave the complete discussions in the Appendix C.1.
3.1. T YPICAL ENCODING EXAMPLES FROM THE CANONICAL FROM
Additive.
In (Huang et al., 2020), the relative positional encoding is formulated as an extra additive term to the query-key inner-product: [Equation]
[Equation]
which can be derived by including an extra identity term as a primitive, formally denoted as:
[Equation]
Multiplicative.
In RoPE (Su et al., 2021), the relative positional encoding works in the form of the weighted inner product:
[Equation]
[Equation] which can be denoted as:
[Equation] [Equation]
3.1. S IMPLIFICATION
For the ease of remaining discussion, we introduce necessary notations and simplify Eq. 6.
[Equation]
with these notations, we can rewrite Eq.
6 into the matrix form: f rel ( q s , k t ) =
ˆ q H s ˆ W t − s ˆ k t .
Since every component of ˆ q s and ˆ k t are handled with no difference, without losing generality, we only discuss cases where m = 1 :
[Equation]
3. L INEARIZED RELATIVE POSITION ENCODING
Eq.
6 is a canonical form of relative positional encoding, meaning that its variants are applicable to vanilla transformers but not necessarily for linear ones.
To design relative encoding compatible with linear transformers, the attention computation has to respect the decomposibilty condition.
This additional condition leads to the linearized relative position encoding (LRPE) family, defined as follows.
Definition 3.1.
A relative position encoding is called linearized relative position encoding ( LRPE ), when the following holds:
[Equation]
where[-q s , k t ∈ C d , W s , M s ∈ C d × d , W 
-]
[Equation]
The assumption of W 0 = I d implies that the interaction between tokens from the same position only depends on the content, which is reasonable enough that most encoding methods respect.
In its essence, Eq. 13 ensures the positional matrix is decomposable.
In this way, the query-key innerproduct can be avoided in the attention computation.
Consequently, complexity of computing LRPE is O ( nd 2 ) , where n is sequence length, d is embedding dimension as Appendix C.2 shows in detail.
We prove that Eq.
13 can be simplified based on the following proposition:
Proposition 3.2.
Eq. 13 is equivalent to Eq.
14 and W t is Unitary matrix, [Equation]
Proof of Proposition 3.2.
According to the arbitrariness of q s , k t , Eq. 13 is equivalent to
[Equation]
Take s = t in Eq 13, we get (since we assume
that W 0 = I d ): [Equation] Thus, M s is a unitary matrix.
On the other hand, note that for any unitary matrix P , we always have [Equation] This means that left multiplying M t by a unitary matrix P does not change Eq. 13. Since M s and
M H 0 are also unitary matrices, we can perform the following transformation: [Equation] With M s , Eq. 15 becomes
[Equation] Take s = 0 , we have [Equation] Thus Eq.
19 becomes [Equation] [Equation] Since M s is a unitary matrix, W s is also a unitary matrix, i .
e ., [Equation]
[Equation]
The detailed proof can be found in Appendix 3.2.
In the following section, we derive some particular solutions for Eq.
3.2. P ARTICULAR SOLUTIONS
In this section, we discuss Eq.
and give a family of solutions.
It is worth noting that the solutions we provide are all in the form of W s =
P H Λ ( s ) P , where P , Λ ( s ) are unitary matrix.
The complete derivation can be found in Appendix C.4, C.5, C.6.
[Table]
Unitary (Solution 1)
The first case is discussed in the complex domain, which is not common in transformer models yet exhibiting an elegant solution.
Proposition 3.1.
The following form of W s ∈ C d × d satisfies Eq.
[Equation]
[Equation]
where P ∈ C
d × d is unitary matrix, α k , k = 1 , . . .
, d are parameters.
Orthogonal (Solution 2) Now we consider the real domain, a more general case in transformers.
Proposition 3.2.
The following form of W s ∈ R d × d satisfies Eq.
[Equation]
where P ∈ R d × d is orthogonal matrix, α k , k = 1 , . .
, d are parameters.
Permutation (Solution 3)
The last case is inspired by PermuteFormer (Chen, 2021), which is associated with the permutation matrix:
Proposition 3.3.
The following form of W k ∈ R d × d satisfies Eq.
[Equation]
where P ∈ R d × d is the orthogonal matrix.
3. T HE LRPE FAMILY
LRPE ( W s = P H Λ ( s ) P ) contains two components, i .
e ., a fixed unitary matrix P and a unitary matrix family Λ ( s ) as mentioned in proposition 3.1, 3.2, and 3.3.
The P can be seen as a rotation matrix that rotates the token feature to a particular coordinate system and the Λ ( s ) derives the positional information from the rotated feature.
In this paper
[Table]
we select three types of commonly used orthogonal matrix,
i
[Table]
e ., (1) householder matrix (Golub & Van Loan, 2013), (2) permutation matrix, and (3) FFT matrix (Bracewell & Bracewell, 1986).
We combine the selected Matrix P with three LRPE solutions above to obtain 8 practical LRPE types as shown in Table 1.
Detailed information can be found in Appendix C.7.
E XPERIMENTS
[Table]
L L
4. E XPERIMENTAL SETTINGS
Primary tasks.
We validate the effectiveness of the proposed LRPE on various NLP tasks that resort to different Transformer architectures.
Specifically, we first study the autoregressive language model (Radford et al., 2018) with a GPT-like decoder-only structure.
This is followed by the bidirectional language model ( encoder-only ), which adopts the Roberta architecture (Liu et al., 2020) and is pretrained and then fine-tuned on several downstream tasks from the GLUE benchmark (Wang et al., 2018).
Then, we evaluate LRPE on machine translation ( encoder-decoder ).
Competing methods.
Our baseline is the Transformer model (Vaswani et al., 2017) without relative positional encoding.
For comparison, we also choose four state-of-the-art methods , i .
e .,
RoPE
(Su et al., 2021), SPE (Liutkus et al., 2021), PermutateFormer (abbreviated as “PER”)
(Chen, 2021), T (Raffel et al., 2019) and Transformer-XL (Dai et al., 2019) (abbreviated as “XL”), and test them in both linear attention and vanilla attention.
In the linear attention, we employ 1+elu( · ) (Katharopoulos et al., 2020) as the kernel function.
Training configuration.
Our experiments are implemented in the Fairseq framework (Ott et al., 2019) with V100 GPUs.
All the methods share the same configurations which are listed in Appendix
D.1.
[Table]
4. R ESULTS IN LINEAR SETTING
Autoregressive language model.
The autoregressive language model has 6 decoder layers and is trained on the WikiText-103 dataset (Merity et al., 2017).
We use the Perplexity (PPL) as the evaluation metric and report the results in Table 2.
We observe that under the linear setting, most variants of LRPE present performance gain over the baseline.
Our best model, i .
e ., Type2, outperforms
RoPE and SPE on both validation and test sets to a large margin, and achieves comparable results
to
PER with minor difference.
Clearly, the proposed method is effective in encoding causal data.
Bidirectional language model.
The bidirectional model follows an encoder-only structure, i .
e .,
Roberta (Liu et al., 2020), with 12 layers.
We first pretrain it on the WikiText-103 dataset, and present the results in Fig.
2 and Appendix D.2.
Generally, LRPE (Type2 in this case) has better performance, i .
e ., smaller validation PPL in all evaluation steps, than competing methods.
Notably, it surpasses RoPE, SPE and PER by nearly 10%, 27% and 37% in terms of the final PPL, indicating its superiority in bidirectional language modeling.
We then fine-tune the pretrained model on the GLUE dataset (Wang et al., 2018).
We use different learning rates among 1e-5, 3e-5, 6e-5, 1e-4 and choosing the best result after fine-tuning forepochs.
From Table 3, we find that the two representative LRPE variants, i .
e ., Type5 , perform consistently better than other methods in all metrics.
The average score of Type5 beats RoPE, SPE, and PER by more than 4.4%.
Machine translation.
For this task, we adopt the base transformer model which consists of 6 encoder layers and 6 decoder layers, and train it on the WMT’14 En-De dataset (Bojar et al., 2014).
We ran each experiment 5 times and report the averaged results.
Note that in practical, we only embed the linear attention and its corresponding relative positional encoding in encoders, since we empirically find that the model cannot converge appropriately when the linear attention appears in decoders.
We measure the accuracy with BLEU, and the quantitative results on both validation and test sets are displayed in Table 4.
Most variants in the LRPE family have comparable performances to the competing methods on the validation set, and Type4 ranks first on the test data, which demonstrates again the validity of our LRPE.
However, a few variants present less competitive results than the state-of-the-arts on the test data.
Empirically, this is caused by the relatively high sensitivity of parameter tuning on the machine translation performance while all listed methods share the identical parameter setting.
We will concentrate on how to further improve their accuracy by specializing the parameters for each variant in the future work.
[Table] [Table]
4. M ODEL ANALYSIS
An explanation of LRPE.
According to the discussion in Section.
3.3, The LRPE rotates the token feature through P , and encodes the positional information through Λ ( s ) .
In Table 5, we ablate the effectiveness of the P matrix on the autoregressive language modeling task.
Our approach with the Householder matrix and the Permutation matrix achieve marginally better results than the one without rotation (Identity matrix).
It indicates that we can get better performance by carefully selecting the projection of the positional encoding.
Complexity and efficiency.
The implementation of the proposed LRPE does not affect the computational complexity of the linear transformer, i
[Table]
e ., preserving the linear complexity as O ( n ) .
We also measure the training speed of the bidirectional language model on the same local machine ( i
[Table] e ., a GeForce GTX 1060 card), and observe that the speed after using LRPE is only 9% slower than the baseline on average.
The detailed comparison of the efficiency can be found in Appendix D.3.
In general, UPRE does not incur significant computational burden to the transformer, and can fulfill the practical needs by maintaining comparable efficiency.
Generalization to vanilla attention.
Finally, we investigate the generalization of LRPE towards the vanilla attention.
The results are reported in Fig.
2, Table 2, 3, and 4.
The conclusion is consistent with that of the linear setting, i .
e ., improving the vanilla transformer baseline and achieving competitive performance to the competing methods.
It indicates the good flexibility of LRPE as it can be seamlessly applied to any attention type.
C ONCLUSION
In this paper, we standardize the form of relative positional encoding in both linear and vanilla transformers, and focus the case in the linear attention.
The unitary transformation is employed as a special solution to the linearized relative positional encoding, and the solutions as per various constraints constitute the unitary relative positional encoding (LRPE) family.
We validate the effectiveness of LRPE through extensive experiments on several NLP tasks with different transformer architectures.
It outperforms state-of-the-art methods under both linear and vanilla settings.
Appendix A M ATHEMATICAL N OTATIONS
[Table]
[Equation]
B C OMPUTATION OF V ANILLA /L INEAR A TTENTION B. B ASIC N OTATIONS
Both vanilla and linear attention blocks involve three matrices, i .
e ., Q ( Query ), K ( Key ) and V ( Value ).
All of them are linear projections of input X ∈ C n × d , i .
e .,
[Equation]
where W Q , W K ,
[Equation]
The vector form is organized as
[Equation]
The attention output is
[Equation]
B. V ANILLA A TTENTION
In vanilla attention, the output is computed using the Softmax weighted sum, i
[Equation]
e .,
[Equation]
B. L INEAR A TTENTION
The linear attention is formulated as follows,
[Equation]
C P ROOF OF T HEOREM C. M ORE E XAMPLES
In the following
[Equation]
we provide two additional examples of relative positional encoding with the canonical form.
RPR (Shaw et al., 2018):
[Equation]
The canonical form is
[Equation]
DeBERTa (Huang et al., 2020):
[Equation]
[Equation]
The canonical form is
[Equation]
cosFormer (Qin et al., 2022):
[Equation]
36) which indicates that the relative positional encoding is effectively a coefficient term in the attention matrix, as such, it can be derived via a positional matrix primitive with the coefficients.
[Equation]
C. L INEARIZED R ELATIVE P OSITIONAL E NCODING
Proof of 3.2.
For this, we only need to prove that the time complexity is linear with respect to n .
To this end, we first give basic notations as follows,
[Equation]
The time complexity of transforming Q , K to ˜ Q , ˜ K is O ( nd 2 ) .
The next step is to calculate the output, i .
e .,
[Equation]
Clearly, Eq. 39 is a standard formulation for the linear attention with the time complexity as O ( nd 2 ) .
Combing it with the first step, we have the total time complexity as O ( nd 2 ) , which is unchanged.
C. L INEARIZED R ELATIVE P OSITIONAL E NCODING
Before the proof, we first give the following theorems (Yao & Algebra, 2015):
Theorem C.1.
If matrix W ∈ C d × d is a unitary matrix, there exists another unitary matrix P ∈ C
d × d , such that
[Equation]
Theorem C.2.
If matrix W ∈ R d × d is an orthogonal matrix, there exists another orthogonal matrix
P ∈ R d × d , such that
[Equation]
C. U NITARY (S OLUTION 1)
Proof of Proposition 3.1.
According to Theorem C.1, we can assume that W s has the following form (
P ∈ C d × d is a unitary matrix),
[Equation]
Hence, Eq. 14 is equivalent to
[Equation]
(cid:110)
[Equation]
[Equation] ( t ) [Equation] θ ( s ) 1 )
[Equation]
j ( θ ( t ) 2 − θ ( s ) 2 )
[Equation]
[Equation]
( θ ( t ) d − θ ( s ) d ) (cid:111) =
diag (cid:110) jθ
( t − s ) [Equation] jθ ( t − s ) [Equation]
[Equation]
jθ ( t −
[Equation]
[Equation]
In this case, ∀ k = 1 , . . .
, d , we have
[Equation]
Note that 2 tπ does not affect the result, so we can assume t = 0 , i .
e .,
[Equation] [Equation]
Taking t =
s + 1 , we get
[Equation]
[Equation]
C. O RTHOGONAL (S OLUTION 2)
Proof of Proposition 3.2.
According to Theorem C.2, we can assume that W s has the following form ( P ∈ R d × d is an orthogonal matrix),
[Equation]
Hence, Eq. 14 is equivalent to
[Equation]
48) where
[Equation]
For A ( s ) , considering the k -th component, we get
[Equation]
Hence, ∀ k = 1 , . . .
, d , we have
[Equation]
Note that 2 tπ does not affect the result, so we can assume t = 0 , i .
e .,
[Equation]
Taking t =
s + 1 , we have
[Equation]
Next, for B ( s ) , the conclusion is more obvious, i .
e .,
[Equation]
Finally, for C ( s ) , we have
[Equation]
In that case, we must have r = 0 .
C. P ERMUTATION (S OLUTION 3)
Prior to the proof, we first provide some relevant definitions and propositions.
Definition C.3.
Permutation π is a bijection defined on the integer set:
[Equation]
[Equation]
Definition C.4.
For matrix
[Equation] [Equation]
M π is defined as
[Equation]
Definition C.5.
For identity matrix I d ∈ R d × d and permutation π , we define
[Equation]
For Λ k , we have the following important properties:
Lemma C.6.
For permutation π , matrix M ∈ R d × d and Λ k ∈ R d × d defined in C.5, we have
[Equation]
Proof.
We first organize I d ∈ R d
× d in the following form, where e k ∈ R d , k = 1 , . . .
, d represents the one-hot vector with the k -th element as one, i .
e .,
[Equation]
Notice that
[Equation]
so we get
[Equation]
Theorem C.7.
For Λ k defined in C.5, we have:
[Equation]
Proof.
We use induction for the proof.
For k = 1 , the conclusion is obvious.
Now assuming that the conclusion holds for k
= s − 1 , when k = s , we have
[Equation]
The next step is to prove
[Equation]
The above conclusion follows from C.6.
Theorem C.8.
Λ k ∈ R d × d defined in C.5 are orthogonal matrices, i.e.,
[Equation]
Proof.
We first prove that the conclusion holds for k = 1 :
[Equation]
Since Λ 1 is a square matrix, we also have
[Equation]
In general cases, we only use C.7, i .
e .,
[Equation]
With the same proof, we get
[Equation]
Based on the above conclusions, we can prove Proposition 3.3 below.
Proof of Proposition 3.3.
According to Theorem C.8 and the production of the orthogonal matrix is an orthogonal matrix, we can assume that W k has the following form (
P ∈ R d × d is an orthogonal matrix), i .
e .,
[Equation]
The next step is to verify that it satisfies Eq. 14, which follows Theorem C.7 and C.8:
[Equation]
C. I MPLEMENTATION
LRPE( W s = P H Λ ( s ) P ) contains two components, i .
e ., the fixed unitary matrix P and the unitary matrix family Λ ( s ) mentioned in proposition 3.1, 3.2, and 3.3.
We first introduce the choice of matrices P / Λ ( s ) , and then illustrate some implementation tricks.
Choice of matrices
For matrix P , we employ three types as follows,
• Householder matrix: denoted as a vector v ∈ R d ,
i
e .
[Equation]
In our implementation, we sample v from standard normal distribution, and make it deterministic or learnable .
• Permutation matrix: formulated as per the following permutation (inspired by Flash (Hua et al., 2022)), i .
e .,
[Equation]
• FFT matrix: a matrix form of FFT (Fast Fourier Transform).
For matrix family Λ ( s ) , we use the following settings:
• For unitary (Solution 1) (3.1), we use the same method in (Su et al., 2021) with initialized α t = 10000 − 2 t/d , and make it deterministic.
Since this method involves complex numbers, we only use the FFT matrix for the choice of P .
• For orthogonal (Solution 2) (3.2), we test with two versions.
In the first version, we set the dimension of the identity submatrix q =
⌊ d/ 2 ⌋ , initialized α t = 10000 − 2 t/d as in (Su et al., 2021) and make it deterministic .
In the second version, we choose the dimension of identity submatrix q = 0 with initialized α t = 10000 − 2 t/d as in (Su et al., 2021), and make it learnable .
– Another notable version to choose the dimension of the identity submatrix q = 0 with initialized α t = 10000 − 2 t/d as in (Su et al., 2021), and make it deterministic .
When using this version along with the identity matrix, we can get RoPE
(Su et al., 2021).
• For permutation (Solution 3) (3.3), we randomly choose the permutation and make it deterministic .
– Notice that when combing this method with identity matrix, we can get a version of PermutateFormer (Chen, 2021).
Implementation tricks
According to the following facts, we can simplify the computation, i .
e .,
[Equation]
Hence, in practice, we can use W s = P H Λ ( s ) instead of W s = P H Λ ( s ) P to reduce the computational costs.
C. P SEUDOCODE
In this section, we provide pseudocodes for LRPE in Python: import torch import torch.nn as nn import numpy as np class Lrpe(nn.
Module): def __init__(self, core_matrix, p_matrix, max_positions=512, embedding_dim=768,theta_type="a", theta_learned=
False, householder_learned=
False): super().__init__() self.core_matrix= core_matrix
self.p_matrix = p_matrix self.theta_type = theta_type self.theta_learned = theta_learned self.householder_learned =
householder_learned
# Lambda matrix if self.core_matrix =
= 1: if self.theta_learned:print("Learntheta!")
self.theta
=nn.
Parameter(
** (-2 / embedding_dim * torch.arange(embedding_dim // 2)).reshape(1, 1, -1)) else: print(f"Theta_type {self.theta_type}") elif self.core_matrix =
= 2: print("Mixed") elif self.core_matrix =
= 3: print("Permutation") permutation=self.get_permutation(max_positions, embedding_dim) self.register_buffer("permutation", permutation) elif self.core_matrix =
= 4: print("Complex exp") if self.theta_learned:print("Learntheta!")
self.theta
=nn.
Parameter(
** (-2 / embedding_dim * torch.arange(embedding_dim)).reshape(1, 1, -1)) else:print(f"Theta_type {self.theta_type}")
# P matrix if self.p_matrix == 1: print("Identity") elif self.p_matrix == 2: print("Householder") if self.householder_learned: print("learnhouseholder!")self.v=nn.Parameter(torch.randn(1, embedding_dim, 1)) else:v= torch.randn(1, embedding_dim, 1) v = v / torch.norm(v) print(f"Householder norm is {torch.norm(v)}") self.v = nn.
Parameter(v, requires_grad=
False) elif self.p_matrix == 3: print("Fourier") elif self.p_matrix == 4: print("Odd_even") self.p = self.get_p() self.core_transform = self.get_core_transform() def forward(self, x): ’’’input shape: (b, l, e), b stands for batch size, l stands for sequence length, e stands for embedding dimension.
’’’x=
self.p(x) x = self.core_transform(x) return x def get_p(self):ifself.p_matrix == 1: def f(x):return x return f elif self.p do permutation elif self
4: return self exp def get_permutation embedding permutation torch.randperm( dim reshape(1,
) expanded =
.arange
_dim unsqueeze( for
: previous = expanded gather
, permutation) expanded.append(current) expanded = torch.stack(expanded, dim return expanded def odd_even_permutation(self, x): # 2k->k, 2k+1->d+k e = x.shape[-1] d = e - e //permutation = torch.arange(e) index = torch.arange(e) permutation[::2] = index[::2] // 2 permutation[1::2] =
(index[1::2] - 1) //
2 + d permutation = permutation.to(x.device) x
= x.gather(
, permutation.expand_as(x
) return x def do_permutation(self, x): b, l, e = x.shape x = x.gather( permutation[:, :l, :] expand_as(x
) return x def reflect(self, x): b, l, d = x.shape e = d - 1 if d % 2 =
= 1 else d return self.transform(x, e) def mix_reflect(self, x): b, l, d = x.shape assert d >= 3 # split e = d // 2 # to even if e % 2: e +=return self.transform(x, e) def transform(self, x, e): assert e % 2 =
=b, l, d = x.shape # do identity transformation x1 = x
[:, :, e:] # do reflection x = x[:, :, :e] if self.theta_learned: theta=self.theta else: if self.theta_type =
= "a": theta = 10000 ** (-2 / e * torch.arange(e // 2))
elif self.theta_type == "b": theta =
np.pi / 2 / l / (e // 2) * torch.arange(1, e // 2 + 1) elif self.theta_type == "c": theta =
np.pi / 2 / l / torch.arange(1, e // 2 + 1) theta = theta.reshape(1, 1, -1).to(x) theta =
torch.stack([theta, theta], dim=-1).reshape(1, 1, e) theta = theta * torch.arange(l).reshape(1, -1, 1).to(x)
# (-q1, -q3), (q0, q2) -> (-q1, q0, -q3, q2) x_half
=
torch.stack([-x[..., 1::2], x[..., ::2]], dim=-1).reshape_as(x) x_transform =
x
* torch.cos(theta) + x_half * torch.sin(theta) # merge if e !
=
d: x_transform = torch.cat([x_transform, x1], dim=-1) return x_transform def complex_exp(self, x): b, l, e = x.shape if self.theta_learned:theta=self.theta else:if self.theta_type == "a": theta = 10000 ** (-2 / e * torch.arange(e)) theta
= theta.reshape(1, 1, -1).to(x.device)
matrix = theta * torch.arange(l).reshape(1, -1, 1).to(x.device) sin_cos = torch.complex(torch.cos(matrix),torch.sin(matrix)).to(x.device)
x = self.element_wise_complex(x, sin_cos) return x def element_wise_complex(self, t1, t2): return torch.complex(t1.real * t2.real - t1.imag * t2.imag, t1.real * t2.imag + t1.imag * t2.real) def householder(self, x, eps=1e-6): if self.householder_learned: v=self.v/(torch.norm self.v) + eps) else: v= self.v # (b, n, e), (1, e, 1) -> (1, n, 1) y = torch.matmul(x, v) # (1, n, 1), (1, 1, e) -> (1, n, e) y = torch.matmul(y, v.transpose(1, 2))
D E XPERIMENT D. C ONFIGURATION
We provide detailed data, model and training configurations in Table 7.
For published datasets,
WikiText-103 is obtained from https://www.salesforce.com/products/einstein/ai-research/thewikitext-dependency-language-modeling-dataset/, with Creative Commons Attribution-ShareAlike
License.
The GLUE dataset is obtained from https://gluebenchmark.com/. The WMT’14 EN-DE dataset is downloaded from https://www.statmt.org/wmt14/:
D. R ESULTS OF B IDIRECTIONAL L ANGUAGE M ODEL
We report the pretrained results of the bidirectional language model in Table 8.
Our LRPE achieves competitive performance in both linear attention and vanilla attention.
Notably, the UPPE variant
UPRE ol h has the best quantitative results in all evaluation metrics.
[Table]
[Table] [Table]
D. E FFICIENCY
We compare the training speed of LRPE with other methods in Table D
[Table] [Table]
which indicates that our method maintains good efficiency without incurring too much computational burden.
[Table]