Transformer Quality in Linear Time

来自QVL

https://github.com/lucidrains/FLASH-pytorch

https://arxiv.org/abs/2202.10447

苏老师的解读：FLASH：可能是近来最有意思的高效Transformer设计 - 科学空间|Scientific Spaces (kexue.fm)

1、提出了一种新的Transformer变体，它依然具有二次的复杂度，但是相比标准的Transformer，它有着更快的速度、更低的显存占用以及更好的效果；

2、提出一种新的线性化Transformer方案，它不但提升了原有线性Attention的效果，还保持了做Decoder的可能性，并且做Decoder时还能保持高效的训练并行性。

苏老师的其它相关分析：

GAU中的relu^2 和softmax与长度泛化的关系：https://kexue.fm/archives/9019

GAU还需要热身训练吗？https://kexue.fm/archives/8990

Pre-norm和post-norm：https://kexue.fm/archives/9009

摘要：我们重新审视了Transformer的设计选择，并提出了解决其在处理长序列时弱点的方法。首先，我们提出了一个名为门控注意单元的简单层，它允许在质量损失最小的情况下使用较弱的单头注意。然后，我们提出了一种线性近似方法来补充这个新层，它对加速器友好，在质量上具有很强的竞争力。最终生成的模型名为FLASH（Fast Linear Attention with a Single Head），在短（512）和长（8K）上下文长度上匹配了改进的Transformers的复杂性，在Wiki-40B上实现了4.9倍的训练加速，在PG-19上实现了12.1倍加速的自回归语言建模，在C4上实现了4.8倍加速的屏蔽语言建模。

1.引言

Transformer（Vaswani et al.，2017）已经成为最先进的深度学习系统的新引擎，在语言（Devlin et al.，2018；Brown et al.，2020）和视觉（Dosovitskiy et al.，2020）方面取得了许多最新突破。尽管它们的模型尺寸一直在增长，但由于在输入长度上的二次复杂性，大多数Transformer仍然限于较短的上下文尺寸。这种限制使Transformer模型无法处理长期信息，这是许多应用的一个关键特性。

许多技术已经被提出，通过更有效的注意机制，在扩展上下文中加速Transformer（Child等，2019；DAI等，2019；Re等人，2019；CuroMangSkyet等人，2020；王等人，2020；KATARPopoul等人，2020；BelTaGy等人，2020；ZaHeer-等人，2020；KiaEV等人，2020；罗伊等人，2021；Jayle等人，2021）。尽管其中一些方法具有线性理论复杂性，但在最先进的系统中，vanilla Transformer仍然是主要的选择。在这里，我们从实践的角度研究这个问题，发现现有的有效注意方法至少存在以下一个缺点：

•    质量低劣。

我们的研究表明，当添加几个简单的调整后，vanilla Transformers可能比文献中使用的常用基线更强（参见图1中的Transformer vs.Transformer++）。与增强型Transformer相比，现有的高效注意方法通常会导致显著的质量下降，这种下降超过了它们的效率优势。

•    Overhead in Practice。

由于有效的注意方法通常会使Transformer层复杂化，并需要大量的数据/内存格式化操作，因此在GPU或TPU等加速器上，它们的理论复杂度和实际速度之间可能存在着不可忽视的差距。

•    效率低下的自回归训练。

大多数注意力线性化技术在推理过程中享受快速解码，但在自动回归任务（如语言建模）上的训练可能非常缓慢。这主要是由于他们的RNN式连续状态更新在大量步骤中进行，使得在训练期间不可能充分利用现代加速器的力量。我们通过开发一个新的模型系列来解决上述问题，该系列第一次不仅在质量上与完全增强的Transformer不相上下，而且在现代加速器上真正享受到相对于上下文大小的线性可扩展性。与现有的直接针对Transformer中的多头部自注意（MHSA）近似的有效注意方法不同，我们从一个新的层设计开始，该层设计自然能够实现更高质量的近似。具体来说，我们的名为FLASH的模型分两步开发：

