Reformer 模型 – 突破语言建模的极限

Reformer 如安在不到 8GB 的内存上练习 50 万个词元

Kitaev、Kaiser 等人于 20202 年引进的 Reformer 模型 是迄今为止长序列建模领域内存功率最高的 transformer 模型之一。

最近，人们对长序列建模的兴趣激增，仅今年一年，就出现出了很多的作业，如 Beltagy 等人的作业 (2020) 、Roy 等人的作业 (2020) 、Tay 等人的作业 以及 Wang 等人的作业 等等。长序列建模背后的动机是，N​​LP 中的许多使命 (例如 摘要、问答 ) 要求模型处理更长的序列，这些序列长度超出了 BERT 等模型的处理才干。在需求模型处理长输入序列的使命中，长序列模型无需对输入序列进行裁剪以防止内存溢出，因而已被证明优于规范的 BERT 类模型 ( 见 Beltagy 等人 2020 年的作业)。

Reformer 能够一次处理多达 50 万个词元，然后打破了长序列建模的极限 (详细可参见本 笔记本)。相形之下，传统的 bert-base-uncased 模型最长仅支撑 512 个词元。在 Reformer 中，规范 transformer 架构的每个部分都经过从头规划，以最小化内存需求，并防止显着下降功能。

内存的改进来自于 Reformer 作者向 transformer 世界引进的 4 大特性:

1. Reformer 自留意力层 – 如安在不受限于本地上下文的状况下高效地完成自留意力机制？
2. 分块前馈层 – 怎么更好地对大型前馈层的时刻和内存进行权衡？
3. 可逆残差层 – 怎么聪明地规划残差架构以大幅削减练习中的内存耗费？
4. 轴向方位编码 (Axial Positional Encodings) – 怎么使方位编码可用于超长输入序列？

本文的目的是 深化 论述 Reformer 的上述四大特性。尽管这四个特性目前是用在 Reformer 上的，但其办法是通用的。因而，读者不应被此捆绑，而应该多思考在哪些状况下能够把这四个特性中的某一个或某几个运用于其他的 transformer 模型，以处理其问题。

下文四个部分之间的联系很松散，因而能够单独阅览。

Reformer 已集成入 Transformers 库。关于想运用 Reformer 的用户，主张咱们阅览本文，以更好地了解该模型的作业原理以及怎么正确装备它。文中一切公式都附有其在 transformers 中对应的 Reformer 装备项 ( 例如 config.<param_name> )，以便读者能够快速相关到官方文档和装备文件。

留意: 轴向方位编码 在官方 Reformer 论文中没有解说，但在官方代码库中广泛运用。本文初次深化阐释了轴向方位编码。

1. Reformer 自留意力层

Reformer 运用了两种特殊的自留意力层: 部分 自留意力层和 LSH (Locality Sensitive Hashing，部分灵敏哈希， LSH ) 自留意力层。

在介绍新的自留意力层之前，咱们先简要回忆一下传统的自留意力，其由 Vaswani 等人在其 2017 年的论文 中引进。

本文的符号及配色与 《图解 transformer》 一文一致，因而激烈主张读者在阅览本文之前，先阅览《图解 transformer》一文。

重要: 尽管 Reformer 最初是为了因果自留意力而引进的，但它也能够很好地用于双向自留意力。本文在解说 Reformer 的自留意力时，将其用于 双向 自留意力。

大局自留意力回忆

Transformer 模型的核心是 自留意力 层。现在，咱们回忆一下传统的自留意力层，这儿称为 大局自留意力 层。首要咱们假定对嵌入向量序列 X=x1,…,xnmathbf{X} = mathbf{x}_1, ldots, mathbf{x}_nX=x1​,…,xn​ 执行一个 transformer 层，该序列中的每个向量 ximathbf{x}_{i}xi​ 的维度为 config.hidden_​​size ， 即 dhd_hdh​。

简而言之，大局自留意力层将 $\mathbf{X}$ 投影到查询矩阵、键矩阵和值矩阵: $\mathbf{Q}$、$\mathbf{K}$、$\mathbf{V}$ 并运用 softmax 核算终究输出 $\mathbf{Z}$，如下所示:

Z=SelfAttn(X)=softmax(QKT)Vmathbf{Z} = text{SelfAttn}(mathbf{X}) = text{softmax}(mathbf{Q}mathbf{K}^T) mathbf{V}Z=SelfAttn(X)=softmax(QKT)V，其间 $\mathbf{Z}$ 的维度为 dhnd_h times ndh​n (为简略起见，此处省掉了键归一化因子和输出映射权重 WOmathbf{W}^{O}WO)。有关完整 transformer 操作的更多详细信息，请参看 《图解 transformer》 一文。

下图给出了 n=16，dh=3n=16，d_h=3n=16，dh​=3 状况下的操作:

请留意，本文一切示意图都假定 batch_size 和 config.num_attention_heads 为 1。为了便于稍后更好地解说 LSH 自留意力 ，咱们还在图中符号出了一些向量， 如 x3mathbf{x_3}x3​ 及其相应的输出向量 z3mathbf{z_3}z3​。图中的逻辑能够轻易扩展至多头自留意力 ( config.num_attention_heads > 1)。如需了解多头留意力，主张读者参看 《图解 transformer》。

敲个要点，关于每个输出向量 zimathbf{z}_{i}zi​，整个输入序列 $\mathbf{X}$ 都需求参与其核算。内积张量 QKTmathbf{Q}mathbf{K}^TQKT 的内存杂乱度为 O(n2)mathcal{O}(n^2)O(n2)，这事实上使得 transformer 模型的瓶颈在内存。

这也是为什么 bert-base-cased 的 config.max_position_embedding_size 只要 512 的原因。

部分自留意力

部分自留意力 是缓解 O(n2)mathcal{O}(n^2)O(n2) 内存瓶颈的一个显着的处理方案，它使咱们能够以更低的核算成本建模更长的序列。在部分自留意力中，输入 X=X1:n=x1,…,xnmathbf{X} = mathbf{X}_{1:n} = mathbf{x}_{1}, ldots, mathbf{x}_{n}X=X1:n​=x1​,…,xn​ 被切成 ncn_{c}nc​ 个块: X=[X1:lc,…,X(nc−1)∗lc:nc∗lc]mathbf{X} = left[mathbf{X}_{1:l_{c}}, ldots, mathbf{X} _{(n_{c} – 1) * l_{c} : n_{c} * l_{c}}right]X=[X1:lc​​,…,X(nc​−1)∗lc​:nc​∗lc​​]，每块长度为 config.local_chunk_length ， 即 lcl_{c}lc​，随后，对每个块别离运用大局自留意力。

持续以 n=16，dh=3n=16，d_h=3n=16，dh​=3 为例:

假定 lc=4，nc=4l_{c} = 4，n_{c} = 4lc​=4，nc​=4，此刻，咱们将分块留意力图示如下:

能够看出，咱们对每个块别离执行了留意力操作 X1:4，X5:8，X9:12，X13:16mathbf{X} _{1:4}，mathbf{X}_ {5:8}，mathbf{X} _{9:12 }，mathbf{X}_ {13:16}X1:4​，X5:8​，X9:12​，X13:16​。 该架构的一个显着的缺点是: 一些输入向量无法拜访其直接上下文， 例如 ，咱们的比方中的 x9mathbf{x} _9x9​ 无法拜访 x8mathbf{x}_ {8}x8​，反之亦然。这是有问题的，由于这些词元无法在学习其向量表征时将其直接上下文的归入考量。

一个简略的补救措施是用 config.local_num_chunks_before ( 即 npn_{p}np​) 以及 config.local_num_chunks_after ( 即 nan_{a}na​) 来扩充每个块，以便每个输入向量至少能够拜访 npn_{p}np​ 个从前输入块及 nan_{a}na​ 个后续输入块。咱们可将其了解为重叠分块，其间 npn_{p}np​ 和 nan_{a}na​ 界说了每个块与其从前块和后续块的重叠量。咱们将这种扩展的部分自留意力表明如下:

