w~Transformer~合集11

我自己的原文哦~ https://blog.51cto.com/whaosoft/12472192

#LightSeq

最高加速9倍！字节跳动开源8比特混合精度Transformer引擎,近年来，Transformer 已经成为了 NLP 和 CV 等领域的主流模型，但庞大的模型参数限制了它的高效训练和推理。于是字节跳动在 2019 年 12 月和 2021 年 6 月分别推出了高效推理和训练引擎 LightSeq，大大加速了 Transformer 系列模型的训练和推理，也打通了 Transformer 从训练到推理的整个流程，极大优化了用户使用体验。最近，LightSeq 训练引擎相关论文[1]，被录用难度极高的超算领域国际顶会 SC22 接收，得到了学术界的广泛认可！

SC22 接收论文：https://sc22.supercomputing.org/presentation/?id=pap211&sess=sess154
代码地址：https://github.com/bytedance/lightseq

如何继续提升速度？降低计算精度是比较直接的方法。2017 年以来，fp16 混合精度技术 [2] 获得了广泛应用。在对模型效果无损的前提下，将模型训练和推理的速度提升了 50% 以上。而为了维持模型效果，更低精度的方法（例如 int8）通常需要使用如下传统方案：

首先使用 fp16 混合精度将模型训练至收敛；
然后在模型计算密集型算子的权重、输入和输出位置处，插入伪量化结点，进行量化感知训练；
最后将带有伪量化结点的模型计算图转换到专用的 int8 推理引擎中，进行服务部署和模型推理。

虽然在多数任务上，上述方案可以实现模型效果无损，但还是存在以下问题：

使用方法复杂。例如要多一次量化感知训练 [4] 的过程，并且带有伪量化节点的计算图转换复杂。
训练速度慢。由于目前流行的深度学习框架不支持 int8 精度，所以量化感知训练需要插入 fp16 的伪量化结点来模拟 int8 量化，导致量化感知训练反而比 fp16 混合精度训练慢 2-3 倍。
推理部署难且加速比低。对比 fp32、fp16 等类型，int8 硬件和底层软件库优化相对滞后。例如在 NVIDIA GPU 上，int8 矩阵乘法加速受限于硬件架构和特定 shape，实际加速比远远低于理论值。

在下文中，如无特殊说明，量化都是指的 int8 精度的量化。

针对这些问题，字节跳动推出了全新版本的 LightSeq GPU 量化训练与推理引擎。支持 Transformer 系列模型的量化训练与推理，并做到了开箱即用，用户友好。LightSeq 快准狠地实现了 int8 精度的量化训练和推理：

快：A100 多卡训练最高加速 5.2 倍，T4 单卡推理最高加速 8.9 倍。
准：训练和推理效果基本无损。
狠：相同数据量下，显存占用最高减少 68%，模型存储空间减少 75%。

总体来说，LightSeq 新版量化训练与推理引擎具有如下几个优点：

1. 丰富的支持

支持完整的 Transformer 模块和多种解码算法，支持 Transformer、BERT、GPT、BART、ViT 等多种模型结构，支持 Fairseq、Hugging Face、NeurST 等多种训练框架接入量化训练、导出模型以及量化推理，提供了丰富的样例供用户参考。

2. 卓越的性能

相比于 fp16 精度的 LightSeq 推理引擎，int8 量化还可以进一步加速最高 70%，相比于 PyTorch 推理更是达到了最高 8.9 倍的加速比。同时显存占用相比 fp16 推理引擎降低了 30% 左右，模型存储空间只需要原来的四分之一。最后经过多个任务的验证，推理效果几乎无损。

3. 便捷的使用

LightSeq 已经针对多个训练库进行了量化支持，可以一键开启量化训练，然后轻松导出为 LightSeq 支持的模型格式，最后实现量化推理。除此之外，LightSeq 还支持训练后量化，无需额外训练即可体验量化推理。

如上图所示，为了最大程度减小量化带来的损失，首先需要用 fp16 精度训练一个浮点数模型，将模型效果训到最好。然后开启量化进行 finetune，得到微调过的量化模型，此时模型效果已经基本恢复到浮点数模型的水平。接着将量化模型转换为 LightSeq 支持的 PB 或者 HDF5 模型格式，最后用 LightSeq 进行量化推理。

安装方法

LightSeq 安装非常简单，只需要一行命令即可：

pip install lightseq

量化训练

LightSeq 支持 Fairseq、Hugging Face、NeurST 等训练框架的量化接入，同时也可以自定义模型并开启量化训练。以 encoder 层为例，只需要先定义浮点数模型，然后开启量化即可：

from lightseq.training import LSTransformerEncoderLayer
from lightseq.training.ops.pytorch.quantization import enable_quant
config = LSTransformerEncoderLayer.get_config(model="bert-base",max_batch_tokens=4096,max_seq_len=512,fp16=True,local_rank=0,
)
layer = LSTransformerEncoderLayer(config)
# 开启量化
layer.apply(enable_quant)

量化推理

LightSeq 提供了便捷的 python 推理接口，只需要三行代码即可实现快速的量化推理：

import lightseq.inference as lsi
model = lsi.QuantTransformer(pb_path, batch_size)
result = model.infer(input)

此外 LightSeq 还提供了 BERT、GPT、ViT 等模型的 python 接口，分别调用 QuantBert、QuantGpt 和 QuanVit 即可体验。

梯度通信量化

LightSeq 支持 Transformer 模型的梯度通信量化[5]，使用 Fairseq 或者 Hugging Face 即可轻松开启分布式量化训练，并同时支持浮点数模型和量化模型。在构建模型后，只需要为模型注册一个 communication hook 即可开启梯度通信量化，再开始训练过程。

from lightseq.training.gradient_comm_quantization import encode_and_decode, GCQState
from torch.nn.parallel import DistributedDataParallel 
# model could be from Fairseq or Hugging Face, wrapped by DDP
model = DistributedDataParallel(model)
state =  GCQState(process_group)
# register hook
model.register_comm_hook(state=state, hook=encode_and_decode)

性能测试

LightSeq 在多个任务上测试了量化训练、量化推理和梯度通信量化的速度，并且分析了显存占用情况和量化模型的效果。

量化训练速度

LightSeq 在 8 张 A100 显卡上进行了训练实验，主要对比对象是 Fairseq 的 Transformer、Hugging Face 的 BERT、GPT2 和 ViT。

可以看出，四种模型结构加速趋势都是类似的，加速比都会随着数据量的增大而减小，原因有三点：

随着数据量的增大，矩阵乘法 GEMM 的占比会明显增加，因此 PyTorch QAT 增加的额外的伪量化结点时间占比会逐渐减小，最后速度会和 PyTorch fp16 无限接近。
与此同时，随着 GEMM 占比升高，LightSeq fp16 自定义算子的提速效果也逐渐减小，因此时间上也会和 PyTorch fp16 无限接近。
由于 Ampere 架构显卡上 int8 GEMM 在 shape 较小时甚至不如 fp16 GEMM 快，在大 shape 下才能稍快一点，因此随着数据量增大，LightSeq int8 也会无限接近 LightSeq fp16 的速度。

量化推理速度

LightSeq 在单张 T4 显卡上进行了推理实验，主要对比对象是 Hugging Face 的 Transformer、BERT、GPT2 和 ViT。

可以看出，随着输入数据量的增大，LightSeq 与 PyTorch 的差距会逐渐减小，这也是 GEMM 占比升高造成的。比较 LightSeq fp16 和 LightSeq int8，可以看出随着数据量的增大，LightSeq int8 越来越快。这是因为在 T4 显卡上，int8 GEMM 的加速会随着 shape 的增大而有明显增加。因此在 T4 显卡上进行量化推理时，输入数据量越大，加速效果越好。

LightSeq 还针对机器翻译多个语向和多个测试集，测试了不同 batch size 下，LightSeq int8 推理相对于 LightSeq fp16 推理的加速比，实验同样是在单张 T4 显卡上进行的，采用的模型都是标准的 Transformer-Big。

可以得到和上文中相同的结论，随着 batch size 的增大，量化推理的加速比会逐渐升高。相比于 LightSeq fp16，最高还可以再加速近 70%，这极大地缩短了线上翻译模型的推理延时。

最后如上图所示，为了展示自动 GEMM 调优技术的效果，LightSeq 测试对比了 A100 显卡上 Transformer 和 BERT 模型 fp16、int8 调优前和 int8 调优后的延时。可以看出调优前某些 shape 的 int8 GEMM 速度甚至比 fp16 还要慢，而调优后全面超越了 fp16。

显存占用

LightSeq 分析了不同 batch size 下，量化模型相对于浮点数模型显存占用的加速比。可以看出随着 batch size 的增大，量化模型的显存占用优势更明显，最高可以减少 30% 左右。而 LightSeq fp16 引擎相对于 PyTorch 模型也极大程度减少了显存占用，因此 LightSeq int8 引擎最终能够减少最多 68% 左右的显存。

量化模型效果

针对机器翻译多个语向和多个测试集，LightSeq 测试了量化模型推理相对于浮点数模型 BLEU 的损失，采用的模型都是标准的 Transformer-Big。

在数据量较大的语向 en2zh 上，LightSeq int8 相对 BLEU 损失较大些，最大达到了 - 0.4。而在数据量较小的语向 en2es 上，LightSeq int8 不仅没有任何效果损失，反而比浮点数模型更好。总体而言，int8 量化模型的平均 BLEU 相比浮点数模型基本无损。在 GLUE 和 SQuAD 等多个任务上，LightSeq 也验证了量化模型的效果。

梯度通信量化

由于在多机多卡场景下通信瓶颈更加明显，所以梯度通信量化主要应用在分布式训练场景。因此 LightSeq 在 2 机 8 卡的 A100 上进行了分布式训练的速度测试。

可以看出，梯度通信量化的训练加速效果整体上随着输入数据的增大而减弱。这主要是因为随着输入数据的增大，计算时间占比升高，梯度通信时间占比减少，梯度量化的收益也随之减小。

LightSeq 还额外增加了不同数量网卡（NIC）下的训练速度测试。可以看到使用梯度通信量化的分布式训练速度相比原始的 LightSeq fp16 有大幅度提升。

量化技术

int8 量化的加速收益主要来自如下几个方面：

GEMM 精度从 fp16 降低到 int8 后，计算时间缩短；
自定义算子采用 int8 输入输出后，数据读写时间缩短；
梯度采用 int8 存储后，多机之间通信时间缩短。

以 Transformer 模型为例，经过 LightSeq fp16 引擎加速后，自定义算子时间大大缩短，而 GEMM 时间占比提升到了 90% 左右，因此优化的重点转移到了 GEMM 提速。将 fp16 GEMM 替换为 int8 GEMM 不仅可以缩短 GEMM 时间，还可以减小前后算子的输入输出位宽，从而减小读写数据的时间。最后多机训练的瓶颈主要在梯度的通信，将梯度量化为 int8 精度可以大大加快分布式训练的速度。

量化原理

为了弥补量化带来的精度损失，通常需要用量化感知训练来模拟量化过程。如上图所示，量化感知训练就是将 float GEMM 的两个 float 输入分别做一遍量化和反量化（称之为伪量化结点），离散化成分段的浮点数输入，然后进行 float GEMM 运算。得到结果后再次进行量化与反量化，得到最终的浮点数结果。而量化的过程是不可导的，因此需要用 STE 方法来估计量化参数的梯度。之所以量化感知训练中需要插入伪量化结点，然后用 float GEMM 去模拟量化过程，是因为 TensorFlow 和 PyTorch 等训练框架不支持 int8 GEMM。

而 LightSeq 量化训练直接采用 int8 GEMM 来真实还原量化过程，因此相比传统的实现要更快，且更加节省显存。在推理的时候，同样采用离散化后的整数进行 int8 GEMM 运算，最后再反量化回浮点数结果。量化推理过程和量化训练完全一致，并且和传统的量化感知训练是完全等价的。

量化位置

整个量化 Transformer 的网络结构如上图所示，红色箭头表示需要加上量化和反量化结点的位置。

首先所有 int8 GEMM 的输入和输出都需要进行量化。由于 int8 GEMM 的 shape 限制，部分 GEMM（例如注意力分数的计算）仍然采用 float GEMM。此外第二层 FFN 的 GEMM 采用的是 int32 的输出，因为它的 GEMM 输入是 ReLU 激活函数的输出结果，只包含正数，非对称，因此如果采用 int8 输出的 GEMM，将无法反量化为正确的浮点数结果。

然后所有的模型权重 weight 都需要存储为 int8 类型，因此需要对 weight 做量化。而权重 bias 参数量较小，无需量化，保留 float 精度反而可以提升模型效果。

最后需要对 decoder 端的 cache 进行量化。因为在推理时，decoder 端的 cache 需要频繁进行读写，因此将 cache 量化为 int8 可以大大加快解码的速度。

量化策略

将一个浮点数矩阵量化为 int8 整数矩阵有很多方法，LightSeq 采用的是对称量化，即将正负数范围对称的浮点数区间等比例地映射到整数区间 [-127, 127] 上。

而实际上浮点数矩阵的数值范围通常并不对称，存在极少的离群值。如果直接按照离群值的范围来量化矩阵，会影响到量化后的精度，所以需要先对矩阵进行数值截断。

LightSeq 采用 PACT 方法进行截断[6]，将截断的范围当作模型可学习的参数，然后利用 STE 算法去估计参数的梯度，并进行反向传播优化。根据实践经验，权重 weight 的初始截断范围设为[-1, 1]，中间结果的初始截断范围设为[-16, 16]，可以在大部分任务上达到最好的效果。最后经过截断范围和其他模型参数的联合优化，量化模型的效果可以达到基本无损。

梯度通信量化

针对分布式训练场景，LightSeq 推出了梯度量化压缩技术。即对浮点精度的梯度进行 int8 量化，以减少梯度通信的时间消耗，从而加速训练，这就是梯度通信量化（GCQ）。

如上图所示，梯度通信量化的主要流程如下：

计算每张卡上各自梯度的截断范围；
对截断范围执行 all-reduce max 操作；
每张卡使用统一的截断范围对各自梯度进行 int8 量化；
对 int8 梯度执行 all-reduce sum 操作；
每张卡对 all-reduce 后的梯度进行反量化，还原为浮点数梯度，并进行参数更新。

为了解决 int8 梯度在 all-reduce 过程中溢出的问题，LightSeq 首先将每张卡上的浮点数梯度除以卡数，再使用除之前的截断范围进行量化，最后进行 all-reduce 操作。这样每张卡上量化后的 int8 整数 all-reduce 完就不会溢出，但是单卡实际用于量化的比特数也因此而减少，所以目前方案在 2 机 8 卡效果几乎无损，但随着卡数的上涨，训练效果会有所下降。以 en2de 和 en2fr 翻译任务为例，在 4 机 8 卡上进行分布式量化训练，BLEU 值分别会下降 0.4 和 1.5 左右。未来 LightSeq 将会持续探索更好的方法来解决这一问题。

通用技术

除了上一章节中提到的量化技术以外，此次更新 LightSeq 还提出了几种通用的优化技术，不仅可以应用在量化模型中，也适用于其它所有精度模型的训练与推理。

算子融合

上图是 encoder 模块量化训练的计算图，LightSeq 将两次 GEMM 运算之间的所有操作融合成一个算子[7]，减少了 kernel 调用的次数，因此减少了总的计算时间。

图中黄色矩形表示 int8 GEMM，绿色矩形表示 float GEMM。这里采用 float GEMM 是由于 shape 的限制，不适合使用 int8 GEMM 加速。红色箭头表示流动数据的类型是 int8，绿色箭头表示第二层 FFN 的 GEMM 输出是 int32 数据类型。int8 GEMM 输入输出的量化与反量化操作都被融合到了前后 kernel 里，这不仅可以减少数据搬运，还可以减小显存占用。

在推理时，LightSeq 还针对 decoder 做了优化。如上图所示，在计算 self-attention 时，注意力得分的维度是(batch size, 1, sequence length)。因此在计算 value 乘积时，可以不采用 GEMM 运算，而直接手写加权求和的算子，从而将图中虚线框中的计算融合成一个 kernel。

自动显存管理

模型量化引入了更复杂的张量类型和张量依赖关系，这给显存管理带来新的挑战。为此，LightSeq 设计了新的显存管理机制。如上图所示，主要包括以下过程：

训练启动前，根据每个算子的拓扑依赖关系，自动计算每个张量的生命周期及显存空间大小。其中，包含动态维度的张量按照此维度的最大量进行计算，例如机器翻译任务中的最大句长和最大 batch 句子数量。这些最大量在训练前已被指定；
张量确定生命周期和大小后，分析显存复用关系。其中，无生命周期重合的张量可以共用一片显存空间，所有显存空间都是无数据类型的，可以被分配到任意数据类型的张量上；
根据张量显存复用关系，申请多段显存空间，为每个张量分配实际的显存起止地址。

张量显存复用的分析，LightSeq 借鉴了论文 [3] 中提出的 Greedy by Size for Offset Calculation 方法，做了三个改进：

支持了整个训练过程的显存复用（forward/backward）；
不同数据类型能做到显存复用（int8/fp16/fp32）；
在多段显存空间上容纳所有张量，而非一段非常大的显存空间，这样能有效提升显存利用率。

自动 GEMM 调优

LightSeq 的 int8 GEMM 采用了 NVIDIA 的 cuBLASLt 库，这也是目前 NVIDIA 显卡上最为高效的矩阵运算库。但是输入数据的 shape 或者显卡不同的话，GEMM 所采用的最优配置（例如数据排布、GEMM 算法等等）也可能不同，因此需要进行自动选取。LightSeq 采取的自动调优方案如下：

