DCA，不需训练让Llama上下文扩大48倍的方法

论文标题：Training-Free Long-Context Scaling of Large Language Models

论文地址：https://arxiv.org/pdf/2402.17463

最近研究Qwen2和Qwen2.5论文的时候，发现都有用到DCA（Dual Chunk Attention），transformers库里面的Qwen2Model只有RoPE和FlashAttention2的实现，没找到DCA代码。更多DCA的评估和使用教程可以去看Qwen团队的博客：https://qwenlm.github.io/zh/blog/qwen2.5-1m/

DCA讲的是在不进行训练、微调的情况下，大大增加模型的输入长度且生成质量不会降低太多。DCA的可行性显然已经被qwen-long这类长文本模型证明了。

DCA主要解决的是RoPE内相对位置编码的差值矩阵问题，其实很简单。但是看到文章把一个矩阵拆解成三个维度介绍的样子，有点好笑，想起了我研究生时期用TFIDF做文本分类的时候，天天写“类间权重”、“类内权重”这些名词的狼狈~

Abstract

大型模型在输入token数量超过其预训练长度时，生成能力会断崖式下降，但如果直接训练长序列的大模型成本非常高。因此作者提出了双块注意力（DCA），使 LLAMA2 70B 能够支持超过 10 万个token的上下文输入，而且不需要持续训练。

DCA可以与 Flash Attention 无缝集成。除此之外，DCA 在实际的长上下文任务上的性能与微调模型相当甚至更好。与专有模型相比，DCA的免训练 70B 模型达到了 gpt-3.5-16k 性能的 94%。

【注】大模型的外推能力上，还有苏剑林提出的RoPE，旋转位置编码，巧妙地采用复数替代正弦余弦，被如今大模型广泛采用，苏神牛！

Introduction

让LLM拥有长上下文处理能力很重要，具体体现在：

• 长文档解析场景
• 保留更多更长的历史对话聊天

通过微调来提高上下文长度的近期研究成果主要有：

1）短上下文模型训练：通过在长文本序列上进一步训练原本为短上下文设计的模型，可以提升这些模型处理长上下文的能力。

2）Llama2 Long：模型通过在长文本数据和原始Llama2预训练语料库的混合上进行训练，也可以提升上下文长度。

但是这些方法最大的缺点就是需要微调，微调的难点在于模型数据集很难拿到（通常不公开），而且微调训练有成本。

因此，更多的方法是探究不微调不训练而使得模型有更长输入的算法。

• LMinfinite和StreamingLLM等免训练方法通过选择性地保留基本的局部信息来处理扩展序列，从而在处理长文本时保持较低的Perplexity。但是失去了长距离文本的上下文信息。

【注】也就是类似于CNN和RNN的区别。这些方法就像是CNN，用卷积提取信息；而decoder-only与attention更多的需要全局序列依赖。

• 还有一些基于LlaMa与RoPE开发出来的技术，如位置插值（Position Interpolation）、NTK感知旋转编码（NTK-Aware RoPE）。

虽然这些方法可以降低训练长上下文输入模型的成本，但是作者发现，一旦输入长度超过训练长度两倍，这些方法就会导致Perplexity显著增加。

【注】Perplexity（PPL）在NLP领域叫做困惑度，用于衡量语言模型性能，公式为：

P

P

L

=

2

−

1

N

∑

i

=

1

N

log

⁡

2

P

(

w

i

∣

w

1

,

w

2

,

.

.

.

,

w

i

−

1

)

PPL = 2^{-frac{1}{N} sum_{i=1}^{N} log_2 P(w_i | w_1, w_2, …, w_{i-1})}

PPL=2−N1​∑i=1N​log2​P(wi​∣w1​,w2​,…,wi−1​)

这是个很古早的性能指标了。我记得以前学马尔可夫链的时候就看到过这个指标。

不同于以上历史方法，DCA作出如下改进：

1）重用原始位置索引以及其嵌入向量，但是重新设计了相对位置矩阵的构造。

2）在块注意力上，设置每个块都小于训练窗口长度，并且提出intra-chunk attention、inter-chunk attention、succes-sive chunk attention，分别负责chunk内注意力、chunk间注意力、连续chunk注意力的提取。

Background

位置编码

论文在这里介绍了Transformer模型使用的正余弦位置编码和旋转位置编码。

这里就不要看论文的了。目前大模型的位置编码分为绝对位置编码和相对位置编码。

绝对位置编码也就是输入该token的索引（i）来获得嵌入，这种写死的方法不可以外推，其实现在已经没人在用了。

相对位置编码是在计算attention时，输入为两个token的索引（i和j），因为注意力其实就是两两token之间计算，因此可行。并且无论输入长度怎么拓展，相对位置编码都能支持。RoPE就是利用复数推导，同时具有绝对位置编码和相对位置编码的技术，被业界广泛采用。

RoPE外推不足之处

RoPE虽然是很好的位置编码技术，但是仍然有不足之处：假设训练上下文长度为6，那么模型就从未训练过超过 6 的相对位置，当输入为10时，超过6的相对位置的效果就不会太好。

因此，PI和NTK通过减少整个位置矩阵的大小来缓解这类问题，以确保它在训练期间处于观察到的上下文长度范围内。

【注】也就是说，RoPE在理论上解决了位置编码无限外推的可行性，但是真要外推的时候，还是会受限制于训练窗口的大小。

Method

【注】DCA就是在改索引矩阵的下标，其目的就是把图1的矩阵值，改为图2©中的矩阵值。

