算法进阶 | 深度解析Transformer、RNN和Mamba的联系！

本文来源公众号“算法进阶”，仅用于学术分享，侵权删，干货满满。

原文链接：深度解析Transformer、RNN和Mamba的联系！

通过探索看似不相关的大语言模型(LLM)架构之间的潜在联系,我们可能为促进不同模型间的思想交流和提高整体效率开辟新的途径。

尽管Mamba等线性循环神经网络(RNN)和状态空间模型(SSM)近来备受关注,Transformer架构仍然是LLM的主要支柱。这种格局可能即将发生变化:像Jamba、Samba和Griffin这样的混合架构展现出了巨大的潜力。这些模型在时间和内存效率方面明显优于Transformer,同时在能力上与基于注意力的LLM相比并未显著下降。

近期研究揭示了不同架构选择之间的深层联系,包括Transformer、RNN、SSM和matrix mixers，这一发现具有重要意义,因为它为不同架构间的思想迁移提供了可能。本文将深入探讨Transformer、RNN和Mamba 2,通过详细的代数分析来理解以下几点:

1. Transformer在某些情况下可以视为RNN(第2节)

2. 状态空间模型可能隐藏在自注意力机制的掩码中(第4节)

3. Mamba在特定条件下可以重写为掩码自注意力(第5节)

这些联系不仅有趣,还可能对未来的模型设计产生深远影响。

1 LLM中的掩码自注意力机制

首先,让我们回顾一下经典的LLM自注意力层的结构:

更详细的结构如下:

自注意力层的工作流程如下:

1. 将查询矩阵Q和键矩阵K相乘,得到一个L×L的矩阵,包含查询和键的标量积。

2. 对结果矩阵进行归一化。

3. 将归一化后的矩阵与L×L的注意力掩码进行元素级乘法。图中展示了默认的因果掩码——左侧的0-1矩阵。这一步骤将较早查询与较晚键的乘积置零,防止注意力机制"看到未来"。

4. 对结果应用softmax函数。

5. 最后,将注意力权重矩阵A与值矩阵V相乘。输出的第t行可表示为:

这意味着第i个值是通过"第t个查询对第i个键的注意力权重"来加权的。

这种架构中的多个设计选择都可能被修改。接下来我们将探讨一些可能的变体。

2 线性化注意力

注意力公式中的Softmax函数确保了值是以和为1的正系数混合的。这种设计保持了某些统计特性,但同时也带来了限制。例如即使我们希望利用结合律,如(QK^T)V = Q(K^TV),也无法突破Softmax的限制。

为什么结合律如此重要?因为改变乘法顺序可能显著影响计算复杂度:

左侧公式需要计算一个L×L矩阵,如果这个矩阵完全显现在内存中,复杂度为O(L²d),内存消耗为O(L²)。右侧公式需要计算一个d×d矩阵,复杂度为O(Ld²),内存消耗为O(d²)。

随着上下文长度L的增加,左侧公式的计算成本rapidly become prohibitively非常的高。为了解决这个问题,我们可以考虑移除Softmax。详细展开带有Softmax的公式:

其中