Zloc=[Z0:lcloc,…,Z(nc−1)∗lc+1:nc∗lcloc]，mathbf{Z}^{text{loc}} = left[mathbf{Z}_{0:l_{c}}^{text{loc}}, ldots, mathbf{Z}_{(n_{c} – 1) * l_{c} + 1 : n_{c} * l_{c}}^{text{loc}}right]，Zloc=[Z0:lc​loc​,…,Z(nc​−1)∗lc​+1:nc​∗lc​loc​]，

其间

Zlc∗(i−1)+1:lc∗iloc=SelfAttn(Xlc∗(i−1−np)+1:lc∗(i+na))[np∗lc:−na∗lc],∀i∈{1,…,nc}mathbf{Z}_{l_{c} * (i – 1) + 1 : l_{c} * i}^{text{loc}} = text{SelfAttn}(mathbf{X}_ {l_{c} * (i – 1 – n_{p}) + 1: l_{c} * (i + n_{a})})left[n_{p} * l_{c}: -n_{ a} * l_{c}right], forall i in {1, ldots, n_{c} }Zlc​∗(i−1)+1:lc​∗iloc​=SelfAttn(Xlc​∗(i−1−np​)+1:lc​∗(i+na​)​)[np​∗lc​:−na​∗lc​],∀i∈{1,…,nc​}

好吧，这个公式看起来有点杂乱，咱们略微剖析一下。在 Reformer 的自留意力层中，nan_{a}na​ 一般设为 0，npn_{p}np​ 设为 1，咱们据此重写 $i=1$ 时的公式:

Z1:lcloc=SelfAttn(X−lc+1:lc)[lc:]mathbf{Z}_{1:l_{c}}^{text{loc}} = text{SelfAttn}(mathbf{X}_{-l_{c} + 1: l_{c}})left[l_{c}:right]Z1:lc​loc​=SelfAttn(X−lc​+1:lc​​)[lc​:]

咱们留意到这儿有一个循环联系，因而榜首个块也能够重视终究一个块。咱们再次图解一下这种增强的部分重视算法。咱们先按块找到其对应的窗口，并在其上运用自留意力，然后仅保留中心输出段作为本块的输出。

终究，将相应的输出串接到 Zlocmathbf{Z}^{text{loc}}Zloc 中，如下所示:

请留意，在完成部分自留意力时，为了核算功率，咱们并不会像图中相同先核算悉数输出并随后 丢掉 一部分。图中红叉所示的地方仅用于阐明，实际并不会发生核算行为。

这儿需求留意的是，扩展每个分块自留意力函数的输入向量能够使得 每个 输出向量 zimathbf{z}_{i}zi​ 都能够学到更好的向量表征。以图中的向量为例，每个输出向量 z5loc，z6loc，z7loc，z8locmathbf{z}_{5}^{text{loc}}，mathbf{z}_{6}^{text{loc}}，mathbf{z}_{7}^{text{loc}}，mathbf{z}_{8}^{text{loc}}z5loc​，z6loc​，z7loc​，z8loc​ 都能够将 X1:8mathbf{X}_{1:8}X1:8​ 的一切输入向量归入考量以学到更好的表征。

内存耗费上的下降也是清楚明了的: O(n2)mathcal{O}(n^2)O(n2) 的内存杂乱度被分化到段，因而总内存杂乱度削减为 O(nc∗lc2)=O(n∗lc)mathcal{O}(n_{c} * l_{c}^2) = mathcal{O}(n * l_{c})O(nc​∗lc2​)=O(n∗lc​)。

这种增强的部分自留意力比普通的部分自留意力架构更好，但仍然存在一个首要缺点，由于每个输入向量只能重视预界说巨细的部分上下文。关于不需求 transformer 模型学习输入向量之间的长途依靠联系的 NLP 使命 ( 例如 语音辨认、命名实体辨认以及短语句的因果言语建模) 而言，或许不是一个大问题。但还有许多 NLP 使命需求模型学习长途依靠联系，因而部分自留意力在这些使命下或许会导致显着的功能下降， 如 :

• 问答 : 模型有必要学习问题词元和相关答案词元之间的联系，这些词元很或许并不相邻;
• 多项挑选 : 模型有必要将多个答案词元段彼此比较，这些答案词元段一般隔得比较远;
• 摘要 : 模型有必要学习长序列的上下文词元和较短的摘要词元序列之间的联系，而上下文和摘要之间的相关联系很或许无法经过部分自留意力来捕获。
• ……

部分自留意力本身很或许不足以让 transformer 模型学习输入向量 (词元) 彼此之间的相关联系。

因而，Reformer 额外采用了一个近似大局自留意力的高效自留意力层，称为 LSH 自留意力 。

LSH 自留意力

鉴于咱们现已了解了部分自留意力的作业原理，下面咱们持续尝试一下或许是 Reformer 中最具创新性的算法改进: LSH 自留意力。

LSH 自留意力的规划方针是在作用上接近大局自留意力，而在速度与资源耗费上与部分自留意力相同高效。

LSH 自留意力因依靠于 Andoni 等人于 2015 年提出的 LSH 算法 而得名。

LSH 自留意力源于以下洞见: 假定 $n$ 很大，则对每个查询向量而言，其对应的输出向量 zimathbf{z}_{i}zi​ 作为一切 $\mathbf{V}$ 的线性组合，其间应只要极少数几个 vimathbf{v}_{i}vi​ 的权重比其他大得多。也就是说对 QKTmathbf{Q}mathbf{K}^TQKT 留意力点积作 softmax 发生的权重矩阵的每一行应仅有极少数的值远大于 0。

咱们展开讲讲: 设 ki∈K=[k1,…,kn]Tmathbf{k}_{i} in mathbf{K} = left[mathbf{k}_1, ldots, mathbf{k}_n right]^Tki​∈K=[k1​,…,kn​]T 和 qi∈Q=[q1,…,qn]Tmathbf{q}_{i} in mathbf{Q} = left[mathbf{q}_1, ldots, mathbf{q}_nright]^Tqi​∈Q=[q1​,…,qn​]T 别离为键向量和查询向量。关于每个 qimathbf{q}_{i}qi​，能够仅用那些与 qimathbf{q}_{i}qi​ 具有高余弦类似度的 kjmathbf{k}_{j}kj​ 的键向量来近似核算 softmax(qiTKT)text{softmax}(mathbf{q}_{i}^T mathbf{K}^T)softmax(qiT​KT) 。这是由于 softmax 函数对较大输入值的输出会呈指数级添加。听起来没毛病，那么下一个问题就变成了怎么高效地找到每个 qimathbf{q}_{i}qi​ 的高余弦类似度键向量集合。

首要，Reformer 的作者留意到同享查询投影和键投影: $\mathbf{Q}=\mathbf{K}$ 并不会影响 transformer 模型 1{}^11。现在，不必为每个查询向量 qiq_iqi​ 找到其高余弦类似度的键向量，而只需核算查询向量彼此之间的余弦类似度。这一简化很重要，由于查询向量之间的余弦类似度满足传递性: 假定 qimathbf{q}_{i}qi​ 与 qjmathbf{q}_{j}qj​ 和 qkmathbf{q}_{k}qk​ 都具有较高的余弦类似度，则 qjmathbf{q}_{j}qj​ 与 qkmathbf{q}_{k}qk​ 也具有较高的余弦类似度。因而，能够将查询向量聚类至不同的桶中，使得同一桶中的一切查询向量彼此的余弦类似度较高。咱们将 CmC_{m}Cm​ 界说为第 m 组方位索引，其间装的是属于同一个桶的一切查询向量: Cm=i∣qi∈第m簇C_{m} = { i | mathbf{q}_{i} in text{第 m 簇}}Cm​=i∣qi​∈第m簇，一起咱们界说桶的数量 config.num_buckets ， 即 nbn_{b}nb​。