在多种型号显卡上（例如 T4 和 A100）进行不同 shape 的 GEMM 最优配置搜索，并将结果保存到配置文件中，用户只需要下载即可；
模型初始化时，加载对应型号显卡的配置文件，解析并保存到键值对为 (shape, 最优配置) 的字典中。如果没有对应型号显卡的配置文件，或者没有需要的 GEMM shape，那么用户可以选择自己搜索并保存，或者直接使用默认配置；
模型前向或后向计算时，根据输入的 shape 在字典中寻找最优配置，然后进行 GEMM 计算。如果没有找到对应的 shape，那么直接采用默认的配置。

未来工作

未来 LightSeq 还将继续探索移动端的低精度量化、反向传播中梯度的量化、大模型量化等方向。

#SCTNet

SCTNet，一种带有transformer语义信息的单分支CNN用于实时分割。借助于提出的transformer类CNN块CFBlock和语义信息对齐模块，SCTNet可以在训练中从transformer分支捕获丰富的语义信息。 80.5mIoU+62.8FPS! 华科与美团联合提出单分支推理分割架构

最新的实时语义分割方法通常采用额外的语义分支来追求丰富的长距离上下文。然而，额外的分支会带来不必要的计算开销，并减缓推理速度。为了消除这一困境，我们提出了SCTNet，一种带有transformer语义信息的单分支CNN用于实时分割。

https://arxiv.org/abs/2312.17071

https://github.com/xzz777/SCTNet

SCTNet在保留轻量级单分支CNN高效性的同时，还拥有语义分支的丰富语义表示。考虑到transformer提取长距离上下文的卓越能力，SCTNet将transformer作为仅用于训练的语义分支。借助于提出的transformer类CNN块CFBlock和语义信息对齐模块，SCTNet可以在训练中从transformer分支捕获丰富的语义信息。在推理过程中，只需要部署单分支CNN。我们在Cityscapes,ADE20K和COCO-Stuff-10K上进行了广泛的实验，结果表明，我们的方法达到了新的最先进水平。

本文贡献主要包含以下三点：

我们提出了一种新的单支实时分割网络SCTNet。通过学习从Transformer到CNN的语义信息对齐来提取丰富的语义信息，SCTNet在保持轻量级单支CNN快速推理速度的同时，具有Transformer的高准确性。
为了缓解CNN特征和Transformer特征之间的语义鸿沟，我们设计了CFBlock(ConvFormer Block)，它可以仅使用卷积操作捕获长距离上下文。此外，我们提出了SIAM(语义信息对齐模块)，以更有效地对齐特征。
在Cityscapes、ADE20K和COCO-Stuff-10K上的大量实验结果表明，所提的SCTNet在实时语义分割方面优于现有的最新方法. SCTNet为提高实时语义切分的速度和性能提供了一个新的视角

本文方案

降低计算成本的同时，获得丰富的语义信息，我们将现在流行的两个分支架构拆解为：

一个CNN分支进行推断；
一个Transformer分支用于训练阶段语义对齐。

Backbone 为了提高推理速度，SCTNet采用了典型的分层CNN骨干。SCTNet的Stem模块由两个3×3卷积构成；前两个阶段是由堆叠的残积模块组成的；后两个阶段则是由所提CFBlock构成。CFBlock采用了几个精心设计的卷积操作来执行类似于Transformer块的远程上下文捕获功能。

Decoder Head 解码头由DAPPM与分割头构成，为进一步丰富上下文信息，作者在Stage4后面添加了DAPPM。然后，作者将S2和S4输出进行拼接并送入分割头。

Training Phase 众所周知，Transformer在捕获全局语义上下文方面表现出色。另一方面，CNN已被证明比变换器更适合于对分层局部信息进行建模。受Transformer和CNN优点的启发，我们探索配备一个具有这两种优点的实时分割网络。我们提出了一个单分支CNN，它学习将其特征与强大的Transformer的特征对齐。这种特征对齐使单分支CNN能够提取丰富的全局上下文和详细的空间信息。具体而言，SCTNet采用了一个仅作用在训练阶段的Transformer作为语义分支来提取强大的全局语义上下文，语义信息对齐模块监督卷积分支以对齐来自Transformer的高质量全局上下文。

Inference Phase 为了避免两个分支的巨大计算成本，在推理阶段只部署了CNN分支。利用transformer对齐的语义信息，单分支CNN可以生成准确的分割结果，而无需额外的语义或昂贵的密集融合。更具体地说，输入图像被送入到单分支层次卷积主干中，解码器头拾取主干中的特征并进行简单的拼接进行像素分类.

本文实验

上图与表为Cityscapes语义分割上不同方案的性能对比，从中可以看到：

所提SCTNet以大幅优势优于其他实时分割方案，取得了最佳的速度-精度均衡；
所提SCTNet-B-Seg100去的了80.5%mIoU且速度达62.8FPS，达成实时分割新SOTA；
所提SCTNet-B-Seg75取得了79.8%mIoU，比RTFormer-B与DDRnet-23精度更高，同时速度快两倍；
在所有输入分辨率下，所提SCTNet-B均比其他方案指标更优；此外，SCTNet-S同样取得了比STDC2、RTFormer-S、SeaFormer-B、TopFormer-B更优的性能均衡。

上表为ADE20K与COCO-Stuff-10K两个数据集上不同分割方案的性能对比，很明显：所提SCTNet同样取得了更优的速度-精度均衡。

#STKET

作者提出了一种基于时空知识嵌入的 Transformer（STKET）将先验时空知识纳入多头交叉注意机制中，从而学习更多有代表性的视觉关系表示。基于时空知识的视频场景图生成

视频场景图生成（VidSGG）旨在识别视觉场景中的对象并推断它们之间的视觉关系。该任务不仅需要全面了解分散在整个场景中的每个对象，还需要深入研究它们在时序上的运动和交互。为此，我们进行了相关的探索，并发现：每对物体组合及其它们之间的关系在每个图像内具有空间共现相关性，并且在不同图像之间具有时间一致性/转换相关性。基于这些先验知识，我们提出了一种基于时空知识嵌入的 Transformer（STKET）将先验时空知识纳入多头交叉注意机制中，从而学习更多有代表性的视觉关系表示。具体来说，我们首先以统计方式学习空间共现和时间转换相关性。然后，我们设计了时空知识嵌入层对视觉表示与知识之间的交互进行充分探索，分别生成空间和时间知识嵌入的视觉关系表示。最后，我们聚合这些特征，以预测最终的语义标签及其视觉关系。大量实验表明，我们所提出的框架大幅优于当前竞争算法。目前，该论文已经被人工智能顶级期刊 IEEE T-IP接收。

论文链接：https://arxiv.org/abs/2309.13237

代码链接：https://github.com/HCPLab-SYSU/STKET

1. 概述

随着场景理解领域的快速发展，许多研究者们开始尝试利用各种框架解决场景图生成（Scene Graph Generation, SGG）任务，并已取得了不俗的进展。但是，这些方法往往只考虑单张图像的情况，忽略了时序中存在着的大量的上下文信息，导致现有大部分场景图生成算法在无法准确地识别所给定的视频中包含的动态视觉关系。因此，许多研究者致力于开发视频场景图生成（Video Scene Graph Generation, VidSGG）算法来解决这个问题。

目前的工作主要关注从空间和时间角度聚合对象级视觉信息，以学习对应的视觉关系表示。然而，由于各类物体与交互动作的视觉外表方差大以及视频收集所导致的视觉关系显著的长尾分布，单纯的仅用视觉信息容易导致模型预测错误的视觉关系。

针对上述问题，我们做了以下两方面的工作：

首先，我们提出挖掘训练样本中包含的先验时空知识用以促进视频场景图生成领域。其中，先验时空知识包括：1）空间共现相关性：某些对象类别之间的关系倾向于特定的交互。2）时间一致性/转换相关性：给定对的关系在连续视频剪辑中往往是一致的，或者很有可能转换到另一个特定关系。
其次，我们提出了一种新颖的基于时空知识嵌入的 Transformer (Spatial-Temporal Knowledge-Embedded Transformer, STKET) 框架。该框架将先验时空知识纳入多头交叉注意机制中，从而学习更多有代表性的视觉关系表示。根据在测试基准上得到的比较结果，我们发现我们所提出的 STKET 框架优于以前的最先进方法。

图 1. 由于视觉外表多变和视觉关系的长尾分布，导致视频场景图生成充满挑战

2. 基于时空知识嵌入的 Transformer

2.1. 时空知识表示

在推断视觉关系时，人类不仅利用视觉线索，还利用积累的先验知识 [1, 2]。受此启发，我们提出直接从训练集中提取先验时空知识，以促进视频场景图生成任务。其中，空间共现相关性具体表现为当给定物体组合后其视觉关系分布将高度倾斜（例如，“人”与“杯子”之间的视觉关系的分布明显不同于“狗”与“玩具”之间的分布）和时间转移相关性具体表现为当给定前一时刻的视觉关系后各个视觉关系的转换概率将大幅变化（例如，当已知前一时刻的视觉关系为“吃”时，下一时刻视觉关系转移为“书写”的概率大幅下降）。如图 2 所示，我们可以直观地感受到给定物体组合或之前的视觉关系后，预测空间可以被大幅的缩减。

图 2. 视觉关系的空间共现概率 [3] 与时间转移概率

图 3. 学习空间 (a) 和时间 (b) 知识表示的过程

2.2. 知识嵌入注意力层

空间知识通常包含有关实体之间的位置、距离和关系的信息。另一方面，时间知识涉及动作之间的顺序、持续时间和间隔。鉴于它们独特的属性，单独处理它们可以允许专门的建模更准确地捕获固有模式。因此，我们设计了时空知识嵌入层，彻底探索视觉表示与时空知识之间的相互作用。

图 4. 空间 (左侧) 和时间 (右侧) 知识嵌入层

2.3. 时空聚合模块

如前所述，空间知识嵌入层探索每个图像内的空间共现相关性，时间知识嵌入层探索不同图像之间的时间转移相关性，以此充分探索了视觉表示和时空知识之间的相互作用。尽管如此，这两层忽略了长时序的上下文信息，而这对于识别大部分动态变化的视觉关系具有帮助。为此，我们进一步设计了时空聚合（STA）模块来聚合每个对象对的这些表示，以预测最终的语义标签及其关系。它将不同帧中相同主客体对的空间和时间嵌入关系表示作为输入。具体来说，我们将同一对象对的这些表示连接起来以生成上下文表示。然后，为了在不同帧中找到相同的主客体对，我们采用预测的对象标签和 IoU（即并集交集）来匹配帧中检测到的相同主客体对。最后，考虑到帧中的关系在不同批次中有不同的表示，我们选择滑动窗口中最早出现的表示。

3. 实验结果

为了全面评估所提出的框架的性能，我们除了对比现有的视频场景图生成方法（STTran, TPI, APT）外，我们也选取了先进的图像场景图生成方法（KERN, VCTREE, ReIDN, GPS-Net）进行比较。其中，为确保对比的公平，图像场景图生成方法通过对每一帧图像进行识别，从而达到对所给定视频生成对应场景图的目标。

#SeTformer

这里提出了SeTformer，一种新的transformer，其中DPSA完全被Self-optimal Transport (SeT)取代，以实现更好的性能和计算效率。在小型和基准尺寸模型下，SeTformer在ImageNet-1K上实现了令人印象深刻的84.7%和86.2%的top-1准确率。

论文链接：https://arxiv.org/pdf/2401.03540.pdf

Transformer（变压器）最初是用于自然语言处理（NLP）的技术，在视觉领域得到了显著的流行，这要归功于Vision Transformer（ViT）的开创性工作它的优势已经在各种视觉任务中得到了证明，包括图像分类、目标检测、分割等。对于捕获长距离依赖关系，点积自注意力（DPSA）与softmax归一化在transformer中起着至关重要的作用。然而，该模型的计算导致了二次时间和内存复杂度，使得训练长序列模型变得困难。

简介

本文提出了SeTformer，一种新的transformer，其中DPSA完全被Self-optimal Transport (SeT)取代，以实现更好的性能和计算效率。SeT基于两个基本的softmax属性：保持非负的注意力矩阵和使用非线性的重新加权机制来强调输入序列中重要的标记。通过引入一个用于最优传输的核成本函数，SeTformer有效地满足了这些属性。特别是，在小型和基准尺寸模型下，SeTformer在ImageNet-1K上实现了令人印象深刻的84.7%和86.2%的top-1准确率。在目标检测方面， SeTformer-base相比FocalNet同类产品超出2.2 mAP，使用的参数和浮点运算数分别减少了38%和29%。在语义分割方面，我们的基准模型相比NAT超出了3.5 mIoU，并且参数减少了33%。SeTformer在GLUE基准测试中也取得了最先进的语言建模结果。这些发现凸显了SeTformer在视觉和语言任务中的适用性。

方法与模型

我们的目标是开发一种强大而高效的自注意力模型，尤其注重简单性。我们不添加任何复杂模块，如卷积、平移窗口或注意力偏置，以提高视觉任务的性能。事实上，我们采用了不同的策略。SeT利用了softmax的重要性质，包括非负性和重新加权机制，同时在设计中也注重了效率。使用具有正定（PD）核的RKHS避免了聚合负相关信息。SeT通过OT引入了非线性的重新加权方案。这涉及在RKHS中计算输入和参考集之间的对齐得分。这个过程引入了对齐得分的非线性，给元素分配权重以突出它们的重要性。这有助于模型捕捉复杂关系并强调局部相关性。

SeTformer 架构首先是一个下采样的卷积层，然后是包含多个 SeT 块的四个序列阶段。连续的阶段通过降采样层相连，降低空间尺寸同时加倍深度。在右边，我们展示了我们的注意力计算：将 x 和 y 元素映射到RKHS，然后通过 x 和 y 之间的 OT 计算聚合 x 特征，如果它们与相应的参考对齐良好。

我们使用Swin作为我们的基线模型，用我们的SeT模块替换其自注意力。我们的模型由四个阶段组成，每个阶段具有不同的空间分辨率，结果是输入图像的1/4大小。输入使用两层3×3卷积和2×2步幅进行嵌入。在每个阶段之后，除了最后一层外，都有一个通过3×3卷积和2×2步幅进行下采样的模块。这与Swin不同，Swin使用的是非重叠的2×2卷积。

1 Representing local image neighborhoods in an RKHS

为了保持线性计算，我们将输入特征向量嵌入到一个RKHS中，其中点评估采用线性函数的形式。核方法使我们能够通过一个正定核函数K，将数据从其原始空间X映射到一个高维希尔伯特空间（特征空间）F中。对于函数u：X → F（特征映射），正定核函数表示为K(x, x′) = 〈u(x), u(x′)〉F。鉴于u(x)可以是无穷维的，核技术允许从Rk中导出一个有限维度的表示v(x)，其中内积〈v(xi), v(x′j)〉表示K(x, x′)。正如所示，如果K是正定的，对于任意的x和x′，我们有K(x, x′) ≥ 0，这与softmax算子的非负性质一致。

2 Optimal transport (OT)

我们模型中的一个基本作用是通过学习它们之间的映射将相关令牌进行聚合。我们的加权聚合依赖于被视为不同测度或加权点云的元素x和x′之间的输运计划。OT在对齐问题中得到了广泛应用，并且具有捕捉数据几何形状的出色能力。在本文中，我们专注于Kantorovich形式的OT ´ ，其中使用熵正则化来平滑输运计划

3 Self-optimal Transport (SeT)

对于一个输入特征向量x和一个位于X中的参考ym，我们进行以下步骤：(i)将特征向量x和y表示为RKHSF中的元素， (ii)使用OT将x的元素与y对齐，(iii)对x的元素进行加权聚合，得到一个对齐矩阵A。我们使用参考y来实现高效的元素聚合。参考集合中的每个元素都作为一个”对齐单元”，输入特征通过加权求和在这些单元中进行聚合。这些权重指示了输入和参考之间的对应关系，通过OT计算得出。假设我们有一个输入特征向量x = {x1, . . . , xn}，其中x属于X ∈ Rd，是从输入图像中随机提取的。在Nystrom¨ 近似方法的背景下，y的样本是通过对训练集X中的特征向量进行K-means聚类来获得的质心，从而我们得到y = {y1, . . . , ym}，其中m ≤ n。使用参考集合有助于优化计算过程，并使模型能够有效地处理更长的输入序列。设k是一个正定的核函数，如定义在RKHS上的高斯核函数，以及映射u : Rd → F。我们创建一个大小为n × m的矩阵k，用于存储比较k(xi, yj )的结果。接下来，我们根据公式(2)计算x和y之间的传输计划，得到大小为n × m的矩阵T(x, y)。传输计划找到将输入特征与参考元素对齐的最佳方法，同时最小化对齐成本。

4 Projecting onto a linear subspace

当处理有限维度的u(x)时， Ay(x)可以直接计算，而不会引起重大的计算开销。对于无限维或高维的u(x)，Nystrom¨ 算法提供了一种有效的近似方法来嵌入væRd → Rk。Nystrom¨ 算法通过对列和行进行采样，并将输入从特征空间F投影到线性子空间F1上来近似计算传输计划，从而得到嵌入〈v(xi), v(x′j)〉F1。子空间F1由k个中心u(z1), . . . , u(zk)张成。显式公式v(xi) =k(z, z)−1/2k(z, xi)表示将z = z1, . . . , zk作为中心来进行新的嵌入。这种高效的方法只需要执行K-means聚类并计算逆平方根矩阵。

5 Linear positional encoding

为了将位置信息融入我们的模型中，我们采用了的方法，在输入集和参考集之间的相似性上应用了指数惩罚，基于它们的位置距离。这涉及到对T(v(x), y)与一个距离矩阵M进行乘法运算，其中Mij = e(− 1τ2)(α−β)2，其中α = i/n，β = j/m，τ表示平滑参数。我们考虑了内容和位置信息的相似性权重与其他位置编码方法相比取得了优秀的性能。

实验与结果

我们在图像和语言领域进行了实验，包括ImageNet、COCO和ADE20K，以及GLUE，以展示我们的模型的影响。我们对超参数进行了微调，例如参考数量（m），OT中的熵正则化ϵ，以及位置嵌入中的τ。我们观察到ϵ和τ在任务之间表现稳定，但对于值m的选

