论文笔记-ULTRA-SPARSE-MEMORY-NETWORK

JAY.LIN 收录于论文 Nlp

2025-03-12 约 1717 字预计阅读 4 分钟

论文笔记 - ULTRA-SPARSE MEMORY NETWORK

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

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

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 层之间，以实现更高效的计算和更好的性能。