对每个索引 CmC_{m}Cm​ 对应的查询向量桶内的查询向量 qimathbf{q}_{i}qi​，咱们能够用 softmax 函数 softmax(Qi∈CmQi∈CmT)text{softmax}(mathbf{Q}_{i in C_{m}} mathbf{Q}^T_{i in C_{m}})softmax(Qi∈Cm​​Qi∈Cm​T​) 经过同享查询和键投影来近似大局自留意力的 softmax 函数 softmax(qiTQT)text{softmax}(mathbf{q}_{i}^T mathbf{Q}^T)softmax(qiT​QT)。

其次，作者利用 LSH 算法将查询向量聚类到预界说的 nbn_{b}nb​ 个桶 中。这儿，LSH 算法是理想之选，由于它非常高效，且可用于近似依据余弦类似度的最近邻算法。对 LSH 进行解说超出了本文的规模，咱们只要记住，对向量 qimathbf{q}_{i}qi​，LSH 算法将其索引至 nbn_{b}nb​ 个预界说桶中的某个桶， 即 LSH(qi)=mtext{LSH}(mathbf{q}_{i}) = mLSH(qi​)=m 其间 $i\in 1,\dots ,n$，m∈1,…,nbm in {1, ldots, n_{b}}m∈1,…,nb​。

还用前面的比方，咱们有:

接着，能够留意到，将一切查询向量聚类至 nbn_{b}nb​ 个桶中后，咱们能够将输入向量 x1,…,xnmathbf{x}_1, ldots, mathbf{x}_nx1​,…,xn​ 按其对应的索引 CmC_{m}Cm​ 进行重排 2{}^22，以便同享查询 – 键自留意力能够像部分留意力相同分段运用。

咱们用比方再解说一下，假定在 config.num_buckets=4 ， config.lsh_chunk_length=4 时重排输入向量 X=x1,…,x16mathbf{X} = mathbf{x}_1, …, mathbf{x}_{16}X=x1​,…,x16​。上图已将每个查询向量 q1,…,q16mathbf{q}_1, ldots, mathbf{q}_{16}q1​,…,q16​ 分配给簇 C1、C2、C3、C4mathcal{C}_{1}、mathcal{C}_{2}、mathcal{C}_{3}、mathcal{C}_{4}C1​、C2​、C3​、C4​ 中的某一个。现在，对其对应的输入向量 x1,…,x16mathbf{x}_1, ldots, mathbf{x}_{16}x1​,…,x16​ 进行重排，并将重排后的输入记为 X′mathbf{X’}X′:

对每个输入向量，仅需在簇内进行自留意力核算即可，因而每个输入向量对应的输出向量可核算如下: Zi∈CmLSH=SelfAttnQ=K(Xi∈Cm)mathbf{Z}^{text{LSH}}_{i in mathcal{C}_m} = text{SelfAttn}_{mathbf{Q}=mathbf{K}}(mathbf{X}_{i in mathcal{C}_m})Zi∈Cm​LSH​=SelfAttnQ=K​(Xi∈Cm​​)。

咱们再次图解一下该进程:

能够看出，自留意力函数的运算矩阵巨细各不相同，这种状况比较费事，由于 GPU 和 TPU 无法高效并行处理不同尺度的矩阵运算。

为了进一步处理高效核算的问题，能够借鉴部分留意力的办法，对重排后的输入进行分块，以使每个块的巨细均为 config.lsh_chunk_length 。经过对重排后的输入进行分块，一个桶或许会被分红两个不同的块。为了处理这个问题，与部分自留意力相同，在 LSH 自留意力中，每个块除了本身之外还重视其前一个块 config.lsh_num_chunks_before=1 ( config.lsh_num_chunks_after 一般设置为 0)。这样，咱们就能够大概率保证桶中的一切向量彼此重视 3{}^33。

总而言之，关于一切块 k∈1,…,nck in {1, ldots, n_{c}}k∈1,…,nc​，LSH 自留意力能够如下表明:

Z’lc∗k+1:lc∗(k+1)LSH=SelfAttnQ=K(X’lc∗(k+1):lc∗(k+1))[lc:]mathbf{Z’}_{l_ {c} * k + 1:l_{c} *(k + 1)}^{text{LSH}} = text{SelfAttn}_{mathbf{Q} = mathbf{K}}(mathbf{X’}_{l_{c} * (k + 1): l_{c} *(k + 1)})left[l_{c}:right]Z’lc​∗k+1:lc​∗(k+1)LSH​=SelfAttnQ=K​(X’lc​∗(k+1):lc​∗(k+1)​)[lc​:]

其间 X′mathbf{X’}X′ 和 Z′mathbf{Z’}Z′ 是按照 LSH 分桶进行重排后的输入和输出向量。公式有点杂乱，咱们还是画个图以协助咱们了解。

这儿，咱们对上图中的重排向量 X′mathbf{X’}X′ 进行分块，并别离核算每块的同享查询 – 键自留意力。

终究，将输出 Z′LSHmathbf{Z’}^{text{LSH}}Z′LSH 重排回原次序。

这儿还要说到的一个重要特征是，能够经过并行运转 LSH 自留意力 config.num_hashes (即 nhn_{h}nh​) 次来进步 LSH 自留意力的精确性，其间每次运用不同的随机 LSH 哈希。经过设置 config.num_hashes > 1 ，关于每个 $i$，会核算多个输出向量 ziLSH,1,…,ziLSH,nhmathbf{z}^{text{LSH}, 1}_{i}, ldots , mathbf{z}^{text{LSH}, n_{h}}_{i}ziLSH,1​,…,ziLSH,nh​​。随后，能够对它们进行加权求和: ziLSH=∑knhZiLSH,k∗weightikmathbf{z}^{text{LSH}}_{i} = sum_k^{n_{h}} mathbf{Z}^{text{LSH}, k}_{i} * text{weight}^k_iziLSH​=∑knh​​ZiLSH,k​∗weightik​，这儿 weightiktext{weight}^k_iweightik​ 表明第 $k$ 轮哈希的输出向量 ziLSH,kmathbf{z}^{text{LSH}, k}_{i}ziLSH,k​ 与其他哈希轮次比较的重要度，其应与其对应输出的 softmax 归一化系数呈指数正比联系。这一规划背后的直觉是，假定查询向量 qikmathbf{q}_{i}^{k}qik​ 与其对应块中的一切其他查询向量具有较高的余弦类似度，则该块的 softmax 归一化系数往往很大，因而相应的输出向量 qikmathbf{q}_{i}^{k}qik​ 应该能更好地近似大局留意力，因而其理应比 softmax 归一化系数较小的哈希轮次所发生的输出向量获得更高的权重。更多详细信息，请参看 该论文 的附录 A。在咱们的比方中，多轮 LSH 自留意力示意图如下。

打完收工！至此，咱们了解了 LSH 自留意力在 Reformer 中是怎么作业的。

说回内存杂乱度，该办法有两个或许的瓶颈点: 点积所需的内存: O(nh∗nc∗lc2)=O(n∗nh∗lc)mathcal{O}(n_{h} * n_{c} * l_{c}^2) = mathcal{O}(n * n_{h} * l_{c})O(nh​∗nc​∗lc2​)=O(n∗nh​∗lc​) 以及 LSH 分桶所需的内存: O(n∗nh∗nb2)mathcal{O}(n * n_{h} * frac{n_{b}}{2})O(n∗nh​∗2nb​​) 其间 lcl_{c}lc​ 是块长度。由于关于大的 $n$ 而言，桶的数量 nb2frac{n_{b}}{2}2nb​​ 的增长速度远远快于块长度 lcl_{c}lc​，因而用户能够持续对存储桶的数量 config.num_buckets 进行分化，详见 此处。