224x224分辨率ImageNet-1K的分类准确率

SeTformer模型以较小的模型大小、Flops和吞吐量稳定优于ConvNeXt。我们的Mini模型的准确率超过Swin-T模型0.4%，参数量减少40%（28M → 16M），Flops减少37%。我们的Tiny模型（83.9%）在性能上超过CSWin 1.2%，并具有类似的模型大小，速度提升12%（从701/s到785/s）。与FocalNet-T模型相比，它在性能上表现更优，提高了1.6%。使用更大的模型，我们在较少的参数和较低的计算成本下实现了最先进的性能。例如，SeTformer-B模型在超过24%和36%的Flops和参数减少的情况下，将NAT-B模型（84.4%）的准确率提高了1.8%。我们还注意到，吞吐量是在V100 GPU上测量的。

COCO数据集上Mask R-CNN目标检测结果

SeTformer在卷积神经网络（如ResNet）和Transformer骨干网络（如CSWin、 NAT、 MViTv2）方面表现优异。例如， SeTformer-T的APb为49.3，APm为44.0，相较于NAT-T增加了1.6和1.4个百分点，同时计算量更小，模型尺寸更小。在扩展规模方面，SeTformer-B的APb为51.9，相比于CSWin-B的50.8，增加了1.1个百分点，同时参数减少28%，计算量减少33%。

ADE20K数据集上的语义分割结果

语义分割任务上我们的模型优于现有最先进的方法；例如，相比于CSWin的对应模型，SeTformer-T和SeTformer-S的mIoU（SS/MS）分别提高了+1.3 / +0.7和+0.7 / +0.4，同时具有更轻、更低复杂度的优势。

#Find+Replace Transforme

论文新提出了一种名为“Find+Replace Transformer”的多 Transformer 架构，并证明了通过集成多个Transformer，能够解决单一 Transformer 无法胜任的任务。

ICLR 匿名研究：单一 Transformer 不具备图灵完备性，但多 Transformer 可以。

Transformer 自 2017 年出世以来就在 AI 领域高举高打，ChatGPT 引发全球大型语言模型热潮后更是在 NLP 领域被赋予了神话般的地位。

但近日，一篇正在审核中的 ICLR 2023 投稿论文（如下）经研究后提出一个观点：单一 Transformer 并不具备图灵完备性，其计算能力存在理论上的局限性，在圈内引起关注。

由于该论文正在审核中，作者信息没有被公开。

论文链接：https://openreview.net/pdf?id=MGWsPGogLH

与此同时，该论文新提出了一种名为“Find+Replace Transformer”的多 Transformer 架构，并证明了通过集成多个Transformer，能够解决单一 Transformer 无法胜任的任务。

这项研究直接对标并超越了当前最先进的GPT-4模型，在一系列极具挑战性的基准测试中展现了显著的优势和潜力。

1 被神化的 Transformer 局限在哪里？

图灵完备性是评判一个计算系统强大与否的关键指标。如果一个系统被确认为图灵完备，则理论上只要赋予其充足的运行时间和内存资源，即可以执行任何可计算的算法。

在实际应用中，尽管 Transformer 模型在诸多自然语言处理任务上表现卓越，但其能力受到设计上的固有限制，例如固定的上下文窗口长度和有限的词汇表大小。这意味着 Transformer 模型并不具备解决所有类型计算问题的能力，特别是那些需要无限存储空间或无限制迭代过程的问题。

在论文中，研究团队特别指出，基础的语言模型工作原理在于根据前 k 个词语的概率来预测下一个词语。在 NLP 领域，通常会构建一些专门针对固定长度输入输出序列设计的模型集合或框架，并将这类模型归入 MF_SMF 类别。

Transformer 作为 MF_SMF 这一框架下的具体实例，其图灵完备性的缺失得到了该研究团队的理论论证。他们基于以下逻辑：

首先，回顾计算理论的基础：图灵停机问题是不可判定的，意味着不存在一个通用的方法来判断任意给定程序何时终止运行，就如同无法找到一把万能钥匙预测每一场棋局结束时间一样。这一原理同样适用于评估模型是否会在执行过程中陷入无尽循环而无法自拔。

研究者进而分析了 MF_S（这里假设 MF_S 代表 MF_SMF 中的子集）集合中的模型：

假设可以构建一个算法H，它可以准确判断MF_S中任意模型m是否终止。
假设MF_S集合中存在一个模型m’，它足够强大以至于能够模拟任何图灵机的计算过程，包括那些永远不会停止的图灵机。
根据算法H的假设能力，如果MF_S集合中的模型m’能够模拟那些不会停止的图灵机，那么算法H应该能够预测m’在模拟这些图灵机时是否会停止。
然而，根据图灵的停机问题不可判定定理，我们知道实际上不可能存在这样一个算法H，因为它会与图灵的定理相矛盾。
因此，MF_S集合中不可能存在能够模拟所有图灵机行为的模型m’，也就是说，MF_S中没有任何模型是图灵完备的。

Transformer便属于 MF_SMF，所以 Transformer 不具备图灵完备性。

研究人员指出，Transformer在处理自然语言任务，尤其是在机器翻译方面，有明显的优势。这类模型能够通过递归的方式输入序列并生成更新后的序列，从而逐个预测下一个符号。

但是，尽管Transformer模型能够基于之前的字符序列连续生成新的字符序列，每次接收一段输入字符后产出相应的输出字符，并利用新产生的字符序列进行迭代计算，它还是受到了上下文长度k和词汇表大小v的限制。这意味着它能够处理的不同字符组合的数量不会超过v^k种。

例如，当 Transformer 遇到重复输入时，由于它的无状态特性（这有利于并行训练多个序列），模型必须保证对同一输入产生一致的输出结果。这可能导致在某些情况下，模型陷入无限循环的模式，即只能生成有限数量的、最多为v^k种不同的输出序列，或者在自我复制的过程中无法停止。

与Transformer相比，图灵在1936年提出的图灵机概念具有无限的计算潜力，不受这些结构性的限制，能够模拟任何可计算的过程，确保不会陷入类似的有限循环困境。

2 如何超越 GPT-4？

实验结果显示，单个 Transformer 架构并不具备图灵完备性，而多 Transformer 则有能力实现图灵完备（如论文中所提出的 Find+Replace Transformer）、并执行如 GPT-4 等最先进的 Transformer 模型所无法解决的问题。

论文中创新性地将 Find Transformer 与 Replace Transformer 相结合，构建了Find+Replace Transformer体系结构——这是一个能在任意长度序列上运行的多Transformer系统，在论文中被形象地比喻为“磁带”（Tape）。

该系统由 Find Transformer、Replace Transformer 以及 Map 三部分组成，其中 Map 是一个从 Replace Transformer 到 Find Transformer 所涉及的有序集合的函数映射关系。

具体运作时，Find Transformer 会在输入序列中定位并标识出需要由 Replace Transformer 处理的部分内容。这两个组件各自具有固定的上下文长度 k，并依次对“磁带”上的每个长度为k的子序列进行分析，Find Transformer 会选择那些在最终层产生最高激活值的特定子序列。

随后，Replace Transformer 会接收 Find Transformer 标识出的子序列作为输入，并基于此生成一个新的长度为k的输出序列，这个过程利用了 Map 关联的 f∈Map(r) 规则，确保了两个 Transformer 之间的协同工作及信息传递。

那这个 Find+Replace Transformer 的多 Transformer 系统是如何可以实现图灵完备的呢？

简单来说，Find+Replace Transformer 是一个学习简化的机器。在编程语言的基石 λ 演算中，有三条被称为“归约”（Reduction）的规则:

Alpha Reduction：这是一个绑定变量的重命名。它被用来避免命名冲突。例如，在λ 演算的项 λx.x，我们可以化简成 λy.y，且不改变其意思。
Beta Reduction：这是将函数应用于其参数的过程。例如，在λ项(λx.x)y（表示将函数λx.x作用于参数y），我们可以化简成 y。
Eta Reduction：这是对函数和参数的简化。如果你有一个函数比如λx.(fx)，而x不出现在f中，那么这个就可以化简为f。

Find+Replace Transformer 的多Transformer 系统之所以能够实现图灵完备性，关键在于其架构设计和训练方式允许模型通过一系列组合操作模拟类似于 λ 演算中的归约规则。尽管单个 Transformer 受限于上下文长度、词汇表大小等因素，但通过构建一个多 Transformer 协作的框架，并结合特定的学习机制，这些简单且局部的“查找与替换”操作得以在更复杂的计算任务中累积并形成强大的综合效应。

具体来说，在Find+Replace Transformer中，多个 Transformer 可能被专门设计来分别或协同地处理不同类型的简化（归约）任务，例如模拟 Alpha Reduction 进行变量重命名、模拟 Beta Reduction 执行函数应用以及模拟 Eta Reduction进行函数简化等。每个 Transformer 可能专注于理解和学习如何执行这类简单的转换操作，并将结果传递给下一个Transformer，从而逐步构建起复杂问题的解决方案。

虽然单个 Transformer 不具备图灵完备性，但当它们以特定的方式组织起来并协同工作时，可以模拟通用图灵机的逻辑行为，进而实现对任意可计算问题的解决能力。这样的体系结构让Find+Replace Transformer在处理大规模、多层次的复杂问题时展现出超越传统单一Transformer的性能表现，实现了更高阶的计算能力。

2023年当OpenAI 发布GPT-4时，微软研究院的研究人员发表了一篇题为“Sparks of Artificial General Intelligence(Bubeck et al., 2023)”的论文，阐述了早期AGI所面临的局限性。

研究者们以汉诺塔问题为例进行了说明。汉诺塔是一个经典的递归问题，要求玩家将按照大小顺序堆叠的圆盘从一根柱子移动到另一根柱子上，期间只能移动一个圆盘且任何时候大盘不能位于小盘之上，借助第三根柱子作为中转。

GPT-4无法解决这个复杂的推理问题，从而突显了当前Transformer在推理过程中缺乏规划能力。

研究者对比了几种模型在解决完整汉诺塔问题上的表现。随着问题规模增大，其难度呈指数级上升：规模为n的问题其解决方案需要2^n - 1步操作。Find+Replace Transformer在此任务上表现出色，甚至能生成比GPT-4至少长18倍的正确解决方案。

除了在汉诺塔这个GPT-4都难以解决的问题上表现优越之外，在其他AI任务，如创作满足特定条件的诗歌等，Find+Replace Transformer都能超越GPT-4，这反映了其在泛化能力上的优势。

3 结语

Find+Replace Transformer模型通过创新性地结合多个Transformer单元，并模拟λ演算中的归约规则，在处理如汉诺塔问题等复杂组合任务时展现出了超越传统单个Transformer的优越性能。

这一研究成果揭示了多Transformer系统在实现图灵完备性方面的潜力，也证明了在面对特定计算难题时，提高模型的逻辑推理和抽象表达能力的重要性。

而纵观整个人工智能技术的发展，从深度学习兴起到大模型浪潮来袭，每一次技术迭代，人们都对于新技术报以极大的热情与崇拜。

然而，无论是深度学习还是Transformer架构，亦或是如今新出现Find+Replace Transformer架构，所带给我们的启示是，在研究和应用深度学习技术时，都需要避免过分神化任何技术，应该理性地看待每一项技术，关注其优势和局限，并结合实际问题来选择和调整合适的技术。只有这样，才能不断地在通往人工通用智能（AGI）的道路上迈进。

#Soft MoE

本文提出了一种可微的稀疏混合专家 Transformer 模型 (fully-differentiable sparse Transformer) Soft MoE 来解决端到端训练困难的问题，同时也能够保持 MoE 方法的优势，即以较低的推理成本更大的模型容量。

Soft MoE 提出了一种新的可微稀疏混合专家模型，稀疏混合专家 (Sparse Mixture of Experts, MoE) 是一种在保证模型训练和推理的成本不显著增加的情况下，大幅度提升模型容量的方法。

MoE 方法已经有很长的一段历史了，是一种扩大模型容量的经典高效的做法，但是它的缺点是：

训练不稳定
Token Dropping 的问题
较难扩展 Expert 的数量
低效率的微调

造成以上问题的一个原因是 MoE 的端到端训练困难，因此，本文提出了一种可微的稀疏混合专家 Transformer 模型 (fully-differentiable sparse Transformer) Soft MoE 来解决端到端训练困难的问题，同时也能够保持 MoE 方法的优势，即以较低的推理成本更大的模型容量。Soft MoE 的特点是给每个专家输入不同 token 的权重混合。

视觉实验结果证明，Soft MoE 大大优于标准 ViT 和流行的 MoE 方法，比如 128 个 Expert，16 个 MoE 层的 Soft MoE-Huge/14 模型参数比 ViT-Huge/14 多 40 倍，但推理时间成本仅增长 2%，同时性能要好得多。

1 Soft MoE：一种完全可微的稀疏 Transformer

论文名称： From Sparse to Soft Mixtures of Experts

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

1 Soft MoE 论文解读：

1.1 背景：把离散优化问题变为可微的优化问题

稀疏混合专家 (Sparse Mixture of Experts, MoE) 是一种在保证模型训练和推理的成本不显著增加的情况下，大幅度提升模型容量的方法。在视觉，语言和多模态任务中都取得了成功，代表像视觉的 V-MoE[1]，文本的 Switch Transformer[2]和多模态的 LIMoE[3]。

如下图1左所示，稀疏 MoE Transformer 的核心是一个离散优化问题，即：模型需要决定每个输入 token 应该输入哪些 Expert 里面，这些 Expert 一般是 MLP 模块。输入 token 和 Expert 之间的匹配 (token-to-expert match) 是 MoE 中要考虑的很重要的问题之一，之前也有各种各样的方法尝试解决此问题，比如基于线性规划的[4]，比如基于 RL 算法的[5]，比如基于固定规则的[6]，比如基于最优传输理论的[7]，和基于贪婪匹配的[8]。总之，解决好稀疏 MoE 的这个离散优化问题的确是件不容易的事情。稀疏 MoE 的缺点有：

训练不稳定
Token Dropping 的问题
较难扩展 Expert 的数量
低效率的微调

图1：Sparse MoE 和 Soft MoE 的区别：左：Sparse MoE，给每个 Expert 分配一定的输入 token。右：Soft MoE，给每个 Expert 分配的是所有输入 token 的加权平均值

如下图1右所示，Soft MoE 把稀疏 MoE Transformer 的这个离散优化问题变成了可微的优化问题。Soft MoE 觉得不必一定要 "hard" 地找到输入 token 和 Expert 之间的一一匹配，而是可以 "Soft" 地混合输入 token 并且分给每一个 Expert。Soft MoE 给每个 Expert 分配的不是某几个输入 token，而是所有输入 token 的加权平均值 (权重取决于 token 和 Expert)，然后由这个对应的 Expert 去处理这个加权平均值。

1.2 变为可微的优化问题之后，解决了之前稀疏 MoE 的什么问题？

问题1： 精心设计的 Expert-to-token 的路由机制通常并不比随机固定路由好。

Soft MoE 可以避免这个问题，因为每个路由的参数都是基于每个输入 token 直接更新的。

问题2： 训练不稳定 (LIMoE[3]这个工作观察到在训练期间，可能有大部分 token 改变路由，给训练带来一定挑战) 导致很多稀疏 MoE 方法的 Expert 都不可以设置得很多。

Soft MoE 可以避免这个问题，扩展到数千个 Expert。

1.3 Soft MoE 算法描述

参数配置：

整个过程如下图2所示。

图2：Soft MoE 算法流程图

遵循稀疏 MoE 的常用设计思想，作者用 Soft MoE 块替换了 Transformer 的一部分 MoE 块。slot 的总数是 Soft MoE 的关键超参数，因为时间复杂度取决于 slot 的数量，而不是 Expert 的数量。比如，可以设置等于输入序列长度的 slot 数以匹配等效密集 Transformer 的 FLOP。

Soft MoE 的 JAX 代码：

def soft_moe_layer(X, Phi, experts):# Compute the dispatch and combine weights.logits = jnp.einsum('md,dnp->mnp', X, Phi)D = jax.nn.softmax(logits, axis=(0,))C = jax.nn.softmax(logits, axis=(1, 2))# The input slots are a weighted average of all the input tokens,# given by the dispatch weights.Xs = jnp.einsum('md,mnp->npd', X, D)# Apply the corresponding expert function to each input slot.Ys = jnp.stack([f_i(Xs[i, :, :]) for i, f_i in enumerate(experts)],axis=0)# The output tokens are a weighted average of all the output slots,# given by the combine weights.Y = jnp.einsum('npd,mnp->md', Ys, C)return Y

全部代码：

https://github.com/google-research/vmoegithub.com/google-research/vmoe

1.4 Soft MoE 的一些关键性质

1) 完全可微：

Sparse MoE 算法的通病是 token 和 Expert 之间存在的分配问题，有时精心设计的 Expert-to-token 的路由机制通常并不比随机固定路由好。输入 token 和 Expert 之间的匹配 (token-to-expert match) 是 MoE 中要考虑的很重要的问题之一，之前也有各种各样的方法尝试解决此问题，比如基于线性规划的[4]，比如基于 RL 算法的[5]，比如基于固定规则的[6]，比如基于最优传输理论的[7]，和基于贪婪匹配的[8][9]。所有这些方法本质上都是离散，不可微的。

Soft MoE 可以避免这个问题，因为每个路由的参数都是基于每个输入 token 直接更新的。

2) 可以避免掉 Token Dropping 和 Expert Unbalance 的问题

MoE 算法里面每个 Expert 都会处理一些 token，很自然地就会带来 Token Dropping (有的 token 不会分配给任何一个 Expert) 和 Expert Unbalance (一些 Expert 会比另一些 Expert 分配到更多 token) 的问题。

Soft MoE 可以避免这个问题，因为每个 slot 的输入都是所有 token 的加权平均值。