如图2所示，三个图分别代表chunk内注意力、chunk间注意力、连续chunk注意力对应的相对位置编码的索引矩阵的值。

设模型预训练的长度为

c

c

c，本次请求模型给模型的输入长度为

l

l

l，chunk的长度为

s

s

s，那么我们可以把本次输入拆分为

l

s

frac{l}{s}

sl​块。

在论文中，

c

=

10

c=10

c=10，

l

=

12

l=12

l=12，

s

=

6

s=6

s=6，并且此时键key的下标

P

k

=

{

0

,

1

,

2

,

.

.

.

,

l

−

1

}

P_k={0,1,2,…,l-1}

Pk​={0,1,2,…,l−1}，查询query的下标

P

q

=

{

0

,

1

,

2

,

.

.

.

,

l

−

1

}

P_q={0,1,2,…,l-1}

Pq​={0,1,2,…,l−1}。

Intra-Chunk Attention

Intra-Chunk Attention 用于计算同一 chunk 中 queries 和 key 的内积。对于长度为

l

l

l 的长序列，将序列划分为为

l

s

frac{l}{s}

sl​块长度为

s

s

s的chunk。现在定义新的键key的下标

P

k

P_k

Pk​与查询的下标

P

q

I

n

t

r

a

P_q^{Intra}

PqIntra​：

P

k

=

P

q

I

n

t

r

a

=

{

0

,

1

,

2

,

.

.

.

,

l

−

1

}

m

o

d

s

=

{

0

,

1

,

2

,

3

,

4

,

5

,

0

,

1

,

2

,

3

,

4

,

5

}

P_k = P_q^{Intra} = {0,1,2,…,l-1} mod s = {0,1,2,3,4,5,0,1,2,3,4,5}

Pk​=PqIntra​={0,1,2,…,l−1} mod s={0,1,2,3,4,5,0,1,2,3,4,5}

然后计算相对距离

M

M

M：

M

[

i

]

[

j

]

=

P

q

I

n

t

r

a

[

i

]

−

P

k

[

j

]

M[i][j]=P_{q}^{Intra }[i]-P_{k}[j]

M[i][j]=PqIntra​[i]−Pk​[j]

所以就得到了图3中的矩阵值。

Inter-Chunk Attention

为了聚合来自其他 chunk 的信息，论文引入了 Inter-Chunk Attention。同时由于因果注意力矩阵的特点（参考图1），Inter-Chunk Attention生成出来的矩阵值在块内要满足如下特点：

1. 最大值不能超过预训练输入长度$c$，超过了不就变成原始RoPE了嘛；
2. 从x轴方向延申，其值要递减，表示距离越来越近；
3. 从x轴方向延申，其值虽然要递减，但是要保证尽可能比Intra-Chunk Attention 生成的值要大，不然没法体现块间长距离。

因此，论文提出让索引从能接受的最大距离

c

−

1

c-1

c−1开始下降，现在定义查询的下标

P

q

I

n

t

e

r

P_q^{Inter}

PqInter​：

P

q

I

n

t

e

r

=

[

c

−

1

,

c

−

1

,

.

.

.

c

−

1

⏟

l

e

l

e

m

e

n

t

s

]

=

{

9

,

9

,

9…

,

9

}

P_{q}^{Inter }=[underbrace{c-1, c-1, … c-1}_{l elements }] = {9,9,9…,9}

PqInter​=[lelements


c−1,c−1,…c−1​​]={9,9,9…,9}

然后计算相对距离

M

M

M：

M

[

i

]

[

j

]

=

P

q

I

n

t

r

a

[

i

]

−

P

k

[

j

]

=

c

−

1

−

P

k

[

j

]

M[i][j]=P_{q}^{Intra }[i]-P_{k}[j]=c-1-P_{k}[j]

M[i][j]=PqIntra​[i]−Pk​[j]=c−1−Pk​[j]

所以就得到了图4中的矩阵值。

Successive-Chunk Attention

我们将图3与图4直接相加，然后和最终效果图5对比，发现中间不一样的地方，就是Successive-Chunk Attention要解决的问题。这个地方其实是chunk与chunk之间的交汇处，但是相对距离这么大的话不太合适，因此定义查询的下标

P

q

S

u

c

c

P_q^{Succ}

PqSucc​：

P

q

S

u

c

c

=

s

,

s

+

1

,

.

.

.

,

s

+

w

−

1

⏟

w

e

l

e

m

e

n

t

s

,

c

−

1

,

.

.

.

,

c

−

1

⏞

t

h

e

s

a

m

e

f

o

r

a

l

l

c

h

u

n

k

s

]

=

{

6

,

7

,

8

,

9

,

9

,

9

,

6

,

7

,

8

,

9

,

9

,

9

}

P_{q}^{Succ }=overbrace{underbrace{s, s+1, …, s+w-1}^{w elements }, c-1, …, c-1}_{the same for all chunks }] = {6,7,8,9,9,9,6,7,8,9,9,9}

PqSucc​=


s,s+1,…,s+w−1​welements,c−1,…,c−1


​thesameforallchunks​]={6,7,8,9,9,9,6,7,8,9,9,9}

其中

w

w

w 表示本地窗口大小，可以直接设置为预训练长度和块大小之间的差值

c

−

s

c-s

c−s。然后计算相对距离

M

M

M，所以就得到了图5中的矩阵值。

【注】以上就是DCA的原理，各个参数应该有经验值，可能就要自行摸索了。

Experiments

实验与总结略。阿里的qwen-long已经证明。

文章来源于互联网:DCA，不需训练让Llama上下文扩大48倍的方法