首先，我们提出了一个能更有效近似的新的层（we propose a new layer that is more desirable for effective approximation）。我们引入了一种门控机制来减轻自注意的负担，从而形成了图2中的门控注意单元（GAU）。与Transformer层相比，每个GAU层都更便宜，更重要的是，其质量不太依赖于自注意力的精度。事实上，拥有一个小的单头、无softmax注意力的GAU和Transformer一样出色。虽然GAU在上下文大小上仍然存在二次复杂性，但它削弱了注意力的作用，因此允许我们以后以最小的质量损失进行近似。

然后，我们提出了一种有效的方法来近似GAU中的二次注意，从而在上下文大小上产生具有线性复杂度的层变量。关键思想是首先将标记分组成块，然后在块内使用精确的二次注意，并在块之间使用快速线性注意，如图4所示。我们进一步描述了如何从这个公式中自然地导出一个高效的加速器实现，在实践中只需几行代码更改即可实现线性可伸缩性。

我们进行了大量实验，以证明FLASH在各种任务（掩蔽和自回归语言建模）、数据集（C4、Wiki-40B、PG19）和模型量表（110M到500M）上的有效性。值得注意的是，FLASH与完全增强的Transformers（Transformer++）在质量上具有竞争力，适用于各种实际感兴趣的上下文大小（512–8K），同时在现代硬件加速器上实现线性可扩展性。例如，在质量相当的情况下，FLASH在Wiki-40B上实现了1.2×到4.9×的语言建模加速，在Transformer++上实现了C4上屏蔽语言建模1.0×到4.8×的加速比。随着我们进一步扩展到PG-19（Rae et al.，2019），FLASH将Transformer++的训练成本降低了12.1倍，并在质量上取得了显著的提高。

2.    门控注意力单元

在这里，我们介绍了门控注意单元（GAU），一个比Transformer更简单但性能更好的层。虽然GAU在上下文长度上仍具有二次复杂性，但它是一种更可取的近似方法，在第3节中会阐述。我们从介绍相关层开始：

Vanilla MLP。令是T个tokens的表示。Transformer的MLP的输出可以表示为，其中，。这里表示特征维度，表示扩展的中间维度，是元素激活函数。

门控线性单元（GLU）。    这是一种改进的MLP变体，增加了选通功能（Shazeer，2020）。在许多情况下，GLU已经被证明是有效的（NARAN等人，2021），并用于最先进的Transformer语言模型（DU等人，2021；ToopPILAN等人，2022）。

        （1）

        （2）

其中代表元素乘法。在GLU中，每个表示由与同一token关联的另一个表示进行选通。

门控注意力单元（GAU）。    关键思想是将注意力和GLU描述为一个统一的层，并尽可能多地共享它们的计算（图2）。这不仅可以提高参数/计算效率，而且自然可以实现强大的注意选通机制。具体而言，GAU将GLU中的等式（2）概括如下：

    （3）

这里包含token-token的注意权重。不同于GLU总是使用来选通（两者都与同一token关联）不同，我们的GAU用一个可能更相关的表示替换，该表示使用注意从所有可用token中“检索”。当是一个单位矩阵时，上述内容将简化为GLU。

与Liu等人（2021，Pay attention to mlps）的研究结果一致，门控的存在允许使用比MHSA更简单/更弱的注意机制，而不会造成质量损失：

        （4）

        （5）

其中是一个共享表示（）（注释：除非特别说明，本本工作中设置s=128），和是两个廉价的变换，它们将每维缩放和偏移量(per-dim scalars and offsets)应用于（类似于LayerForms中的可学习变量），是相对位置偏差。我们还发现，在GAU的情况下，MHSA中的softmax可以简化为规则激活函数。（注释：我们使用平方ReLU，来自Primer，在NLP的实验上表现得好）GAU层及其伪代码如图2所示。