咱们快速总结一下:

1. 咱们期望利用 softmax 运算仅对极少数键向量赋予重要权重的先验常识来对大局留意力进行近似。
2. 假定键向量等于查询向量，这意味着 关于每个 查询向量 qimathbf{q}_{i}qi​，softmax 只需给与其余弦类似度高的其他查询向量赋予重要权重就行了。
3. 这种联系是对称的，也就是说，假定 qjmathbf{q}_{j}qj​ 与 qimathbf{q}_{i}qi​ 类似，则 qjmathbf{q}_{j}qj​ 也与 qimathbf{q}_{i}qi​ 类似，因而咱们能够在核算自留意力之前对输入进行大局聚类。
4. 咱们对输入按簇进行重排，并对重排后的输入核算部分自留意力，终究将输出从头恢复为原次序。

1{}^{1}1 作者进行了一些开端试验，确认同享查询 – 键自留意力的体现与规范自留意力大体一致。

2{}^{2}2 更精确地说，对存储桶中的查询向量依据其原始次序进行排序。举个比方， 假定 向量 q1,q3,q7mathbf{q}_1, mathbf{q}_3, mathbf{q}_7q1​,q3​,q7​ 悉数散列到存储桶 2，则存储桶 2 中向量的次序仍应是先 q1mathbf{q}_1q1​，后跟 q3mathbf{q}_3q3​ 和 q7mathbf{q}_7q7​。

3{}^33 顺带阐明一下，作者在查询向量 qimathbf{q}_{i}qi​ 上放了一个掩码，以防止向量重视本身。由于向量与其本身的余弦类似度总是大于等于其与其他向量的余弦类似度，所以激烈不主张同享查询 – 键自留意力中的查询向量重视本身。

基准测验

Transformers 最近添加了基准测验相关的代码，你可参看 此处 以获取更详细的阐明。

为了展示部分 LSH 自留意力能够节约多少内存，咱们在不同的 local_attn_chunk_length 和 lsh_attn_chunk_length 上对 Reformer 模型 google/reformer-enwik8 进步行了基准测验。你能够从 此处 找到更详细的有关 google/reformer-enwik8 模型的默许装备和用法信息。

咱们先进行一些必要的导入和装置。

#@title Installs and Imports
# pip installs
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments

首要，咱们测验一下在 Reformer 模型上运用 大局 自留意力的内存运用状况。这能够经过设置 lsh_attn_chunk_length = local_attn_chunk_length = 8192 来达成，此刻，关于一切小于或等于 8192 的输入序列，模型事实上就回退成大局自留意力了。

config = ReformerConfig.from_pretrained("google/reformer-enwik8", lsh_attn_chunk_length=16386, local_attn_chunk_length=16386, lsh_num_chunks_before=0, local_num_chunks_before=0)
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[2048, 4096, 8192, 16386], batch_sizes=[1], models=["Reformer"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config], args=benchmark_args)
result = benchmark.run()

HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1279.0, style=ProgressStyle(description…
1 / 1
Doesn't fit on GPU. CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 11.17 GiB total capacity; 8.87 GiB already allocated; 1.92 GiB free; 8.88 GiB reserved in total by PyTorch)
====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
           Reformer                  1              2048            1465     
           Reformer                  1              4096            2757     
           Reformer                  1              8192            7893     
           Reformer                  1             16386            N/A      
--------------------------------------------------------------------------------

输入序列越长，输入序列和峰值内存运用之间的平方联系 O(n2)mathcal{O}(n^2)O(n2) 越显着。能够看出，实际上，需求更长的输入序列才干清楚地观察到输入序列翻倍会导致峰值内存运用量添加四倍。

对运用大局留意力的 google/reformer-enwik8 模型而言，序列长度超越 16K 内存就溢出了。

现在，咱们运用模型的默许参数以使能 部分 LSH 自留意力。

  config = ReformerConfig.from_pretrained("google/reformer-enwik8")
  benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[2048, 4096, 8192, 16384, 32768, 65436], batch_sizes=[1], models=["Reformer"], no_speed=True, no_env_print=True)
  benchmark = PyTorchBenchmark(configs=[config], args=benchmark_args)
  result = benchmark.run()

1 / 1
Doesn't fit on GPU. CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 11.17 GiB total capacity; 7.85 GiB already allocated; 1.74 GiB free; 9.06 GiB reserved in total by PyTorch)
Doesn't fit on GPU. CUDA out of memory. Tried to allocate 4.00 GiB (GPU 0; 11.17 GiB total capacity; 6.56 GiB already allocated; 3.99 GiB free; 6.81 GiB reserved in total by PyTorch)
====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
           Reformer                  1              2048            1785     
           Reformer                  1              4096            2621     
           Reformer                  1              8192            4281     
           Reformer                  1             16384            7607     
           Reformer                  1             32768            N/A      
           Reformer                  1             65436            N/A      
--------------------------------------------------------------------------------

不出所料，关于较长的输入序列，运用部分 LSH 自留意力机制的内存功率更高，关于本文运用的 11GB 显存 GPU 而言，模型直到序列长度为 32K 时，内存才耗尽。

2. 分块前馈层

依据 transformer 的模型一般在自留意力层之后会有一个非常大的前馈层。该层或许会占用很多内存，有时乃至成为模型首要的内存瓶颈。Reformer 论文中初次引进了前馈分块技能，以用时刻交换内存。

Reformer 中的分块前馈层

在 Reformer 中， LSH 自留意力层或部分自留意力层一般后边跟着一个残差连接，咱们可将其界说为 transformer 块 的榜首部分。更多相关常识，可参看此 博文。

Transformer 块 榜首部分的输出，称为 归范化自留意力 输出，能够记为 Z‾=Z+Xmathbf{overline{Z}} = mathbf{Z} + mathbf{X}Z=Z+X。在 Reformer 模型中，$\mathbf{Z}$ 为 ZLSHmathbf{Z}^{text{LSH}}ZLSH 或 Zlocmathbf{Z}^text{loc}Zloc。

在咱们的比方中，输入 x1,…,x16mathbf{x}_1, ldots, mathbf{x}_{16}x1​,…,x16​ 的规范化自留意力输出图示如下:

Transformer 块 的第二部分一般由两个前馈层 1^{1}1 组成，其间 Linearint(…)text{Linear}_{text{int}}(ldots)Linearint​(…) 用于将 Z‾mathbf{overline{Z}}Z 映射到中心输出 Yintmathbf{Y}_{text{int}}Yint​，Linearout(…)text{Linear}_{text{out}}(ldots)Linearout​(…) 用于将中心输出映射为终究输出 Youtmathbf{Y}_{text{out}}Yout​。咱们将两个前馈层界说如下:

Yout=Linearout(Yint)=Linearout(Linearint(Z‾))mathbf{Y}_{text{out}} = text{Linear}_{text{out}}(mathbf{Y} _text{int}) = text{Linear}_{text{out}}(text{Linear}_{text{int}}(mathbf{overline{Z}}))Yout​=Linearout​(Yint​)=Linearout​(Linearint​(Z))

敲要点！在数学上，前馈层在方位 $i$ 处的输出 yout,imathbf{y}_{text{out}, i}yout,i​ 仅取决于该方位的输入 y‾imathbf{overline{y}}_{i}y​i​。与自留意力层相反，每个输出 yout,imathbf{y}_{text{out}, i}yout,i​ 与其他方位的输入 y‾j≠imathbf{overline{y}}_{j ne i}y​j=i​ 完全独立。

z‾1,…,z‾16mathbf{overline{z}}_1, ldots, mathbf{overline{z}}_{16}z1​,…,z16​ 的前馈层图示如下:

从图中能够看出，一切输入向量 z‾imathbf{overline{z}}_{i}zi​ 均由同一前馈层并行处理。