3) 运算速度快

Soft MoE 的主要优点是完全避免了之前算法中的 token 排序或 top-k 操作，因为这些操作的速度慢，而且不太适合硬件加速器。因此，Soft MoE 明显快于大多数 Sparse MoE 算法。

4) Soft MoE 算法是密集的 MoE 算法还是稀疏的 MoE 算法？

要回答这个问题我们需要首先搞明白为什么 Sparse MoE 算法是稀疏的。Sparse MoE 是稀疏的这件事的根本原因是每个 Expert 的输入特征仅仅是一部分的 token，而 Soft MoE 的输入是所有输入 token 的加权平均值，因此不能算作是稀疏的。

Soft MoE 也不能算作是 Dense MoE 算法，因为每个 Expert 仅仅会处理输入 token 的子集。

5) Soft MoE 算法需要归一化

Transformers 中，MoE 层通常用于替换每个编码器块中的 FFN 层，因此如果去遵循大部分 Transformer 架构的 Pre-Normalization 方法，就需要使用归一化，这里 Soft MoE 针对的操作是：

l2_normalize(X, axis=1)

scale * l2_normalize(Phi, axis=0)

其中，scale 是可学习的参数，l2_normalize 的定义是：

def l2_normalize(x, axis, eps=1e-6):norm = jnp.sqrt(jnp.square(x).sum(axis=axis, keepdims=True))return x * jnp.reciprocal(norm + eps)

6) Soft MoE 算法和注意力机制 (Multi-Head Self-Attention) 的区别和联系？

1.5 Soft MoE 算法的局限性

自回归解码 (Auto-regressive decoding)：

因为 Soft MoE 算法要在运行过程中合并所有的输入 token，因此很难实现自回归。因为自回归必须在训练期间保留过去的 token 和未来 token 之间的因果关系 (Causality)。

Self-Attention 解决这个问题的手段是依赖于注意力的掩码 (Mask) 机制。如果想在 Soft MoE 中实现这一点就需要特别小心 token 之间的依赖和相关关系。总之研究 Soft MoE 算法的自回归解码是个很有价值的方向。

内存消耗

Soft MoE 倾向于利用大量 Expert，而其成本和 Dense Backbone 类似，使得模型的内存需求可能变大。

1.6 图像分类实验结果

训练数据集

预训练数据集： JFT-4B：一个私有数据集，其最新版本包含超过 4B 张图像，涵盖超过 29k 个类。预训练的过程中评价指标是 JFT-4B 上的上游验证精度 Precision-at-1 和 ImageNet 10-shot 精度 (冻结模型权重，并用一个新的权重来计算的，该数据集仅在包含来自 ImageNet-1K 的每个类包含 10 张图像的数据集上进行训练)。

微调数据集： ImageNet-1K 训练集。

验证集： ImageNet-1K 验证集。

模型尺寸：

ViT-S/8, ViT-S/16, ViT-S/32, ViT-B/16, ViT-B/32, ViT-L/16, ViT-L/32, ViT-H/14。

方法：

Token Choice, Expert Choice 和本文的 Soft MoE。

训练策略：

300k steps, Batch Size 4096

Pareto Model 实验结果：

如下图3所示是四种方法 Soft MoE, Experts Choice, Tokens Choice, Dense 在预训练过程中的 JFT-4B Precision-at-1 的结果和 ImageNet 10-shot 的精度的训练成本/性能帕累托边界。Soft MoE 算法在这两个指标上都优于之前的方法。

图3：四种方法在预训练过程中的 JFT-4B Precision-at-1 的结果和 ImageNet 10-shot 的精度的训练成本/性能帕累托边界

更长的训练结果：

本文还测试在更长的训练 step 下模型的性能如何，把从 Small 到 Huge 的模型训练了 4K steps，用 128 个 Expert 的 Soft MoE 替换 ViT S/16、B/16、L/16 和 H/14 中的最后一半 Block 中的 FFN，每个 Expert 使用一个 slot。

由于模型并行性所需的额外数据传输，Large Soft MoE 模型产生的 wall-clock time overhead 很小。所有变体都训练了 4M 步，除了 H/14，出于成本原因训练了 2M 步，实验结果如下图4和5所示。

如下图4所示是 Soft MoE 和 ViT 的 JFT-4B 精度、ImageNet 10-shot 精度和 ImageNet 微调精度与 ExaFLOPS 的训练成本。

图4：不同模型更长的训练 step 下的 JFT-4B 精度

如下图5所示是所有结果。对于给定的计算预算，Soft MoE 模型大大优于 Vision Transformer 模型。比如 Soft MoE-S/16 在 JFT-4B 和 ImageNet 10-shot 上的表现优于 ViT-B/16，它还提高了完整 ImageNet 数据的微调分数，即使它的训练 (和推理) 成本要小得多。同样，Soft MoE-B/16 在上游任务 JFT-4B 和 ImageNet 10 shot 的表现优于 ViT-L/16，微调后仅落后 0.5，同时速度快 3 倍，所需的 FLOP 减少了近 4 倍。最后，Soft MoE-L/16 模型优于 Dense H/14 模型，同时在推理速度又快 3 倍左右。

图5：不同模型更长的训练 step 下的实验结果

根据前面的实验结果，较小的 Soft MoE 的性能可以匹配较大的视觉 Transformer，作者因此继续训练小模型 Backbone，希望以非常低的训练成本获得更高质量的模型。

作者观察到对于 Soft MoE 方法而言，较长的 cooldown (学习率线性减小到零的时期) 可以很好地适用于 Soft MoE，因此将 cooldown 从 50k steps 增加到 500k steps。

实验结果如下图6和7所示。Soft MoE-B/16 训练了 1k TPUv3 Days，优于在相似时间预算上训练的 ViT-H/14，而 Soft MoE-B 模型的 FLOPs 要低 10 倍，wall-clock time 低 5.7 倍。即使将 ViT-H/14 的训练代价加倍，Soft MoE-B 模型的性能也可以与之相匹配。Soft MoE-L/16 模型的在推断上比 ViT H/14 快近 2 倍的同时性能大大优于所有模型。

图6：不同训练代价和尺寸的 Soft MoE 模型和 ViT 的 JFT-4B Precision-at-1 性能和 ImageNet 10-shot 性能

图7：Soft MoE 模型和 ViT 的实验结果

视觉-文本对比学习实验结果

作者还验证了 Soft MoE 得到的模型在其他任务的性能。具体而言作者探索了一种流行的范式，即图像语言对比学习，这里遵循的是 LiT[10] 方法，其中图像塔在图像分类任务上进行了预训练，然后在在图像-文本对数据集上训练文本编码器时冻结。

视觉编码器作者重用了在 JFT 上训练的模型，对比学习在 WebLI 上训练，这是一个专有数据集，由 10B 图像和从互联网上抓取的 ALT 文本组成。图像编码器被冻结，而文本编码器从头开始训练。实验结果如下图8所示，Soft MoE -L/16 在 Imagenet 和 Cifar-100 零样本上的性能分别比 ViT-L/16 高出 1% 和 2% 以上。

图8：对比学习实验结果

#Transformers18~ Diffusion

还是Transformers,来自 UC 伯克利的 William Peebles 以及纽约大学的谢赛宁撰文揭秘扩散模型中架构选择的意义，并为未来的生成模型研究提供经验基线。

近几年，在 Transformer 的推动下，机器学习正在经历复兴。过去五年中，用于自然语言处理、计算机视觉以及其他领域的神经架构在很大程度上已被 transformer 所占据。

不过还有许多图像级生成模型仍然不受这一趋势的影响，例如过去一年扩散模型在图像生成方面取得了惊人的成果，几乎所有这些模型都使用卷积 U-Net 作为主干。这有点令人惊讶！在过去的几年中，深度学习的大事件一直是跨领域的 Transformer 的主导地位。U-Net 或卷积是否有什么特别之处使它们在扩散模型中表现得如此出色？

将 U-Net 主干网络首次引入扩散模型的研究可追溯到 Ho 等人，这种设计模式继承了自回归生成模型 PixelCNN++，只是稍微进行了一些改动。而 PixelCNN++ 由卷积层组成，其包含许多的 ResNet 块。其与标准的 U-Net 相比，PixelCNN++ 附加的空间自注意力块成为 transformer 中的基本组件。不同于其他人的研究，Dhariwal 和 Nichol 等人消除了 U-Net 的几种架构选择，例如使用自适应归一化层为卷积层注入条件信息和通道计数。

本文中来自 UC 伯克利的 William Peebles 以及纽约大学的谢赛宁撰文《 Scalable Diffusion Models with Transformers 》，目标是揭开扩散模型中架构选择的意义，并为未来的生成模型研究提供经验基线。该研究表明，U-Net 归纳偏置对扩散模型的性能不是至关重要的，并且可以很容易地用标准设计（如 transformer）取代。

这一发现表明，扩散模型可以从架构统一趋势中受益，例如，扩散模型可以继承其他领域的最佳实践和训练方法，保留这些模型的可扩展性、鲁棒性和效率等有利特性。标准化架构也将为跨领域研究开辟新的可能性。

论文地址：https://arxiv.org/pdf/2212.09748.pdf
项目地址：https://github.com/facebookresearch/DiT
论文主页：https://www.wpeebles.com/DiT

该研究专注于一类新的基于 Transformer 的扩散模型：Diffusion Transformers（简称 DiTs）。DiTs 遵循 Vision Transformers (ViTs) 的最佳实践，有一些小但重要的调整。DiT 已被证明比传统的卷积网络（例如 ResNet ）具有更有效地扩展性。

具体而言，本文研究了 Transformer 在网络复杂度与样本质量方面的扩展行为。研究表明，通过在潜在扩散模型 (LDM) 框架下构建 DiT 设计空间并对其进行基准测试，其中扩散模型在 VAE 的潜在空间内进行训练，可以成功地用 transformer 替换 U-Net 主干。本文进一步表明 DiT 是扩散模型的可扩展架构：网络复杂性（由 Gflops 测量）与样本质量（由 FID 测量）之间存在很强的相关性。通过简单地扩展 DiT 并训练具有高容量主干（118.6 Gflops）的 LDM，可以在类条件 256 × 256 ImageNet 生成基准上实现 2.27 FID 的最新结果。

Diffusion Transformers

DiTs 是一种用于扩散模型的新架构，目标是尽可能忠实于标准 transformer 架构，以保留其可扩展性。DiT 保留了 ViT 的许多最佳实践，图 3 显示了完整 DiT 体系架构。

DiT 的输入为空间表示 z（对于 256 × 256 × 3 图像，z 的形状为 32 × 32 × 4）。DiT 的第一层是 patchify，该层通过将每个 patch 线性嵌入到输入中，以此将空间输入转换为一个 T token 序列。patchify 之后，本文将标准的基于 ViT 频率的位置嵌入应用于所有输入 token。

patchify 创建的 token T 的数量由 patch 大小超参数 p 决定。如图 4 所示，将 p 减半将使 T 翻四倍，因此至少能使 transformer Gflops 翻四倍。本文将 p = 2,4,8 添加到 DiT 设计空间。

DiT 块设计：在 patchify 之后，输入 token 由一系列 transformer 块处理。除了噪声图像输入之外，扩散模型有时还会处理额外的条件信息，例如噪声时间步长 t、类标签 c、自然语言等。本文探索了四种以不同方式处理条件输入的 transformer 块变体。这些设计对标准 ViT 块设计进行了微小但重要的修改。所有模块的设计如图 3 所示。

本文尝试了四种因模型深度和宽度而异的配置：DiT-S、DiT-B、DiT-L 和 DiT-XL。这些模型配置范围从 33M 到 675M 参数，Gflops 从 0.4 到 119 。

实验

研究者训练了四个最高 Gflop 的 DiT-XL/2 模型，每个模型使用不同的 block 设计 ——in-context（119.4Gflops）、cross-attention（137.6Gflops）、adaptive layer norm（adaLN，118.6Gflops）或 adaLN-zero（118.6Gflops）。然后在训练过程中测量 FID，图 5 为结果。

扩展模型大小和 patch 大小。图 2（左）给出了每个模型的 Gflops 和它们在 400K 训练迭代时的 FID 概况。可以发现，增加模型大小和减少 patch 大小会对扩散模型产生相当大的改进。

图 6（顶部）展示了 FID 是如何随着模型大小的增加和 patch 大小保持不变而变化的。在四种设置中，通过使 Transformer 更深、更宽，训练的所有阶段都获得了 FID 的明显提升。同样，图 6（底部）展示了 patch 大小减少和模型大小保持不变时的 FID。研究者再次观察到，在整个训练过程中，通过简单地扩大 DiT 处理的 token 数量，并保持参数的大致固定，FID 会得到相当大的改善。

图 8 中展示了 FID-50K 在 400K 训练步数下与模型 Gflops 的对比：

SOTA 扩散模型 256×256 ImageNet。在对扩展分析之后，研究者继续训练最高 Gflop 模型 DiT-XL/2，步数为 7M。图 1 展示了该模型的样本，并与类别条件生成 SOTA 模型进行比较，表 2 中展示了结果。

当使用无分类器指导时，DiT-XL/2 优于之前所有的扩散模型，将之前由 LDM 实现的 3.60 的最佳 FID-50K 降至 2.27。如图 2（右）所示，相对于 LDM-4（103.6 Gflops）这样的潜在空间 U-Net 模型来说，DiT-XL/2（118.6 Gflops）计算效率高得多，也比 ADM（1120 Gflops）或 ADM-U（742 Gflops）这样的像素空间 U-Net 模型效率高很多。

表 3 展示了与 SOTA 方法的比较。XL/2 在这一分辨率下再次胜过之前的所有扩散模型，将 ADM 之前取得的 3.85 的最佳 FID 提高到 3.04。

#Diffusion Transformers (DiTs)

本文探索了一类新的基于 Transformer 的扩散模型 Diffusion Transformers (DiTs)。本文训练 latent diffusion models 时，使用 Transformer 架构替换常用的 UNet 架构，且 Transformer 作用于 latent patches 上。

作者探索了 DiT 的缩放性，发现具有较高 GFLOPs 的 DiT 模型，通过增加 Transformer 宽度或者深度或者输入 token 数量，始终有更好的 FID 值。最大的 DiT-XL/2 模型在 ImageNet 512×512 和 256×256 的测试中优于所有先前的扩散模型，实现了 2.27 的 FID 值。

做了什么工作

探索了一类新的基于 Transformer 的 Diffusion Model，称为 Diffusion Transformers (DiTs)。
研究了 DiT 对于模型复杂度 (GFLOPs) 和样本质量 (FID) 的缩放性。
证明了通过使用 Latent Diffusion Models (LDMs)[1]框架，Diffusion Model 中的 U-Net 架构可以被 Transformer 替换。

1 DiT：Transformer 构建扩散模型

论文名称：Scalable Diffusion Models with Transformers (ICCV 2023, Oral)

论文地址：https//arxiv.org/pdf/2212.09748.pdf

论文主页：https//www.wpeebles.com/DiT.html

1 DiT 论文解读：

1.1 把 Transformer 引入 Diffusion Models

机器学习正经历着 Transformer 架构带来的复兴：NLP，CV 等许多领域正在被 Transformer 模型覆盖。尽管 Transformer 在 Autoregressive Model 中得到广泛应用[2][3][4][5]，但是这种架构在生成式模型中较少采用。比如，作为图像领域生成模型的经典方法，Diffusion Models[6][7]却一直使用基于卷积的 U-Net 架构作为骨干网络。

