[论文解读] Block-Recurrent Transformers

DeLesley Hutchins, Imanol Schlag+3 more|arXiv (Cornell University)|Mar 11, 2022

Topic Modeling被引用 27

一句话总结

Block-Recurrent Transformer 通过对块循环应用 Transformer 层来处理长序列，在序列长度上实现线性复杂度，并且在类似或较低的计算量下对长文档的困惑度显著更好，优于非循环基线。

ABSTRACT

We introduce the Block-Recurrent Transformer, which applies a transformer layer in a recurrent fashion along a sequence, and has linear complexity with respect to sequence length. Our recurrent cell operates on blocks of tokens rather than single tokens during training, and leverages parallel computation within a block in order to make efficient use of accelerator hardware. The cell itself is strikingly simple. It is merely a transformer layer: it uses self-attention and cross-attention to efficiently compute a recurrent function over a large set of state vectors and tokens. Our design was inspired in part by LSTM cells, and it uses LSTM-style gates, but it scales the typical LSTM cell up by several orders of magnitude. Our implementation of recurrence has the same cost in both computation time and parameter count as a conventional transformer layer, but offers dramatically improved perplexity in language modeling tasks over very long sequences. Our model out-performs a long-range Transformer XL baseline by a wide margin, while running twice as fast. We demonstrate its effectiveness on PG19 (books), arXiv papers, and GitHub source code. Our code has been released as open source.

研究动机与目标

• 激励并开发一种在保持长程记忆能力的同时避免二次注意力成本的 Transformer 架构。
• 引入一种在块级令牌与状态向量上运行的递归单元，以实现长上下文建模。
• 在与标准 Transformer 相当的计算量下，在长序列语言建模任务上实现显著的困惑度改进。

提出的方法

• 使用滑动窗口（因果）注意力，并通过区块级递归处理长度为 N 的段落（块大小为 W）。
• 实现一个递归单元，结合垂直方向（自注意力 + 对递归状态的交叉注意力）和水平方向（状态对状态自注意力并对输入有交叉注意力）的信息流，并配备遗忘/记忆信息的门控。
• 在垂直和水平方向之间共享键和值，同时保持单独的查询。
• 引入学习得到的状态 ID 和相对位置偏置，以在长序列上稳定训练。
• 尝试不同的门控配置（固定门、LSTM 风格门控），并评估稳定性、扩展性和消融效果。

实验结果

研究问题

• RQ1 区块级递归在 Transformer 中是否能对极长上下文（例如 60k 标记）保持记忆，并在困惑度方面与 Transformer-XL 基线相比具有竞争力甚至更优？
• RQ2 递归门类型和状态大小对困惑度、训练稳定性和计算成本的影响如何？
• RQ3 在长文档数据集上，单层递归与增加多层非递归结构相比，性能和成本有何差异？

主要发现

模型	segment length	window length	step time	PG19 (字节/标记)	arXiv (标记)	GitHub (标记)
XL:512	512	512	0.88	1.01	3.62±0.01	1.45	1.21
XL:1024	1024	1024	1.20	0.997	3.59±0.01	1.37	1.08
XL:2048	2048	2048	2.11	0.990	3.58±0.01	1.31	1.01
Slide:12L	4096	512	0.93	0.989	3.60	1.43	1.19
Slide:13L	-	-	1.00	0.989	3.58±0.01	1.42	1.17
Rec:lstm:dual	4096	512	1.06	0.985	3.54±0.01	1.26	1.01
Rec:lstm:single	-	-	1.05	0.962	3.54±0.01	1.29	1.03
Rec:lstm:skip	-	-	1.00	0.969	3.56±0.01	1.31	1.10
Rec:fixed:dual	-	-	1.01	0.957	3.52±0.01	1.27	0.991
Rec:fixed:single	-	-	1.02	0.966	3.58±0.01	1.25	1.00
Rec:fixed:skip	-	-	0.99	0.952	3.53±0.01	1.24	0.976
Feedback:lstm:single	4096	512	1.40	0.977	3.50	1.22	-
Feedback:fixed:skip	-	-	1.35	0.935	3.49	1.24	-
Memorizing Trans. 64k	512	512	1.94	0.950	3.53	1.22	-

• 单层递归在计算量相似的情况下，比增加非递归层带来更大的困惑度提升。
• 在 PG19、arXiv 和 GitHub 上，Block-Recurrent 变体在类似或更低的步时下优于 Transformer-XL 基线，并且在比更强的广_CONTEXT 基线更快的情况下也可运行更快。
• 最佳配置（Rec:fixed:skip）通常优于其他设置，且在 XL:2048 上也能击败并以两倍以上的速度运行；它还在 1.3B 参数模型的 PG19 上实现了词级困惑度的最新状态（26.50）。
• 缩放结果显示，递归在不同模型规模上均有收益，且随着模型规模扩大，收益进一步提升（参数翻倍后）。
• 通常一层递归就足以获得大部分收益；再加入超过一层递归层在性能提升上有限或取决于数据集。
• 区块递归模型可以在没有专门的 CUDA 内核且开源发布的情况下，达到甚至超过基线的性能。