咱们再观察一下前馈层的输出维度，看看有没有啥有意思的作业。在 Reformer 中，Linearinttext{Linear}_{text{int}}Linearint​ 的输出维度为 config.feed_forward_size ， 即 dfd_ {f}df​; 而 Linearouttext{Linear}_{text{out}}Linearout​ 的输出维度为 config.hidden_​​size ， 即 dhd_ {h}dh​。

Reformer 作者观察到 2^{2}2，在 transformer 模型中，中心维度 dfd_{f}df​ 一般往往比输出维度 dhd_{h}dh​ 大许多。这意味着尺度为 dfnd_{f} times ndf​n 的张量 Yintmathbf{mathbf{Y}}_text{int}Yint​ 占有了很多的内存，乃至或许成为内存瓶颈。

为了更好地感受维度的差异，咱们将本文比方中的矩阵 Yintmathbf{Y}_text{int}Yint​ 和 Youtmathbf{Y}_text{out}Yout​ 图示如下:

很显着，张量 Yintmathbf{Y} _text{int}Yint​ 比 Youtmathbf{Y}_{text{out}}Yout​ 占用了更多的内存 (精确地说，多占 dfdhnfrac{d_{f}}{d_{h}} times ndh​df​​n 字节的内存)。可是，是否有必要存储完整的中心矩阵 Yintmathbf{Y}_text{int}Yint​ ？并非如此，由于咱们关心的实际上只要输出矩阵 Youtmathbf{Y}_ text{out}Yout​。为了以速度换内存，咱们能够对线性层核算进行分块，一次只处理一个块。界说 config.chunk_size_feed_forward 为 cfc_{f}cf​，则分块线性层界说为 Yout=[Yout,1:cf,…,Yout,(n−cf):n]mathbf{Y}_{text{out}} = left[mathbf{Y}_{text{out}, 1: c_{f}}, ldots, mathbf{Y}_{text{out}, (n – c_{f}): n}right]Yout​=[Yout,1:cf​​,…,Yout,(n−cf​):n​] 即 Yout,(cf∗i):(i∗cf+i)=Linearout(Linearint(Z‾(cf∗i):(i∗cf+i)))mathbf{Y}_{text{out}, (c_{f} * i):(i * c_{f} + i)} = text{Linear}_{text{out}}( text{Linear}_{text{int}}(mathbf{overline{Z}}_{(c_{f} * i):(i * c_{f} + i)}))Yout,(cf​∗i):(i∗cf​+i)​=Linearout​(Linearint​(Z(cf​∗i):(i∗cf​+i)​))。这么做意味着咱们能够增量核算输出终究再串接在一起，这样能够防止将整个中心张量 Yintmathbf{Y}_{text{int}}Yint​ 存储在内存中。

假定 cf=1c_{f}=1cf​=1，咱们把增量核算 $i=9$ 的进程图示如下:

当块巨细为 1 时，有必要完整存储在内存中的仅有张量是巨细为 16dh16 times d_{h}16dh​ 的输入张量 Z‾mathbf{overline{Z}}Z，其间 dhd_{h}dh​ 为 config.hidden_size 。而中心张量只需求存储巨细为 dfd_{f}df​ 的 yint,imathbf{y}_{text{int}, i}yint,i​ 就能够了 3^{3}3。

终究，重要的是要记住， 分块线性层 与传统的完整线性层比较，其输出在数学上是等效的，因而能够运用于一切 transformer 线性层。因而，在某些场景下，能够考虑运用 config.chunk_size_feed_forward 在内存和速度之间进行更好的权衡。

1{}^11 为了简略起见，咱们省掉了前馈层之前的层归一化操作。

2{}^22 以 bert-base-uncased 为例，其间间维度 dfd_{f}df​ 是 3072，为输出维度 dhd_{h}dh​ 的 4 倍。

3{}^33 提醒一下，为明晰阐明起见，本文假定输出 config.num_attention_heads 为 1，因而假定自留意力层的输出巨细为 config.hidden_​​size 。

读者也能够在 Transformers 的 相应文档 中找到有关分块线性/前馈层的更多信息。

基准测验

咱们测验一下运用分块前馈层能够节约多少内存。

#@title Installs and Imports
# pip installs
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments

  Building wheel for transformers (setup.py) ... [?25l[?25hdone

首要，咱们将没有分块前馈层的默许 google/reformer-enwik8 模型与有分块前馈层的模型进行比较。

config_no_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8") # no chunk
config_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8", chunk_size_feed_forward=1) # feed forward chunk
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[1024, 2048, 4096], batch_sizes=[8], models=["Reformer-No-Chunk", "Reformer-Chunk"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_chunk, config_chunk], args=benchmark_args)
result = benchmark.run()

1 / 2
Doesn't fit on GPU. CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 11.17 GiB total capacity; 7.85 GiB already allocated; 1.74 GiB free; 9.06 GiB reserved in total by PyTorch)
2 / 2
Doesn't fit on GPU. CUDA out of memory. Tried to allocate 2.00 GiB (GPU 0; 11.17 GiB total capacity; 7.85 GiB already allocated; 1.24 GiB free; 9.56 GiB reserved in total by PyTorch)
====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
      Reformer-No-Chunk              8              1024            4281     
      Reformer-No-Chunk              8              2048            7607     
      Reformer-No-Chunk              8              4096            N/A      
        Reformer-Chunk               8              1024            4309     
        Reformer-Chunk               8              2048            7669     
        Reformer-Chunk               8              4096            N/A      
--------------------------------------------------------------------------------

风趣的是，分块前馈层似乎在这儿底子没有协助。原因是 config.feed_forward_size 不够大，所以作用不显着。仅当序列长度较长 (4096) 时，才干看到内存运用量略有下降。

咱们再看看假定将前馈层的巨细添加 4 倍，并将留意力头的数量一起削减 4 倍，然后使前馈层成为内存瓶颈，此刻峰值内存景象怎么。

config_no_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8", chunk_size_feed_forward=0, num_attention_{h}eads=2, feed_forward_size=16384) # no chuck
config_chunk = ReformerConfig.from_pretrained("google/reformer-enwik8", chunk_size_feed_forward=1, num_attention_{h}eads=2, feed_forward_size=16384) # feed forward chunk
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[1024, 2048, 4096], batch_sizes=[8], models=["Reformer-No-Chunk", "Reformer-Chunk"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_chunk, config_chunk], args=benchmark_args)
result = benchmark.run()

1 / 2
2 / 2
====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
      Reformer-No-Chunk              8              1024            3743     
      Reformer-No-Chunk              8              2048            5539     
      Reformer-No-Chunk              8              4096            9087     
        Reformer-Chunk               8              1024            2973     
        Reformer-Chunk               8              2048            3999     
        Reformer-Chunk               8              4096            6011     
--------------------------------------------------------------------------------

现在，关于较长的输入序列，能够看到峰值内存运用量显着削减。总之，应该留意的是，分块前馈层仅关于具有很少留意力头和较大前馈层的模型才有含义。

3. 可逆残差层

可逆残差层由 N. Gomez 等人 首要提出并运用在 ResNet 模型的练习上以削减内存耗费。从数学上讲，可逆残差层与 真实的 残差层略有不同，其不需求在前向传达期间保存激活，因而能够大大削减练习的内存耗费。

Reformer 中的可逆残差层

咱们首要研究为什么模型练习比推理需求更多的内存。

在模型推理时，所需的内存差不多等于核算模型中 单个 最大张量所需的内存。而在练习模型时，所需的内存差不多等于一切可微张量的 总和。

假定读者现已了解了深度学习框架中的主动微分的作业原理，对此就比较简略了解了。多伦多大学 Roger Grosse 的这些 幻灯片 对咱们了解主动微分很有协助。

