好资源源码网站东莞网站推广策划

1、目前Transformer模型现状

• dense模型相同激活参数下，性能远低于MOE模型，因此大家倾向于训练MOE模型
• 虽然同激活参数下，MOE性能比dense好，但MOE模型内存访问高，因此推理速度相比dense要慢不少。比如top2的moe，推理速度要慢2倍多。
  因此，本文提出Ultra-Mem架构，拟解决MOE模型内存访问成本高，推理速度慢的问题。

2、结论说在前面

• 对比同参数、同激活参数的MOE
  Ultra-Mem的valid loss 更低，推理速度快了1.7 ~ 6倍不等
• 对比4倍激活参数的dense 模型
  Ultra-Mem的valid loss 更低，推理速度仅慢一点。
  

3、怎么做的？

3.1 相关工作

介绍之前，先介绍几种常见的架构：

• Dense模型：MLP采用致密的矩阵向量
• MOE架构，MLP采用多个稀疏的矩阵向量，选择topn个专家进行推理

• large memory layer，PKM（Lample et al. (2019) ）
  mlp采用类似Attention q、k、v。每个头的qi从k中进行检索，基于检索的索引很分数从value中取值，公式如下：
  

PKM为了减少计算，k采用了分块的思想。qi每次检索可以获得2个分数，一个是kj块的分数，另一个是kj种每个key的分数，基于这2个分数相加得到一个二维逻辑查询表，从value中抽取出相应的向量进行加权求和。公式如下：

PKM存在问题：每个token在推理时仅仅激活极少的value，推理时不会遇到访存瓶颈，但效果较差，且scaling能力差

原因：
1、随着值的数量 N 显著增加，query更难找到正确的值：
当记忆表value的规模变得非常大时，查询向量在检索过程中会面临更高的难度，因为需要在庞大的候选集中找到最相关的值，这可能导致检索效率下降。
2、乘积键分解引入了检索拓扑的偏差：
例如，假设 (i,j) 是网格中得分最高的逻辑地址，那么得分第二高的值必须位于第 i 行或第 j 列。这种限制显著降低了 top-m 选择的多样性，因为检索结果被限制在特定的行或列中。
（因为qk分数是基于行和列的分数相加得到的）
3、大规模参数训练时存在多 GPU 计算和通信不平衡的问题：
在大规模训练中，由于模型参数过多，无法将整个模型放置在单个 GPU 上，这导致了计算和通信的不平衡。例如，当一个层的参数量超过单个 GPU 的容量时，传统的并行策略可能无法有效处理。

3.2 Ultra-Mem，改进版PKM

小trick

• 移除qk之后的softmax

• qk进行nom

• value常量学习率 -> 逐渐变小学习率

• q proj之前增加一个对q的单向深度卷积
  大trick

• PKM的部分采用share query

• TDQKR： 针对PKM 原因2
  引入可学习参数C ∈ r x r。r表示Tucker 分解中的秩。并且C也可以是多个，即C∈hxrxr。这就是论文后面用到的MCS方法。论文的h采用的是2.
  

  但同时也带来了新问题：Srow X C X Scol的比较低效。因此，将C拆分为2个低秩矩阵：
  

为了确保u和t不会被训练为全1矩阵，因此加上了一个辅助loss：

其中λi表示u，t。
τ：常量，论文采用0.15
r：表示Tucker 分解中的秩，论文采用2，更到r如3和4效果不大
α：常量，论文采用0.001

• value向量的维度减半，数量加倍，并将value分块。最后减半的维度加一个线性层转为model dim。针对PKM 原因1
  

• 针对PKM 原因3
  UltraMem 层被设计为替代或与 Transformer 中的 MLP 层并行运行
  1、UltraMem 层可以完全替代 Transformer 中的 MLP 层，尤其是在模型的深层部分。
  2、或者，UltraMem 层可以与 MLP 层并行运行，共享输入并合并输出。这种方式允许模型同时利用 MLP 的灵活性和 UltraMem 的高效性。
  在论文的实验中，作者采用了以下配置：

• 对于较小的模型（如 151M 参数模型），UltraMem 层被插入到 Transformer 的特定层之间，例如在第 3 层插入，得到的输出添加到第3层的输出和第 5 层的输入里面。…

• 对于较大的模型（如 680M 和 1.6B 参数模型），UltraMem 层被分布在多个 Transformer 层之间，以实现更高效的计算和更好的性能。
  

4、实验结果

• 模型对比
  

• Scaling、稀疏性对loss和推理速度的影响
  bc图横坐标表示稀疏参数和激活参数的比例
  20k、40k…表示value的大小
  越稀疏loss越高，推理速度越快
  

• 消融实验

  • 效果提升最多的是half vdim + proj 和value lr decay