与Transformer的MHSA不同，它带有参数，GAU的注意力只在GLU的基础上引入了单个小密集矩阵，该矩阵参数量为（和中的缩放和偏移可以忽略不计）。通过为GAU设置，这种紧凑的设计允许我们用两个GAUs替换每个Transformer块（MLP/GLU+MHSA），同时保持相似的模型尺寸和训练速度。

GAU对Transformer。

图3显示，不同型号模型下，GAUs在TPU上与Transformer（MSHA+MLP/GLU）竞争。注：这些实验是在相对较短的上下文大小（512）上进行的。我们将在后面的第4节中看到，事实上，当上下文长度较长时，GAUs的表现更出色，这是因为它们的注意力能力降低了(thanks to their reduced capacity in attention)。

层烧蚀。    在表1和表2中，我们表明高斯和Transformer本身都是局部最优的(locally optimal on their own)。

3. 使用GAU的快速线性注意

第2节中的两个观察结果促使我们将GAU扩展到长序列建模：

•    首先，GAU中的门控机制允许使用较弱（单头、无softmax）的注意力，而不会造成质量损失。如果我们进一步将这种直觉应用于长时间的注意序列建模，GAU还可以提高近似（弱）注意机制的有效性，例如局部、稀疏和线性化注意。

•    此外，使用GAU，注意力模块的数量可以增倍，在成本上，MLP+MHSA≈2×GAU（第2节）。由于近似注意通常需要更多的层来捕捉完全依赖性（Dai等人，2019；Child等人，2019），该特性也使得GAU在建模长序列时更具吸引力

考虑到这种直觉，我们首先回顾了一些使用注意力做长序列建模的相关工作，然后展示了我们如何使GAU在长序列的线性时间内实现 Transformer级质量。

3.1 现有的线性复杂度变量

部分注意力。

一类流行的方法试图用不同的部分/稀疏模式来近似全注意矩阵，包括局部窗口（Dai等人，2019；Rae等人，2019）、局部+稀疏（Child等人，2019；Li等人，2019；Beltagy等人，2020；Zaheer等人，2020）、轴向（Ho等人，2019；Huang等人，2019），通过哈希（Kitaev等人，2020年）或聚类（Roy等人，2021年）可学习的模式。虽然没有完全注意力那么有效，但这些变体通常能够从扩展到更长序列中获得质量收益。然而，这类方法的关键问题是，它们涉及大量不规则或规则的内存重新格式化操作，如gather，scatter，slice and concatenation，这对大规模并行的现代加速器，尤其是像TPU这样的专用ASICs，是不友好的。因此，它们的实际效益（速度和RAM效率）在很大程度上取决于加速器的选择，通常落后于理论分析。因此，在这项工作中，我们故意最小化模型中内存重新格式化操作的数量。

线性注意力。

或者，另一种流行的研究路线通过分解注意力矩阵，然后重新安排矩阵乘法的顺序，将注意力计算线性化（Choromanski等人，2020年；Wang等人，2020年；Katharopoulos等人，2020年）。从示意图上看，线性注意可以表示为：

其中分别表示查询、键和值。重新安排计算将复杂性w.r.t从二次降低到线性。

线性注意的另一个理想特性是，在推理时，它对每个自回归解码步骤的计算和内存是恒定的（注释：恒定是相对于序列长度而言）。

具体地说，我们定义，并注意Mt的计算为增量的（incremental）形式：

    （6）

这意味着我们只需要维护一个具有恒定内存的高速缓存，每当新输入到达时间戳时，只需要进行恒定计算即可将累积到并得到。相反，在每个解码步，完全二次注意需要线性计算和内存，因为每个新输入必须注意前面的所有步骤。