Diffusion Models 的开创性工作 DDPM [8]首次引入了基于 U-Net 骨干网络的扩散模型。U-Net 继承自 PixelCNN++[9][10]，变化很少。与标准 U-Net[11]相比，额外的空间 Self-Attention 块 (Transformer 中必不可少的组件) 以较低分辨率穿插。[12]这个工作探索了 U-Net 的几种架构选择，例如自适应归一化层 (Adaptive Normalization Layer[13]为卷积层注入条件信息和通道计数。然而，DDPM 里面 U-Net 的高级设计在很大程度上都保持不变。

本文的目的是探索 Diffusion Models 架构选择的重要性，并为未来生成式模型的研究提供基线。本文的结论表明 U-Net 架构设计对 Diffusion Models 的性能并不重要，并且它们可以很容易地替换为 Transformers。

本文证明了 Diffusion Models 也可以受益于 Transformer 架构，受益于其训练方案，受益于其可扩展性，受益于其鲁棒性和效率等等。标准化架构还将为跨域研究开辟了新的可能性。

1.2 Diffusion Models 简介

DDPM

高斯扩散模型假设有一个前向的加噪过程 (Forward Noising Process)，在这个过程中逐渐将噪声应用于真实数据：

这个优化的目标函数比较复杂，最后通过 variational lower bound 方法得到的结论是优化下式 (此处详细推导可以参考开创性工作 DDPM[8])：

1.3 DiT 架构介绍

1.3.1 Patchify 过程

图1：图片的 Patchify 操作。当 Patch 的大小 p 越小时，token 的数量 T 越大

1.3.2 DiT Block 设计

在 Patchify 之后，输入的 tokens 开始进入一系列 Transformer Block 中。除了噪声图像输入之外，Diffusion Model 有时会处理额外的条件信息，比如噪声时间步长 ttt , 类标签 ccc , 自然语言。

作者探索了4种不同类型的 Transformer Block，以不同的方式处理条件输入。这些设计都对标准 ViT Block 进行了微小的修改，所有 Block 的设计如下图2所示。

图2：Diffusion Transformer (DiT) 架构

In-Context Conditioning

作者将以上几种方法 In-Context Conditioning，Cross-Attention Block，Adaptive Layer Norm (adaLN) Block，adaLN-Zero Block 的做法列入 DiT 的设计空间中。

1.3.3 模型尺寸

图3：DiT 模型的详细配置

作者将以上几种配置列入了 DiT 的设计空间中。

1.3.4 Transformer Decoder

在最后一个 DiT Block 之后，需要将 image tokens 的序列解码为输出噪声以及对角的协方差矩阵的预测结果。

最终，完整 DiT 的设计空间是 Patch Size、DiT Block 的架构和模型大小。

1.4 DiT 训练策略

1.4.1 训练配方

作者在 ImageNet 数据集上训练了 class-conditional latent DiT 模型，标准的实验设置。

数据增强技术只使用 horizontal flips。

作者发现 learning rate warmup 和 regularization，对训练 DiT 模型而言不是必须的。

作者使用了 exponential moving average (EMA)，参数为 0.9999 。

训练超参数基本都来自 ADM，不调学习率, decay/warm-up schedules, Adam 参数以及 weight decay.

1.4.2 扩散模型配置

作者保留了 ADM 中使用的超参数。

1.5.1 DiT 架构设计

作者首先探索的是不同 Conditioning 策略的对比。对于一个 DiT-XL/2 模型，其计算复杂度分别是：in-context (119.4 Gflops), cross-attention (137.6 Gflops), adaptive layer norm (adaLN, 118.6 Gflops), adaLN-zero (118.6 Gflops)。实验结果如下图4所示。

adaLN-Zero 的 Block 架构设计取得了最低的 FID 结果，同时在计算量上也是最高效的。在 400K 训练迭代中，adaLN-Zero Block 架构得到的 FID 几乎是 In-Context 的一半，表明 Condition 策略会严重影响模型的质量。

初始化同样也重要：adaLN-Zero Block 架构在初始化时相当于恒等映射，其性能也大大优于 adaLN Block 架构。

因此，在后续实验中，DiT 将一直使用 adaLN-Zero Block 架构。

图4：不同 Conditioning 策略对比

1.5.2 缩放模型尺寸和 Patch Size

作者训练了12个 DiT 模型 (尺寸为 S, B, L, XL，Patch Size 为 8,4,2)。下图是不同 DiT 模型的尺寸和 FID-50K 性能。如下图5所示是不同大小 DiT 模型的 GFLOPs 以及在 400K 训练迭代中的 FID 值。可以发现，在增加模型大小或者减小 Batch Size 时可以显著改善 DiT 的性能。

图5：不同尺寸 DiT 模型的 GFLOPs 以及它们在 400K 训练迭代中的 FID

下图6上方是 Patch Size 不变，增加模型规模时 FID 的变化。当模型变深变宽时，FID 会下降。

下方是模型规模不变，减小 Patch Size 时 FID 的变化。当 Patch Size 下降时，FID 出现显著改善。

图5：缩放 DiT 模型可以改善训练各个阶段的 FID

1.5.3 GFLOPs 对性能很重要

上图5的结果表明，参数量并不能唯一确定 DiT 模型的质量。当 Patch Size 减小时，参数量仅仅是略有下降，只有 GFLOPs 明显增加。这些结果都表明了缩放模型的 GFLOPs 才是性能提升的关键。为了印证这一点，作者在下图6中绘制了不同 GFLOPs 模型在 400K 训练步骤时候的 FID-50K 结果。这些结果表明，当不同 DiT 模型的总 GFLOPs 相似时，它们的 FID 值也相似，比如 DiT-S/2 和 DiT-B/4。

作者还发现 DiT 模型的 GFLOPs 和 FID-50K 之间存在很强的负相关关系。

图6：GFLOPs 与 FID 密切相关

1.5.4 大模型更加计算高效

图7：大模型更加计算高效

1.5.5 缩放结果可视化

图8：缩放对于视觉质量的影响

1.6 DiT 实验结果

作者将 DiT 与最先进的生成模型进行了比较，结果如图9所示。DiT-XL/2 优于所有先前的扩散模型，将 LDM 实现的先前最佳 FID-50K 降低到 2.27。图5右侧显示 DiT-XL/2 (118.6 GFLOPs) 相对于 LDM-4 (103.6 GFLOPs) 等Latent Space U-Net 模型的计算效率很高，并且比 Pixel Space U-Net 模型更高效，例如 ADM (1120 GFLOPs) 或 ADM-U (742 GFLOPs)。

图9：ImageNet 256×256 图像生成结果

作者在 ImageNet 上训练了一个新的 DiT-XL/2，这次分辨率是 512×512，3M training iterations，超参数与 256×256 模型相同。这个模型 latent 的维度是 64×64×4，然后 Patch Size 为2，这样 Transformer 模型需要处理的 token 的数量就是 1024。如下图10所示是比较结果。DiT-XL/2 在此分辨率下再次优于所有先前的扩散模型，将 ADM 实现的先前最佳 FID 提高了 3.85 到 3.04。即使 token 的数量增加了，DiT-XL/2 的计算效率依然很高，比如 ADM 使用 1983 GFLOPs，ADM-U 使用 2813 GFLOPs，DiT-XL/2 仅仅使用 524.6 GFLOPs。

图10：ImageNet 512×512 图像生成结果

缩放模型大小还是采样次数？

Diffusion Model 的一个独特之处是它们可以通过在生成图像时增加采样步骤的数量来在训练期间使用额外的计算。也就是扩散模型的计算量既可以来自模型本身的缩放，也可以来自采样次数的增加。因此，作者在这里研究了通过使用更多的采样计算，较小的 DiT 模型是否可以胜过更大的模型。

作者计算了所有的 12 个 DiT 模型在 400K training iteration 时候的 FID 值，每张图分别使用 [16, 32, 64, 128, 256, 1000] sampling steps。

实验结果如下图11所示，考虑使用 1000 个采样步骤的 DiT-L/2 和使用 128 步的 DiT-XL/2。在这种情况下：

DiT-L/2 使用 80.7 TFLOPs 对每张图像进行采样。
DiT-XL/2 使用 15.2 TFLOPs 对每张图像进行采样。

但尽管如此，DiT-XL/2 具有更好的 FID-10K 结果。说明增加采样的计算量也无法弥补模型本身计算量的缺失。

图11：增加采样的计算量也无法弥补模型本身计算量的缺失

#eventful-transformers

如何降低视觉Transformer计算成本？时间冗余方法让人大吃一惊

在为语言领域带来变革之后，Transformer 正在进军视觉领域，但其也有着高计算成本的问题。近日，威斯康星大学麦迪逊分校一个研究团队提出了 Eventful Transformer，可通过在视觉 Transformer 中利用时间冗余来节省成本。

Transformer 一开始是为自然语言处理任务设计的，但现在却已经被广泛用于视觉任务。视觉 Transformer 在一系列视觉识别任务上实现了出色的准确度，并在图像分类、视频分类和目标检测等任务上取得了当前最优的表现。

视觉 Transformer 的一大缺点是计算成本高。典型的卷积网络（CNN）处理每张图像需要数十 GFlops，而视觉 Transformer 所需的往往会多上一个数量级，达到每张图像数百 GFlops。在处理视频时，由于数据量巨大，这个问题更为严重。高昂的计算成本让视觉 Transformer 难以被部署到资源有限或有严格延迟需求的设备上，这就限制了这项技术的应用场景，否则我们已经有一些激动人心的应用了。

在近期一篇论文中，威斯康星大学麦迪逊分校的三位研究者 Matthew Dutson、Yin Li 和 Mohit Gupta 首先提出可以在后续输入之间使用时间冗余来降低视觉 Transformer 在视频应用中的成本。他们也发布了模型代码，其中包含用于构建 Eventful Transformer 的 PyTorch 模块。

论文地址：https://arxiv.org/pdf/2308.13494.pdf
项目地址：http://wisionlab.com/project/eventful-transformers

时间冗余：首先假设有一个视觉 Transformer，其可以逐帧或逐视频片段地处理视频序列。这个 Transformer 可能是简单的逐帧处理的模型（如目标检测器）或是某个时空模型的中间步骤（如 ViViT 的分解式模型的第一步）。不同于一个输入就是一个完整序列的语言处理 Transformer，在这里，研究者的做法是随时间为 Transformer 提供多个不同的输入（帧或视频片段）。

自然视频包含显著的时间冗余，即后续帧之间的差异很小。尽管如此，包括 Transformer 在内的深度网络通常都会「从头开始」计算每一帧。该方法会丢弃之前推理获得的潜在相关信息，浪费极大。故而这三位研究者设想：是否可以复用之前计算步骤的中间计算结果来提升处理冗余序列的效率？

自适应推理：对于视觉 Transformer 以及一般意义上的深度网络而言，推理成本通常由架构决定。然而在现实应用中，可用的资源可能会随时间而变化，比如可能因为存在相竞争的进程或电源发生变化。如此一来，可能就存在运行时修改模型计算成本的需求。在这项新成果中，研究者设定的一大主要设计目标便是适应性 —— 其方法可实现对计算成本的实时控制。下图 1（底部）给出了在视频处理过程中修改计算预算的示例。

Eventful Transformer：本文提出了 Eventful Transformer，这类 Transformer 能利用输入之间的时间冗余来实现高效且自适应的推理。Eventful 这一术语的灵感来自事件相机（event camera），这种传感器能在场景变化时离散地记录影像。Eventful Transformer 会跟踪随时间发生的 token 层面的变化情况，并在每个时间步骤有选择性地更新 token 表征和自注意力映射图。Eventful Transformer 的模块中包含一种门控模块，用于控制运行时间被更新 token 的数量。

该方法可用于现成的模型（通常无需再训练）并且兼容许多视频处理任务。研究者也进行了实验论证，结果表明 Eventful Transformer 可用于现有的当前最佳模型，在极大降低它们的计算成本的同时还能维持其原有的准确度。

Eventful Transformer

这项研究的目标加速用于视频识别的视觉 Transformer。在这个场景中，视觉 Transformer 需要反复处理视频帧或视频片段，具体的任务包括视频目标检测和视频动作识别等。这里提出的关键思想是利用时间冗余，即复用之前时间步骤的计算结果。下面将详细描述如何通过修改 Transformer 模块来使其具备感知时间冗余的能力。

token 门控：检测冗余

这一小节将介绍研究者提出的两种新模块：token 门和 token 缓冲器。这些模块让模型可以识别和更新自上次更新后有明显变化的 token。

门模块：该门会从输入 token N 中选择一部分 M 发送给下游层执行计算。其记忆中维护着一个参照 token 集，记为 u。这种参照向量包含每个 token 在其最近一次更新时的值。在每个时间步骤，比较各个 token 与其对应的参照值，其中与参照值相差较大的 token 获得更新。

现在将该门的当前输入记为 c。在每个时间步骤，按照以下流程更新门的状态并决定其输出（见下图 2）：

构建可感知冗余的 Transformer

为了利用上述时间冗余，研究者提出了一种对 Transformer 模块的修改方案。下图 4 展示了 Eventful Transformer 模块的设计。该方法可以加速针对各个 token 的运算（如 MLP）以及查询 - 键值和注意力 - 值乘法。

在针对各个 token 的运算 Transformer 模块中，很多运算都是针对各个 token 的，也就是说它们不涉及到 token 之间的信息交换，其中包括 MLP 和 MSA 中的线性变换。为了节省计算成本，研究者表示可以跳过未被门选取的 token 的面向 token 的运算。由于 token 之间的独立性，这不会改变对所选 token 的运算结果。参见图 3。

具体来说，针对各个 token 的运算（包括 W_qkv 变换、W_p 变换和 MLP）的连续序列，研究者使用了一对门 - 缓冲器。注意，他们还在 skip 连接之前添加了缓冲器以确保两个加法操作数的 token 是正确对齐的。

针对各个 token 的运算的成本正比于 token 的数量。门可将这个数量从 N 降至 M，也就将下游的针对各个 token 的运算的计算成本降低了 N/M 倍。

查询 - 键值的积：现在来看看查询 - 键值积 B = q k^T。

下图 5 展示了稀疏地更新查询 - 键值积 B 中一部分元素的方法。

这些更新的总体成本为 2NMD，相较而言，从头开始计算 B 的成本为 N^2D。注意，新方法的成本正比于 M，即门选取的 token 的数量。当 M < N/2 时（此时更新的 token 不到总量一半），可节省计算量。

注意力 - 值的积：研究者为此提出了一种基于增量 ∆ 的更新策略。

下图 6 展示了新提出的高效计算三个增量项的方法。

同样当 M < N/2 时，可节省计算量。

token 选取策略

Eventful Transformer 的一大重要设计是其 token 选取策略。给定一个门误差张量 e，这样一个策略的目标是生成一个掩码 m，其中指示了应当被更新的 token。具体的策略包括：

Top-r 策略：该策略选取 r 个误差 e 有最大范数的 token（这里使用的是 L2 范数）。

阈值策略：该策略选取误差 e 的范数超过一个阈值 h 的所有 token。

其它策略：更复杂精细的 token 选取策略可实现更好的准确度 - 成本权衡，比如可以使用一个轻量级策略网络来学习一个策略。但是，训练策略的决策机制的难度可能很大，因为二元掩码 m 一般是不可微分的。另一个思路是使用重要度分数作为选取的参考信息。但这些想法都还有待未来研究。

实验

研究者用实验评估了新提出的方法，具体使用的任务是视频目标检测和视频动作识别。

下图 7 展示了视频目标检测的实验结果。其中正轴是计算节省率，负轴是新方法的 mAP50 分数的相对减少量。可以看到，新方法用少量的准确度牺牲换来了显著的计算量节省。

下图 8 给出了在视频目标检测任务上的方法比较和消融实验结果。

下图 9 给出了视频动作识别的实验结果。

下表 2 给出了在一台 CPU（Xeon Silver 4214, 2.2 GHz）和一台 GPU（NVIDIA RTX3090）上运行时间（毫秒）结果，可以看到时间冗余在 GPU 上带来的速度提升可达 1.74 倍，在 CPU 上带来的提升可达 2.47 倍。

#Llama~transformers搭建

本例从零开始基于transformers库逐模块搭建和解读Llama模型源码(中文可以翻译成羊驼)。

并且训练它来实现一个有趣的实例：两数之和。

输入输出类似如下：

输入："12345+54321="

输出："66666"

我们把这个任务当做一个文本生成任务来进行。输入是一个序列的上半部分，输出其下半部分.

这和文本生成的输入输出结构是类似的，所以可以用Llama来做。

目前大部分开源LLM模型都是基于transformers库来做的，它们的结构大部分都和Llama大同小异。

俗话说，魔鬼隐藏在细节中，深入理解Llama模型的的源码细节，将会帮助你打通和开源LLM模型相关的基础原理(如旋转位置编码以及长度外推)，并让你熟悉各种参数的配置和使用(如past_key_value，attention_mask的使用等等)。

一，准备数据

import math  
from typing import List, Optional, Tuple, Union  import torch  
import torch.nn.functional as F  
import torch.utils.checkpoint  
from torch import nn  
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss  from transformers.activations import ACT2FN  
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast  
from transformers.modeling_utils import PreTrainedModel  
from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings  from transformers.models.llama.configuration_llama  import LlamaConfig  
from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING,LLAMA_START_DOCSTRING  logger = logging.get_logger('llama')  config = LlamaConfig(  vocab_size=len(vocab),  hidden_size=512,  intermediate_size=2752,  num_hidden_layers=8,  num_attention_heads=16,  hidden_act='silu',  max_position_embeddings=128,  initializer_range=0.02,  rms_norm_eps=1e-06,  use_cache=True,  pad_token_id=0,  bos_token_id=1,  eos_token_id=2,  tie_word_embeddings=False  
)

<BOS>3914835626735057733+318829464988=3914835945564522721<EOS>

# 定义数据集  
class TwoSumDataset(torch.utils.data.Dataset):  def __init__(self,size = 100000, min_length=10,max_length=20):  super(Dataset, self).__init__()  self.size = size  self.min_length=min_length  self.max_length=max_length  def __len__(self):  return self.size  def __getitem__(self, i):  x,y = self.get(i)  # 编码成token  context_ids = [vocab[i] for i in x]  target_ids = [vocab[i] for i in y]  input_ids = context_ids + target_ids  #-100标志位后面会在计算loss时会被忽略不贡献损失，我们集中优化target部分生成的loss  labels = [-100]*len(context_ids)+ target_ids  masks = [0 if t==vocab['<PAD>'] else 1 for t in input_ids]  example = {'input_ids':input_ids,  'labels':labels,'attention_mask':masks}  return example  def get(self,i):  return get_data(self.min_length,self.max_length)  def show_example(self,example):  input_ids,labels = example['input_ids'],example['labels']  x = ''.join([vocab_r[a] for a,b in zip(input_ids,labels) if b==-100])  y = ''.join([vocab_r[a] for a,b in zip(input_ids,labels) if b!=-100])  print(x+y)  ds_train = TwoSumDataset(size = 100000,min_length=10,max_length=20)  
ds_val = TwoSumDataset(size = 10000,min_length=10,max_length=20)  
example = ds_train[0]  
ds_train.show_example(example)

<BOS>12878683929048906366+11274414130675477=12889958343179581843<EOS>

def data_collator(examples: list):  len_ids = [len(example["input_ids"]) for example in examples]  longest = max(len_ids) #之后按照batch中最长的input_ids进行padding  input_ids = []  labels_list = []  masks_list = []  for length, example in sorted(zip(len_ids, examples), key=lambda x: -x[0]):  ids = example["input_ids"]  labs = example["labels"]  masks = example['attention_mask']  ids = [vocab['<PAD>']] * (longest - length)+ids   labs = [-100] * (longest - length)+labs  masks = [0]*(longest - length)+masks  input_ids.append(torch.LongTensor(ids))  labels_list.append(torch.LongTensor(labs))  masks_list.append(torch.LongTensor(masks))  input_ids = torch.stack(input_ids)  labels = torch.stack(labels_list)  attention_mask = torch.stack(masks_list)  return {  "input_ids": input_ids,  "labels": labels,  "attention_mask":attention_mask  }  # 数据加载器  
dl_train = DataLoader(dataset=ds_train,  batch_size=200,  drop_last=True,  shuffle=True,  collate_fn = data_collator          )  dl_val = DataLoader(dataset=ds_val,  batch_size=200,  drop_last=True,  shuffle=False,  collate_fn = data_collator    )  for batch in dl_train:  break

batch   {'input_ids': tensor([[ 1, 11,  6,  ...,  7, 11,  2],         [ 0,  1,  6,  ...,  5,  4,  2],         [ 0,  1,  7,  ...,  8,  8,  2],         ...,         [ 0,  0,  0,  ..., 10, 11,  2],         [ 0,  0,  0,  ..., 12,  3,  2],         [ 0,  0,  0,  ..., 11, 12,  2]]), 'labels': tensor([[-100, -100, -100,  ...,    7,   11,    2],         [-100, -100, -100,  ...,    5,    4,    2],         [-100, -100, -100,  ...,    8,    8,    2],         ...,         [-100, -100, -100,  ...,   10,   11,    2],         [-100, -100, -100,  ...,   12,    3,    2],         [-100, -100, -100,  ...,   11,   12,    2]]), 'attention_mask': tensor([[1, 1, 1,  ..., 1, 1, 1],         [0, 1, 1,  ..., 1, 1, 1],         [0, 1, 1,  ..., 1, 1, 1],         ...,         [0, 0, 0,  ..., 1, 1, 1],         [0, 0, 0,  ..., 1, 1, 1],         [0, 0, 0,  ..., 1, 1, 1]])}

二，定义模型

下面，我们会像搭积木建城堡那样从低往高地构建LLaMA模型。先构建4个基础组件：旋转位置编码，多头注意力、前馈网络、层归一化。类似用最基础的积木块搭建了墙壁，房顶，房门，窗户这样的模块。然后用这4个基础组件构建中间成品: 解码层。类似用基础组件构建了房间。接着用多个中间成品解码层的堆叠组装成了LlamaModel完整模型，相当于通过构建多个房间建成了城堡的主体结构。最后我们在LlamaModel基础上设计了两种不同的输出head，一种是语言模型Head，得到了LlamaForCausalLM，可用于文本生成。另外一种是分类head，得到了LlamaForSequenceClassification，可用于文本分类。相当于我们在城堡主体结构完成的基础上设计了两种不同的装修风格，一种是加装了一些游乐设施以便用于商业活动，另一种则是加装了一些武器以便用于军事活动。

1, 旋转位置编码: RoPE (使用旋转矩阵实现的绝对位置编码，可以起到相对位置编码的效果)

2, 多头注意力: LlamaAttention (用于融合不同token之间的信息)

3, 前馈网络: LlamaMLP (用于逐位置将多头注意力融合后的信息进行高维映射变换)

4, 层归一化: LlamaRMSNorm (用于稳定输入，相当于保持每个词向量的方向不变，但对模长标准化。)

5, Llama解码层: LlamaDecoderLayer (同时具备信息融合，信息转换功能的基本结构单元)

6, Llama解码器: LlamaModel (多个解码层的堆叠)7，Llama语言模型: LlamaForCausalLM (解码器加上语言模型head，可用于文本生成)8，Llama分类模型: LlamaForSequenceClassification (解码器加上分类head，可用于文本分类)

import math  
from typing import List, Optional, Tuple, Union  import torch  
import torch.nn.functional as F  
import torch.utils.checkpoint  
from torch import nn  
from torch.nn import BCEWithLogitsLoss, CrossEntropyLoss, MSELoss  from transformers.activations import ACT2FN  
from transformers.modeling_outputs import BaseModelOutputWithPast, CausalLMOutputWithPast, SequenceClassifierOutputWithPast  
from transformers.modeling_utils import PreTrainedModel  
from transformers.utils import add_start_docstrings, add_start_docstrings_to_model_forward, logging, replace_return_docstrings  from transformers.models.llama.configuration_llama  import LlamaConfig  
from transformers.models.llama.modeling_llama import LLAMA_INPUTS_DOCSTRING,LLAMA_START_DOCSTRING  logger = logging.get_logger('llama')  config = LlamaConfig(  vocab_size=len(vocab),  hidden_size=512,  intermediate_size=2752,  num_hidden_layers=8,  num_attention_heads=16,  hidden_act='silu',  max_position_embeddings=128,  initializer_range=0.02,  rms_norm_eps=1e-06,  use_cache=True,  pad_token_id=0,  bos_token_id=1,  eos_token_id=2,  tie_word_embeddings=False  
)

1，旋转位置编码 RoPE

旋转位置编码即使用旋转矩阵表示位置编码(Rotary Position Encoding),简称RoPE。

关于RoPE的3个核心要点知识如下：

RoPE的设计思想是使用绝对位置编码来达到相对位置编码的效果。

RoPE的实现方式是使用旋转矩阵来表示绝对位置编码。

使用NTK扩展方法可以让RoPE在短文本上训练并在长文本上做预测。

参考文章：

《博采众长的旋转式位置编码》https://kexue.fm/archives/8265

《RoPE是一种进制编码》https://kexue.fm/archives/9675

（1）绝对位置编码和相对位置编码

位置编码一般可以分成绝对位置编码和相对位置编码。

绝对位置编码的优点是计算简单高效，缺点是一般效果不如相对位置编码。

相对位置编码的优点是效果较好，缺点是计算效率不如绝对位置编码。

在相对位置编码中，注意力权重的结果仅仅和参与注意力计算的token向量的相对位置有关，不和绝对位置直接关联。

这符合NLP领域在序列长度方向上具有平移不变性的特点，所以相对位置编码一般效果会优于绝对位置编码。

不过绝对位置编码并非一无是处，绝对位置编码只需要初始化时对序列的每个位置(数量正比于序列长度)赋予位置编码即可，后续无需干预。

而相对位置编码要在计算过程中获取许多个(数量正比于序列长度平方)相对位置。

因此绝对位置编码更加简单高效。

（2）使用旋转矩阵表示位置编码

上述讨论可以看到，绝对位置编码和相对位置编码互有优劣，那么有没有什么办法能够对二者进行取长补短呢？

有的，这个方法就是RoPE，它的设计思想就是使用绝对位置编码来达到相对位置编码的效果。

那么旋转位置编码如何使用绝对位置编码来达到相对位置编码的效果的呢？答案是使用旋转矩阵来表示位置编码。

由于旋转矩阵是稀疏矩阵，直接使用乘法计算会很浪费算力，可以将旋转位置编码过程由矩阵乘法运算简化成两次向量的哈达玛积求和。

（3）旋转位置编码的长度扩展

在LLM的应用中，有一个非常重要的参数，叫做LLM支持的上下文长度(max context length)。

更长的上下文长度允许我们进行更多轮次的对话，允许我们对更长的本文进行总结分析，也允许我们生成更长的文章。

但是在训练LLM的时候，我们的训练语料大部分是不够长的，许多LLM训练时候设计的最大文本长度都是只有2k，也就是最长2048个token。

那么，能否在训练的时候使用较短的文本，而在推理的时候扩展到长文本上呢？

是有可能的，我们可以对RoPE进行长度扩展。

我们介绍3种扩展方案。

第一种是直接外推：直接外推其实就是继续沿用现有的位置编码公式，不做任何修改。

在扩展长度不太长的时候，例如由2k扩展到2.5k时，这种方法可能对性能的影响并不大。

因为旋转位置编码只和相对位置m-n的大小有关，一般具有远程衰减性，即相对距离越大的两个token，其相关性一般越弱。

因此如果我们的模型已经从训练数据那里学习到了token之间的相关性相对于相对距离在0-2k的一个合适的衰减规律的时候，可以设想把这个规律应用到0-2.5k也是没有太大的问题的。

但是如果我们要扩展到更长的长度，例如从2k扩展到32k，这种直接外推的方案通常会严重地影响性能。因为我们学习到的衰减规律有可能在5k的那里就完全衰减截断基本降为0了，这样我们就无法捕捉相对距离长于5k的两个token之间的相互作用，外推就会导致性能下降。

总结一下，直接外推对衰减规律在长距离情况下的使用容易出现问题，导致性能下降。

为了减少长度外推对性能的影响，我们可以让训练好的模型在更长的上下文上做少许步骤的微调。

第二种是线性内插：线性内插需要改变位置编码公式，等效于将位置序号等比例缩小。

线性内插没有改变模型学习到的衰减规律的应用范围，不考虑微调的话，其效果一般好于直接外推方案。

但是，扩展倍数非常大的时候，例如从2k扩展到32k，其性能也会明显的受到影响。

因为在这种情况下，衰减规律在短距离情况下的使用会受到较严重的影响，本来距离为1的两个token，长度扩展后相当于变成了距离为1/16，衰减规律在短距离时可能具有非常大的变化率，因此对相关性的评估可能会极端地偏离合理值。

应用线性内插时，在长文本上做少许步骤的微调也能够明显地改善性能。

第三种是NTK扩展方式：这种方式综合了外推和内插的优点，做长度扩展后即使不微调也能够保持较好的性能。

前面的分析我们知道直接外推对衰减规律在长距离情况下的使用容易出问题，在短距离情况下的使用不受影响。

而线性内插对衰减规律在短距离情况下的使用容易出现问题，在长距离的情况下影响较小。

我们能否将它们综合起来，在短距离情况下具有外推特性(与扩展前基本一致)，在长距离情况下具有内插特性(缩放到扩展前的范围)，从而使得长距离情况下和短距离情况下衰减规律的使用都不太受到影响呢。

NTK扩展方式的要点是高频外推，低频内插，实现方法是直接对底数base进行缩放，类似进制编码转换。

采用NTK扩展到长文本，即使不做微调，性能会只会略有下降。

下面是RoPE以及三种长度扩展方式的实现。

class LlamaRotaryEmbedding(torch.nn.Module):  def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None):  super().__init__()  self.dim = dim  self.max_position_embeddings = max_position_embeddings  self.base = base  inv_freq = 1.0 / (self.base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))  self.register_buffer("inv_freq", inv_freq, persistent=False) #persistent=False将不会作为state_dict  # Build here to make `torch.jit.trace` work.  self._set_cos_sin_cache(  seq_len=max_position_embeddings, device=self.inv_freq.device, dtype=torch.get_default_dtype()  )  def _set_cos_sin_cache(self, seq_len, device, dtype):  self.max_seq_len_cached = seq_len  t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)  freqs = torch.einsum("i,j->ij", t, self.inv_freq)  # Different from paper, but it uses a different permutation in order to obtain the same calculation  emb = torch.cat((freqs, freqs), dim=-1)  self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)  self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)  def forward(self, x, seq_len=None):  # x: [bs, num_attention_heads, seq_len, head_size]  #超过预设的max_position_embeddings则重新计算更大的Rope缓存，否则直接在缓存上切片  if seq_len > self.max_seq_len_cached:   self._set_cos_sin_cache(seq_len=seq_len, device=x.device, dtype=x.dtype)  return (  self.cos_cached[:, :, :seq_len, ...].to(dtype=x.dtype),  self.sin_cached[:, :, :seq_len, ...].to(dtype=x.dtype),  )  class LlamaLinearScalingRotaryEmbedding(LlamaRotaryEmbedding):  """LlamaRotaryEmbedding extended with linear scaling. Credits to the Reddit user /u/kaiokendev"""  def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):  self.scaling_factor = scaling_factor  super().__init__(dim, max_position_embeddings, base, device)  def _set_cos_sin_cache(self, seq_len, device, dtype):  self.max_seq_len_cached = seq_len  t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)  t = t / self.scaling_factor #线性内插相当于将位置序号等比例缩小  freqs = torch.einsum("i,j->ij", t, self.inv_freq)  # Different from paper, but it uses a different permutation in order to obtain the same calculation  emb = torch.cat((freqs, freqs), dim=-1)  self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)  self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)  class LlamaDynamicNTKScalingRotaryEmbedding(LlamaRotaryEmbedding):  """LlamaRotaryEmbedding extended with Dynamic NTK scaling. Credits to the Reddit users /u/bloc97 and /u/emozilla"""  def __init__(self, dim, max_position_embeddings=2048, base=10000, device=None, scaling_factor=1.0):  self.scaling_factor = scaling_factor  super().__init__(dim, max_position_embeddings, base, device)  def _set_cos_sin_cache(self, seq_len, device, dtype):  self.max_seq_len_cached = seq_len  if seq_len > self.max_position_embeddings:  base = self.base * (  (self.scaling_factor * seq_len / self.max_position_embeddings) - (self.scaling_factor - 1)  ) ** (self.dim / (self.dim - 2))  #NTK扩展方式直接对base进行缩放  inv_freq = 1.0 / (base ** (torch.arange(0, self.dim, 2).float().to(device) / self.dim))  self.register_buffer("inv_freq", inv_freq, persistent=False)  t = torch.arange(self.max_seq_len_cached, device=device, dtype=self.inv_freq.dtype)  freqs = torch.einsum("i,j->ij", t, self.inv_freq)  #此处处理逻辑与原始的ROPE有差异，原始逻辑如下  #emb = torch.cat((freqs, freqs), dim=-1)  #emb[...,0::2]=freqs  #emb[...,1::2]=freqs  # Different from paper, but it uses a different permutation in order to obtain the same calculation  emb = torch.cat((freqs, freqs), dim=-1)  self.register_buffer("cos_cached", emb.cos()[None, None, :, :].to(dtype), persistent=False)  self.register_buffer("sin_cached", emb.sin()[None, None, :, :].to(dtype), persistent=False)  def rotate_half(x):  """Rotates half the hidden dims of the input."""  #此处逻辑与原始的ROPE有所差异，原始逻辑如下  #x1 = x[..., 0::2]   #x2 = x[..., 1::2]  #res = torch.cat((x1, x2), dim=-1)  #res[...,0::2]=-x2  #res[...,1::2]=x1  #return res  x1 = x[..., : x.shape[-1] // 2]   x2 = x[..., x.shape[-1] // 2 :]  return torch.cat((-x2, x1), dim=-1)  def apply_rotary_pos_emb(q, k, cos, sin, position_ids):  # The first two dimensions of cos and sin are always 1, so we can `squeeze` them.  cos = cos.squeeze(1).squeeze(0)  # [seq_len, dim]  sin = sin.squeeze(1).squeeze(0)  # [seq_len, dim]  cos = cos[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]  sin = sin[position_ids].unsqueeze(1)  # [bs, 1, seq_len, dim]  q_embed = (q * cos) + (rotate_half(q) * sin)  k_embed = (k * cos) + (rotate_half(k) * sin)  return q_embed, k_embed  x = torch.randn(1,8,4,2)  
rope = LlamaRotaryEmbedding(dim=8)  
cos,sin = rope.forward(x,seq_len=4)  
print(cos.shape)   
print(cos)  
torch.Size([1, 1, 4, 8])tensor([[[[ 1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,  1.0000,            1.0000],          [ 0.5403,  0.9950,  0.9999,  1.0000,  0.5403,  0.9950,  0.9999,            1.0000],          [-0.4161,  0.9801,  0.9998,  1.0000, -0.4161,  0.9801,  0.9998,            1.0000],          [-0.9900,  0.9553,  0.9996,  1.0000, -0.9900,  0.9553,  0.9996,            1.0000]]]])

