Transformer升级之路:12、无限外推的ReRoPE?
By 苏剑林 | 2023-08-07 | 62633位读者 |自从在《Transformer升级之路:11、将β进制位置进行到底》中引入混合进制的思路进一步推广了NTK-aware Scaled RoPE后,笔者感觉类似思路的效果已经达到了上限,想要更大幅度的提升就必须另辟蹊径了。这时候笔者想起了此前构思过的一个思路,该思路由于复杂度较高所以被搁置下了,既然现在已经遇到了瓶颈,那么“唯一的办法就是最好的办法”,于是便将它重拾起来。
万万没想到的是,尽管该方法增加了一些推理复杂度,但它的实验效果却惊人地好——甚至隐约有无限的长度外推能力!因此,笔者迫不及待地撰写了本文来分享该方法。由于形式上跟ReLU激活函数的相似性,所以笔者将该方法命名为“ReRoPE (Rectified Rotary Position Embeddings)”。
重温 #
我们知道,RoPE形式上是一种绝对位置编码,但实际上给Attention带来的是相对位置信息,即如下的Toeplitz矩阵:
\begin{equation}\begin{pmatrix}0 & \\
1 & 0 & \\
2 & 1 & 0 &\\
3 & 2 & 1 & 0 & \\
\ddots & 3 & 2 & 1 & 0 & \\
\ddots & \ddots & 3 & 2 & 1 & 0 & \\
\ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{L - 2} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{L - 1} & \small{L - 2} & \ddots & \ddots & \ddots & 3 & 2 & 1 & 0 & \\
\end{pmatrix}\label{eq:rope}\end{equation}
这里的$L$是当前样本长度。当$L$明显超出了训练长度时,多出来的位置由于没有被训练过,所以无法保证效果,这就是直接外推(Length Extrapolation)表现通常比较差的原因。
后来,研究人员提出了位置内插(Position Interpolation),它相当于将相对位置矩阵改为:
\begin{equation}\begin{pmatrix}0 & \\
\frac{1}{k} & 0 & \\
\frac{2}{k} & \frac{1}{k} & 0 &\\
\frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{\frac{L-2}{k}} & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots & \ddots \\
\small{\frac{L-1}{k}} & \small{\frac{L-1}{k}} & \ddots & \ddots & \ddots & \frac{3}{k} & \frac{2}{k} & \frac{1}{k} & 0 & \\
\end{pmatrix}\end{equation}
这样一来,只要调整$k$,就可以保证最大的相对位置也不超过训练长度,因此避免了外推。然而,它使得位置信息更加“拥挤”了,所以还需要进行一定步数的微调才能让模型重新工作。而也正因为避免了外推,所以它所需要的微调步数相比直接外推要少得多(神经网络往往更擅长内插而不是外推)。
至于后面提出的NTK-aware Scaled RoPE,则是“剑走偏锋”,巧妙地将外推压力平摊到每一个维度上,所以它不微调也能有不错的效果,但它终究还是依赖外推,这是神经网络不擅长的事情,所以效果存在上限,在笔者的实验中,它的Long Context表现还无法很接近训练效果。
融合 #
我们也可以从语言模型的局域性来考察这些方法。所谓局域性,是指语言模型在推断下一个token时,明显更依赖于邻近的token。直接外推保持了局域性(0附近位置编码不变),效果差是因为引入了超出训练长度的位置编码;位置内插虽然没有外推位置编码,但扰乱了局域性(0附近位置编码被压缩为$1/k$),所以不微调效果也不好;而NTK-aware Scaled RoPE通过“高频外推、低频内插”隐含了两者优点,保证了局域性,又没有明显外推位置编码,所以不微调也有不错的效果。
有没有能更直接地结合外推和内插的方法呢?有,我们可以设定一个窗口大小$w$,在窗口内我们使用大小为$1$的位置间隔,在窗口外我们使用大小为$1/k$的位置间隔,整个相对位置矩阵如下:
\begin{equation}\begin{pmatrix}
\color{red}{0} & \\
\color{red}{1} & \color{red}{0} & \\
\color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\small{w + \frac{L-1-w}{k}}} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\small{w + \frac{2}{k}}} & \color{green}{\small{w + \frac{1}{k}}} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\end{pmatrix}\label{eq:leaky-rerope}\end{equation}
只要$w$小于训练长度,那么通过控制$k$,我们就可以在精确保持了局域性的前提下,使得所有位置编码不超过训练长度,简单直接地结合了直接外推和位置内插。
特别地,矩阵$\eqref{eq:leaky-rerope}$还有一个特别的case:当$k\to\infty$时,它简化为
\begin{equation}\begin{pmatrix}
\color{red}{0} & \\
\color{red}{1} & \color{red}{0} & \\
\color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{green}{w} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{green}{w} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \color{red}{\ddots} & \\
\color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{green}{w} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\color{green}{w} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{\ddots} & \color{green}{w} & \color{green}{w} & \color{green}{w} & \color{red}{\small{w - 1}} & \color{red}{\ddots} & \color{red}{2} & \color{red}{1} & \color{red}{0} & \\
\end{pmatrix}\label{eq:rerope}\end{equation}
在这个case下,不管输入长度是多少,它的位置编码范围都不超过$w$,所以这是一种有可能支持任意长度的Context的方案!
形式上,矩阵$\eqref{eq:rerope}$、$\eqref{eq:leaky-rerope}$与标准RoPE矩阵$\eqref{eq:rope}$的关系,就相当于ReLU、Leaky ReLU与Linear的关系,所以笔者将$\eqref{eq:rerope}$称为“ReRoPE(Rectified RoPE)”,将$\eqref{eq:leaky-rerope}$称为“Leaky ReRoPE”。
计算 #
其实,类似的思路并不难想到,以往基于Attention Bias的相对位置编码(比如经典相对位置编码、T5位置编码)经常会出现这样的分块运算。然而跟这些相对位置编码不同,在RoPE中实现这样的分块运算会明显增加计算量,这也是该思路会被笔者搁置的主要原因。
怎么理解增加计算量呢?我们知道RoPE是“通过绝对位置实现相对位置”,这样只能得到线性的相对位置,而矩阵$\eqref{eq:leaky-rerope}$、$\eqref{eq:rerope}$是非线性的(或者说分段线性的),要实现它只能算两次Attention矩阵,然后组合起来。具体来说,首先用标准的RoPE计算一次Attention矩阵(Softmax之前)
\begin{equation}a_{i,j}^{(1)} = \left(\boldsymbol{\mathcal{R}}^i\boldsymbol{q}_i\right)^{\top}\left(\boldsymbol{\mathcal{R}}^j\boldsymbol{k}_j\right) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^{j-i}\boldsymbol{k}_j\end{equation}
这里第一个等号是实现方式,第二个等号是等效结果,其中$\boldsymbol{\mathcal{R}}$就是RoPE的旋转矩阵,简单起见我们省略了Attention的scale因子。接着,我们需要计算间隔为$1/k$的RoPE的Attention矩阵(Leaky ReRoPE):
\begin{equation}a_{i,j}^{(2)} = \left(\boldsymbol{\mathcal{R}}^{(i-w)/k+w}\boldsymbol{q}_i\right)^{\top}\left(\boldsymbol{\mathcal{R}}^{j/k}\boldsymbol{k}_j\right) = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^{(j-i+w)/k-w}\boldsymbol{k}_j\end{equation}
如果是ReRoPE,那么简单一些:
\begin{equation}a_{i,j}^{(2)} = \left(\boldsymbol{\mathcal{R}}^w\boldsymbol{q}_i\right)^{\top}\boldsymbol{k}_j = \boldsymbol{q}_i^{\top}\boldsymbol{\mathcal{R}}^w\boldsymbol{k}_j\end{equation}
最后,根据$i - j < w$这个条件,将它们合并起来:
\begin{equation}a_{i,j} = \left\{\begin{aligned}
&a_{i,j}^{(1)},\quad (i - j < w) \\[8pt] &a_{i,j}^{(2)}, \quad (i - j \geq w)
\end{aligned}\right.\end{equation}
不管是ReRoPE还是Leaky ReRoPE,都不可避免地计算两次Attention矩阵(如果有更高效的实现方法,请赐教),这便是增加的计算量之一。此外,需要自定义计算Attention矩阵也导致了不能直接套用现成的flash attention实现,因此相对之下又增加了计算成本。
另一方面,同样是由于非线性的相对位置,所以在自回归解码时,Key序列的cache只能存RoPE之前的,然后在每步解码时给整个Key序列补上对应的RoPE,这样的改动也会增加推理计算量。唯一的好消息是,在token by token解码时,从第二步开始Query序列的长度就为1,此时只需要为Key序列定制RoPE,那么可以只算一次Attention矩阵:
\begin{equation}a_{i,j} = \left\{\begin{aligned}
&\boldsymbol{q}_i^{\top}\left(\boldsymbol{\mathcal{R}}^{\max(j-i,-w)}\boldsymbol{k}_j\right), \quad(\text{ReRoPE})\\[8pt]
&\boldsymbol{q}_i^{\top}\left(\boldsymbol{\mathcal{R}}^{\max(j-i,(j-i+w)/k-w)}\boldsymbol{k}_j\right), \quad(\text{Leaky ReRoPE})
\end{aligned}\right.\end{equation}
实验 #
继续沿着《Transformer升级之路:11、将β进制位置进行到底》的设置,我们对ReRoPE进行了实验,效果如下表:
\begin{array}{c|cc}
\hline
\text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\
\hline
\text{Baseline} & 49.41\% & 24.17\% & 23.16\% \\
\text{Baseline-}\log n & 49.40\% & 24.60\% & 24.02\% \\
\hline
\text{PI-RoPE} & 49.41\% & 15.04\% & 13.54\% \\
\text{PI-RoPE-}\log n & 49.40\% & 14.99\% & 16.51\% \\
\hline
\text{NTK-RoPE-old} & 49.41\% & 51.28\% & 39.27\% \\
\text{NTK-RoPE-}\log n\text{-old} & 49.40\% & 61.71\% & 43.75\% \\
\hline
\text{NTK-RoPE-fixed} & 49.41\% & 51.86\% & 39.61\% \\
\text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-fixed} & 49.41\% & 55.94\% & 41.11\% \\
\text{NTK-RoPE-}\log n\text{-fixed} & 49.40\% & 62.85\% & 44.14\% \\
\text{NTK-RoPE-mixed} & 49.41\% & 53.09\% & 40.12\% \\
\text{NTK-RoPE-}\log n^{\color{red}{\dagger}}\text{-mixed} & 49.41\% & 59.11\% & 42.38\% \\
\text{NTK-RoPE-}\log n\text{-mixed} & 49.40\% & 68.91\% & 45.41\% \\
\hline
\text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\
\text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\
\text{ReRoPE-w256-}\log n & 49.40\% & \boldsymbol{85.12\%} & \boldsymbol{49.07\%} \\
\hline
\text{HFWA} & 48.70\% & 80.84\% & 48.15\% \\
\hline
\end{array}
正如文章开头所说,ReRoPE不微调外推的效果可谓出奇地好,不仅明显超越了此前最优的NTK-RoPE-mixed,还明显超过了从零预训练的HFWA!这里的$\text{w256}$指的$w=256$,$\log n^{\color{red}{\dagger}}$是指预训练没有加入$\log n$缩放(比如LLAMA),测试阶段每个$\boldsymbol{q}_n$都乘上$\max(1, \log_{\text{maxlen}} n)$,$\log n$则是指预训练就加入了$\log n$缩放因子。
以下是一些消融实验,显示出ReRoPE关于$w$还是很鲁棒的,最优值大致是训练长度的$1/4\sim 1/2$左右:
\begin{array}{c|cc}
\hline
\text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\
\hline
\text{ReRoPE-w64} & 49.41\% & 69.39\% & 45.19\% \\
\text{ReRoPE-w64-}\log n^{\color{red}{\dagger}} & 49.41\% & 78.58\% & 47.42\% \\
\text{ReRoPE-w64-}\log n & 49.40\% & 84.38\% & 48.14\% \\
\hline
\text{ReRoPE-w128} & 49.41\% & 76.11\% & 47.82\% \\
\text{ReRoPE-w128-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.28\% & 48.78\% \\
\text{ReRoPE-w128-}\log n & 49.40\% & \boldsymbol{85.47\%} & 48.87\% \\
\hline
\text{ReRoPE-w256} & 49.41\% & 77.90\% & 48.48\% \\
\text{ReRoPE-w256-}\log n^{\color{red}{\dagger}} & 49.41\% & 82.40\% & 48.85\% \\
\text{ReRoPE-w256-}\log n & 49.40\% & 85.12\% & \boldsymbol{49.07\%} \\
\hline
\text{ReRoPE-w384} & 49.41\% & 70.72\% & 48.15\% \\
\text{ReRoPE-w384-}\log n^{\color{red}{\dagger}} & 49.41\% & 76.42\% & 48.31\% \\
\text{ReRoPE-w384-}\log n & 49.40\% & 83.24\% & 48.62\% \\
\hline
\text{ReRoPE-w512} & 49.41\% & 7.09\% & 8.25\% \\
\text{ReRoPE-w512-}\log n^{\color{red}{\dagger}} & 49.41\% & 7.08\% & 8.25\% \\
\text{ReRoPE-w512-}\log n & 49.40\% & 15.84\% & 10.83\% \\
\hline
\end{array}
下表则对比了ReRoPE和Leaky ReRoPE:
\begin{array}{c|cc}
\hline
\text{测试长度} & 512(\text{训练}) & 4096(\text{重复}) & 4096(\text{不重复})\\
\hline
\text{ReRoPE-w128-}\log n & 49.40\% & \boldsymbol{85.47\%} & 48.87\% \\
\text{Leaky ReRoPE-w128-k64-}\log n & 49.40\% & 85.29\% & 48.96\% \\
\text{Leaky ReRoPE-w128-k32-}\log n & 49.40\% & 85.31\% & 49.03\% \\
\text{Leaky ReRoPE-w128-k16-}\log n & 49.40\% & 85.15\% & \boldsymbol{49.10\%} \\
\text{Leaky ReRoPE-w128-k8-}\log n & 49.40\% & 80.00\% & 48.11\% \\
\hline
\text{ReRoPE-w256-}\log n & 49.40\% & 85.12\% & 49.07\% \\
\text{Leaky ReRoPE-w256-k64-}\log n & 49.40\% & 84.60\% & 49.03\% \\
\text{Leaky ReRoPE-w256-k32-}\log n & 49.40\% & 84.30\% & 48.97\% \\
\text{Leaky ReRoPE-w256-k16-}\log n & 49.40\% & 83.59\% & 48.87\% \\
\text{Leaky ReRoPE-w256-k8-}\log n & 49.40\% & 69.80\% & 45.72\% \\
\hline
\end{array}
作为ReRoPE的一般化,经过精调的Leaky ReRoPE是有机会超过ReRoPE的,但提升很微弱。此外,当$k$取有限值时,能处理的最大长度也是有限的,因为我们不能提前知道要生成的总长度,所以只能预设一个足够大的$k$,但设定为有限值之后,当输入足够长时,就会因为位置编码超出训练长度而效果大幅下降,相比之下ReRoPE则不会有这个风险。总的来说,精调Leaky ReRoPE相比ReRoPE的价值似乎不大。
以上实验结果都只是在1亿参数的GAU模型上测试的,下面给出基于llama2-13b的测试结果(指标是loss,越小越好),它代表了在真正的LLM表现:
\begin{array}{c|cc}
\hline
\text{测试长度} & 4096(\text{训练}) & 8192 & 16384\\
\hline
\text{RoPE} & 1.4967 & 8.8615 & \text{-} \\
\text{NTK-RoPE} & 1.6081 & 1.5417 & 1.5163 \\
\text{ReRoPE} & 1.4996 & 1.4267 & 1.4001 \\
\hline
\end{array}
可以看到,ReRoPE真正做到了几乎不损训练效果(RoPE-4096代表训练效果),并且满足“longer context, lower loss”的理想特点(更多的context应该更加有助于预测)。此外,笔者也在OpenBuddy开源的LLAMA2-13b微调模型上测试了chat的效果,自我感觉还不错(最多测试过20k tokens的Context)。
最后,分享笔者在transformers的LLAMA模型基础上实现ReRoPE和Leaky ReRoPE的代码,读者也可以自行加载LLAMA系列模型进行测试:
小结 #
在这篇文章中,笔者提出了ReRoPE (Rectified RoPE),它同样是一种RoPE的后处理方案,实验结果显示它的不微调长度外推能力不仅明显超过了此前的NTK-aware Scaled RoPE,甚至还超过了之前专门设计的需要从零训练的HFWA。此外,不同于NTK-aware Scaled RoPE在超过某个长度后能力会大幅下降,ReRoPE似乎在任意长度下都表现良好。除了对比实验外,文章还给出了基于transformers-llama的参考实现,有兴趣的读者可以自行测试。
转载到请包括本文地址:https://spaces.ac.cn/archives/9708
更详细的转载事宜请参考:《科学空间FAQ》
如果您还有什么疑惑或建议,欢迎在下方评论区继续讨论。
如果您觉得本文还不错,欢迎分享/打赏本文。打赏并非要从中获得收益,而是希望知道科学空间获得了多少读者的真心关注。当然,如果你无视它,也不会影响你的阅读。再次表示欢迎和感谢!
如果您需要引用本文,请参考:
苏剑林. (Aug. 07, 2023). 《Transformer升级之路:12、无限外推的ReRoPE? 》[Blog post]. Retrieved from https://spaces.ac.cn/archives/9708
@online{kexuefm-9708,
title={Transformer升级之路:12、无限外推的ReRoPE?},
author={苏剑林},
year={2023},
month={Aug},
url={\url{https://spaces.ac.cn/archives/9708}},
}
August 14th, 2023
为什么不能用一次计算拟合ReRePE呢?类似$a_{i,j}= q^T_i\Re^{w(1-e^{i-j})}k_j$
RoPE只能通过“绝对位置编码来实现相对位置编码”,因为直接按照相对位置的结果硬算的话成本更加大。
August 14th, 2023
苏老师考虑过用softplus这种形状的函数来模拟这样的区分么,这样似乎转变会更加平滑一点,而且也可以在趋近无限的同时一直保持递增(而不是都是 w)
同这里@苏剑林|comment-22509。如果不去从cuda层面优化实现(参考这里@袁境阳|comment-22463)的话,RoPE只能通过“绝对位置编码来实现相对位置编码”来实现线性的相对位置,以及通过多次重复计算来实现分段线性的相对位置,重复计算的次数取决于分段线性的段数,所以完全光滑的非线性相对位置基本上是做不到的。事实上我也试过分更多段,效果并没有优势,所以最终只保留了ReRoPE和Leaky ReRoPE。
August 14th, 2023
在不微调的情况下,ReRoPE的最优窗口长度为啥不是设置和训练长度一样呢?
不好意思,打扰了,明白了
这个大概是因为相对位置的训练不充分问题,在 https://kexue.fm/archives/9706 中我们也提过,越大的相对位置训练次数越少。
August 15th, 2023
学习了苏神的位置编码系列博文,感觉收获很大。关于提高计算两次Attention矩阵的效率的问题,参考LLAMA开源的代码,个人有一点点愚见,不知是否有用。
1.q、k中,使用相同频率编码的相邻各两维转换为复数的实部与虚部;
2.位置编码中对应上述q和k中各两维的2*2旋转矩阵直接写为复数$cos(\alpha)+isin(\alpha)$,所有这些位置编码计为位置编码矩阵。其计算可以先算好各$\alpha$,然后整体算$\exp(i\alpha)$;。
3.先乘q和kT,然后按位乘位置编码矩阵。
按这个思路ReRoPE和LeakRoPE的计算中所需计算两次的无非是$\alpha$组成的矩阵,代价不太大。FlashAttention应该也可以直接用。
水平有限,想法不一定对,欢迎大家批评
有点不大懂。你是说$Q,K$先相乘但不求和吗?如果$Q,K\in\mathbb{R}^{b\times h\times n\times d}$,相乘不求和得到的矩阵大小是$b\times h\times n\times n\times d$,显存受不了呀...
其实应用Flash Attention的困难不是算两次Attention矩阵,算两次Attention矩阵简单,放到multi-head那里算也行,问题是Flash Attention是一个包含$Q,K,V$完整的Attention运算,而ReRoPE需要将两个Attention矩阵合并为一个,然后再softmax再乘以$V$。
抱歉我写的不太对。确实有这个问题。感谢苏神拨冗指导
August 29th, 2023
非常感谢 ~ 后面测试下~
August 31st, 2023
感觉这种方法和flash attention的算法可以很容易结合,而且flash attention分块计算的方法也恰恰可以用来减少冗余的矩阵计算。用triton简单试了下前向传播的实现供参考(https://gist.github.com/chu-tianxiang/4307937fd94b49c75b61a6967716bae9)。
感谢感谢,确实是任何分block实现attention的思路都有可能用来改进它。可惜我很少用torch,不然就好好学习了。
还有一个改进的空间,就是直接以block size为单位,实现“锯齿状”的rerope,这样就完全不存在重复计算了。当block size明显小于window时,结果几乎是没区别的。
September 6th, 2023
苏神,刚开始学这一块,请教个比较小白的问题,就是目前已有的长度外推方法,有一些是需要微调的,有一些不需要是吗?
是的。
September 20th, 2023
苏神您好!非常感谢您的分享!有一句内容我没有太懂,"由于非线性的相对位置,所以在自回归解码时,Key序列的cache只能存RoPE之前的,然后在每步解码时给整个Key序列补上对应的RoPE,这样的改动也会增加推理计算量。"这个可否展开讲一下呢?或者有其他文章的引用。没有太明白为什么只能存RoPE之前的,以及为什么补上对应的RoPE就可以了呢?
我不知道你理解实现ReRoPE要算两次Attention矩阵没有?如果理解了,那就好办。因为要算两次Attention矩阵,所以要不你就把算这两次Attention矩阵的K(带有不同的位置信号)都缓存起来,要不你就只缓存一份没有位置信号的K,然后在解码的时候动态把位置信息加进去。
September 25th, 2023
你好,请问苏神,如何做一个博客系统?需要全栈知识吗
一般情况下,就是选一个框架,然后买一台对应环境的云服务器,然后像安装软件那样安装上去,简单配置一下就可以了。
October 8th, 2023
苏神,RoPE的一系列外推策略最终是要达到什么目的呢?
是不是想要让模型在输入任意长度下,都能full attention的进行计算和有效生成?
但是这样的话,总是会面对显存不足的问题吧?真实来讲,永远做不到任意长度?
不知道这样理解对不对... ...