然而，另一方面，在自回归训练中，线性注意的 re-arranging计算会导致严重的效率低下。如图4（mid）所示，由于自回归训练的因果约束，每个时间步处的查询向量对应于不同的缓存值。这要求模型计算并缓存个不同的值，而不是在非自回归情况下仅一个值。理论上，序列可以以获得，通过首先计算，然后对T个tokens执行大的累积和（cumsum）来获得。但在实践中，cumsum引入了一种RNN风格的T步顺序依赖关系，其中每个步都需要处理一个状态。顺序依赖不仅限制了并行度，而且更重要的是需要循环的内存访问，这通常比在现代加速器上计算元素加法要花费更多的时间。因此，理论复杂度与实际运行时间之间存在相当大的差距。在实践中，我们发现在TPU（图6（a））和GPU（附录C.1）上，直接计算完整的二次注意矩阵甚至比重新排列（线性化）的版本更快。

3.2. 我们的方法：混合块注意Mixed Chunk Attention

基于现有线性复合注意的优缺点，我们提出了混合组块注意，它融合了部分注意和线性注意的优点。图4展示了高级思想。下面我们将重新制定GAU以纳入这一理念。

准备    首先将输入序列分为size为的个非重叠块，即。然后，据式（1）和式（4）中的GAU公式，为每个chunk生成，和。接下来，通过应用per-dim scaling and offset（这是非常便宜的），从产生四种类型的注意头、、、。

我们将描述如何使用局部注意力加上全局注意力有效地近似GAUs注意力。注：所有主要张量、和在两个组件之间共享。在原始GAU上引入的唯一附加参数是用于生成和的每维缩放和偏移（参数）。

Note all the major tensors Ug, Vg and Zg are shared between the two components. The only additional parameters introduced over the original GAU are the per-dim scalars and offsets for generating Q^lin_g and K^lin_g (4×s parameters).

每个块的局部注意。首先，将局部二次注意独立地应用于长度C的每个块，以产生预选通状态的部分(part of the pre-gating state)：

这部分的复杂度是，如果保持不变，这相对于是线性的。

块间的全局注意。

此外，一个全局线性注意机制被用来捕捉块间的远程交互

非因果：（7）

因果：（8）

注：式（7）和式（8）中的求和是在块级别执行的。对于因果（自回归）的情况，这将使token-level的线性注意力的cumsum计算的元素数减少了C倍（在我们的实验中，典型的C为256），从而导致显著的训练加速。

最后，和相加，然后是门控和类似于等式（3）的后注意投影(post-attention projection)：

混合块注意易于实现，代码1给出了相应的伪代码。

3.2.1. 讨论

快速自回归训练。

重要的是，如图4（底部）所示，由于分块，自回归情况下的顺序依赖性从标准线性注意中的步减少到等式（8）中分块版本中的步。因此，我们观察到，当块大小（chunk size）在{128，256，512}时，自回归训练会显著加快。在消除了自回归训练的低效性的情况下，所提出的模型仍然享有每步解码的恒定的存储和计算，其中额外的常数C来自局部二次注意。

不重叠的局部注意。    我们的方法中的块不会相互重叠。理论上，在保持分块线性注意不变的情况下，任何部分注意变体都可以作为非重叠局部注意的替代。作为一个具体的例子，我们探索了允许每个块额外地关注其附近的块，这基本上使局部注意力变成重叠的，类似于Longformer（Beltagy et al.，2020）和BigBird（Zaheer et al.，2020）。虽然重叠的局部注意力持续地提高了质量，但它也引入了许多明显损害实际运行速度的内存重新格式化操作。在我们对TPU上的语言建模进行的初步实验中，我们发现，在内存和速度方面，使用重叠局部注意的成本效益权衡可能不如增加更多层。总的来说，我们认为最佳的部分注意变体是取决于具体任务的(optimal partial attention variant is task-specific)，而非重叠的局部注意与分块线性注意的选择相结合时，总是一个很好的选择。

与Combiner的联系。

