KERPLE: Kernelized Relative Positional Embedding for Length Extrapolation
Anonymous Author(s)
AffiliationAddress email
Abstract
Relative positional embeddings (RPE) have received considerable attention since
RPEs effectively model the relative distance among tokens and enable length extrapolation.
We propose KERPLE, a framework that generalizes relative positionembedding for extrapolation by kernelizing positional differences.
We achieve thisgoal using conditionally positive definite (CPD) kernels, a class of functions knownfor generalizing distance metrics.
To maintain the inner product interpretation ofself-attention, we show that a CPD kernel can be transformed into a PD kernelby adding a constant offset.
This offset is implicitly absorbed in the Softmaxnormalization during self-attention.
The diversity of CPD kernels allows us toderive various RPEs that enable length extrapolation in a principled way.
Experiments demonstrate that the logarithmic variant achieves excellent extrapolationperformance on three large language modeling datasets
Introduction
Transformer-based models have excelled in various natural language processing tasks such as chatbot
[Roller et al., 2021], code completion [Chen et al., 2021a], and paper abstract summarization [Zhang et al., 2020].
These sequence modeling tasks often require the model to operate well onsignificantly longer text sequences than the fixed maximum length L used at training time.
Training(or retraining) the model using a substantially larger value of L is often infeasible since the transformer training cost is O ( L 2 ) .
Hence, one desires a transformer that continues to perform well onlonger sequences than those used during training; i.e., perform length extrapolation at inference time.
Most transformer designs do not have this property [Press et al., 2022].
While recent work on absolutepositional embeddings demonstrated the extrapolation ability [Kiyono et al., 2021, Likhomanenkoet al., 2021], it is believed that relative positional embeddings are more robust to input length change[Likhomanenko et al., 2021], for example, ALiBi
[Press et al., 2022] and T5 [Raffel et al., 2020].
Hence, we are motivated to study the inner workings of relative positional embeddings.
Relative positional embeddings (RPE) encode the idea of shift-invariance: for any shift p , ( m + p ) −( n + p ) = m − n .
It is often added directly to the self-attention matrix before Softmax normalization[Chen et al., 2021b].
Inspired by shift-invariance and the ability of a kernel to define a similarityfunction, there have been studies on shift-invariant kernels for RPE [Wennberg and Henter, 2021]with a focus on Gaussian kernel.
However, in our preliminary experiments, the Gaussian kerneldemonstrates limited length extrapolation ability (see Appendix A.3).
Hence, a distinct class ofshift-invariant kernels is needed to achieve adequate length extrapolation.
To this end, we note a set of well-established conditionally positive definite (CPD) kernels suitablefor modeling distance metrics [Schölkopf, 2000].
However, CPD kernels do not conform to an inner
[Figure]
product.
We can remedy this issue by transforming a CPD kernel into a PD kernel by adding asufficiently large constant.
This constant offset is subsequently absorbed implicitly in the Softmaxnormalization (see the discussion below Eq.
For example, ALiBi implicitly admits a PD kernelof the form c − | m − n | (see the end of section 4), which is reduced to a CPD kernel −| m − n | .
The
CPD kernel and Softmax normalization combination opens the door to a sea of possible CPD kernels.
We investigate structures from this class that exhibit a strong length extrapolation ability, like ALiBi.
Our main result is a framework for KE rnelize R elative P ositional Embedding for L ength
E xtrapolation ( KERPLE ).
The framework elucidates key principles that encourage the lengthextrapolation property.
We show that ALiBi is a particular instance within our framework.
Oursubsequent experiments suggest that the proposed method yields better length extrapolation on largedatasets such as OpenWebText2, GitHub, and ArXiv.
Background and Related Work2.1 Preliminary
Let { w m } Lm
=1 be the input tokens to a transformer model, where L is the total number of tokens.
Each w m is a scalar and is used to index the embedding vector e
m ∈ R d as the input to the transformer.
A transformer converts each e m into query, key, and value vectors in R d : q m =
W q e m , k m =
W k e m , v m =
W v e m , where W q , W k , W v ∈ R
d
×
d are learnable matrices.
Then, the self-attention module computes the scaled attention scores and generates the output vector o m at position m as:
[Equation]
Since the operation is position-agnostic, it is believed that positional information helps the tokeninteractions modeling [Vaswani et al., 2017], which we survey in the next subsection.
2.2 Positional Embedding
Absolute.
Absolute positional embeddings assign a positional vector p m to each position m andadds p m to the embedding vector e m .
The very first version of which is the predefined sinusoidalfunction [Vaswani et al., 2017].
Followed by the success of BERT [Devlin et al., 2019], learnableabsolute positional embeddings have been applied on the task of masked language modeling [Devlinet al., 2019, Liu et al., 2019, Clark et al., 2020, Lan et al., 2020], Autoregressive-decoding [Radfordet al., 2018, 2019], and sequence-to-sequence [Gehring et al., 2017, Lewis et al., 2019] settings.
Recent work studied ways to extrapolate sinusoidal positional embeddings to longer sequencesby randomly shifting absolute positions during training [Kiyono et al., 2021] or augmenting withcontinuous signals [Likhomanenko et al., 2021].
Relative.
As opposed to the modeling of absolute position m , relative positional embeddings (
)that model the positional difference m − n has become popular in the literature [Shaw et al., 2018,
Huang et al., 2019, Dai et al., 2019, Yang et al., 2019, Huang et al., 2020, He et al., 2021, Ke et al.,
2021, Chen et al., 2021b].
In particular, the T5 model that considers bucketed relative distances andlog-binning has been shown to perform well on various transformer architectures [Raffel et al., 2020].
Rotary positional embedding [Su et al., 2021]
encodes the position with rotations: f
( q m , m ) =
R m q m where R m is a rotation matrix with angles proportional to m .
With the rotation’s property,the query-key product exhibits a positional difference: f ( q m , m ) ⊤ f ( k n , n ) =
q ⊤ m R n − m k n .
We note that the overview above focuses on the NLP domain.
Recent work has applied positionalembeddings to other domains such as vision [Wu et al., 2021a] and speech [Likhomanenko et al.
A survey can be found in [Dufter et al., 2022].
2.3 Kernel and its Application in Transformer
The kernel trick is a classic approach to generalize the inner product to high dimensional spaces[Mika et al., 1998, Schölkopf, 2000, Leslie et al., 2001, Dhillon et al., 2004, Takeda et al., 2007].
Inthe context of transformers, there have been interests in applying kernels to the self-attention structureto enhance the performance.
Examples of such work include kernel for positional embeddings [Tsaiet al., 2019, Wu et al., 2021b, Wennberg and Henter, 2021,
Luo et al., 2021].
Another line of researchleverages the kernel’s feature map [Rahimi and Recht, 2007] to linearize the self-attention moduleand reduce the computational cost [Katharopoulos et al., 2020, Chen et al., 2021c, Xiong et al., 2021,
Peng et al., 2021, Choromanski et al., 2021, Qin et al., 2022].
Theoretical Foundations of CPD Kernels3.1 PD and CPD Kernels
In this work, we use shift-invariant conditionally positive definite (CPD) kernels to model the effectof relative positional differences.
We propose this formulation because the notion of relative ismodeled by a shift-invariant function: a bivariate function k over two positions ( m, n ) such thatk ( m, n ) =
f ( m − n ) for some univariate f .
The notion of positional difference m − n is generalizedby the CPD kernel.
To elaborate, we review the definitions of PD and CPD kernels as below.
Definition 1 (PD Kernel) .
A (real) symmetric function k : X × X
→ R is a positive definitekernel if for any integer N and any { x i ∈ X} Ni =1 , { c i ∈ R } Ni =1 , the quadratic form is nonnegative: 88 (cid:80) Ni =1 (cid:80) Nj =1 c i c j k ( x i , x j ) ≥ 0 .
Definition 2 (CPD Kernel) .
A (real) symmetric function ˜
k
X × X
→ R is a conditionally positivedefinite kernel if for any integer N and any { x i ∈ X} Ni =1 , the quadratic form is conditionallynonnegative:
(cid:80) Ni =1 (cid:80) Nj =1 c i c j ˜ k
( x i , x j ) ≥ 0 for { c i ∈ R } Ni =1 with (cid:80) Ni =1 c
i = 0 .
Fact 1 (Berg et al.
[1984] and Prop. 5 of Schölkopf [2000]) .
Let ˜ k
X × X → ( −∞ , 0] be a CPDkernel with ˜ k
( x, x ) =
∀ x ∈ X .
Then, there exists a Hilbert space H and a mapping ϕ :
X → Hsuch that ∥ ϕ ( x ) − ϕ ( x ′ )
∥ 2 = − ˜ k
( x, x ′ ) .
Fact 1 suggests that CPD kernels generalize the distance metric to high dimensional spaces.
Since weare interested in the positional difference, it is of our interest to model the distance between positionsusing CPD kernels.
However, Fact 1 also implies that CPD kernels do not encode inner products as required by selfattention for pairwise relations computation.
It is PD kernels that represent inner products.
To betterunderstand the effect of CPD kernels on self-attention, we need to establish relations between CPDand PD kernels.
As noted in Schölkopf [2000], if one takes any PD kernel and offsets it by a constant,the result is at least a CPD kernel.
In the next subsection, we show that the converse is nearly true:if ˜ k is CPD, so is c + ˜ k for large enough c ∈ R (Lemma 1).
Therefore, we may generate the CPDkernels of our interest and transform them into PD kernels if needed.
3.2 Constructing PD Kernels From CPD Kernels via Constant Shifts
In this subsection, we review a few properties of CPD kernels and use these to generate a variety of
CPD kernels.
Then, we present a lemma that transforms CPD kernels into PD kernels via constantshifts.
This enables the production of a family of PD kernels from CPD kernels.
Finally, we present our critical observation that the exact value of the constant shift is not needed, thanks to a niceproperty of Softmax normalization.
Below are some important facts about CPD kernels.
Fact 2 (Scaling and Summation) .
If ˜ k 1 and ˜ k 2 are CPD, then so are a · ˜ k 1 (for a > 0 ) and ˜ k 1 + ˜ k 2 .
Fact 3 (Berg et al.
[1984] and Prop. 4 of Schölkopf [2000]) .
If ˜ k : X × X → ( −∞ , 0] is CPD, thenso are − ( − ˜ k ) α for 0 < α < 1 and − log(1 − ˜ k ) .
Fact 4 (Page 3 of Schölkopf [2000]) .
The negative squared distance −∥ x − x ′ ∥ 2 is CPD.
The three Facts above jointly yield a rich family of CPD kernels as shown below.
Corollary 1.
The following are CPD kernels.
(a) ˜ k
( x, x ′ )
= − a ∥ x − x ′ ∥
p
with 0 < p ≤ 2 and a > 0 .
(b) ˜ k
( x, x ′ )
= − b · log(1 + a ∥ x − x ′ ∥ p ) with 0 < p ≤ 2 and a, b > 0 .
We note that it is possible to keep iterating between Fact 2 and 3 and generate more complicated ones,for example, − a ∥ x − x ′ ∥ p − b · log(1 + a ∥ x − x ′ ∥ p ) or − b · log(1 + a ∥ x − x ′ ∥ p ) c for 0 < c < 1 .
Considering kernel is a means to model relative positional embeddings, we only present CPD kernelswith reasonable complication.
Those with complicated forms can be deferred to our future work.
Now that Corollary 1 has presented a few class of CPD kernels, we prove a lemma (in Appendix A.1)that constructs PD kernels from CPD kernels through shifting.
Later in Eq.
(2), we will see that theshifting construction is combined neatly with the Softmax normalization of self-attention.
Lemma 1 (CPD Shift Lemma) .
Let ˜ k
X × X → R be a CPD kernel.
There exists c ≥ 0 such thatc + ˜ k is a PD kernel.
Lemma 1 implies the CPD kernels in Corollary 1 can be made PD if they are added with large enoughconstants;
for example, c − ∥ x − x ′ ∥ p for large enough c .
Although Lemma 1 does not have anexplicit construction of c , thanks to the shift-invariant property of the Softmax normalization, wecan leave it as an under-determined constant in our positional embedding design, which is Eq.
(1) insection 4.
Still, given a set of test points { x i } Ni =1 , one can do a geometric sequence search 1 to searchfor a c such that the N × N matrix
c
+ ˜ k
( x i , x j )]
Ni,j =1 ⪰ 0 .
Hence, in this work, we do not need thevalue of c , but we can compute it if we do need its value, e.g., deriving the feature map of c + ˜ k .
Alternative Proof of c − ∥ x − x ′
∥
p c − ∥ x − x ′ ∥
p c − ∥ x − x ′ ∥ p .
While the CPD shift lemma is convenient, one can provec − ∥ x − x ′ ∥
p is PD for large enough c using a kernel representation theorem in Schoenberg [1938].
See Appendix A.2 for details.
Kernelized Relative Positional Embedding
Let { q m } Lm =1 and { k n } Ln =1 be the input queries and keys.
Let ( r 1 , ..., r ℓ ) be learnable parameters.
We propose a kernelized relative positional embedding as follows.
[Equation]
where ˜ k r 1 ,...,r ℓ ( m, n ) is any shift-invariant CPD kernel with ℓ number of parameters.
Due to Lemma 1,the above can be reformulated into its kernel form as follows.
[Equation] [Equation]
(*) is due to the shift-invariant property of the Softmax normalization:
[Equation]
for any c ∈ R .
The second equality defines a bias kernel which is positive definite using Lemma 1:
[Equation]
The last equality introduces a composite kernel k comp : R d +1 × R d +1 → R as
[Equation]
Interpretation.
The proposed method can be interpreted as applying a composite kernel to selfattention.
The composite kernel combines the information from query q m , key k n , and positions( m, n ) in a way that augments the original self-attention structure by multiplicative and additiveposition embeddings.
The augmentation allows k comp to not only retain the original q ⊤ m k n but alsoinclude positional information from the bias kernel k r 1 ,...,r ℓ .
Practical Choice.
In section 5.2, we fix ℓ = 2 and experiment on two variants of the compositekernel, Eq.
(4), where we call these the power variant and the logarithmic variant of our proposed
KERPLE framework, Eq.
These are from a combination of Corollary 1 and Eq.
(3).(power) k comp ([ q m , m ] , [ k n , n ]) =
q ⊤ m k n + c − r 1 | m − n | r 2 with r 1 > 0 and 0 < r 2 ≤ 2 .(logarithmic)
k comp ([ q m , m ] , [ k n , n ]) =
q ⊤ m k n + c − r 1 · log(1 + r 2 | m − n | ) with r 1 , r 2 > 0 .
We note that these are not the only variants of the composite kernel.
In section 5.3, we experimentwith two more complicated variants, but only find lower training speeds and marginal improvementin perplexities (e.g., logarithmic variant vs. 3-para-log).
Thus, based on our study, the choices abovehold advantages in both performance and speed.
Connection to Prior Work.
When the bias kernel, Eq.
(3), is a triangle kernel: c − | m − n | , ourmodel reduces to ALiBi
[Press et al., 2022].
Wennberg and Henter [2021] discuss the situation wherethe bias kernel is a Gaussian kernel.
Tsai et al.
[2019] is the case where there is no bias kernel andthe attention product q ⊤ m k n is multiplied by an exponentiated inner product kernel, exp( x ⊤ y ) .
Since
ALiBi is the state-of-the-art and has great input length extrapolation, we will focus on the comparisonwith ALiBi in our experiments.
The logarithmic variant has an implicit connection to T5 positional bias [Raffel et al., 2020].
According to the official GitHub repository https://github.com/google-research/ text-to-text-transfer-transformer and the HuggingFace Transformer [Wolf et al., 2020],
T5 bias is implemented with a log-binning strategy.
For each head of the transformer, they maintain a bucket of 32 learnable parameters and assign the relative positional bias b m − n to these parameters as b m − n =   bucket [0] if m − n <bucket [ m − n ] if 0 ≤ m − n <bucket [min(31 , ⌊ log(( m − n ) / 16) log(128 / 16) ) · 16 ⌋ ] if m − n ≥ 16 , where ⌊·⌋ is the floor function.
Note that the log factor is approximately 7 .
7 log m − n 16 .
Therefore, T5is using a logarithmic bucket assignment, which turns out to extrapolate to different input lengths.
Compared with T5, our logarithmic variant uses less parameters (2x12 vs. 32x12) but cannot learnnon-monotonic relations (the log function is monotonic).
We will conduct more comparisons with T5bias in our experiments.
Experiments
5.1 Dataset and Implementation Description
Dataset.
We conduct experiments on OpenWebText2, GitHub, and ArXiv datasets gathered in
Gao et al.
OpenWebText2 includes recent content from Reddit submissions until 2020,content from multiple languages, document metadata, multiple dataset versions, and open-sourcereplication code.
GitHub includes open-source repositories written in primary coding languages suchas Java, C/C++, Python, and Go.
ArXiv includes papers written in LaTex in Math, Computer Science,
Physics, and some related fields.
These tasks are motivated by the downstream applications suchas online chatting [Roller et al., 2021], code completion [Chen et al., 2021a], and academic papersummarization [Zhang et al., 2020].
[Table]
Implementation.
We adapt our model from GPT-NeoX
[Black et al., 2021], a transformer implementation by the EleutherAI team.
The codebase is based on NVIDIA Megatron Language Model[Shoeybi et al., 2019] and further accelerated using Microsoft DeepSpeed library [Rasley et al., 2020].
Our model is trained on a machine with one NVIDIA A100 GPU with 40 GB of memory.
We adoptalmost all configurations of small GPT-NeoX 2 , except that we change the train-micro-batch-sizeto 32, seq-length to 512, and max-position-embeddings to 512.
Table 2 summarizes the importantconfigurations fixed throughout our experiments.
In particular, the floating-point encoding is set as
[Table]
189 bfloat16 (Brain Floating Point, developed by Google Brain) so that the training can be accelerated byhalf-precision computation with reliable stability [Kalamkar et al., 2019].
Hidden size 64 means thatd = 64 in Eq.
5.2 Experimental Results (Also c.f. Appendix A.7)
We conduct experiments to cover aspects such as input length extrapolation, application on differentdomains, and comparison with the prior work.
These are elaborated as below.
(i)
Motivated by theinput length extrapolation demonstrated in [Press et al., 2022], we train our model with length 512and test on lengths ranging from 512 to 16384.
We hope that the emphasis on extrapolation enablesthe application of transformers to longer sequences.
(ii) To evaluate the applicability of the modelin different domains, we conduct experiments on OpenWebText2, GitHub, and ArXiv datasets.
(iii)
To validate the effectiveness of our method, we compare KERPLE with Sinusoidal [Vaswani et al.
,2017], Rotary [Su et al., 2021], T5 [Raffel et al., 2020], and ALiBi
[Press et al., 2022].
Table 3 reports the perplexity at different extrapolation lengths.
We perform non-overlappingevaluation: Suppose text is segmented in different manner for 512 and 1024 tokens, we have Nsentences and N/2 correspondingly to evaluate.
We also perform a paired two-sided t-test to validatethe statistical significance (significance level=0.05).
We compare each candidate RPE with ourproposed logarithmic variant and mark the candidate with a † if the log variant is statisticallysignificantly better .
Table 4 reports the training speeds.
These tables yield three conclusions.
First,within the KERPLE framework, the logarithmic variant is better than the power variant.
Secondly, thelogarithmic variant is 9.7% faster than T5.
In terms of extrapolation, the logarithmic variant generallydoes better than T5 but could be slightly worse than T5 at shorter lengths.
Third, the logarithmicvariant is slightly slower than two of the prior work (ALiBi, Rotary, and Sinusoidal) but consistentlyoutperform them at all extrapolation lengths.
These are elaborated as follows.
Logarithmic Variant vs. Power Variant.
In our proposed KERPLE framework, the logarithmicvariant is better than the power variant.
Precisely, the logarithmic variant is 4.4% faster and has lowerperplexities across all extrapolation lengths and all tasks.
Logarithmic Variant vs. T5.
In terms of speed, the logarithmic variant is 9.7% faster than T5.
In terms of extrapolation perplexity, the logarithmic variant is close to or slightly worse than T5when the extrapolation length is shorter than 2048, and consistently excels T5 at longer extrapolationlengths.
The tendency of extrapolation holds for all datasets evaluated in this work.
Logarithmic Variant vs. ALiBi, Rotary, and Sinusoidal.
The logarithmic variant is 1.6% slower,7.5% faster, and 3.0% slower than ALiBi, Rotary, and Sinusoidal.
The speed comparison makes sensebecause we require only a limited amount of learnable parameters for RPEs (at most 3 · H ).
Also, thelogarithmic variant consistently outperforms prior work at all extrapolation lengths and tasks.
[Table] [Table]
5.3 Experiments on Complicated Kernels
In addition to the practical variants (power & logarithmic) in section 4, we consider two complicatedversions of the composite kernel, Eq.
(4), as follows.(bias+wht) bias + weight:
[Equation]
with r 1 , r 3 > 0 and 0 < r 2 , r 4 ≤ 2 .
3-para-log) 3-parameter-logarithmic
:k
[Equation]
m ]
[Equation]
k [Equation] [Equation] )with r 1 , r 2 > 0 and 0 < r 3 ≤ 2 .
Recall the tensor product property of a kernel: if k 1 is a kernel on X and k 2 is a kernel on Y , thenk (( x, y ) , ( x ′ , y ′ ))
=
k 1 ( x, x ′ ) k 2 ( y, y ′ ) is a kernel on X × Y .
Therefore, (bias+wht) is the setting where we train a weight exp( − r 3 | m − n | r 4 ) and a bias kernel c − r 1 | m − n | r 2 .
q ⊤ m k n is multiplied bythe weight kernel and then added with the bias kernel.
(3-para-log) is the setting where we consider| m − n | r 3 in the log.
When r 3 = 1 , it is reduced to the logarithmic variant proposed in section 4.
We plug in these composite kernel k comp into our KERPLE framework, Eq.
(2), and test the performance of these RPE.
Compared with section 5.2, Table 5 suggests that these variants do not haveclear advantage in extrapolation performance, e.g., 3-para-log is slightly better in perplexity than the(two-parameter) logarithmic variant.
Thus, enlarging the complexity of kernels does not necessarilygive better performance in the context of RPE.
[Table]
Additional Analyses
Since the power and logarithmic variants derived from KERPLE achieve superior performance onlength extrapolation across various datasets, we investigate the underlying reason by visualizing theeffective length as shown in Figure 2.
The visualization works in the following procedure.
For each training dataset, the learnable parameters ( r ( h ) 1 , ..., r ( h ) ℓ ) associated with each head h ( in total) are extracted from the model checkpoint.
The CPD kernel at head h is ˜ k
( h ) =
˜
k r ( h ) 1 ,...,r ( h ) ℓ . Both the power and the logarithmic variants in corollary 1 undergo a similar procedure.
The onlydifference is that their ˜ k ’s are different.
For each head h , we compute the effective length of ˜ k
( h ) as eff ( h ) =
[Equation]
That is, the relative positional difference | m − n | such that ˜ k
( m, n ) shift-inv.
=
˜ k
(0 , | m − n | ) justbecomes smaller than -2.
Note ˜ k
(0 , | m − n | ) strictly decreases in | m − n | , so there is only onepossible value.
We pick − 2 here because ˜ k is a bias and is followed by the Softmax normalization.
A bias of − 2 or smaller can make a great impact on the output of Softmax 3 .
eff ( h ) is interpretedas the effective length because, when | m − n | < eff ( h ) , the attenuation due to ˜ k
( h ) is not strong.
When | m − n | > eff ( h ) , the attenuation is strong and has a large impact on q ⊤ m k n + ˜ k ( h ) ( m, n ) .3.
Then, for each | m − n | ∈ [0 , ..., 20480] , we count the number of heads that satisfies eff ( h ) ≤| m − n | .
This gives a cumulative plot as shown in Figure 2, where the x-axis is | m − n | and they-axis is Count ( { h : h ∈ [1 , ..., 12] , eff ( h ) ≤ | m − n |} ) .4.
Repeat the above steps for other datasets and kernels.
Interpretation of Curves.
For a point ( x, y ) on a curve, it means that there are y heads with atleast − 2 bias when the token distance | m − n | is greater than x .
In other words, the slower the yconverges to 12, the longer the inter-token range that the model focuses on.
The Advantage of Learnable Parameters.
We observe that ALiBi [Press et al., 2022] producesthe same curve no matter which dataset is used.
The reason is that ALiBi selects a fixed parameter [Figure] r = 2 − 8 hH at head h for its linear bias − r | m − n | ( H heads in total) regardless of the dataset.
Whilethis strategy is useful for extrapolation, we hypothesize that different datasets might have differentcharacteristics, e.g., the average distance of highly related tokens should differ among the datasets asshown in Figure 2.
These characteristics are easier adapted by learnable parameters.
Thus, we believethat learnable parameters have more advantages in capturing the dataset-dependent characteristics.
Trends Across Datasets.
We notice that both kernels trained on OpenWebText2 tend to focus moreon distant relations.
This makes sense because OpenWebText2 has the highest perplexity scoresamong all datasets, implying that more context is needed to disambiguate the next predicted token.
The opposite trend holds for Arxiv and GitHub datasets, which is reasonable considering their lowerperplexity scores.
Characteristics Learned by Kernels.
Under any dataset, the logarithmic variant tends to focusmore on distant relations than the power variant does.
We can explain it through their functionalforms.
Because logarithm ( − a log(1 + b | m − n | ) ) decays much slower than power ( − a | m − n | p )does, the log variant might encourage the focus on distant relations.
Conclusion and Future Work
A general framework, KERPLE, is proposed to kernelize relative positional embeddings for lengthextrapolation.
At the core of this framework is the application of CPD kernels and the derivation ofpractical variants.
We show that these CPD kernels can be implicitly converted to PD kernels, whichkeep the inner product interpretation of self-attention.
We also demonstrate that the logarithmicvariant achieves exceptional extrapolation performance on three large language modeling datasets.
We believe our work paves the way for some interesting future directions that resolve our limitations.
For instance, we can consider general kernel families and model non-monotonic effects due topositional differences.
In addition, the use of learnable parameters in KERPLE might enable bettergeneralization to inputs higher than 1-dimensional.
Last but not least, there is always room forimproving memory efficiency by adjusting the model architecture and training procedure.
Broader Impact
Our work develops a better understanding of relative positional embedding for transformers basedon expressive kernel classes that adapt well to various datasets.
The results apply to domains wherethe positional information is helpful in the modeling, e.g., natural language, programming language,and DNA/protein sequences for biology/medicine.
The studies of transformers may have positiveeconomic effects by enabling new tasks which cannot be done by humans or enhancing accuracy andefficiency.
But inappropriate use can have negative societal impacts.
These include job loss due toautomation, the ethical challenges from improper text generation, and the privacy issues in the datacollection process.
These implications apply to any research on natural language processing and arenot associated with any specific work.
Checklist
For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’scontributions and scope?
[Yes]
(b) Did you describe the limitations of your work?
[Yes] See the end of sections 4 and 7.
(c) Did you discuss any potential negative societal impacts of your work?
[Yes] Seesection 8.
(d) Have you read the ethics review guidelines and ensured that your paper conforms to
them?
[Yes]
If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results?
[Yes]
(b) Did you include complete proofs of all theoretical results?
[Yes] See Appendix A.1.3.
If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)?
[Yes] See section 5and the supplemental material.
(b) Did you specify all the training details (e
g., data splits, hyperparameters, how theywere chosen)?
[Yes] See Table 2.
(c) Did you report error bars (e.g., with respect to the random seed after running experi-
ments multiple times)?
[Yes] See Table 3.
(d) Did you include the total amount of compute and the type of resources
(e
g
type
cluster
cloud provider
(b) Did you mention the license of the assets?
[Yes] We use publicly released assets from
Black et al.
[2021], Gao et al.
and the licenses are under Apache License 2.0.
(c) Did you include any new assets either in the supplemental material or as a URL?
[Yes]
See section 5.
(d) Did you discuss whether and how consent was obtained from people whose data you’re
using/curating?
[Yes] The data are publicly released.
See section 5.
(e) Did you discuss whether the data you are using/curating contains personally identifiableinformation or offensive content?
[N/A]
The data comes from Gao et al.
[2020] where
the provider has already anonymized all user-contributed content.
See section 5.
If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable?
[N/A]
(b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable?
[N/A]
(c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation?
[N/A]
A Appendix
A.1 Proof of CPD Shift Lemma
Lemma 1 (CPD Shift Lemma) .
Let k : X
×
X
→ R be a conditionally positive definite (CPD)kernel.
Then, there exists c ≥ 0 such that c + k ( x, y ) is a positive definite kernel.
Proof.
Let K =
[ k ( x i , x j )]
Ni,j =1 be the matrix generated by { x 1 , ..., x N } with N ∈ N .
Consider f c ( v ) =
v ⊤ ( c 11 ⊤ + K ) v
= c ( v ⊤ 1 ) 2 + v ⊤ Kv.
We want to show there exists a large enough c such that f c ( v ) ≥ 0 for all v ∈ { v : ∥ v ∥ = 1 } .(i)
It is sufficient to consider
[Equation]
Let a ∗ be the solution to the minimization:
[Equation]
Since v ⊤ Kv is continuous in v and { v : ∥ v ∥ = 1 } is compact (i.e., closed and bounded), a ∗must exist.
If a ∗ ≥ 0 , K is positive semidefinite and f c ( v ) ≥ 0 for c ≥ 0 .
Thus, without loss ofgenerality, we assume a ∗ < 0 .
(ii) It is sufficient to consider KKK without zero eigenvalues (
i
e., full rank)
If there exists v 0 such that Kv
=
0 , then c ≥ 0 is enough to satisfy f c ( v 0 ) ≥ 0 .
For any v 1satisfying v ⊤ 1 v 0 = 0 , we have ( v 1 + v 0 )
⊤ K ( v 1 + v 0 )
=
v ⊤ 1 Kv 1 .
Therefore, whether there existsc to have f c ( v ) ≥ 0 doesn’t depend on the eigenvector corresponding to zero eigenvalue (if thereis such a vector).
This means it is enough to consider K without zero eigenvalues.
(iii)
It is sufficient to consider strict CPD
By definition of conditional positive definiteness (CPD), we know v ⊤ Kv ≥ 0 when v ⊤ 1 = 0 .
Since K has no zero eigenvalue, we cannot have v ⊤ Kv = 0 when v ⊤ 1 = 0 4 .
This means theinequality is strict here: v ⊤ Kv > 0 when v ⊤ 1 = 0 , and it is enough to consider strict CPD. iv)
It is sufficient to show there exists (small enough) δ > 0 δ > 0 δ > [-
-]0 such that v ′ ∈
T δ ⇒ v ′⊤ Kv ′ > 0 v ′ ∈ T δ ⇒ v ′⊤ Kv ′ > 0 v ′ ∈ T δ ⇒ v ′⊤ Kv ′ > 0 .
Note T δ is defined as
[Equation]
For any v ′ ∈ T δ , if v ′ satisfies v ′⊤ Kv ′ > 0 , then c ≥ 0 is enough to have f c ( v ′ ) ≥ 0 .
Conversely, when | v ⊤ 1 | ≥ δ and ∥ v ∥ = 1 , observe that
[Equation]
Then, f c ( v ) ≥ 0 when c ≥ − a ∗ δ 2 .
Therefore, we need to prove v ′ ∈ T δ ⇒ v ′⊤ Kv ′ > 0 for smallenough δ .
v)
Leveraging Continuity.
Consider v ′ satisfying ∥ v ′ − v ∥ <
δ 2 with v ∈ S
= { v : ∥ v ∥ =1 , v ⊤ 1 = 0 } .
Since v ⊤ Kv is continuous in v and ∥ K ∥ < ∞ , for any ϵ > 0 and any v ∈ S , asmall enough δ 2 > 0 gives | v ′⊤ Kv
′ − v ⊤ Kv | < ϵ .
To see this, taking v ′ = v + p with ∥ p ∥ <
δ 2 , we have | v ′⊤ Kv [Equation] − v ⊤ Kv | = | p ⊤
[Equation]
⊤
[Equation]
p ⊤
[Equation]
|
[Equation]
Therefore, 0 < δ 2 < (cid:112) 1 + ϵ/ ∥ K ∥ − 1 is enough to have | v ′⊤ Kv ′ − v ⊤ Kv | < ϵ .
By definition of strict CPD, we know min v ∈ S v ⊤ Kv = λ > 0 .
Thus, take ϵ < λ , a small enoughδ 2 gives v ′⊤ Kv ′ > v ⊤ Kv − ϵ > = λ − ϵ > 0 .
In other words, there exists a small enough δ 2such that v ′⊤ Kv ′ > 0 for v ′ ∈ S
δ 2 = { v ′ :
∥ v ′ − v ∥ < δ 2 , v ∈ S } .
(vi)
Proving
∃ δ > 0 s.t. v ′ ∈ T
δ
⇒ v ′⊤
Kv ′
> 0 ∃ δ > 0 s.t.
v ′ ∈ T δ ⇒ v ′⊤ Kv ′ > 0 ∃ δ > 0 s.t. v ′ ∈
T δ ⇒ v ′⊤ Kv ′ > 0 .
Due to (iv), we want to show ∃ δ > 0 s.t.
v ′ ∈
T δ ⇒ v ′⊤ Kv ′ > 0 .
We will prove by the conclusion of (v).
Let ∥ v ′ ∥
= 1 , v ′⊤ 1 = r with | r | < δ and v ′′ = v ′ − rn 1 .
We have
[Equation] [Equation]
Take v = v ′′ ∥ v ′′ ∥ = v ′′ q , where q = (cid:113) 1 − r 2 n .
We have v ⊤ 1 = 0 , ∥ v ∥ = 1 and∥
[Equation]
| x | ≪ 1)
Thus, ∥ v ′ − v ∥ =
O ( | r |√ n ) ≤ O ( δ √ n ) < δ 2 for small enough δ .
This implies that, with a smallenough δ , for any v ′ ∈ T δ , we can find v ∈ S such that ∥ v ′ − v ∥ < δ 2 .
Thus, v ′ ∈ S δ 2 , and by(v), we arrive at v ′⊤ Kv ′ > 0 .
A.2 Shift-invariant Kernels with Bounded and Unbounded Ranges
Definition 1 implies a shift-invariant kernel is generated by a univariate function f : X → R .
Tocharacterize the set of valid univariate functions, we introduce the positive definite functions as below.
Definition 3 (Positive definite function) .
A (real) positive definite function is a function f : X → Rsuch that for any integer N and any set of N points { x i ∈ X} Ni =1 , the N × N matrix A =[ f
( x i − x j )]
Ni,j =1 is positive semidefinite.
We will interchange the ideas of shift-invariant kernels and positive definite functions because theyare equivalent by definition.
Any statement in positive definite functions can be translated intoshift-invariant kernels, and vice versa.
Because of this, we will use some facts about the positivedefinite functions to derive the shift-invariant kernels of our interest.
Generalizing Classical Bounded Shift-invariant Kernel.
In the literature, there has been studies on applying kernels in the attention mechanism [Tsai et al., 2019, Choromanski et al., 2021, Peng et al., 2021].
One of the most common approach is to consider the Gaussian kernel:
[Equation]
Note the Gaussian kernel is bounded ( k ( m, n ) ∈ (0 , 1] for the case above).
To generalize it to abroader class of bounded shift-invariant kernels, observe that the Gaussian kernel generated by apositive definite function of the form f ( x ) =
exp( −| x | 2 ) .
Since there is no strong reason to stick tothe power of 2, one may generalize it to a broader class of positive definite functions as below.
Fact 5 (Corollary 3 of Schoenberg [1938]) .
exp( −| x | p ) is positive definite if 0 < p ≤ 2 and notpositive definite if p > 2 .
Fact 5 implies that, if one wants to find a class of bounded shift-invariant kernel (i.e., k ( m, n ) iswithin some fixed interval for any m, n ), then k ( m, n ) =
exp( − a | m − n | p ) with a > 0 and p ∈ (0 , 2]may be of interest.
Constructing Unbounded Shift-invariant Kernels.
A limitation of Fact 5 is that it only generateskernels with a bounded range (here, the range is bounded in (0 , 1] ).
In situations where there are noexplicit bounds, one might want to consider kernels with unbounded range.
To construct such kernels,we utilize a kernel representation theorem presented in Schoenberg [1938]:
Fact 6 (Theorem 4 of Schoenberg [1938]) .
f ( x ) is bounded away from zero ( f ( x ) > 0 ) and its positive powers f ( x ) λ ( λ > 0 ) are all positive definite if and only if f ( x ) is of the form f ( x ) =
exp( c + ψ ( x )) , where ψ ( x ) is positive definite and c is a real constant.
Since Fact 6 works for non-negative kernels, we combine it with Fact 5 and show the following classof shift-invariant kernel with an unbounded range.
Proposition 1 (Kernel from Distance Powers) .
For any p ∈ (0 , 2] , there exists c min ∈ R such that for any c ≥ c min ,
[Equation]
, is a positive definite kernel.
When p > 2 , there is no c to make k ( m, n ) positive definite.
Proof.
Due to Fact 5, we know exp( −| x | p ) is positive definite when p ∈ (0 , 2] .
Since exp( −| x | p ) > 0 and exp( − λ | x | p ) is positive definite for any λ > 0 5 , Fact 6 implies there exists a c ′ ∈ R and a positive definite ψ ( x ) such that
[Equation]
In other words, − c ′ −| x | p is a positive definite function.
Take c min = − c ′ .
We see that c min −| m − n | p is a shift-invariant positive definite kernel.
Finally, let k ( m, n ) =
c − | m − n | p with c ≥ c min .
The
N × N matrix
[ k ( x i , x j )]
Ni,j =1 generated by k ( m, n ) on points { x i ∈ R } Ni =1 obeys [ {+k ( x i , x j )] Ni,j =
+}
= [ c min − | x i − x j | p ] Ni,j =1 + ( c − c min ) 11 ⊤ ⪰ ( c − c min ) 11 ⊤ ⪰ 0 , where ⪰ is the Loewner order and 1 =
[1 , ..., 1] ⊤ in the N -dimensional vector with all ones.
Thisshows k
( m, n ) is a shift-invariant positive definite kernel when 0 < p ≤ 2 .
The conclusion on p > 2is proved by contradiction.
When p > 2 , if there exists a c such that k ( m, n ) is positive definite, thenexp( k ( m, n )) is positive definite, which contradicts to the case of p > 2 in Fact 5.
Prop. 1 introduces a kernel with unbounded range ( k ( m, n ) ∈ ( ∞ , c ] ) and is adapted from the p-thpower of the distance | m − n | .
Since the distance is a notion of "dissimilarity", −| m − n | p becomesa notion of "similarity", which gives a sense of kernel.
Thereby, we can interpret the constant c as therequired value to shift −| m − n | p such that c − | m − n | p becomes a positive definite kernel.
In fact, −| m − n | p is a conditional positive definite kernel for p ∈ (0 , 2] Schölkopf [2000].
Therefore,the fact that −| m − n | p can become a positive definite kernel by shifting is not a coincidence, as ithas already had an intimate relation to positive definite kernels.
A.3 Experiments on Gaussian-like Kernels
Since the prior work on shift-invariant kernels mainly focuses on Gaussian kernels, we presentpreliminary experiments on Gaussian-like kernels.
Compared with section 5.2, the perplexities ofthese kernels are large at every extrapolation length.
This verifies our previous assertion
that the
Gaussian-like kernels have limited extrapolation ability.
Because the kernel can be used as a weight or a bias, we consider four kinds of the composite kernel(see section 4) as follows.
(2-para-bias) r 1 , r 2 > 0 .
[Equation]
(3-para-bias) r 1 , r 2 > 0 and 0 < r 3 ≤
[Equation]
[Equation]
(1-para-wht) r 1 >
[Equation] [Equation]
, n ]) =
q ⊤ m k n · exp( − r 1 | m − n | 2 ) .(2-para-wht) r 1 > 0 and 0 < r 2 ≤ 2 .k comp ([ q m , m ] , [ k n , n ]) =
q ⊤ m k n · exp( − r 1 | m − n | r 2 ) .(2-para-bias) and (1-para-wht) are the settings where we put the Gaussian kernel as a bias anda weight, respectively.
(3-para-bias) and (2-para-wht) generalize these settings by considering alearnable power between 0 and 2.
Note we must constrain the power in (0,2]; otherwise, the functionis not positive definite.
See Fact 5 for details.
These composite kernel k comp are plugged into the KERPLE framework, Eq.
(2), and are evaluated on
OpenWebText2, GitHub, and ArXiv datasets.
Table 6 shows the Gaussian-like kernel is better to be aweight instead of a bias.
As discussed in Appendix A.2, the Gaussian-like kernels are bounded.
Tosome extent, this implies that the bounded positive kernel can model a weight.
However, comparedwith section 5.2 the Gaussian-like kernels have limited advantages in extrapolation.
Although theperformance might be improved if the power of exp( −| x | p ) is relaxed from p = 2 to p ∈ (0 , 2] , stillit cannot be as good as the logarithmic variant as we demonstrate in section 5.2.
Therefore, while the
Gaussian kernel is frequently used in the literature, we need a better class of shift-invariant kernels totackle the length extrapolation challenge.
A.4 Position-wise Perplexity Evaluation
We follow the reviewer’s suggestion and plot the position-wise perplexity with different evaluationlengths in Figure 3.
The evaluation is done by measuring the loss at each position in each sequenceand averaging over the sequences to get the position-wise perplexity.
First, we notice that as a general trend, tokens appear at the beginning of an evaluation sequencehave much higher PPL; This is expected as these tokens have access to only a limited amount ofcontext (early token curse).
Second, we focus on Figure 3(a) and observe that all methods exceptrotary can extrapolate smoothly to length=4096.
We can also observe that our proposed KERPLE-logvariant achieves the lowest PPL with a continuously decreasing trend toward length=4096.
We wantto highlight the decreasing trend indicates successful extrapolation as our model is using distant tokeninformation not seen during training time (training length=512).
Finally, we push the evaluationsequence length to 16384 and plot the result in Figure 3(b); This time we can see that the PPL of
T5 bias starts to explode after length=5000, which also corroborates the finding (that T5 bias cannotextrapolate well) in the original ALiBi paper.
In addition, we can see that the decreasing trend of our
PPL is not as prominent as Figure 3(a).
This is also expected as there is very likely a upper limit ofhow distant two relavant tokens could be.
In summary, using larger context and embedding is reallyhelpful to our model up to length=4096 on the GitHub dataset.
[Table]
[Table] Following reviewer’s further suggestion, we provide the zoom-in plots in Figure 4 and 5 and simulatea model with window attention using the existing numbers. To do so, we make one small changeto the original Figure 3: Instead of plotting the horizontal line signifying the minimum value of
KERPLE-log until length=4096, we plot another curve, which is the KERPLE-log-windowed@512variant. This curve is constructed by sharing the same first 512 PPLs (y-values) with KERPLE-logfollowed by a horizontal line equal to PPL@512 of KERPLE-log (i.e. for pos. diff.>512, y-value= KERPLE-log PPL@512). 512 is chosen because it was used as the train length.
We note that
PPL@512 of KERPLE-log is the lowest PPL among the first 512 positions of all model variants. Wecan derive several critical observations for evaluation length=4096 in Figure 4:1. KERPLE-log lies below KERPLE-log-windowed@512, indicating its usage of more distantinformation than window attention: If our model is not using more information other thana fixed-window=512 tokens , the y-values after position=512 should also overlap with the
new horizontal line (window attention). This is clearly not the case. In addition, the PPL of KERPLE-log continues to decrease till the end of 4096 positions (Not plateauing).
T5 lies below KERPLE-log-windowed@512 for most of the time and fluctuates around KERPLE-log-windowed@512 after length=3000. It is still worse than KERPLE-log. ALiBi lies above KERPLE-log-windowed@512 for almost all the positions, indicating that
window attention might be a better choice than ALiBi. As for the long run comparison (evaluation length=16384) in Figure 5: KERPLE-log lies below KERPLE-log-windowed@512 for most of the time, indicating its usage of more distant information than window attention.
The PPL of T5 explodes.
The PPL of ALiBi does not explode, but it is still worse than window attention,
i e. lies
above KERPLE-log-windowed@512.
Although window attention is a strong baseline, our KERPLE-log is almost like a free upgradecompared to window attention: With only 24 additional learnable parameters (2 para. for each head),the almost same training speed, and the same train length=512 as window attention, our model isable to achieve lower PPLs across different positions.
[Figure]
[Figure] [Figure]
Position A v e r a g e P e r p l e x i t y
[Figure]
[Figure] [Figure]
KERPLE-log-windowed@
A.5 Experiments on Large Model, Longer Training {+Length, and Wikitext-
+}
[-and Wikitext-
-] In this subsection, we present additional experiments on (a) large models, (b) longer training length,and (c) Wikitext-103.
Below is the summary of the experiments.
(a) The 1.3B large model is trained on a machine with two NVIDIA A100 GPU with 40 GB ofmemory.
We adopt almost all configurations of XL GPT-NeoX 6 , except that we change the train- micro-batch-size to 16, model-parallel-size to 2, seq-length to 512, and max-position-embeddingsto 512.
Table 8 summarizes the configurations of the 1.3B model. [Figure] Position
A v e r a g e P
e r p l e x i t y Position-wise Perplexity Length
Rotary
ALiBi
T
KERPLE-log
KERPLE-log-windowed@
b) The 162M Model with training sequence length=1024 follows the same configurations as the
ones in Table 2 except that the train seq.
length is changed to 1024.
(c) The Wikitext-103 model is implemented on ALiBi’s
GitHub 7 with exactly the same config-
urations (247M parameters), except that the function buffered_future_mask() at line 1011 of attention_with_linear_biases/fairseq/models/transformer.py is adapted to our KERPLE-log. [Table]
[Table]
Table 9 shows the results on the large model (1.3B).
Compared with the small model results in
Table 3, we see that T5 bias becomes weaker than KERPLE-log and ALiBi, and KERPLE-logremains stronger than ALiBi.
This makes sense because T5 is more susceptible to overfitting and thelarge model has a higher risk of overfitting.
Again, Table 9 shows the results on long training length (1024).
compared with the short traininglength (512) in Table 3, KERPLE-log remains better than ALiBi and T5 bias, especially on longerevaluation length.
This shows the robustness of KERPLE-log over different training lengths.
Table 10 compares KERPLE-log with ALiBi using ALiBi’s implementation and configurations.
Theresults show that KERPLE-log is superior to ALiBi on Wikitext-103.
A.6 Plots of Kernel Functions
We follow reviewer’s suggestion to add the plots of kernel functions including the power, log variants,and ALiBi for different heads (not doing thresholding and selecting exp(-2)) to see the differencesin tail and contribution to softmax.
We again use the GitHub dataset for demonstration.
Please see
Figure 6, 7, and 8 for details.
We can see that both the ALiBi and its generalized power variant quicklyreach a very negative value.
In contrast, the log variant successfully discovers several kernels that
[Table]
[Table] is trained on ALiBi’s codebase with exactly the same configurations except for the positional embeddings.
The results show that KERPLE-log is superior to ALiBi on Wikitext-103.
[Table]
remain flat, effectively extending the window attention.
This corroborates our previous observationthat KERPLE-log can utilize more distant token information: Flat kernel functions permit the usageof more distant token information in that distant information would not be masked out by the implicitwindow attention effect in ALiBi.
A.7 The Choice of codebase and Hyperparameters
We adopt almost all the hyperparameters (except batch size to fit in our GPU) and all implementationsof the T5 bias, ALiBi, Rotary, and Sinusoidal baselines from the GPT-NeoX codebase. To ensure faircomparisons, we did not fine-tune hyper-parameters for KERPLE.
The datasets we used are exactlythe same as the ones released with the GPT-NeoX codebase.
We just ran their prepare_data.py scriptto automatically download and parse the datasets.
All our code was uploaded with the submission onopenreview, and https://github.com/EleutherAI/gpt-neox is the original GitHub repository.
As a side note, we chose this codebase and adopt their parameter settings because it is built by
EleutherAI, which is a well-known and truly non-profit group of researchers publishing variousfamous pretrained models for the academia including GPT-J-6B and GPT-NeoX-20B.