简而言之，为了核算可微函数 ( 如 一层) 的梯度，主动微分需求函数输出的梯度以及函数的输入、输出张量。尽管梯度是能够动态核算并随后丢掉的，但函数的输入和输出张量 ( 又叫 激活) 需求在前向传达进程中被保存下来，以供反向传达时运用。

咱们详细看下 transformer 模型中的状况。Transformer 模型是由多个 transformer 层堆叠起来的。每多一个 transformer 层都会迫使模型在前向传达进程中保存更多的激活，然后添加练习所需的内存。 咱们细看一下 transformer 层。Transformer 层本质上由两个残差层组成。榜首个残差层是第 1) 节中解说的 自留意力 机制，第二个残差层是第 2) 节中解说的 线性层 (或前馈层)。

运用与之前相同的符号，transformer 层的输入 即 $\mathbf{X}$ 首要被归一化 1^{1}1，然后经过自留意力层获得输出 $\mathbf{Z}=\text{SelfAttn}(\text{LayerNorm}(\mathbf{X}))$。为便利评论，咱们将这两层缩写为 $G$，即 $\mathbf{Z}=G(\mathbf{X})$。 接下来，将残差 $\mathbf{Z}$ 与输入相加 Z‾=Z+Xmathbf{overline{Z}} = mathbf{Z} + mathbf{X}Z=Z+X，得到张量输入到第二个残差层 —— 两个线性层。Z‾mathbf{overline{Z}}Z 经过第二个归一化层处理后，再经过两个线性层，得到 $\mathbf{Y}=\text{Linear}(\text{LayerNorm}(\mathbf{Z}+\mathbf{X}))$。咱们将第二个归一化层和两个线性层缩写为 $F$ ，得到 Y=F(Z‾)mathbf{Y} = F(mathbf{overline{Z}})Y=F(Z)。终究，将残差 $\mathbf{Y}$ 加到 Z‾mathbf{overline{Z}}Z 上得到 transformer 层的输出 Y‾=Y+Z‾mathbf{overline{Y}} = mathbf{Y} + mathbf{overline{Z}}Y=Y+Z。

咱们仍以 x1,…,x16mathbf{x}_1, ldots, mathbf{x}_{16}x1​,…,x16​ 为例对完整的 transformer 层进行图解。

比方 ，要核算自留意力块 $G$ 的梯度，有必要事先知道三个张量: 梯度 $\partial \mathbf{Z}$、输出 $\mathbf{Z}$ 以及输入 $\mathbf{X}$。尽管 $\partial \mathbf{Z}$ 能够即时核算并随后丢掉，但 $\mathbf{Z}$ 和 $\mathbf{X}$ 有必要在前向传达期间核算并保存下来，由于在反向传达期间比较难轻松地即时从头核算它们。因而，在前向传达进程中，大张量输出 (如查询 – 键点积矩阵 QKTmathbf{Q}mathbf{K}^TQKT 或线性层的中心输出 Yintmathbf{Y}^{text{int}}Yint) 有必要保存在内存中 2^{2}2。

此刻，可逆残差层就有用了。它的想法相对简略: 残差块的规划方法使得不必保存函数的输入和输出张量，而在反向传达期间就轻松地对二者进行从头核算，这样的话在前向传达期间就无需将这些张量保存在内存中了。

这是经过两个输入流 X(1)、X(2)mathbf{X}^{(1)}、mathbf{X}^{(2)}X(1)、X(2) 及两个输出流 Y‾(1)、Y‾(2)mathbf{overline {Y}}^{(1)}、mathbf{overline{Y}}^{(2)}Y(1)、Y(2) 来完成的。榜首个残差 $\mathbf{Z}$ 由榜首个输出流 Z=G(X(1))mathbf{Z} = G(mathbf{X}^{(1)})Z=G(X(1)) 算得，然后其加到第二个输入流的输入上，即 Z‾=Z+X(2)mathbf{overline{Z}} = mathbf{Z} + mathbf{X}^{(2)}Z=Z+X(2)。类似地，再将残差 Y=F(Z‾)mathbf{Y} = F(mathbf{overline{Z}})Y=F(Z) 与榜首个输入流相加。终究，两个输出流即为 Y(1)=Y+X(1)mathbf{Y}^{(1)} = mathbf{Y} + mathbf{X}^{(1)}Y(1)=Y+X(1)、Y(2)=X(2)+Z=Z‾mathbf{Y}^{(2)} = mathbf{ X}^{(2)} + mathbf{Z} = mathbf{overline{Z}}Y(2)=X(2)+Z=Z。

以 x1,…,x16mathbf{x}_1, ldots, mathbf{x}_{16}x1​,…,x16​ 为例来图示可逆 transformer 层，如下:

能够看出，输出 Y‾(1)、Y‾(2)mathbf{overline{Y}}^{(1)}、mathbf{overline{Y}}^{(2)}Y(1)、Y(2) 的核算方法与不可逆层 Y‾mathbf{overline{Y}}Y 的核算方法非常类似，但在数学上又不同。Reformer 的作者在一些开端试验中观察到，可逆 transformer 模型的功能与规范 transformer 模型的功能相当。与规范 transformer 层的一个显着区别是有两个输入流和输出流 3^{3}3，这一开端反而略微添加了前向传达所需的内存。但即使如此，咱们还是着重双流架构至关重要，由于其在前向传达进程中无需保存任何激活。咱们解说一下: 关于反向传达，可逆 treansformer 层有必要核算梯度 $\partial G$ 和 $\partial F$。除了可即时核算的梯度 $\partial \mathbf{Y}$ 和 $\partial \mathbf{Z}$ 之外，为了核算 $\partial F$ 有必要已知张量值 $\mathbf{Y}$、Z‾mathbf{overline{Z}}Z，为了核算 $\partial G$ 有必要已知 $\mathbf{Z}$ 和 X(1)mathbf{X}^{(1)}X(1)。

假定咱们知道 Y‾(1)，Y‾(2)mathbf{overline{Y}}^{(1)}，mathbf{overline{Y}}^{(2)}Y(1)，Y(2)，则从图中能够很简略看出，咱们能够如下核算出 X(1)，X(2)mathbf{X}^{(1)}，mathbf{X}^{(2)}X(1)，X(2) 。X(1)=F(Y‾(1))−Y‾(1)mathbf{X}^{(1)} = F(mathbf{overline{Y}}^{(1)}) – mathbf{overline{Y}}^{(1)}X(1)=F(Y(1))−Y(1)。X(1)mathbf{X}^{(1)}X(1) 核算出来了！然后，X(2)mathbf{X}^{(2)}X(2) 能够经过 X(2)=Y‾(1)−G(X(1))mathbf {X}^{(2)} = mathbf{overline{Y}}^{(1)} – G(mathbf{X}^{(1)})X(2)=Y(1)−G(X(1)) 算出。之后，$\mathbf{Z}$ 和 $\mathbf{Y}$ 的核算就简略了，能够经过 Y=Y‾(1)−X(1)mathbf{Y} = mathbf{overline{Y}}^{(1)} – mathbf{X}^{(1)}Y=Y(1)−X(1) 和 Z=Y‾(2)−X(2)算出mathbf{Z} = mathbf{overline{Y}}^{(2)} – mathbf{X }^{(2)} 算出Z=Y(2)−X(2)算出。总结一下，仅需在前向传达期间存储 终究一个 可逆 transformer 层的输出 Y‾(1)，Y‾(2)mathbf{overline{Y}}^{(1)}，mathbf{overline{Y}}^{(2)}Y(1)，Y(2)，一切其他层的激活就能够经过在反向传达期间运用 $G$ 和 $F$ 以及 X(1)mathbf {X}^{(1)}X(1) 和 X(2)mathbf{X}^{(2)}X(2) 推导而得。在反向传达期间，每个可逆 transformer 层用两次前向传达 $G$ 和 $F$ 的核算开支交换前向传达时不必保存任何激活。好生意！