与我们的方法类似，Combiner（Ren et al.，2021）也将序列分割成非重叠的块，并在每个块中利用二次局部注意。关键的区别在于如何汇总远程信息，并将其与本地信息结合起来（例如，我们的混合块注意允许每个块拥有更大的有效内存，从而提高质量）。详细讨论见附录A。

4.实验

我们关注我们模型的两个版本，它们在上下文长度方面具有不同的复杂性。二次复杂度模型FLASH Quad指的是堆叠的GAUs，而线性复杂度模型FLASH由GAUs和提出的混合块注意组成。为了证明它们的有效性和普遍适用性，我们在多个大规模数据集的双向和自回归序列建模任务中对它们进行了评估。

基线。

首先，带有GELU激活功能（Hendrycks&Gimpel，2016）的vanilla Transformer（Vaswani et al.，2017）被列为校准的标准基线。尽管在文献中是一个流行的基线，但我们发现RoPE（Su et al.，2021）和GLU（Shazeer，2020）可以显著提高性能。因此，我们还将Transformer+RoPE（Transformer+）和Transformer+RoPE+GLU（Transformer++）作为两个二次复杂度更高的基线。

为了证明我们的模型在长序列上的优势，我们进一步将我们的模型与两个显著的线性复杂度变型Performer（Choromanski et al.，2020）和Combiner（Ren et al.，2021）进行了比较，其中Performer是一种具有代表性的线性注意力方法，而Combiner（使用类似于我们的分块注意力设计）与许多其他方法相比表现出了更高的成本效益权衡（Ren等人，2021年）。为了获得最佳性能，我们使用了Combiner的rowmajor-axial变体（CombinerAxial）和Performer的ReLU核变体。这两种型号都配有RoPE。

为了进行公平比较，所有模型都在同一个代码库中实现，以确保用于训练和评估的相同tokenizer和超参数。使用TensorFlow Profiler测量所有模型的每一步训练潜延迟。有关详细设置和型号规格，请参见附录B。

4.1. 双向语言建模

在BERT（Devlin等人，2018年）中，屏蔽语言建模（MLM）在输入序列中重建随机屏蔽的标记。我们对C4数据集上的所有模型进行预训练和评估（Raffel等人，2020年）。我们始终使用每批2^18个token对每个模型进行125K步的训练，同时在很大范围内改变上下文长度，包括512、1024、2048、4096和8192。每个模型的质量都作为下游任务性能的代理指标在perplexity中报告。每个模型的训练速度（即每一步的训练延迟）用64个TPU-v4核来衡量，总训练成本以TPU-v4-core-days来报告。

图5（a）显示了所有模型在不同上下文长度下每个训练步骤的延迟。为简洁起见，Transformer+的结果被省略，因为它位于Transformer和Transformer++之间。在所有六种模型中，随着上下文长度的增加，Combiner、Performer和FLASH的延迟大致保持不变，这表明了上下文长度的线性复杂性。FLASH Quad在所有上下文长度上都始终比Transformer和Transformer++更快。特别是，当上下文长度增加到8192时，FLASH Quad的速度是Transformer++的2倍。更重要的是，如图5（b）-5（f）所示，对于从512到8192的所有序列长度，我们的模型总是在相同的计算资源下获得最佳质量（即最低的复杂度）。特别是，如果目标是匹配Transformer++在步骤125K的最终困惑，FLASH Quad和FLASH可以分别将训练成本降低1.1×-2.5×和1.0×-4.8×。还值得一提的是，FLASH是唯一一个与二次复杂度对应的线性复杂度模型实现竞争复杂度的模型。有关所有型号的详细质量和速度比较，请参见附录C.2。

4.2. 自回归语言建模