2，多头注意力 LlamaAttention

这里的LlamaAttention 基本上和《Attention Is All You Need》论文里的是一致的，主要差异有以下一些。

1，k和v的head数量可以是q的head数量的几分之一，类似分组卷积的思想，可以减少参数规模。

2，rope位置编码是每次做多头注意力时都进行一次，而不是原论文只在输入的时候进行一次。

3，允许传入key和value的states的缓存past_key_value，这在多轮对话中可以减少重复计算，起到加速效果。

4，attention_mask是通过加法形式作用到softmax之前的attention矩阵上的。

def repeat_kv(hidden_states: torch.Tensor, n_rep: int) -> torch.Tensor:  """  This is the equivalent of torch.repeat_interleave(x, dim=1, repeats=n_rep). The hidden states go from (batch,  num_key_value_heads, seqlen, head_dim) to (batch, num_attention_heads, seqlen, head_dim)  """  batch, num_key_value_heads, slen, head_dim = hidden_states.shape  if n_rep == 1:  return hidden_states  hidden_states = hidden_states[:, :, None, :, :].expand(batch, num_key_value_heads, n_rep, slen, head_dim)  return hidden_states.reshape(batch, num_key_value_heads * n_rep, slen, head_dim)  class LlamaAttention(nn.Module):  """Multi-headed attention from 'Attention Is All You Need' paper"""  def __init__(self, config: LlamaConfig):  super().__init__()  self.config = config  self.hidden_size = config.hidden_size  self.num_heads = config.num_attention_heads  self.head_dim = self.hidden_size // self.num_heads  self.num_key_value_heads = config.num_key_value_heads  self.num_key_value_groups = self.num_heads // self.num_key_value_heads  self.max_position_embeddings = config.max_position_embeddings  if (self.head_dim * self.num_heads) != self.hidden_size:  raise ValueError(  f"hidden_size must be divisible by num_heads (got `hidden_size`: {self.hidden_size}"  f" and `num_heads`: {self.num_heads})."  )  self.q_proj = nn.Linear(self.hidden_size, self.num_heads * self.head_dim, bias=False)  self.k_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)  self.v_proj = nn.Linear(self.hidden_size, self.num_key_value_heads * self.head_dim, bias=False)  self.o_proj = nn.Linear(self.num_heads * self.head_dim, self.hidden_size, bias=False)  self._init_rope()  def _init_rope(self):  if self.config.rope_scaling is None:  self.rotary_emb = LlamaRotaryEmbedding(self.head_dim, max_position_embeddings=self.max_position_embeddings)  else:  scaling_type = self.config.rope_scaling["type"]  scaling_factor = self.config.rope_scaling["factor"]  if scaling_type == "linear":  self.rotary_emb = LlamaLinearScalingRotaryEmbedding(  self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor  )  elif scaling_type == "dynamic":  self.rotary_emb = LlamaDynamicNTKScalingRotaryEmbedding(  self.head_dim, max_position_embeddings=self.max_position_embeddings, scaling_factor=scaling_factor  )  else:  raise ValueError(f"Unknown RoPE scaling type {scaling_type}")  def _shape(self, tensor: torch.Tensor, seq_len: int, bsz: int):  return tensor.view(bsz, seq_len, self.num_heads, self.head_dim).transpose(1, 2).contiguous()  def forward(  self,  hidden_states: torch.Tensor,  attention_mask: Optional[torch.Tensor] = None,  position_ids: Optional[torch.LongTensor] = None,  past_key_value: Optional[Tuple[torch.Tensor]] = None,  output_attentions: bool = False,  use_cache: bool = False,  ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:  bsz, q_len, _ = hidden_states.size()  if self.config.pretraining_tp > 1:  key_value_slicing = (self.num_key_value_heads * self.head_dim) // self.config.pretraining_tp  query_slices = self.q_proj.weight.split(  (self.num_heads * self.head_dim) // self.config.pretraining_tp, dim=0  )  key_slices = self.k_proj.weight.split(key_value_slicing, dim=0)  value_slices = self.v_proj.weight.split(key_value_slicing, dim=0)  query_states = [F.linear(hidden_states, query_slices[i]) for i in range(self.config.pretraining_tp)]  query_states = torch.cat(query_states, dim=-1)  key_states = [F.linear(hidden_states, key_slices[i]) for i in range(self.config.pretraining_tp)]  key_states = torch.cat(key_states, dim=-1)  value_states = [F.linear(hidden_states, value_slices[i]) for i in range(self.config.pretraining_tp)]  value_states = torch.cat(value_states, dim=-1)  else:  query_states = self.q_proj(hidden_states)  key_states = self.k_proj(hidden_states)  value_states = self.v_proj(hidden_states)  query_states = query_states.view(bsz, q_len, self.num_heads, self.head_dim).transpose(1, 2)  key_states = key_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)  value_states = value_states.view(bsz, q_len, self.num_key_value_heads, self.head_dim).transpose(1, 2)  kv_seq_len = key_states.shape[-2]  if past_key_value is not None:  kv_seq_len += past_key_value[0].shape[-2]  cos, sin = self.rotary_emb(value_states, seq_len=kv_seq_len)  query_states, key_states = apply_rotary_pos_emb(query_states, key_states, cos, sin, position_ids)  if past_key_value is not None:  # reuse k, v, self_attention  key_states = torch.cat([past_key_value[0], key_states], dim=2)  value_states = torch.cat([past_key_value[1], value_states], dim=2)  past_key_value = (key_states, value_states) if use_cache else None  # repeat k/v heads if n_kv_heads < n_heads  key_states = repeat_kv(key_states, self.num_key_value_groups)  value_states = repeat_kv(value_states, self.num_key_value_groups)  attn_weights = torch.matmul(query_states, key_states.transpose(2, 3)) / math.sqrt(self.head_dim)  if attn_weights.size() != (bsz, self.num_heads, q_len, kv_seq_len):  raise ValueError(  f"Attention weights should be of size {(bsz, self.num_heads, q_len, kv_seq_len)}, but is"  f" {attn_weights.size()}"  )  if attention_mask is not None:  if attention_mask.size() != (bsz, 1, q_len, kv_seq_len):  raise ValueError(  f"Attention mask should be of size {(bsz, 1, q_len, kv_seq_len)}, but is {attention_mask.size()}"  )  attn_weights = attn_weights + attention_mask  # upcast attention to fp32  attn_weights = nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).to(query_states.dtype)  attn_output = torch.matmul(attn_weights, value_states)  if attn_output.size() != (bsz, self.num_heads, q_len, self.head_dim):  raise ValueError(  f"`attn_output` should be of size {(bsz, self.num_heads, q_len, self.head_dim)}, but is"  f" {attn_output.size()}"  )  attn_output = attn_output.transpose(1, 2).contiguous()  attn_output = attn_output.reshape(bsz, q_len, self.hidden_size)  if self.config.pretraining_tp > 1:  attn_output = attn_output.split(self.hidden_size // self.config.pretraining_tp, dim=2)  o_proj_slices = self.o_proj.weight.split(self.hidden_size // self.config.pretraining_tp, dim=1)  attn_output = sum([F.linear(attn_output[i], o_proj_slices[i]) for i in range(self.config.pretraining_tp)])  else:  attn_output = self.o_proj(attn_output)  if not output_attentions:  attn_weights = None  return attn_output, attn_weights, past_key_value