留意: 最近，首要的深度学习框架都支撑了梯度检查点技能，以答应仅保存某些激活并在反向传达期间重核算尺度较大的激活 (Tensoflow 代码见 此处，PyTorch 代码见 此处)。关于规范可逆层，这仍然意味着有必要为每个 transformer 层保存至少一个激活，但经过界说哪些激活能够动态从头核算，能够节约很多内存。

1^{1}1 在前两节中，咱们省掉了自留意力层和线性层之前的层归一化操作。读者应该知道 $\mathbf{X}$ 和 Z‾mathbf{overline{Z}}Z 在输入自留意力层和线性层之前都别离经过层归一化处理。

2^{2}2 在原始自留意力中，$\mathbf{QK}$ 的维度为 $nn$; 而在 LSH 自留意力 或 部分自留意力 层的维度为 nlcnhn times l_{c} times n_{h}nlc​nh​ 或 nlcn times l_{c}nlc​ 其间 lcl_{c}lc​ 为块长度，nhn_{h}nh​ 为哈希数。

3^{3}3 榜首个可逆 transformer 层的 X(2)mathbf{X}^{(2)}X(2) 等于 X(1)mathbf{X}^{(1)}X(1)。

测验基准

为了丈量可逆残差层的作用，咱们将添加模型层数的一起比较 BERT 和 Reformer 的内存耗费。

#@title Installs and Imports
# pip installs
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, BertConfig, PyTorchBenchmark, PyTorchBenchmarkArguments

咱们把规范 bert-base-uncased BERT 模型的层数从 4 添加到 12 ，一起丈量其所需内存。

config_4_layers_bert = BertConfig.from_pretrained("bert-base-uncased", num_hidden_layers=4)
config_8_layers_bert = BertConfig.from_pretrained("bert-base-uncased", num_hidden_layers=8)
config_12_layers_bert = BertConfig.from_pretrained("bert-base-uncased", num_hidden_layers=12)
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Bert-4-Layers", "Bert-8-Layers", "Bert-12-Layers"], training=True, no_inference=True, no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_4_layers_bert, config_8_layers_bert, config_12_layers_bert], args=benchmark_args)
result = benchmark.run()