对于自回归语言建模，我们重点关注Wiki-40B（郭等人，2020年）和PG-19（Rae等人，2019年）数据集，这两个数据集分别由纯英语维基百科页面和从古腾堡项目中提取的书籍组成。值得注意的是，PG-19中的平均文档长度为69K字，非常适合在长上下文长度上评估模型性能。我们对所有模型进行训练和评估，每批使用2^18个token，共125K个步骤，其中Wiki-40B的上下文长度为512到8K，PG-19的上下文长度为1K到8K。我们报告了Wiki-40B的标记级困惑和PG-19的单词级困惑。

图6（a）显示，在二次复杂度和线性复杂度模型中，FLASH Quad和FLASH分别实现了最低的延迟。在图6（B）-6（f）中，我们比较了Wiki40-B上所有模型在增加上下文长度方面的质量和训练成本权衡。与MLM任务的研究结果类似，我们的模型在所有序列长度的质量训练速度方面主导了所有其他模型。具体来说，FLASH Quad将Transformer++的训练时间减少了1.2倍到2.5倍，FLASH将计算成本减少了1.2倍到4.9倍，同时达到了与Transformer++类似的复杂程度。在我们自己的模型之间，FLASH密切跟踪FLASH Quad的复杂度，并在上下文长度超过2048时开始实现更好的复杂度成本权衡。所有型号的详细质量和速度比较见附录C.2。

对于PG-19，在Rae等人之后，所有模型都使用了大约500M参数（见表9）的增加模型比例进行比较。结果汇总在表3中。与Wiki-40B中的数字相比，FLASH在困惑度和训练时间方面比PG-19上的Transformer+有了更显著的改善。例如，当上下文长度为8K时，FLASH Quad和FLASH能够在55K和55K步中达到Transformer+的最终复杂度（125K步），分别产生5.23倍和12.12倍的加速比。我们假设，相对于Transformer+增加的收益来自PG-19（由书籍组成）的长期性质。与我们之前的实验类似，FLASH比所有全注意Transformer变体的复杂度都要低，这证明了我们高效注意设计的有效性。

4.3. 消融研究

二次和线性分量的重要性。

为了更好地理解FLASH的功效，我们首先研究局部二次注意和全局线性注意分别对性能的贡献。为此，我们创建了FLASH（LocalOnly）和FLASH（GlobalOnly），分别保持FLASH中的局部二次注意和全局线性注意。在FLASH（GlobalOnly）中，我们将块大小从256减少到64，以便为全局线性关注生成更多的局部摘要。在图7中，我们看到完整模型和两种变体之间存在显著差异，这表明线性注意和全局注意是相辅相成的——两者都对拟议的混合组块注意的质量至关重要。

GAUs的意义。

这里我们研究在FLASH中使用GAU的重要性。为了实现这一点，我们对Transformer++应用了相同的混合块注意思想。我们将这种变体称为MC-TFM++（MC代表混合组块），它在每个组块中使用二次MHSA，并在组块之间使用多头线性注意。实际上，MCTFM++与FLASH具有相同的线性复杂度，但MC-TFM++的核心是Transformer++而不是GAU。

图7显示FLASH的性能大大优于MC-TFM+（当序列长度大于2048时，加速比超过2倍），证实了GAU在我们设计中的重要性。我们进一步研究了由于表4中的近似方法而增加的困惑，表明从FLASH Quad到FLASH时，近似导致的质量损失比从TFM++到MC-TFM++时小得多。这表明混合组块注意比MHSA更符合GAU，这与我们的直觉相吻合，即GAU更有利于较弱/近似的注意机制。

块大小的影响。区块大小的选择会影响FLASH的质量和训练成本。我们观察到，总体而言，随着上下文长度的增加，较大的块大小表现得更好。例如，当上下文长度超过1024时，将区块大小设置为512显然比默认区块大小（C=256）更可取。在实践中，可以对区块大小进行超参数搜索，以进一步优化FLASH的性能，尽管我们在实验中没有探索这种选项。更详细的分析见附录C.3。

5.结论