3，前馈网络 LlamaMLP

前馈网络是一个2层的感知机MLP。

先从hidden_size维度up_proj到intermediate_size维度，然后再down_proj还原为hidden_size维度。

这里的主要特色是引入了一个gate_proj配合激活函数来实现一个门控注意力的作用。

class LlamaMLP(nn.Module):  def __init__(self, config):  super().__init__()  self.config = config  self.hidden_size = config.hidden_size  self.intermediate_size = config.intermediate_size  self.gate_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)  self.up_proj = nn.Linear(self.hidden_size, self.intermediate_size, bias=False)  self.down_proj = nn.Linear(self.intermediate_size, self.hidden_size, bias=False)  self.act_fn = ACT2FN[config.hidden_act]  def forward(self, x):  if self.config.pretraining_tp > 1:  slice = self.intermediate_size // self.config.pretraining_tp  gate_proj_slices = self.gate_proj.weight.split(slice, dim=0)  up_proj_slices = self.up_proj.weight.split(slice, dim=0)  down_proj_slices = self.down_proj.weight.split(slice, dim=1)  gate_proj = torch.cat(  [F.linear(x, gate_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1  )  up_proj = torch.cat([F.linear(x, up_proj_slices[i]) for i in range(self.config.pretraining_tp)], dim=-1)  intermediate_states = (self.act_fn(gate_proj) * up_proj).split(slice, dim=2)  down_proj = [  F.linear(intermediate_states[i], down_proj_slices[i]) for i in range(self.config.pretraining_tp)  ]  down_proj = sum(down_proj)  else:  down_proj = self.down_proj(self.act_fn(self.gate_proj(x)) * self.up_proj(x))  return down_proj

4，层归一化 LlamaRMSNorm

这里的层归一化叫做RMSNorm，和标准的LayerNorm有少许差异。

首先是没有移除均值，直接除的RootMeanSquare，然后也没有加上bias。

这两个小的修正可以保证在层归一化不会改变hidden_states对应的词向量的方向，只会改变其模长。

在一定的意义上具有合理性。

class LlamaRMSNorm(nn.Module):  def __init__(self, hidden_size, eps=1e-6):  """  LlamaRMSNorm is equivalent to T5LayerNorm  """  super().__init__()  self.weight = nn.Parameter(torch.ones(hidden_size))  self.variance_epsilon = eps  def forward(self, hidden_states):  input_dtype = hidden_states.dtype  hidden_states = hidden_states.to(torch.float32)  variance = hidden_states.pow(2).mean(-1, keepdim=True)  hidden_states = hidden_states * torch.rsqrt(variance + self.variance_epsilon)  return self.weight * hidden_states.to(input_dtype)

5，Llama解码层

解码层LlamaDecoderLayer由LlamaAttention，LlamaMLP，以及两个LlamaRMSNorm组成，并使用了两次残差结构。

class LlamaDecoderLayer(nn.Module):  def __init__(self, config: LlamaConfig):  super().__init__()  self.hidden_size = config.hidden_size  self.self_attn = LlamaAttention(config=config)  self.mlp = LlamaMLP(config)  self.input_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)  self.post_attention_layernorm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)  def forward(  self,  hidden_states: torch.Tensor,  attention_mask: Optional[torch.Tensor] = None,  position_ids: Optional[torch.LongTensor] = None,  past_key_value: Optional[Tuple[torch.Tensor]] = None,  output_attentions: Optional[bool] = False,  use_cache: Optional[bool] = False,  ) -> Tuple[torch.FloatTensor, Optional[Tuple[torch.FloatTensor, torch.FloatTensor]]]:  """  Args:  hidden_states (`torch.FloatTensor`): input to the layer of shape `(batch, seq_len, embed_dim)`  attention_mask (`torch.FloatTensor`, *optional*): attention mask of size  `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.  output_attentions (`bool`, *optional*):  Whether or not to return the attentions tensors of all attention layers. See `attentions` under  returned tensors for more detail.  use_cache (`bool`, *optional*):  If set to `True`, `past_key_values` key value states are returned and can be used to speed up decoding  (see `past_key_values`).  past_key_value (`Tuple(torch.FloatTensor)`, *optional*): cached past key and value projection states  """  residual = hidden_states  hidden_states = self.input_layernorm(hidden_states)  # Self Attention  hidden_states, self_attn_weights, present_key_value = self.self_attn(  hidden_states=hidden_states,  attention_mask=attention_mask,  position_ids=position_ids,  past_key_value=past_key_value,  output_attentions=output_attentions,  use_cache=use_cache,  )  hidden_states = residual + hidden_states  # Fully Connected  residual = hidden_states  hidden_states = self.post_attention_layernorm(hidden_states)  hidden_states = self.mlp(hidden_states)  hidden_states = residual + hidden_states  outputs = (hidden_states,)  if output_attentions:  outputs += (self_attn_weights,)  if use_cache:  outputs += (present_key_value,)  return outputs

6，Llama解码器

LlamaModel由多个Llama解码层堆叠而成。

有几个理解上的要点：

1，_make_causal_mask用于构造下三角这种mask结构以实现语言模型的单向注意力。

2，_expand_mask用于将传入的等特殊符号相关的mask信息展开成和attention矩阵相同的张量结构。

3，设置gradient_checkpointing=True可以节约显存。其主要应用了torch.utils.checkpoint.checkpoint方法。它的原理非常简单，在对decoder_layer进行forward时不保存中间激活值从而节约显存，backward时重新计算相关值，从而通过时间换取了空间。

4，gradient_checkpointing和use_cache不能同时设置为True，前者是为了节约显存时间换空间的，后者是为了节约时间空间换时间。

# Copied from transformers.models.bart.modeling_bart._make_causal_mask  
def _make_causal_mask(  input_ids_shape: torch.Size, dtype: torch.dtype,   device: torch.device, past_key_values_length: int = 0  
):  """  Make causal mask used for bi-directional self-attention.  """  bsz, tgt_len = input_ids_shape  mask = torch.full((tgt_len, tgt_len), torch.finfo(dtype).min, device=device)  mask_cond = torch.arange(mask.size(-1), device=device)  mask.masked_fill_(mask_cond < (mask_cond + 1).view(mask.size(-1), 1), 0)  mask = mask.to(dtype)  if past_key_values_length > 0:  mask = torch.cat([torch.zeros(tgt_len, past_key_values_length, dtype=dtype, device=device), mask], dim=-1)  return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)  # Copied from transformers.models.bart.modeling_bart._expand_mask  
def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):  """  Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.  """  bsz, src_len = mask.size()  tgt_len = tgt_len if tgt_len is not None else src_len  expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)  inverted_mask = 1.0 - expanded_mask  return inverted_mask.masked_fill(inverted_mask.to(torch.bool), torch.finfo(dtype).min)  @add_start_docstrings(  "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",  LLAMA_START_DOCSTRING,  
)  
class LlamaPreTrainedModel(PreTrainedModel):  config_class = LlamaConfig  base_model_prefix = "model"  supports_gradient_checkpointing = True  _no_split_modules = ["LlamaDecoderLayer"]  _skip_keys_device_placement = "past_key_values"  def _init_weights(self, module):  std = self.config.initializer_range  if isinstance(module, nn.Linear):  module.weight.data.normal_(mean=0.0, std=std)  if module.bias is not None:  module.bias.data.zero_()  elif isinstance(module, nn.Embedding):  module.weight.data.normal_(mean=0.0, std=std)  if module.padding_idx is not None:  module.weight.data[module.padding_idx].zero_()  def _set_gradient_checkpointing(self, module, value=False):  if isinstance(module, LlamaModel):  module.gradient_checkpointing = value  @add_start_docstrings(  "The bare LLaMA Model outputting raw hidden-states without any specific head on top.",  LLAMA_START_DOCSTRING,  
)  
class LlamaModel(LlamaPreTrainedModel):  """  Transformer decoder consisting of *config.num_hidden_layers* layers. Each layer is a [`LlamaDecoderLayer`]  Args:  config: LlamaConfig  """  def __init__(self, config: LlamaConfig):  super().__init__(config)  self.padding_idx = config.pad_token_id  self.vocab_size = config.vocab_size  self.embed_tokens = nn.Embedding(config.vocab_size, config.hidden_size, self.padding_idx)  self.layers = nn.ModuleList([LlamaDecoderLayer(config) for _ in range(config.num_hidden_layers)])  self.norm = LlamaRMSNorm(config.hidden_size, eps=config.rms_norm_eps)  self.gradient_checkpointing = False  # Initialize weights and apply final processing  self.post_init()  def get_input_embeddings(self):  return self.embed_tokens  def set_input_embeddings(self, value):  self.embed_tokens = value  # Copied from transformers.models.bart.modeling_bart.BartDecoder._prepare_decoder_attention_mask  def _prepare_decoder_attention_mask(self, attention_mask, input_shape, inputs_embeds, past_key_values_length):  # create causal mask  # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]  combined_attention_mask = None  if input_shape[-1] > 1:  combined_attention_mask = _make_causal_mask(  input_shape,  inputs_embeds.dtype,  device=inputs_embeds.device,  past_key_values_length=past_key_values_length,  )  if attention_mask is not None:  # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]  expanded_attn_mask = _expand_mask(attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]).to(  inputs_embeds.device  )  combined_attention_mask = (  expanded_attn_mask if combined_attention_mask is None else expanded_attn_mask + combined_attention_mask  )  return combined_attention_mask  @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)  def forward(  self,  input_ids: torch.LongTensor = None,  attention_mask: Optional[torch.Tensor] = None,  position_ids: Optional[torch.LongTensor] = None,  past_key_values: Optional[List[torch.FloatTensor]] = None,  inputs_embeds: Optional[torch.FloatTensor] = None,  use_cache: Optional[bool] = None,  output_attentions: Optional[bool] = None,  output_hidden_states: Optional[bool] = None,  return_dict: Optional[bool] = None,  ) -> Union[Tuple, BaseModelOutputWithPast]:  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions  output_hidden_states = (  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states  )  use_cache = use_cache if use_cache is not None else self.config.use_cache  return_dict = return_dict if return_dict is not None else self.config.use_return_dict  # retrieve input_ids and inputs_embeds  if input_ids is not None and inputs_embeds is not None:  raise ValueError("You cannot specify both decoder_input_ids and decoder_inputs_embeds at the same time")  elif input_ids is not None:  batch_size, seq_length = input_ids.shape  elif inputs_embeds is not None:  batch_size, seq_length, _ = inputs_embeds.shape  else:  raise ValueError("You have to specify either decoder_input_ids or decoder_inputs_embeds")  seq_length_with_past = seq_length  past_key_values_length = 0  if past_key_values is not None:  past_key_values_length = past_key_values[0][0].shape[2]  seq_length_with_past = seq_length_with_past + past_key_values_length  if position_ids is None:  device = input_ids.device if input_ids is not None else inputs_embeds.device  position_ids = torch.arange(  past_key_values_length, seq_length + past_key_values_length, dtype=torch.long, device=device  )  position_ids = position_ids.unsqueeze(0).view(-1, seq_length)  else:  position_ids = position_ids.view(-1, seq_length).long()  if inputs_embeds is None:  inputs_embeds = self.embed_tokens(input_ids)  # embed positions  if attention_mask is None:  attention_mask = torch.ones(  (batch_size, seq_length_with_past), dtype=torch.bool, device=inputs_embeds.device  )  attention_mask = self._prepare_decoder_attention_mask(  attention_mask, (batch_size, seq_length), inputs_embeds, past_key_values_length  )  hidden_states = inputs_embeds  if self.gradient_checkpointing and self.training:  if use_cache:  logger.warning_once(  "`use_cache=True` is incompatible with gradient checkpointing. Setting `use_cache=False`..."  )  use_cache = False  # decoder layers  all_hidden_states = () if output_hidden_states else None  all_self_attns = () if output_attentions else None  next_decoder_cache = () if use_cache else None  for idx, decoder_layer in enumerate(self.layers):  if output_hidden_states:  all_hidden_states += (hidden_states,)  past_key_value = past_key_values[idx] if past_key_values is not None else None  if self.gradient_checkpointing and self.training:  def create_custom_forward(module):  def custom_forward(*inputs):  # None for past_key_value  return module(*inputs, output_attentions, None)  return custom_forward  layer_outputs = torch.utils.checkpoint.checkpoint(  create_custom_forward(decoder_layer),  hidden_states,  attention_mask,  position_ids,  None,  )  else:  layer_outputs = decoder_layer(  hidden_states,  attention_mask=attention_mask,  position_ids=position_ids,  past_key_value=past_key_value,  output_attentions=output_attentions,  use_cache=use_cache,  )  hidden_states = layer_outputs[0]  if use_cache:  next_decoder_cache += (layer_outputs[2 if output_attentions else 1],)  if output_attentions:  all_self_attns += (layer_outputs[1],)  hidden_states = self.norm(hidden_states)  # add hidden states from the last decoder layer  if output_hidden_states:  all_hidden_states += (hidden_states,)  next_cache = next_decoder_cache if use_cache else None  if not return_dict:  return tuple(v for v in [hidden_states, next_cache, all_hidden_states, all_self_attns] if v is not None)  return BaseModelOutputWithPast(  last_hidden_state=hidden_states,  past_key_values=next_cache,  hidden_states=all_hidden_states,  attentions=all_self_attns,  )

7，Llama语言模型

Llama语言模型 LlamaForCausalLM是在Llama解码器LlamaModel的基础上增加了一个lm_head作为Generator。

从而实现了一个完整的语言模型。

除此之外，Llama语言模型还实现了以下重要功能。