HBox(children=(FloatProgress(value=0.0, description='Downloading', max=433.0, style=ProgressStyle(description_…
1 / 3
2 / 3
3 / 3
====================        TRAIN - MEMORY - RESULTS        ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
        Bert-4-Layers                8              512             4103     
        Bert-8-Layers                8              512             5759     
        Bert-12-Layers               8              512             7415     
--------------------------------------------------------------------------------

能够看出，BERT 层数每添加 1，其所需内存就会有超 400MB 的线性增长。

config_4_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8", num_hidden_layers=4, num_hashes=1)
config_8_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8", num_hidden_layers=8, num_hashes=1)
config_12_layers_reformer = ReformerConfig.from_pretrained("google/reformer-enwik8", num_hidden_layers=12, num_hashes=1)
benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Reformer-4-Layers", "Reformer-8-Layers", "Reformer-12-Layers"], training=True, no_inference=True, no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_4_layers_reformer, config_8_layers_reformer, config_12_layers_reformer], args=benchmark_args)
result = benchmark.run()

1 / 3
2 / 3
3 / 3
====================        TRAIN - MEMORY - RESULTS        ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
      Reformer-4-Layers              8              512             4607     
      Reformer-8-Layers              8              512             4987     
      Reformer-12-Layers             8              512             5367     
--------------------------------------------------------------------------------

另一方面，关于 Reformer 而言，每添加一层所带来的内存增量会显着削减，平均不到 100MB。因而 12 层的 reformer-enwik8 模型比 12 层的 bert-base-uncased 模型的内存需求更少。

4. 轴向方位编码

Reformer 使得处理超长输入序列成为或许。然而，关于如此长的输入序列，仅存储规范方位编码权重矩阵就需求超越 1GB 内存。为了防止如此大的方位编码矩阵，官方 Reformer 代码引进了 轴向方位编码 。

重要: 官方论文中没有解说轴向方位编码，但经过阅览代码以及与作者评论咱们很好地了解了它。

Reformer 中的轴向方位编码

Transformer 需求方位编码来对输入序列中的单词次序进行编码，由于自留意力层 没有次序的概念 。方位编码一般由一个简略的查找矩阵 E=[e1,…,enmax]mathbf{E} = left[mathbf{e}_1, ldots, mathbf{e}_{n_text{max}}right]E=[e1​,…,enmax​​] 来界说，然后将方位编码向量 eimathbf{e}_{i}ei​ 简略地加到 第 i 个 输入向量上，即 xi+eimathbf{x}_{i} + mathbf{e}_{i}xi​+ei​，以便模型能够区别输入向量 ( 即 词元) 坐落方位 $i$ 还是方位$j$。关于每个输入方位，模型需求能够查找到相应的方位编码向量，因而 $\mathbf{E}$ 的维度由模型能够处理的最大输入序列长度 config.max_position_embeddings ( 即 nmaxn_text{max}nmax​) 以及输入向量的维度 config.hidden_​​size ( 即 dhd_{h}dh​) 一起决定。

假定 dh=4d_{h}=4dh​=4，nmax=49n_text{max}=49nmax​=49，其方位编码矩阵如下图所示:

此处，咱们仅展示方位编码 e1mathbf{e}_{1}e1​、e2mathbf{e}_{2}e2​ 及 e49mathbf{e}_{49}e49​，其维度 ( 即 高度) 为 4。

想象一下，咱们想要在长度最长为 0.5M 个词元，输入向量维度 config.hidden_​​size 为 1024 的序列上练习 Reformer 模型 (请参看 此笔记本)。其对应的方位嵌入的参数量为 $0.5M1024\sim 512M$，巨细为 2GB。

在将模型加载到内存中或将其保存在硬盘上时，所需求的内存是很大且很没必要的。

Reformer 作者经过将 config.hidden_​​size 维度一分为二，并巧妙地对 nmaxn_text{max}nmax​ 维进行分化，然后成功地大幅缩小了方位编码的巨细。在 transformers 中，用户能够将 config.axis_pos_shape 设置为一个含有两个值的列表: nmax1n_text{max}^ 1nmax1​、nmax2n_text{max}^2nmax2​，其间 nmax1nmax2=nmaxn_text{max}^1 times n_text{max}^2 = n_text{max}nmax1​nmax2​=nmax​，然后对 nmaxn_text{max}nmax​ 维度进行分化。一起，用户能够把 config.axis_pos_embds_dim 设置为一个含有两个值 dh1d_{h}^{1}dh1​ 和 dh2d_{h}^2dh2​ 的列表，其间 dh1+dh2=dhd_{h} ^1 + d_{h}^2 = d_{h}dh1​+dh2​=dh​，然后决定躲藏维度应该怎么切开。下面用图示来直观解说一下。

咱们能够将对 nmaxn_{text{max}}nmax​ 的分化视为将其维度折叠到第三个轴，下图所示为 config.axis_pos_shape = [7, 7] 分化:

三个直立矩形棱柱别离对应于编码向量 e1,e2,e49mathbf{e}_{1}, mathbf{e}_{2}, mathbf{e}_{49}e1​,e2​,e49​，咱们能够看到 49 个编码向量被分为 7 行，每行 7 个向量。现在的想法是仅运用 7 个编码向量中的一行，并将这些向量扩展到其他 6 行。本质上是想让七行重用一行的值，可是又不能让不同方位的编码向量的值相同，所以要将每个维度 ( 或称 高度) 为 config.hidden_size=4 的向量切开成两个部分: 巨细为 $1$ 的低区编码向量 edownmathbf{e}_text{down}edown​ 以及巨细为 $3$ 的高区编码向量 eupmathbf{e}_text{up}eup​，这样低区就能够沿行扩展而高区能够沿列扩展。为了讲清楚，咱们还是画个图。

能够看到，咱们已将嵌入向量切为 edownmathbf{e}_text{down}edown​ ( 蓝色 ) 和 eupmathbf{e}_text{up}eup​ ( 黄色 ) 两个部分。现在对 子 向量 Edown=[edown,1,…,edown,49]mathbf{E} _text{down} = left[mathbf{e}_ {text{down},1}, ldots, mathbf{e} _{text{down},49}right]Edown​=[edown,1​,…,edown,49​] 仅保留榜首行的 7 个子向量， 即 图中宽度，并将其沿列 ( 又叫 深度) 扩展。相反，对 子 向量 Eup=[eup,1,…,eup,49]mathbf{E}_text{up} = left[mathbf{e}_{text{up},1}, ldots, mathbf{e }_{text{up},49}right]Eup​=[eup,1​,…,eup,49​] 仅保留榜首列的 $7$ 个子向量并沿行扩展。此刻，得到的嵌入向量 e′imathbf{e’}_{i}e′i​ 如下:

e′i=[[edown,i%nmax1]T,[eup,⌊inmax2⌋]T]Tmathbf{e’}_{i} = left[ left[mathbf{e}_{text{down, } i % n_text{max}^1}right]^T, left[mathbf{e}_{text{up, } left lfloor{frac{i}{{n}^2_{text{max}}}}right rfloor} right]^T right]^T e′i​=⎣⎡​[edown,i%nmax1​​]T,⎣⎡​eup,⌊nmax2​i​⌋​⎦⎤​T⎦⎤​T

本例中，nmax1=7n_text{max}^1 = 7nmax1​=7，nmax2=7n_text{max}^2 = 7nmax2​=7 。这些新编码 E′=[e′1,…,e′nmax]mathbf{E’} = left[mathbf{e’}_{1}, ldots, mathbf{e’}_{n_text{max}}right]E′=[e′1​,…,e′nmax​​] 称为 轴向方位编码。

下图针对咱们的比方对轴向方位编码进行了更详细的阐明。

现在应该很清楚怎么仅依据维度为 dh1nmax1d_{h}^1 times n_{text{max}^1}dh1​nmax1​ 的 Edownmathbf{E}_{text{down}}Edown​ 及维度为 dh2nmax2d_{h}^2 times n_{text{max}}^2dh2​nmax2​ 的 Eupmathbf{E}_{text{up}}Eup​ 核算终究方位编码向量 E′mathbf{E’}E′ 了。

这儿的关键是，轴向方位编码能够从规划上保证向量 [e′1,…,e′nmax]left[mathbf{e’}_1, ldots, mathbf{e’}_{n_{text{max} }}right][e′1​,…,e′nmax​​] 之间各不相等，并且使编码矩阵的巨细从 nmaxdhn_{text{max}} times d_{h}nmax​dh​ 减小到 nmax1dh1+nmax2dh2n_{text{max}}^1 times d_{h}^1 + n_text{max}^2 times d_{h}^2nmax1​dh1​+nmax2​dh2​。由于规划上答应每个轴向方位编码向量不同，所以一旦模型中的轴向方位编码训出来后，模型就能够灵敏高效地获取方位编码。

为了证明方位编码矩阵的尺度得到了大幅减小，假定咱们为 Reformer 模型设置了参数 config.axis_pos_shape = [1024, 512] 以及 config.axis_pos_embds_dim = [512, 512] ，且该模型支撑的最长输入序列长度为 0.5M 词元。此刻，生成的轴向方位编码矩阵的参数量仅为 $1024512+512512\sim 800K$，即大约 3MB。这个数字与规范方位编码矩阵所需的 2GB 比较，简直是小巫见大巫。

如需更简练、更数学化的解说，请参看 此处 的 Transformers 文档。

基准测验

终究，咱们对传统方位嵌入与 轴向方位嵌入 的峰值内存耗费进行比较。

#@title Installs and Imports
# pip installs
!pip -qq install git+https://github.com/huggingface/transformers.git
!pip install -qq py3nvml
from transformers import ReformerConfig, PyTorchBenchmark, PyTorchBenchmarkArguments, ReformerModel

方位嵌入仅取决于两个装备参数: 输入序列答应的最大长度 config.max_position_embeddings 以及 config.hidden_​​size 。咱们运用一个模型，其支撑的输入序列的最大答应长度为 50 万个词元，即 google/reformer-crime-and-punishment ，来看看运用轴向方位嵌入后的作用。

首要，咱们比较轴向方位编码与规范方位编码的参数形状，及其相应模型的总参数量。

config_no_pos_axial_embeds = ReformerConfig.from_pretrained("google/reformer-crime-and-punishment", axial_pos_embds=False) # disable axial positional embeddings
config_pos_axial_embeds = ReformerConfig.from_pretrained("google/reformer-crime-and-punishment", axial_pos_embds=True, axial_pos_embds_dim=(64, 192), axial_pos_shape=(512, 1024)) # enable axial positional embeddings
print("Default Positional Encodings")
print(20 *'-')
model = ReformerModel(config_no_pos_axial_embeds)
print(f"Positional embeddings shape: {model.embeddings.position_embeddings}")
print(f"Num parameters of model: {model.num_parameters()}")
print(20 *'-' + 'nn')
print("Axial Positional Encodings")
print(20 *'-')
model = ReformerModel(config_pos_axial_embeds)
print(f"Positional embeddings shape: {model.embeddings.position_embeddings}")
print(f"Num parameters of model: {model.num_parameters()}")
print(20 *'-' + 'nn')

HBox(children=(FloatProgress(value=0.0, description='Downloading', max=1151.0, style=ProgressStyle(description…
Default Positional Encodings
--------------------
Positional embeddings shape: PositionEmbeddings(
  (embedding): Embedding(524288, 256)
)
Num parameters of model: 136572416
--------------------
Axial Positional Encodings
--------------------
Positional embeddings shape: AxialPositionEmbeddings(
  (weights): ParameterList(
      (0): Parameter containing: [torch.FloatTensor of size 512x1x64]
      (1): Parameter containing: [torch.FloatTensor of size 1x1024x192]
  )
)
Num parameters of model: 2584064
--------------------

了解了相应的理论后，读者应该不会对轴向方位编码权重的形状感到惊讶。

从结果中能够看出，关于需求处理如此长输入序列的模型，运用规范方位编码是不切实际的。以 google/reformer-crime-and-punishment 为例，仅规范方位编码本身参数量就超越 100M。轴向方位编码能够将这个数字削减到略高于 200K。

终究，咱们比较一下推理所需内存。

benchmark_args = PyTorchBenchmarkArguments(sequence_lengths=[512], batch_sizes=[8], models=["Reformer-No-Axial-Pos-Embeddings", "Reformer-Axial-Pos-Embeddings"], no_speed=True, no_env_print=True)
benchmark = PyTorchBenchmark(configs=[config_no_pos_axial_embeds, config_pos_axial_embeds], args=benchmark_args)
result = benchmark.run()

1 / 2
2 / 2
====================      INFERENCE - MEMORY - RESULT       ====================
--------------------------------------------------------------------------------
          Model Name             Batch Size     Seq Length    Memory in MB 
--------------------------------------------------------------------------------
Reformer-No-Axial-Pos-Embeddin       8              512             959      
Reformer-Axial-Pos-Embeddings        8              512             447      
--------------------------------------------------------------------------------

能够看出，在 google/reformer-crime-and-punishment 模型上，运用轴向方位嵌入可削减大约一半的内存需求。

英文原文: hf.co/blog/reform…

原文作者: Patrick von Platen

译者: Matrix Yao (姚伟峰)，英特尔深度学习工程师，作业方向为 transformer-family 模型在各模态数据上的运用及大规模模型的练习推理。