我们已经介绍了FLASH，这是一种解决现有高效Transformer变体的质量和经验速度问题的实用解决方案。这是通过设计性能层（选通线性单元）并将其与加速器高效近似策略（混合块注意）相结合来实现的。对双向和自回归语言建模任务的实验表明，FLASH在质量（困惑度）上与完全增强的Transformers一样好，但训练速度远远快于最先进的技术。未来的工作是研究这个新模型族的标度律和下游任务的性能。感谢作者感谢Gabriel Bender、John Blitzer、Maarten Bosma、Ed Chi、Dai Hanjun、Pieter Jan Kindermans和David So的反馈

A.与Combiner的连接

为了获取长期信息，Combiner（Ren等人，2021年）还将每个区块汇总为汇总键和值向量Ksum，V sum∈ R T/C×d，并将它们连接成局部二次注意，即。

有效地，Combiner将C向量的每个块压缩为O（d）的单个向量，而我们的分块线性注意部分将每个块压缩为大小为O（sd）的矩阵Klin h>Vh，后者大s倍。换句话说，在分块线性注意中进行的压缩更少，允许增加内存，因此比组合器具有潜在优势。

另一个区别在于来自不同区块的压缩长期信息如何组合，组合器重用二次注意，而我们的区块线性注意只执行（累积）和。然而，通过构造一个额外的[T/C×T/C]注意矩阵来组合组块摘要，很容易将Combiner所做的事情融入我们提出的方法中。

我们确实对这种变体进行了短暂的实验，发现它很有用。但它显然会使整个模型设计复杂化，更重要的是，它需要模型存储并处理所有块摘要。因此，自回归译码复杂度将增加到O（（C+T/C）d2），它与长度有关，不再是常数。因此，我们在默认配置中不包含此功能。

B.实验装置

B.1。超参数

双向语言建模。

表5列出了C4上MLM任务的超参数。所有模型都是使用相同的代码库实现、训练和评估的，以确保公平比较。

自回归语言建模。

表6列出了Wiki-40B和PG-19上LM任务的超参数。所有模型都是使用相同的代码库实现、训练和评估的，以确保公平比较

B.2。型号规格

表7、8和9总结了我们实验中使用的所有模型的详细规格。在实验中，SiLU/Swish（Elfwing等人，2018年；Hendrycks&Gimpel，2016年；Ramachandran等人，2017年）被用作FLASH Quad和FLASH的非线性，因为它在我们的模型中略优于GELU（Hendrycks&Gimpel，2016年）。还值得注意的是，我们对一些屏蔽语言模型使用ScaleNorm，因为ScaleNorm在TPU-v4上的运行速度略快于LayerForm，而不会影响模型的质量。

C.其他实验结果

在这里，我们使用Nvidia V100 GPU（表10）和FLASH块大小的消融研究（图8），提供了不同语言模型的训练速度的完整结果。

C.1。GPU上的自回归训练

我们观察到，自回归训练的低效性不仅限于TPU等硬件加速器。如表10所示，在三种模型中，Performer的延迟最大，因为它需要按顺序对所有token执行累积求和。相比之下，当上下文长度超过1024时，建议的FLASH实现了最低的延迟，这表明了建议的混合块注意机制的有效性。

C.2。MLM和LM结果表

我们在表11和表12中总结了C4上的MLM和Wiki-40B上的LM的实验结果。

C.3。块大小的烧蚀研究

区块大小的选择会影响FLASH的质量和训练成本。在块大小等于上下文长度的极端情况下，FLASH会退回到FLASH Quad，失去对长上下文长度的可伸缩性。在另一种极端情况下，当组块大小等于1时，所提出的注意模块变成了线性注意，这会导致低效的自回归训练。图8显示了上下文长度从1K到8K的四种不同块大小的质量和训练成本之间的权衡。

D. FLASH-Quad和FLASH-Quad的伪代码

我们在代码6和8中展示了FLASH Quad和FLASH的详细实现