1，loss计算功能。当forward方法中传入labels时，会自动计算语言模型的交叉熵损失。注意labels中的-100会被忽略不参与计算。

2，文本生成generate方法。这个方法继承自PreTrainedModel，可以设置model.generation_config.num_beams选择束搜索的束宽度，默认为1即贪心搜索。

_CONFIG_FOR_DOC = "LlamaConfig"  class LlamaForCausalLM(LlamaPreTrainedModel):  _tied_weights_keys = ["lm_head.weight"]  def __init__(self, config):  super().__init__(config)  self.model = LlamaModel(config)  self.vocab_size = config.vocab_size  self.lm_head = nn.Linear(config.hidden_size, config.vocab_size, bias=False)  # Initialize weights and apply final processing  self.post_init()  def get_input_embeddings(self):  return self.model.embed_tokens  def set_input_embeddings(self, value):  self.model.embed_tokens = value  def get_output_embeddings(self):  return self.lm_head  def set_output_embeddings(self, new_embeddings):  self.lm_head = new_embeddings  def set_decoder(self, decoder):  self.model = decoder  def get_decoder(self):  return self.model  @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)  @replace_return_docstrings(output_type=CausalLMOutputWithPast, config_class=_CONFIG_FOR_DOC)  def forward(  self,  input_ids: torch.LongTensor = None,  attention_mask: Optional[torch.Tensor] = None,  position_ids: Optional[torch.LongTensor] = None,  past_key_values: Optional[List[torch.FloatTensor]] = None,  inputs_embeds: Optional[torch.FloatTensor] = None,  labels: Optional[torch.LongTensor] = None,  use_cache: Optional[bool] = None,  output_attentions: Optional[bool] = None,  output_hidden_states: Optional[bool] = None,  return_dict: Optional[bool] = None,  ) -> Union[Tuple, CausalLMOutputWithPast]:  output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions  output_hidden_states = (  output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states  )  return_dict = return_dict if return_dict is not None else self.config.use_return_dict  # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)  outputs = self.model(  input_ids=input_ids,  attention_mask=attention_mask,  position_ids=position_ids,  past_key_values=past_key_values,  inputs_embeds=inputs_embeds,  use_cache=use_cache,  output_attentions=output_attentions,  output_hidden_states=output_hidden_states,  return_dict=return_dict,  )  hidden_states = outputs[0]  if self.config.pretraining_tp > 1:  lm_head_slices = self.lm_head.weight.split(self.vocab_size // self.config.pretraining_tp, dim=0)  logits = [F.linear(hidden_states, lm_head_slices[i]) for i in range(self.config.pretraining_tp)]  logits = torch.cat(logits, dim=-1)  else:  logits = self.lm_head(hidden_states)  logits = logits.float()  loss = None  if labels is not None:  # Shift so that tokens < n predict n  shift_logits = logits[..., :-1, :].contiguous()  shift_labels = labels[..., 1:].contiguous()  # Flatten the tokens  loss_fct = CrossEntropyLoss()  shift_logits = shift_logits.view(-1, self.config.vocab_size)  shift_labels = shift_labels.view(-1)  # Enable model parallelism  shift_labels = shift_labels.to(shift_logits.device)  loss = loss_fct(shift_logits, shift_labels)  if not return_dict:  output = (logits,) + outputs[1:]  return (loss,) + output if loss is not None else output  return CausalLMOutputWithPast(  loss=loss,  logits=logits,  past_key_values=outputs.past_key_values,  hidden_states=outputs.hidden_states,  attentions=outputs.attentions,  )  def prepare_inputs_for_generation(  self, input_ids, past_key_values=None, attention_mask=None, inputs_embeds=None, **kwargs  ):  if past_key_values:  input_ids = input_ids[:, -1:]  position_ids = kwargs.get("position_ids", None)  if attention_mask is not None and position_ids is None:  # create position_ids on the fly for batch generation  position_ids = attention_mask.long().cumsum(-1) - 1  position_ids.masked_fill_(attention_mask == 0, 1)  if past_key_values:  position_ids = position_ids[:, -1].unsqueeze(-1)  # if `inputs_embeds` are passed, we only want to use them in the 1st generation step  if inputs_embeds is not None and past_key_values is None:  model_inputs = {"inputs_embeds": inputs_embeds}  else:  model_inputs = {"input_ids": input_ids}  model_inputs.update(  {  "position_ids": position_ids,  "past_key_values": past_key_values,  "use_cache": kwargs.get("use_cache"),  "attention_mask": attention_mask,  }  )  return model_inputs  @staticmethod  def _reorder_cache(past_key_values, beam_idx):  reordered_past = ()  for layer_past in past_key_values:  reordered_past += (  tuple(past_state.index_select(0, beam_idx.to(past_state.device)) for past_state in layer_past),  )  return reordered_past

8，Llama分类模型

LlamaForSequenceClassification是一个序列分类模型。

这个分类模型可以用来训练RLHF流程中的Reward模型。

@add_start_docstrings(  """  The LLaMa Model transformer with a sequence classification head on top (linear layer).  [`LlamaForSequenceClassification`] uses the last token in order to do the classification, as other causal models  (e.g. GPT-2) do.  Since it does classification on the last token, it requires to know the position of the last token. If a  `pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each row. If  no `pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot guess the  padding tokens when `inputs_embeds` are passed instead of `input_ids`, it does the same (take the last value in  each row of the batch).  """,  LLAMA_START_DOCSTRING,  
)  
class LlamaForSequenceClassification(LlamaPreTrainedModel):  def __init__(self, config):  super().__init__(config)  self.num_labels = config.num_labels  self.model = LlamaModel(config)  self.score = nn.Linear(config.hidden_size, self.num_labels, bias=False)  # Initialize weights and apply final processing  self.post_init()  def get_input_embeddings(self):  return self.model.embed_tokens  def set_input_embeddings(self, value):  self.model.embed_tokens = value  @add_start_docstrings_to_model_forward(LLAMA_INPUTS_DOCSTRING)  def forward(  self,  input_ids: torch.LongTensor = None,  attention_mask: Optional[torch.Tensor] = None,  position_ids: Optional[torch.LongTensor] = None,  past_key_values: Optional[List[torch.FloatTensor]] = None,  inputs_embeds: Optional[torch.FloatTensor] = None,  labels: Optional[torch.LongTensor] = None,  use_cache: Optional[bool] = None,  output_attentions: Optional[bool] = None,  output_hidden_states: Optional[bool] = None,  return_dict: Optional[bool] = None,  ) -> Union[Tuple, SequenceClassifierOutputWithPast]:  r"""  labels (`torch.LongTensor` of shape `(batch_size,)`, *optional*):  Labels for computing the sequence classification/regression loss. Indices should be in `[0, ...,  config.num_labels - 1]`. If `config.num_labels == 1` a regression loss is computed (Mean-Square loss), If  `config.num_labels > 1` a classification loss is computed (Cross-Entropy).  """  return_dict = return_dict if return_dict is not None else self.config.use_return_dict  transformer_outputs = self.model(  input_ids,  attention_mask=attention_mask,  position_ids=position_ids,  past_key_values=past_key_values,  inputs_embeds=inputs_embeds,  use_cache=use_cache,  output_attentions=output_attentions,  output_hidden_states=output_hidden_states,  return_dict=return_dict,  )  hidden_states = transformer_outputs[0]  logits = self.score(hidden_states)  if input_ids is not None:  batch_size = input_ids.shape[0]  else:  batch_size = inputs_embeds.shape[0]  if self.config.pad_token_id is None and batch_size != 1:  raise ValueError("Cannot handle batch sizes > 1 if no padding token is defined.")  if self.config.pad_token_id is None:  sequence_lengths = -1  else:  if input_ids is not None:  sequence_lengths = (torch.eq(input_ids, self.config.pad_token_id).long().argmax(-1) - 1).to(  logits.device  )  else:  sequence_lengths = -1  pooled_logits = logits[torch.arange(batch_size, device=logits.device), sequence_lengths]  loss = None  if labels is not None:  labels = labels.to(logits.device)  if self.config.problem_type is None:  if self.num_labels == 1:  self.config.problem_type = "regression"  elif self.num_labels > 1 and (labels.dtype == torch.long or labels.dtype == torch.int):  self.config.problem_type = "single_label_classification"  else:  self.config.problem_type = "multi_label_classification"  if self.config.problem_type == "regression":  loss_fct = MSELoss()  if self.num_labels == 1:  loss = loss_fct(pooled_logits.squeeze(), labels.squeeze())  else:  loss = loss_fct(pooled_logits, labels)  elif self.config.problem_type == "single_label_classification":  loss_fct = CrossEntropyLoss()  loss = loss_fct(pooled_logits.view(-1, self.num_labels), labels.view(-1))  elif self.config.problem_type == "multi_label_classification":  loss_fct = BCEWithLogitsLoss()  loss = loss_fct(pooled_logits, labels)  if not return_dict:  output = (pooled_logits,) + transformer_outputs[1:]  return ((loss,) + output) if loss is not None else output  return SequenceClassifierOutputWithPast(  loss=loss,  logits=pooled_logits,  past_key_values=transformer_outputs.past_key_values,  hidden_states=transformer_outputs.hidden_states,  attentions=transformer_outputs.attentions,  )

三，训练模型

下面，我们来训练一个LlamaForCausalLM 实现两数之和的任务。

config = LlamaConfig(  vocab_size=len(vocab),  hidden_size=512,  intermediate_size=2752,  num_hidden_layers=8,  num_attention_heads=16,  num_key_value_heads=4,  rope_scaling = None,  hidden_act='silu',  max_position_embeddings=128,  initializer_range=0.02,  rms_norm_eps=1e-06,  use_cache=True,  pad_token_id=0,  bos_token_id=1,  eos_token_id=2,  tie_word_embeddings=False,  pretraining_tp = 1,  max_new_tokens = 100  
)   #试算一下  
model = LlamaForCausalLM(config)  
out = model.forward(**batch)  
print(out.loss)  
tensor(2.7630, grad_fn=)from torchkeras import KerasModel   
from accelerate import Accelerator   class StepRunner:  def __init__(self, net, loss_fn, accelerator=None, stage = "train", metrics_dict = None,   optimizer = None, lr_scheduler = None  ):  self.net,self.loss_fn,self.metrics_dict,self.stage = net,loss_fn,metrics_dict,stage  self.optimizer,self.lr_scheduler = optimizer,lr_scheduler  self.accelerator = accelerator if accelerator is not None else Accelerator()   if self.stage=='train':  self.net.train()   else:  self.net.eval（)  def __call__(self, batch):  #loss  with self.accelerator.autocast():  loss = self.net(**batch).loss  #backward()  if self.stage=="train" and self.optimizer is not None:          self.accelerator.backward(loss)  if self.accelerator.sync_gradients:  self.accelerator.clip_grad_norm_(self.net.parameters(), 1.0)  self.optimizer.step()  if self.lr_scheduler is not None:  self.lr_scheduler.step()  self.optimizer.zero_grad()  all_loss = self.accelerator.gather(loss).sum()  #losses (or plain metrics that can be averaged)  step_losses = {self.stage+"_loss":all_loss.item()}  #metrics (stateful metrics)  step_metrics = {}  if self.stage=="train":  if self.optimizer is not None:  step_metrics['lr'] = self.optimizer.state_dict()['param_groups'][0]['lr']  else:  step_metrics['lr'] = 0.0  return step_losses,step_metrics  KerasModel.StepRunner = StepRunner   
keras_model = KerasModel(model,loss_fn = None,  optimizer=torch.optim.AdamW(model.parameters(),lr=3e-5))  #加载 之前训练过的权重  
ckpt_path = 'llama_twosum'  keras_model.fit(train_data = dl_train,  val_data = dl_val,  epochs=100,patience=5,  monitor='val_loss',mode='min',  ckpt_path = ckpt_path,  mixed_precision='fp16'  )

四，使用模型

from transformers.generation.utils import GenerationConfig  
model.generation_config = GenerationConfig.from_dict({'num_beams':1,  'max_new_tokens':100,  'max_length':200})  
model.generation_config.num_beams=1  
model.generation_config.max_new_tokens = 100   
model.generation_config.max_length=200  
def get_ans(tensor) ->"str":  s = "".join([vocab_r[i] for i in tensor.tolist()])  ans = s[s.find('=')+1:s.find('<EOS>')].replace('<BOS>','').replace('<EOS>','')  return ans  
x,y = get_data()   
print('x: '+''.join(x).replace('<BOS>',''))  
print('y: '+''.join(y).replace('<EOS>',''))

x: 3481340050+90157504501803=

y: 90160985841853

input_ids = torch.tensor([[vocab[i] for i in x]])   
out = model.generate(inputs=input_ids)  
out

tensor([[ 1, 5, 6, 10, 3, 5, 6, 12, 12, 7, 12, 13, 11, 12, 3, 7, 9, 7, 12, 6, 7, 12, 3, 10, 12, 5, 14, 11, 12, 3, 8, 12, 11, 10, 7, 10, 6, 3, 10, 7, 5, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 2, 12, 2, 2, 2, 2, 2, 2, 2, 2, 12, 3, 12, 3]])

get_ans(out[0])

'90160985841853'

五，评估模型

from tqdm import tqdm   
loop = tqdm(range(1,201))  
correct = 0  
for i in loop:  x,y = get_data()   input_ids = torch.tensor([[vocab[i] for i in x]])   out = model.generate(inputs=input_ids)  pred = get_ans(out[0])  gt = ''.join(y).replace('<EOS>','')  if pred==gt:  correct+=1  loop.set_postfix(acc = correct/i)  print("acc=",correct/len(loop))

acc= 0.99漂亮，我们的测试准确率达到了99%！

#Replacing softmax with ReLU in Vision Transformers

对于视觉 Transformer，将其 Self-Attention 中的 Softmax 操作替换为 ReLU/序列长度 (seqlen) 之后，性能的下降问题有所缓解。本文在 ImageNet-21K 上训练了从 Small 级别到 Large 级别的视觉 Transformer，证明了 ReLU-attention 可以在缩放性上接近或者匹配 Softmax-attention 的性能。Google出品，使用ReLU取代Softmax，ViT性能不退化

本文的研究结论是：对于视觉 Transformer，将其 Self-Attention 中的 Softmax 操作替换为 ReLU/序列长度 (seqlen) 之后，性能的下降问题有所缓解。本文在 ImageNet-21K 上训练了从 Small 级别到 Large 级别的视觉 Transformer，证明了 ReLU-attention 可以在缩放性上接近或者匹配 Softmax-attention 的性能。

1 在 ViT 中使用 ReLU 取代 Softmax

论文名称： Replacing softmax with ReLU in Vision Transformers (Arxiv 2023)

论文地址：https//arxiv.org/pdf/2309.08586.pdf

1.1 ReLU-attention 的新发现

Transformer 架构[1]在现代机器学习中无处不在。Attention 是 Transformer 的核心组件，包括一个 Softmax 操作，它在 token 上产生概率分布。Softmax 操作涉及到内部的计算所有输入的指数之和，它的计算代价相当昂贵，使得 Transformer 架构的并行化具有挑战性[2]。

本文作者探索了 Softmax 操作的 Point-wise 的替代方案，该操作不一定输出概率分布。本文的核心贡献是观察到：ReLU/序列长度(seqlen) ，可以在缩放性方面接近或匹配传统的 Softmax 操作。这一结果为并行化提供了新的机会，因为 ReLU-attention 相比传统的 Softmax-attention 可以使用更少的 gather 操作在序列长度维度实现并行化。

1.2 去掉 Softmax 的相关工作

替换 Softmax 的研究：

ReLU 和 squared ReLU：[3][4]把 Softmax 替换成了 ReLU，[5]把 Softmax 替换成了 squared ReLU。但是这些方法不会除以序列长度，本文通过实验发现对于达到与 Softmax 相当的准确度很重要。
[6]仍然需要对序列长度轴进行归一化，以确保注意力权重之和为1，这依然需要 gather。

去掉激活函数的研究：

1.3 ReLU-attention 方法

在进行 Self-attention 的操作时，首先计算注意力权重：

图1：Scaled point-wise attention 实验结果

Sequence length scaling

1.4 实验结果

作者在 ImageNet-21K 上训练了 30 Epochs，在 ImageNet-1K 上训练了 300 Epochs。作者使用了 ViT-22B[10]中提出的 qk-norm 技术，因为这个技术被验证在扩大视觉模型时有益于优化稳定性，但是作者发现在本文量级的模型这一技术没那么重要。

如下图2所示说明了 ReLU-attention 与 ImageNet-21K 训练的 Softmax-attention 的缩放趋势相匹配。x 轴表示实验所需的总 core hours。ReLU-attention 的优势是能够以比 Softmax-attention 以更少的 gather 操作对序列长度维度进行并行化。

图2：Softmax 操作替换为 ReLU/seqlen 的缩放性能与传统带有 qk-layernorm 的 Transformer 的缩放性能匹配

1.5 qk-norm 实验结果

本文主要实验使用了 qk-norm，其中 query 和 key 在计算注意力权重之前通过 LayerNorm 传递，作者发现有必要在扩大模型大小时防止不稳定性。如图3所示是 qk-layernorm 的实验结果。结果表明，qk-norm 对这些模型没有很大的影响。

图3：qk-norm 实验结果

1.6 添加 gate 的影响

[11]这个工作删除了 Softmax 之后，添加了一个门控单元，并且不按序列长度缩放。具体而言，在门控注意力单元中，通过额外的投影层产生输出，该输出在输出映射之前与注意力的结果做 Element-wise 的乘法。

如图4所示是添加 gate 的影响实验结果。作者研究了 gate 的存在是否消除了序列长度缩放的需要。总体而言，作者观察到无论有没有 gate 的存在，使用序列长度缩放都实现了最佳精度。注意到对于带有 ReLU 的 S/8 模型，添加 gate 操作将实验所需的 core hour 增加了大约 9.3%。