【大模型专栏—百日文“新”】天下苦Transformer久矣

大模型专栏介绍

😊你好，我是小航，一个正在变秃、变强的文艺倾年。

🔔本文为大模型专栏子篇，大模型专栏将持续更新，主要讲解大模型从入门到实战打怪升级。如有兴趣，欢迎您的阅读。

💡适合人群：本科生、研究生、大模型爱好者，期待与你一同探索、学习、进步，一起卷起来叭！

🔗篇章一：本篇主要讲解Python基础、数据分析三件套、机器学习、深度学习、CUDA等基础知识、学习使用AutoDL炼丹
🔗篇章二：本篇主要讲解基本的科研知识、认识数据和显卡、语言模型如RNN、LSTM、Attention、Transformer、Bert、T5、GPT、BLOOM、LLama、Baichuan、ChatGLM等系列、强化学习教程、大模型基础知识及微调等
🔗篇章三：本篇主要讲解智能对话、大模型基础实战如Ollama、Agent、QLoar、Deepspeed、RAG、Mobile Agent等、大模型领域前沿论文总结及创新点汇总

Transformer

• 创新点
  • 核心思想：通过自注意力机制处理序列数据。
  • 大白话解释：用大剪子把序列数据按窗口进行截断，然后通过位置嵌入编码和自注意力计算简化了时序关系。【裁缝思想】
• 优点
  • 减少了时间维度
  • 暴力拆解具有通用性，可以通过堆叠模块处理大数据。
• 缺点
  • 没有由于升维，空间复杂度增大。
    
    两个相乘的矩阵大小分别为(N $\times$ d) 和(d $\times$ N)，我们需要拿第一个矩阵的每一行去与第二个矩阵的每一列做点乘，所以总共就需要 N² 次点乘。而每次点乘又需要 d 次乘法，所以总复杂度就为O(N²d)。

改进思路：

• 魔改注意力机制：S4、FlashAttention
• 提出了新的序列模型：Mamba

Mamba1

单层网络

输入是x，经过变换Wx+b和激活函数f，得到输出y。

RNN

因为序列数据就不太好用原始的神经网络处理了，所以RNN引入了隐状态h（hidden state）的概念，隐状态h可以对序列形的数据提取特征，接着再转换为输出。

• $h_1$基于上一个隐藏层的状态$h_0$和当前的输入$x_1$计算得来
• 参数共享：RNN的权值是在同一个向量中，只是不同时刻而已。所以每一步使用的参数U、W、b都是一样的，也就是说每个步骤的参数都是共享的。

• 细品：先根据输入$x_t$和前一时刻的隐藏状态$h_{t-1}$计算出最新的隐藏状态$h_t$，便可以根据最新的隐藏状态$h_t$预测出$y_t$了
• 存在问题：
  • 远距离的梯度消失，不适合长距离依赖；
  • 无法并行训练，导致训练成本较高；

---

SSM（时序状态空间模型 ）

🔗推荐教程：https://www.youtube.com/watch?v=luCBXCErkCs

连续时间表示( continuous-time representation )：

• A就是存储着之前所有历史信息的浓缩精华(可以通过一系列系数组成的矩阵表示之)，以基于A更新下一个时刻的空间状态hidden state
• 与RNN循环结构：$h_t=tanh\left(W{h}_{t-1}+U{x}_t\right)$非常类似：
• A、B、C、D这4个矩阵是参数，是可以学习到的

统一方程后：

• 1.假设我们有一些输入信号x(t)，该信号首先乘以矩阵 B
• 2.上面第一步的结果，加上：上一个状态与矩阵A相乘(矩阵A描述了所有内部状态如何连接)的结果，用来更新状态state
• 3.然后，使用矩阵C来将状态转换为输出
• 4.最后，再利用矩阵D提供从输入到输出的直接信号，这通常也称为跳跃连接skip-connection
• 5.由于矩阵D类似于跳跃连接，因此在没有跳跃连接的情况下，SSM 通常被视为如下

SSM升级到S4

• 离散化：
  • 原因：文本序列是离散输入
  • 目的：将函数到函数$x(t)\to y(t)$，而是序列到序列$x_k\to y_k$
  • 处理方式：利用零阶保持技术(Zero-order hold technique)【在对应变量上面加一个横杠】
    • 1.保留每次收到离散信号，直到收到新的离散信号
    • 2.保持该值的时间由一个新的可学习参数表示，称为步长(size)——$\Delta$ ，它代表输入的阶段性保持(resolution)
    • 3.有了连续的输入信号后，便可以生成连续的输出，并且仅根据输入的时间步长对值进行采样
• 循环结构表示
  • 目的：方便快速推理
  • 推导：
    • 展开每个时间步
    • 对于y2，具体展开
    • 推广：
    • 由此，我们可以采用RNN的结构进行处理，$h_k$始终是$\overline{\mathbf{B}}{\mathbf{x}}_{\mathrm{k}}$和$\overline{\mathbf{A}}{\mathbf{h}}_{\mathrm{k}-1}$的共同作用之下更新的,
• 卷积结构表示
  • 目的：方便并行训练
  • 卷积核（过滤器，kernels）：处理的是文本而不是图像，因此我们需要一维视角
    • 过滤器（$\overline{K}$）的理解：
      • 1.使用 SSM 内核来检查每组token并计算输出
      • 2.内核移动，执行下一步的计算
      • 3.同理，继续移动
      • 4.转化为点积的形式：
        $$y_2=\left[\begin{array}{ccc}C{\bar{A}}^2\bar{B}& C\bar{A}\bar{B}& C\bar{B}\end{array}\right]\cdot \left[\begin{array}{c}{x}_0\\ x_1\\ x_2\end{array}\right]=\bar{K}\cdot X$$
      • 5.由于其中三个离散参数A、B、C都是常数，因此我们可以预先计算左侧向量并将其保存为卷积核

基于卷积与循环的利弊，我们采用两全其美的办法：

推理用RNN结构，训练用CNN结构
💡这里我们用生活的例子，帮你更好理解RNN和SSM的区别，想象你在读一本书

• 时序嵌套的 RNN每次只能读一行，然后把记忆传递到下一行，这种方法只适合处理短故事，故事一长，容易忘记前面的情节。
• 而 SSM 并行处理，同时打开所有页看到每行内容，这样就能快速找到和理解整本书，无需逐行传递记忆

提出Mamaba

Mamaba这个名字来源于黑曼巴蛇(Black Mamba)，以速度和致命著称。这种命名意在传达该架构的速度、灵活性和高效性，反映出它在处理和转换数据方面的强大能力。

• 论文名称：Mamba: Linear-Time Sequence Modeling with Selective State Spaces，202405
• 论文地址：https://arxiv.org/abs/2312.00752
• 论文代码：https://github.com/state-spaces/mamba
• 论文作者：Albert Gu, Tri Dao
• 研究设计和结论
  • 问题：
    • Transformer的注意力机制虽然有效，但效率很低，因为它需要存储整个上下文(storing the entire context，也就是KV缓存)，导致训练和推理消耗算力大。
    • RNN的推理和训练效率高，但性能容易受到对上下文压缩程度的限制。【写的快，忘的快】
    • SSM的问题在于其中的矩阵A B C不随输入不同而不同，即无法针对不同的输入针对性的推理
  • 方法：平衡一下，既要有抓重点的能力（选择性复制任务），也要有上下文联想/推理能力（诱导头任务）。
    • 选择性SSM：保持时间维度，通过不同的阀门（矩阵 ($B_t,$ $C_t$) 和函数 $({\Delta }_t$)）调节记忆融合和特征捕捉。【大白话就是一个流淌的大水管子，通过设计不同的阀门来调节记忆的融合和特征的捕捉【水暖工旋钮思想】
    • 算法改进：左边是原来的算法，可以看出ABC和$\Delta$都是固定的。现在分别用三个S 函数根据输入把 B、C、$\Delta$都变成了时变的。
      • B：批次大小(Batch size)。表示一次输入的数据量的大小。类似于RNN中的输入门。
      • L：序列长度 (Sequence length)。表示每个序列中包含的时间步数。
      • N：特征维度(Feature dimension)。表示每个时间步的特征数量。
      • D：输入特征维度(nput feature dimension。
      • C：类似于 RNN 中的输出门。
      • A：对应这个维度的SSM来说，A在每个hidden state维度上的作用可以不相同，起到multi-scale/fine-grained gating的作用，这也是LSTM网络里面用element-wise product的原因。
      • $\Delta$：就好比放大镜观察窗口，影响信息处理的焦点。较小的步长Δ会忽略当前输入，而更多地使用先前的上文，而较大的步长Δ会更多地关注当前输入而不是上文。
      • 维度变化具体执行过程：
        • $s_B(x)=Linea{r}_N(x)、s_C(x)=Linea{r}_N(x)$都是线性投影，这是种常见的神经网络操作，用于将输入数据转换到一个新的空间或维度。这里的 Linear 表示是用线性层来学习这几个函数。
        • $s_{\Delta }(x)=Broadcas{t}_D(Linea{r}_1(x))$，广播是一个数操作，它使得维度较小的数组能够与维度较大的数组进行算术操作
        • $\tau \Delta =softplus$，这是个平滑的非线性函数，通常用于网络中以添加非线性特征并帮助网络学习复杂的模式。
          $Softplus(x)=log(1+e^x)$
    • 流体力学与李指数映射：
      • Transformer 描述的是粒子运动，通过自注意力机制映射动态调整每个输入的权重，类似粒子间通过牛顿力学相互作用力来动态调整自己的轨迹。训练的过程，就是在用牛顿力学拟合粒子轨迹，每个输入 (粒子)独立计算与其他输入的关系。
      • Mamba 描述的是流体运动，通过李指数映射来建模时空结构。
        • 流体运动：流体运动描述的是连续介质中的分子集体行为，运动是整体的，内部各点之间有强烈的相互关系和依赖。流体的每个部分都受到整体流体运动的影响，通过内部压力、粘性等因素相互作用。
        • 记忆：记忆系统具有连续性、动态变化性和整体关联性，这些特性与流体的性质非常相似。流体模型能够更好地描述记忆中的信息如何相互关联、如何随着时间和新信息的出现进行动态调整和整合。
        • 李指数映射(Lieexponential map)：李指数映射是一种数学工具，用于描述和分析一个向量场如何沿着另一个向量场发生变化，比如流体力学、电磁场、广义相对论的时空结构等解决了动态系统中相互作用的描述。它是群论和微分几何中重要的概念，来源于李群和李代数的理论，是挪威数学家索菲斯李引入的。训练 mamba 的过程就是用李指数映射拟合流体力学动态系统，找到主管道 A，调整阀门和旋钮 $B_t、C_t、{\Delta }_t$，获得最优流体流动路径，让模型能在高维特征空间中进行高效导航和决策。
          • 把记忆的流淌比作一个水流管道系统，可以看做一个“李群”，进行各种复杂变换(比如旋转、推移等)。
          • 固定矩阵A就是主管道(全局演变路径)，类似于流体运动的全局关系，让系统状态更新有固定的全局路径和规则。
          • $B_t、C_t$就是阀门或旋钮，$\Delta$ 这个离散化因子，就像是流体力学中的时间步长，决定流体运动的离散时间点。
    • 内核融合：把离散化和循环在GPU SRAM 内存中实现，然后加载和存储参数ABC矩阵都用HBM高带宽内存。【Flash Attention技术】
    • 并行扫描：
      • 原因：由于A、B、C这些矩阵现在是动态的了，因此无法使用卷积表示来计算它们(CNN需要固定的内核)，因此，我们只能使用循环表示，如此也就而失去了卷积提供的并行训练能力。
      • 扫描操作(scan operation)：每个状态比如$H_1$都是前一个状态比如$H_0$乘以$\bar{A}$，加上当前输入$X_1$乘以$\bar{B}$的总和。这种状态之下想并行化是不可能的(因为只有在获取到前一个状态的情况下才能计算当前的每个状态)
      • 大白话：每个状态H相当于一个人，看它对应的这条处理线程。自己忙活自己的，吃着碗里看着锅里，当别人的饭做好后，就抢过来用一下，所以某种程度上实现了并行计算。
    • 重计算：避免存储反向传播所需的中间状态，在输入从HBM加载到SRAM时在反向通道中重计算。
    • 网络结构：
      • 经过线性投影后，输入嵌入的维度可能会增加，以便让模型能够处理更高维度的特征空间，从而捕获更细致、更复杂的特征。
      • SSM之前的CNN负责提取局部特征(因其擅长捕捉局部的短距离特征)，而SSM则负责处理这些特征并捕捉序列数据中的长期依赖关系，两者算互为补充
      • CNN有助于建立token之间的局部上下文关系，从而防止独立的token计算。如果每个 token 独立计算，那么模型就会丢失序列中 token 之间的上下文信息。若通过先进行卷积操作，可以确保在进入 SSM 之前，序列中的每个 token 已经考虑了其邻居 token 的信息。这样，模型就不会单独地处理每个 token，而是在处理时考虑了整个局部上下文。
  • 实验结果：
    • 任务对比：coping、selective copying、induction heads
      • (左)复制任务：LTI的效果，输出只能对规则的输入特征进行发现，
      • (右上)选择性复制任务：能自己找重点了，带色的尽管开始间隔大小不一，但都能找出来排好队。
      • (右下)归纳头部任务：联想能力的体现再看到黑的后就想到以前后面应该跟着蓝色的。也就是说，对于非线性时变数据，具备了很强的特征捕捉能力。
    • 合成任务验证选择机制：
      
      • 选择性复制任务要求模型能够记住并复制序列中的特定单词。表一Mamba架构与选择性机制结合后的表现优秀，(S6):准确率为99.8。
      • 扩展序列长度任务要求模型看到一个二元组(如“Harry Potter)时，能够记住“Harry并在序列中再次出现时预测“Potter”。表二Mamba架构，也就是最上面的棕色线，比其他方法要好两倍。
    • 语言模型预训练：
      • 左图是较短的序列长度2048，右图是较长的序列长度8192。两图对比了不同模型，横轴为 FLOPs 计算复杂度由低到高，纵轴为困感度。可以看出Mamba 模型困感度最低。
      • 结论：它是第一个无需注意力机制就能在扩展定律scaling Laws 上匹敌强大Transformer++模型的架构。
    • DNA序列：由于大型语言模型的成功，人们开始探索将基础模型范式应用于基因组学。DNA被视为一种由有限词汇组成的离散序列，需要模型处理长程依赖。实验和图表展示了Mamba 架构在 DNA建模任务中的卓越性能，特别是在处理长序列和扩展模型大小方面。
      • 横轴为参数量，纵轴为困惑度。
      • Mamba (橙色线)的困感度：左图随着参数数量的增加(从约200K到约40M)显著下降，右图随着序列长度保持稳定。
      • 下面是物种 DNA分类任务的微调准确率和不同数据集上的扩展定律。随序列长度，Mamba性能更好。
    • 音频例子：音频波形建模和生成问题也是序列建模任务。
    • 训练推理效率分析：
      • 训练效率：Mamba的扫描实现比标准实现快40倍，处理长序列时时间增长最慢(橙线)。对比蓝线，最右边512 时，1ms /0.025ms = 40倍。
      • 推理效率：Mamba在推理阶段的吞吐量比Transformers高5倍，特别是在大批次处理时，显著优于其他模型。蓝色条的高度明显高于其他模型。
    • 消融实验：${\Delta }_t$是最重要的参数，其次是B和C的组合使用，这个好理解，选择性主要就靠它来确定几个函数$S_B,S_C,S_{\Delta }$。随机初始化表现较好，复杂初始化效果较差。增大N(SSM状态维度)显著改善性能，成本增加微乎其微，但只有在B和C也选择性时才有效。这些发现验证了选择性SSM在语言建模任务中的有效性和优势。
• 论文贡献
  • 创新点：
    • 对输入信息有选择性处理(Selection Mechanism)
    • 硬件感知的算法(Hardware-aware Algorithm)：该算法采用“并行扫描算法”而非“卷积”来进行模型的循环计算(使得不用CNN也能并行训练)，但为了减少GPU内存层次结构中不同级别之间的IO访问，它没有具体化扩展状态
      当然，这点也是受到了S5(Simplified State Space Layers for Sequence Modeling)的启发
    • 更简单的架构：将SSM架构的设计与transformer的MLP块合并为一个块(combining the design of prior SSM architectures with the MLP block of Transformers into a single block)，来简化过去的深度序列模型架构，从而得到一个包含selective state space的架构设计
  • 不足：模型堆叠能力差点意思

MambaOut

自注意力机制的类别

• 因果模式：只能看过去，不能看未来，只有记忆没有未卜先知，适合自回归用来生成和预测，以史为鉴。如GPT。
• 全可见模式：过去未来都可见，适合理解，左顾右看瞻前顾后。如Bert。

Mamba的选择性注意力机制属于因果模式，但和Transformer的因果注意力有区别。

• Transformer的因果注意力（左图）是组合(叠加)之前所有的记忆，记忆无损但复杂度增加，越累越长，计算复杂度同样为0(L²);
• Mamba的因果注意力（右图）是融合之前的记忆到新的隐藏状态，记忆有损但复杂度恒定。

Mamba适用的特征任务

• 任务涉及处理长序列，因为复杂度低，更高效。
• 任务需要因果token 混合模式。

💡此时我们反向思考，Mamba不适合短序列任务，且不需要考虑因果。

那么什么样的任务适合满足这种类型呢？视觉识别任务中的图像分类。

视觉识别任务：

• 图像分类：主要关注整体特征空间特征就够了，目标也只是粗犷的类别标号，因此不涉及什么序列信息，而且需要全局信息。
• 目标检测和语义分割：不一定，比如要考虑边缘的连贯性，因此可能有序列问题。

因果模式没必要：

右图显示以ViT 为例，将自注意力机制从全可见模式切换到因果模式后，性能有所下降，说明对于图像分类问题，用因果模式没必要。

---

图像处理任务不属于长序列：

Transformer的浮点运算次数公式：$FLOPs=24D^{2}L+4D{L}^2$

• L是token 长度 (即输入序列的长度)
• D是通道维度(即特征维度)
• $24D^{2}L$代表线性复杂度
• $4D{L}^2$代表二次复杂度

如何判断它是不是需要长序列建模，只需要看计算量是不是对L敏感。

Transformer的浮点运算次数公式推导：

自注意力部分：【X矩阵：L * D，权重矩阵W：D * D】

• 查询矩阵：$Q=X{W}_q$
• 键矩阵：$K=X{W}_k$
• 值矩阵：$V=X{W}_v$

总的计算复杂度为：$3L{D}^2$

注意力分数计算公式：【Q矩阵：L * D，$K^T$矩阵：D * L，V矩阵：L * D】

总的计算复杂度为：$L^{2}D$ + $L^{2}D$ = $2L^{2}D$

前馈神经网络通常包括两个线性层和一个激活函数计算复杂度为：$4L{D}^2+4L{D}^2=8L{D}^2$

再考虑到多头注意力机制、残差连接、归一化以及前馈神经网络的多次计算等等，最后得到：$FLOPs=24D^{2}L+4D{L}^2$

对于FLOPs(L)，二次项与线性项的比率为：$r_L=\frac{24D^{2}L}{4D{L}^2}=\frac{6D}{L}$，大于1说明二次项的计算复杂度超过了线性项，因此任务涉及长序列建模。

---

图像分类任务推导：

以VIT模型为例：

我们假设输入的图像大小为 224 * 224，patch大小为16 * 16

则 L（token数量）= $(\frac{224}{16}{)}^2$ = 14 * 14 = 196。

每个这样的 patch 通道被展平成一个长向量，RGB三通道就是16163=768，然后通过一个线性投影层(粉色的)映射到高维空间，也就是通道维度 D，它是个指定的超参数。对于 ViT-S，常见的通道维度 D 为384。对图像分类任务L=196，远小于6D=6*384=2304，因此不涉及长序列建模。

---

目标检测和实例分割问题：

在 COCO 数据集上推理图像大小为 800x1280，生成的 token 数约为 4000，大于6D=6*384=2304，因此涉及长序列建模。

MambaOut

“What can I say, Mamba out.” — Kobe Bryant, NBA farewell speech, 2016

• 论文名称：MambaOut: Do We Really Need Mamba for Vision?，20240520
• 论文地址：https://arxiv.org/abs/2405.07992
• 论文代码：
• 论文作者：Weihao Yu, Xinchao Wang
• 研究设计和结论
  • 问题： 解决注意力机制的二次复杂性问题。
  • 方法：通过堆叠 Mamba 块，同时移除其核心标记混合器 SSM，构建了一系列名为 MambaOut 的模型。
    • 左图为MambaOut的架构：
      • 输入图像大小为 $H\times W\times 3$，表示图像的高度、宽度和 RGB 三个颜色通道。
      • 采用了分层架构，共有四个阶段，每个阶段进行特征提取和降采样。
      • 每个阶段包含若干个 GatedCNN 块，用于特征提取。
      • 每个阶段之间有降采样操作，将特征图的大小逐渐减小，从而增加特征的抽象层次。通道维度为 $D_1，D_2，D_3，D_4$。
    • 右图为Gated CNN 块，Gated CNN 块包含两个线性层、中间夹一个卷积层和归一化层，通过残差连接实现输入和输出的融合。和Mamba块的区别在于Gated CNN 块没有SSM(状态空间模型)。
    • MambaOut 的架构与 Swin Transformer 和 DenseNet 在分层结构和降采样方面有相似之处，但在特征提取和信息混合机制上有所不同。
      • MambaOut 使用Gated CNN 块。
      • Swin Transformer 使用窗口注意力机制的Transformer块。
      • DenseNet 则使用密集连接的卷积层。
    • Gated CNN与ResNet的区别：
      • Gated CNN 块使用了线性层进行升维操作，能够在特征空间中进行更灵活的变换，这与传统的 ResNet 中主要使用卷积操作进行特征提取有所不同。
      • Gated CNN 块中跳线增加了非线性激活函数可以被看作一种简单的门控机制，根据输入值调整输出信息量。增加了模型的非线性能力，使得模型能够学习更复杂的特征。
    • 核心代码实现：Gated CNN 块通过线性变换、卷积操作和残差连接，实现了对输入特征的扩展、局部特征提取和信息保留。结合了深度卷积网络和残差网络的优点，同时通过门控机制(如激活函数)来控制信息流。
  • 实验结果：
    • 图像分类比较
      • SSM 有没有对图像分类意义不大
      • 不如最新的CAFormer-M36 使用简单的可分离卷积和原始注意力机制，比所有同等大小的视觉Mamba 模型高出超过1%的准确率85.2%。
    • 目标检测和实例分割
      • 数据集：COCO，Mask R-CNN的主干网络：MambaOut
      • Mamba适合处理长序列视觉任务：尽管MambaOut在COCO 上的目标检测和实例分割任务中可以超越一些视觉 Mamba模型，但它仍然落后于最先进的视觉 Mamba 模型，例如 VMamba 和LocalVMamba。
      • Mamba还需努力：与最先进的卷积-注意力混合模型 TransNext相比51.7%%，视觉Mamba仍表现出显著的性能差距49.2%。
    • 语义分割的比较：
      • SSM 模块在语义分割任务上效果很好，同时也验证了MambaOut在某些情况下的有效性。视觉Mamba 需要进一步展示其在长序列建模任务中的强大性能，以在语义分割任务中实现更强的性能
• 论文贡献
  • 创新点：
    • 定量分析论证了图像分类任务不是长序列建模问题，而目标检测和实例分割是。
    • 借鉴Mamba的GatedCNN 结构微调了 ResNet，实现了一种新型全局可见注意力机制下的改进版模型。
  • 不足：
    • 没有利用矩阵乘法单元，而现代加速器(如 GPU 和 TPU)正是为此而专门设计的。

Mamba2

Transformer is RNN

• 文章：
  • Transformers are RNNs: Fast Autoregressive Transformers with Linear Attention，2020
  • COSFORMER : RETHINKING SOFTMAX IN ATTENTION，2022
• 介绍：引入线性注意力机制（Linear Transformer），将transformer的复杂度从O(N²)降低为O(N)。
• softmax自注意力机制：$Y=softmax(Q{K}^{\top })\cdot V$，其中$Q,K,V\in {\mathbb{R}}^{(1,P)}$，注意力机制需要一次次计算两两token之间的注意力，导致了二次方的计算复杂度。
  • 外层是对于每一个Query，我们需要计算它对应token的新表征。
  • 内层for循环是为了计算每一个Query对应的新表征，需要让该Query与每一个Key进行计算。
  • 大白话就是，军训时，甲乙丙丁4个人列成一队，计算注意力机制的过程相当于每个人都需要站出去算和自己在内所有人的相似度。
• 问题：降低二次方的复杂度。
• 方法：将softmax折叠到核特征映射中，并利用矩阵乘法的结合性将注意力计算中的矩阵左乘改成右乘。即$\left(Q{K}^{\top }\right)\cdot V=Q\cdot \left(K^{\top }V\right)$。
• 推导：
  • 1.将自注意力的计算可以分解为向量运算：$O=\mathrm{softmax}\left(\frac{Q{K}^{T^{\prime }}}{\sqrt{d}}\right)V=\frac{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}\odot v_i}{{\sum }_{i=1}^T{e}^{q_t^{\top }{k}_i}}$
  • 2.假设$sim()$为抽象出的计算Query和Key相似度的函数，用下标i来表示矩阵的第 i 行（如 $Q_i$表示矩阵 Q 的第 i 行）$O_i=\frac{{\sum }_{j=1}^N\mathrm{sim}\left(Q_i,K_j\right)}{{\sum }_{j=1}^N\mathrm{sim}\left(Q_i,K_j\right)}{V}_j$
  • 3.Linear Transformer采用了kernel来定义$sim()$，其中 $\varphi$ 是一个特征映射函数：$\mathrm{sim}\left(Q_i,K_j\right)=\varphi \left(Q_i\right)\varphi {\left(K_j\right)}^T$，则$O_i=\frac{{\sum }_{j=1}^N\left(\varphi \left(Q_i\right)\varphi {\left(K_j\right)}^T\right)V_j}{{\sum }_{j=1}^N\left(\varphi \left(Q_i\right)\varphi {\left(K_j\right)}^T\right)}$
  • 4.根据矩阵乘法结合律：【softmax只能左乘，linear可以右乘，而右乘更快】$\left(\varphi (Q)\varphi (K{)}^T\right)V=\varphi (Q)\left(\varphi (K{)}^{T}V\right)$，最终复杂度转化为：$O\left(N{d}^2\right)$。
  • 5.在一般的NLP任务中，一个头d的特征维度总是比输入序列长度$N(d\ll N)$小得多，因此可以忽略d，实现O(N)的计算复杂度
  • 6.所以，$O_i=\frac{\varphi \left(Q_i\right){\sum }_{j=1}^N\varphi {\left(K_j\right)}^T{V}_j}{\varphi \left(Q_i\right){\sum }_{j-1}^N\varphi {\left(K_j\right)}^T}$
  • 7.我们可以再假设$\begin{array}{l}{S}_i={\sum }_{j=1}^i\varphi {\left(K_j\right)}^T{V}_j=\varphi {\left(K_i\right)}^T{V}_i+{\sum }_{j=1}^{i-1}\varphi {\left(K_j\right)}^T{V}_j=\varphi {\left(K_i\right)}^T{V}_i+S_{i-1}\\ Z_i={\sum }_{j=1}^i\varphi {\left(K_j\right)}^T=\varphi {\left(K_i\right)}^T+{\sum }_{j=1}^{i-1}\varphi {\left(K_j\right)}^T=\varphi {\left(K_i\right)}^T+Z_{i-1}\end{array}$，当需要计算第 i 时刻的输出时，我们可以复用之前的状态 $S_{i-1}$ 和 $Z_{i-1}$ ，再额外加上一个与当前时刻相关的计算量即可。而Transformer在计算第 i 时刻的输出时，它在第i-1个时刻的所有计算都无法被i时刻所复用。
• 结论：Transformer is RNN

SSM is RNN

• 左边就是一个简单的线性时不变系统建模
• 中间是离散化后的模型就是个 RNN
• 最右边是并行化用卷积核进行处理，也就是 CNN 化的模型。

这种表示方式是用图模型来建模，强调的是序列数据之间的依赖关系和动态变化。所谓的 SSM 其实可以理解为就是 RNN，只不过更强调通过线性代数方程来描述系统状态的变化，利用状态空间模型中的状态转移矩阵和观测知阵来进行建模。

对偶

• 介绍：在数学、物理学乃至哲学中，“对偶性”是指两种看似不同的理论或模型之间存在的一种深层次的等价关系。通过这种对偶关系，可以将一个复杂的问题转化为另一个相对简单的问题来解决，或者在一种表示形式下无法轻易看到的性质在另一种表示形式下变得显而易见。
• 思考：那么如何在注意力机制和SSM之间建立统一的对偶关联关系？ => 结构化矩阵

结构化矩阵

• 介绍：如果一个m×n阶矩阵只需要少于m×n个参数来描述，就是一个结构化矩阵(Structured Matrices)。如稀疏矩阵、低秩矩阵、Toeplitz矩阵、Cauchy矩阵、Vandermonde矩阵和蝶形矩阵。
  • Toeplitz 矩阵：每条对角线上的元素都相同的矩阵。
  • Cauchy 矩阵：每个元素都由两个向量的元素之间的差的倒数来定义的矩阵。
  • Vandermonde 矩阵：由一个向量的幂组成的矩阵。
  • 低秩矩阵：其秩远小于其行或列数的矩阵。
• 特性：压缩表示、通过快速算法直接操作这种压缩表示。
• 目的：通过压缩表示可以用更少参数和更快算法计算，减少存储需求，加快运算速度。
• SSM本质上也是一种结构化矩阵：
  • SSM可以表示为 y=Mx 的形式，其中M是ABC的表达式。
    • t = 0 时，$h_0=B_0{x}_0={\sum }_{s=0}^0{A}_{0:s}{B}_s{x}_s$ 【$A_{0:0}$是单位矩阵 𝐼 】
    • t = 1 时，$h_1=A_1{h}_0+B_1{x}_1=A_1{B}_0{x}_0+B_1{x}_1={\sum }_{s=0}^1{A}_{0:s}{B}_s{x}_s$
    • t = 2 时，$h_2=A_2{h}_1+B_2{x}_2=A_2(A_1{B}_0{x}_0+B_1{x}_1)+B_2{x}_2=A_2{A}_1{B}_0{x}_0+A_2{B}_1{x}_1+B_2{x}_2={\sum }_{s=0}^2{A}_{2:s}{B}_s{x}_s$
    • 推广：$h_t=A_t\dots A_1{B}_0{x}_0+A_t\dots A_2{B}_1{x}_1+\cdots +A_t{A}_{t-1}{B}_{t-2}{x}_{t-2}+A_t{B}_{t-1}{x}_{t-1}+B_t{x}_t=\underset{s=0}{\sum ^t}{A}_{t:s}{B}_s{x}_s$
    • $y_t=\underset{s=0}{\sum ^t}{C}_t^T{A}_{t:s}{B}_s{x}_s$
    • $y=SSM(A,B,C)(x)=Mx$
    • SSM的矩阵变换形式：$M_{ji}:=C_j^T{A}_j\cdots A_{i+1}{B}_i$
  • M 具有专门设计的半可分结构，能简化运算。
    • 顺序半可分矩阵 （SSS，SequentiallySemiseparable Structure）：
      • “半”：主要关注下三角部分。
      • “可分”：每个蓝色小块的秩较小，不超过N，意味着可以用更少的独立成分表示，从而实现高效计算。
      • 特点：序列化、下三角、低秩。
    • 即 $y=SSS(A,B,C)\cdot x$，状态空间模型SSM，如果状态的维度为N，等价于一个秩为N的SSS。也就是说任何SSM其实都可以转写成一个等价的局部下对角阵M的形式。
• 所以我们在SSM的计算中，特别是矩阵A，引入了类似注意力机制的公式和方法：
  
  • 1.A的结构从对角线进一步简化为标量乘以单位矩阵结构。 在这种情况下，每个$A_t$也可以仅用一个标量来表示
  • 2.类似Transformer 中多头注意力的概念，使用了更大的头维度 P，相比于Mamba1中使用的 P = 1，通常选择P={64,128}
  • 3.使用类似注意力的对偶形式，去除了 softmax，并引入了一个额外的掩码矩阵 L，根据数据生成，控制信息在时间上的传递量。
    其中圆圈表示元素相乘，也就是哈德马积。$a_i$是依赖于输入的标量，范围在 [0, 1]之间。我们可以假设$a=[a_1,a_2,a_3]$，则
    行标是i，列标是j，i<j的部分全是零，意味着只考虑时间上早于或同一时间点的元素之间的关系。换句话说，它是一种类似 GPT 模型中的单向注意力机制，只考虑过去的时间步，而不考虑未来的时间步。通过这种下三角矩阵，可以有效地控制信息在时间上的流动，确保信息只能从过去传递到现在，而不能反向传播。

张量收缩计算的对偶统一

• SSM计算：SSS的计算过程可以被看作是一系列张量收缩操作，借助顺序半可分矩阵的特殊结构能实现高效计算。所以SSM理论上也可以这样做
  第一步将输入矩阵X 与矩阵 B 进行结合，以产生一个中间结果 Z。矩阵A没有出现，它体现在第二步因状态更新中，L的定义依赖于 A。第三步是最终输出。
• 注意力机制计算：
  • 注意力基本形式：
    • S 和 T 表示源和目标序列长度，分别意指：source、target之意
    • N表示特征维度
    • P表示头维度
  • 自注意力：
    • 源序列和目标序列相同(即 S = T）
    • 通常特征维度和头维度相同(即 N = P)
    • 𝑄, 𝐾, 𝑉是通过对同一输入向量的线性投影生成的，即$Q=W_Q\cdot X,K=W_K\cdot X,V=W_V\cdot X$
  • 掩码(核)注意力 ：
    • $y=(L\circ (Q{K}^{\top }))\cdot V$
    • 将其分解为精确的计算序列：先计算相似性矩阵 $G=Q{K}^T$，应用掩码矩阵 L 后得到新的相似性矩阵 $M=G\odot L$，然后再计算最终的输出$Y=MV$，后面对应的是相应的维度。
    • 写成张量收缩的形式：$Y=contract(TN,SN,SP,TS\to TP)(Q,K,V,L)$
    • 进一步拆解成多步收缩：
  • 线性注意力
    • 用张量收缩表达：$Y=(L\circ (Q{K}^{\top }))\cdot V=Q\cdot cumsum(K^{T}V)$
    • 多步收缩变为了：
  • 结构化掩码注意力SMA：
    • 无论是从状态空间模型(SSM)侧，还是从注意力机制侧来看，都可以统一到张量收缩的视角下进行操作。
    • 不同的掩码矩阵（如因果掩码、衰减掩码、Toeplitz 矩阵等）L 定义了不同的序列变换矩阵 M，从而实现不同形式的结构化注意力。

于是，Mamba-2孕育而生。

Mamba2

• 论文名称：Transformers are SSMs: Generalized Models and Efficient Algorithms Through Structured State Space Duality，202405
• 论文地址：https://arxiv.org/abs/2405.21060
• 论文代码：https://github.com/state-spaces/mamba
• 论文作者：Tri Dao, Albert Gu
• 研究设计和结论
  • 问题： 既然SSM和注意力机制两种对偶等价，那就研究如何结合起来进行计算。
  • 方法：整体上是SSM，但是通过块分解把大矩阵拆解成小的子矩阵，每个小问题特别是低秩块，再用注意力机制计算，利用矩乘法上的高效性和并行计算能力使得计算过程更加高效。【分解块，独立运算，并行计算】
    • 如下图所示，一个大的 M矩阵，分解成9 块，其中蓝色块用矩阵乘法。
    • 半可分矩阵 MMM 被分解成多个矩阵块：
      • 对角块(Diagonal Block)表示输入到输出的计算，
      • 低秩块 (Low-Rank Block)分成三类：
        • 从输入到状态 (lnput State)。
        • 从状态到状态 (State - State)。
        • 从状态到输出 (State ->Output)。
    • 这种块分解方法进行计算的流程。
      • 输入序列 X被分解成多个块，每个块对应图中的一个黑色虚线框。
      • 输入块通过低秩块 (绿色头)和对角块(橙色头)进行计算，得到中间的状态块 H。
      • 状态块之间通过低秩块(黄色箭头)进行计算，表示状态间的传递。
      • 最终，状态块通过对角块 (蓝色箭头)计算得到输出块 Y。
    • 核心代码
      • segsum(x)：用于计算分段累加和。
      • ssd(X，A，B，C，block len=64,initial_states=None)：用于计算SSD 模型，其中A，B，C分别表示状态矩阵、扩展矩阵和收缩矩阵
        • 0.先对输入张量 X、A、B和C进行重排，将它们重排成块的形式
        • 1.对每个块内的对角块 (Diagonal Block)进行计算，使用 torch.einsum 计算块内的矩阵乘法。
        • 2.计算每个块内的低秩块 (Lw-Rank Block)，用于生成下一个块的输入状态。
        • 3.生成块间的状态转移，确保在块边界处的状态正确。
        • 4.对块内的低秩块进行计算，将状态转换为输出。
        • 5.最后将块内和块间的输出汇总，得到最终输出Y和最终状态 final state。
      • 总结：通过块分解的方法，将一个大规模的状态空间模型问题分解成多个小规模的块级别运算问题。这种方法利用了半可分矩阵的特性，能够提高计算效率和并行性，适合硬件加速。
    • Mamba-2 架构：
      • Mamba-1基于SSM (状态空间模型)设计，线性投影之后生成SSM 参数 A，B，C。
      • Mamba-2的两种块设计：
        • 顺序 Mamba 块。SSM 层从 A，X，B，C 直接映射到输出Y，在序列变换中并行生成参数 A，B，C。
        • 并行 Mamba块。SSM 层从 A，X，B，C 直接映射到输出Y，在块开始时并行生成，适合更大规模的并行处理。这种方法类似于标准注意力架构中的并行生成 Q,K,V。这种设计能够减少参数数量，适合更大模型的张量并行计算。
        • 每个 Mamba 块中增加额外归一化层，改善模型稳定性。
        • Mamba-2 模型通过并行化和增加归一化层来优化原始 Mamba 模型的计算效率和稳定性。
      • 并行化处理方法
        • 张量并行（Tensor Parallelism，简称TP）
          • 左侧：输入和输出投影矩阵分割，并在单个设备上处理。每个SSM头(即A、B、C、X)到Y 的映射都在单个设备上进行。最终归一化层选择 GroupNorm，以避免额外的通信。
          • 右侧：将序列维度上的计算分配到多个设备上，每个设备负责一部分序列的计算，然后将结果传递给下一个设备。
        • 序列/上下文并行（Sequence/Context Parallelism）：对于非常长的序列，可能需要沿序列长度维度将输入和激活拆分到不同的GPU上。
          • 用于残差和归一化操作的序列并行(SP)。
          • 序列并行用于token混合操作(注意力或SSM)，也称为“上下文并行”(context parallelism，简称CP)。
  • 实验结果：
    • 合成记忆任务：该图展示了不同模型在多查询关联记忆(MQAR)任务中的表现。
      • 三张子图对应不同的序列长度(256、512、1024)、横轴代表模型的维度 (32、64、128、256)、纵轴代表准确率(从0到1)。
      • Mamba-2 系列模型在较大的模型维度下表现优异，特别是当维度达到128和256时其准确率接近1.0。
      • Mamba-2 模型明显优于 Mamba-1 和普通注意力模型，尤其在更大的状态规模(N=256)下表现尤为显著。
    • 语言模型预训练和评估
      • 在The Pile 上进行训练的模型，Mamba-2的性能匹配或超过了Mamba 和强大的“Transformer++”方案。
      • 与Transformer 基线相比，Mamba-2 在性能(困感度)、理论FLOPs和实际壁钟时间上都是帕累托占优的。
    • 零样本评估
      • 在每种模型规模中，Mamba-2 模型的表现普遍优于其他模型。特别是Mamba-2在较大的模型规模下(27B 参数)表现尤为突出，证明其在不同任务上的泛化能力更强。
      • 在不同数量的注意力层下的困感度。大约10%的注意力层比例表现最佳。适量的注意力层可以显著提高模型性能，超过了完全不使用注意力层或完全使用注意力层的情况。
      • 比较了SSD、MLP和注意力层的不同组合方式，在2.7B 规模上进行评估困感度(ppl)和准确率 (acc)。Mamba-2 与注意力层的结合 (后4个)在多个任务上的表现优于其他模型组合，显示出更强的泛化能力和任务适应性。
    • 速度性能
      • 左图：不同方法在处理序列长度(从512到512k)时所需的时间 (以毫秒为单位)。
      • 右图：展示了在处理固定序列长度(4K)时，不同状态维度(从16到56下所需的时间(以毫秒为单位)。
      • SSD 方法在处理大状态扩展时表现优异，比Mamba 的融合扫描快2到8倍(比如64k时紫色线1毫秒，Mamba为10 毫秒)，并且在序列长度超过2k 时也比 FlashAttention-2更快。
    • 消融实验
      • Mamba-2 模块在结合并行处理和额外归一化后，显著提升了模型性能，表现优于传统的Mamba-1 模块。
        
      • 多头结构中，复杂的头组合和状态扩展通常可以提高性能，特别是当模型规模增大时对于核近似，Swish和LayerNorm 方法通常效果较好，且适用于不同规模的模型
      • 增加复杂度和头的数量一般有助于提高模型性能，但需要权衡参数数量的增加。
• 论文贡献
  • 创新点：将 SSM 模型与注意力机制结合，实现了模型并行堆叠，提升了模型的堆叠能力。大白话就是一个流淌的大水管子，通过设计不同的阀门来调节记忆的融合和特征的捕捉。【水暖工旋钮思想】

xLSTM

LSTM

• 大白话理解：因记忆能力有限，记住重要的，忘记无关紧要的。
• RNN —> LSTM：
  
• tanh 激活函数的作用：帮助调节流经网络的值，使得数值始终限制在 -1 和 1 之间。
  
• Sigmoid 激活函数的作用：帮助调节流经网络的值，使得数值始终限制在 0 和 1 之间。
• 公式：
• Forget Gate（遗忘门）：
  • 
• Input Gate（输入门）：
  • 
• Cell State（细胞状态）：
  • 
• OutPut Gate（输出门）：
  • 
• 贡献：一定程度上解决梯度消失和梯度爆炸的问题
• 问题：
  • 在处理长序列时效率低；
  • 记忆容量有限；
  • 不能并行处理数据；

sLSTM

• 公式：
  
  • 输入门和遗忘门的激活函数从 sigmoid 改成了指数函数(红色部分)
    • 指数函数相比于 sigmoid 函数，具有更大的输出范围和更大的梯度(右图黄色，左图红色)，可以减轻梯度消失问题，使得梯度在反向传播过程中不会迅速减小，从而使得模型在训练时能够更有效地更新权重。
    • 指数函数的增长速度比 sigmoid 函数快，对输入变化更加敏感。因此，可以更迅速地强烈的调整输入和遗忘门的输出，使得模型能够更快地捕捉到输入信息的变化，更加选择性地记住或忘记信息，从而提高模型的记忆和遗忘能力。
    • 这种强烈的选择性，让模型能够更准确地保留重要信息和丢弃不重要的信息。在特定任务(如长序列的最近邻搜索或稀有事件预测) 中表现得尤为显著，能够显著提升模型性能。
  • 引入了归一化状态 $n_t$(公式9)，相当于搞了一个大分母，因为指数激活函数可能导致数值过大而溢出。
  • 相应的隐层$h_t$的计算方式变了，改成了$c_t/n_t$;也就是公式 (10)
  • 引入了一个额外状态 m_t 来进一步稳定门控
    • 式子（15）使用了 log，指数函数的逆运算，相当于降一级运算，然后取最大值，意思就是输入门和遗忘门都别太猛。
    • 根据 $m_t$ 再调整输入门和遗忘门，相当于设置了一个缓冲区。【忍一时风平浪静，退一步海阔天空】
• 贡献：解决了敏感度，某种程度上也是长序列处理效率问题。

mLSTM

• 公式：
  • 状态和权重参数都变成了矩阵形式，对应的运算变成了向量矩阵乘法和哈达玛积，公式（21）。
  • 增加了$q_t,k_t,v_t$这种键值对的计算公式（22-24），优化了子注意力机制，多了好几个权重模型增强了模型表达能力。
  • 这种框架可以使用多头模式，头与头之间没有记忆混合，因此可以充分并行，无形中提升了并行能力。
• 贡献：增强LSTM的存储能力。

xLSTM

线性不可分

• 设输入空间为X={(0,0),(1,1),(0,1),(1,0)}，显然输入空间是二维的。对于异或问题，(0,0),(1,1)的值为0（图中用o表示），(0,1),(1,0)的值为1（图中用×表示），显然没有一条直线可以解决异或问题的分类。
• 现在我们将输入空间中的任一点$(x_1,x_2)$，转换为$(x_1,x_2,(x_1-x_2{)}^2)$，称转换后的空间为特征空间，记为Z，则Z={(0,0,0),(1,1,0),(0,1,1),(1,0,1)}。显然特征空间是三维的。(0,0,0),(1,1,0)的值为0（图中用o表示），均在底平面；(0,1,1),(1,0,1)的值为1（图中用×表示），均在顶平面。显然，任意一个在底平面和顶平面之间的平面可以解决异或问题的分类。

核方法 (Kernel Method)

• 动机：解决分类任务中线性不可分的问题。
• 方法：将数据映射到更高维的空间，将原本在低维空间中线性不可分的数据转换为在高维空间中线性可分的数据，以大大降低分类任务的难度。
• 问题：更高维的数据就更加难以训练，也更容易过拟合。
• 应用：Transformer 模型就是通过多头注意力机制在高维空间中进行并行处理，使得不同位置的特征可以相互影响和结合，从而提高了模型的性能。

那么xLSTM如何升维？

• “先干后变”：先在原始空间中总结信息，然后映射到高维空间，再返回原始空间。看图从下往上输入sLSTM，然后向上投影，也就是用一个倒着的梯形矩阵升维，处理后再降维。

  • 内部结构：
    • 输入先LN（LayerNorm）整理，然后一分为二。一部分卷积提取特征，激活非线性变换（左），另一部分直接输入sLSTM（右）。
      • LN（LayerNorm）：层归一化，帮助稳定和加速训练过程。
      • Conv4：卷积层，卷积核大小为 4。提取局部特征。
      • Swish：一种平滑的非线性激活函数，可以帮助模型学习到更复杂的模式。
    • 接着所有运算都采用了4个头的多头并行进制，每个头可以专注于捕捉输入数据的不同特征或模式，从而使模型能够更全面地理解数据。
      • 内部采用块对角线结构，在计算时可以并行处理，从而显著降低计算复杂度和内存需求；
      • 每个子矩阵（块）主要关注输入数据的一部分能够更好地捕捉局部特征；结构化的稀疏性，这有助于减少过拟合。
    • 在 sLSTM 图中的箭头表示信息在不同时间步之间的流动和处理，代表的是与先前时刻状态的混合计算。这部分相当于记忆的重新组合。
    • 然后组内归一化、降维、再激活、再降维，然后与残差相加再输出。
      • GN（GroupNorm）：组归一化（Group Normalization）。在每一组内进行归一化，有助于加速训练和提高模型稳定性，特别是在小批量（batch）训练时。
      • PF=3/4 和 PF=4/3：投影因子（Projection Factor），分别将输入维度缩小为原来的 3/4，将输入维度扩大 4/3 倍。
      • GeLU（Gaussian Error Linear Unit，高斯误差线性单元）也是一种激活函数，数学表达式为：$GELU(x)=x\Phi (x)$，其中Φ(x)是输入x的标准正态累积分布函数（CDF），具体公式为：
        图像为：
        
        特性：
        • 平滑性：与ReLU的尖锐转折点不同，GELU提供了一个平滑的激活曲线，这有助于深度学习模型在训练过程中更加稳定，特别是在处理不连续输入数据时。
        • 非饱和性：GELU与ReLU一样，具有非饱和性质，这意味着它可以缓解梯度消失问题，特别是在训练深层网络时。
        • 自适应门控机制：GELU通过其内部的高斯CDF组件，实现了一种自适应的门控机制。这意味着GELU可以根据输入的属性自动调整激活的量，类似于神经元的开/关切换，这有助于网络自动学习重要的特征。

• “先变后干” ：先映射到高维空间，总结信息后再返回原始空间。也就是输入直接上投影，再用mLSTM处理，然后再降维。

  • 高维空间中的记忆容量更大，因此用有矩阵化记忆单元的mLSTM更合适，而在低维空间处理sLSTM更合适。
  • 内部结构：
    • 整体上都是充分利用了残差堆叠结构，层归一化技术等稳定网络，通过升降维度实现空间变换，激活函数非线性变换，然后利用 LSTM 进行记忆混合，或者说时序上的选择性自注意力机制计算，采用多头和块对角模式实现并行处理，当然也没少了用卷积提取特征。
    • PF=1/2 和 PF=2：投影因子（Projection Factor）。前者将输入维度缩小一半，后者将输入维度扩大两倍。
    • LSkip 是个跳线，类似于残差连接，可以帮助梯度更好地传递，防止梯度消失和爆炸。这里相当于有两种跳线残差。
    • BS=4：块大小为 4 的块对角投影矩阵。
    • mLSTM 单元中的q、k、v分别表示查询（query）、键（key）和值（value），我们刚讲过，都是从输入中生成的，用于计算注意力权重和进行信息检索。

实验论证：

• 合成任务和长程任务
  • 每行表示一种模型，包括Llama、Mamba等7种模型的12中变体，xLSTM[0:1]：主要是sLSTM块；xLSTM[1:0]：主要是mLSTM块；xLSTM[1:1]：均衡使用mLSTM和sLSTM块。
  • 每列表示一种任务，包括上下文敏感、确定性上下文无关、正则，最后是多数任务，也是正则。
  • 使用sLSTM和mLSTM的组合（如xLSTM[1:1]）在大多数任务上表现出色，特别是在
    复杂和状态跟踪任务上。
  • 不同模型在多查询联想记忆任务中的性能对比：
    
    • 图表说明：横轴模型的尺寸，纵轴验证准确率。
    • 实验结果：xLSTM[1:1]（粉色）表现最佳，越难越好，Llama等Transformer模型在较小和中等难度任务中表现优越。Mamba略强。
• 验证集困惑度比较
  • 数据集：15B个Token训练的 SlimPajama。
  • 比较内容：下一词预测性能。
  • 图表说明：横轴为模型参数量，纵轴为验证困惑度。
  • 实验结果：总体趋势都差不多，但xLSTM明显更好。说明其在语言建模任务中的优势。
• 大规模语言建模实验
  • 数据集：300B个Token训练的SlimPajama
  • 比较内容：不同模型在下一词预测任务中的验证困惑度（Validation Perplexity），特别是对长序列的外推性能。
  • 图表说明：横轴为token数量，也就是序列长度。纵轴为验证困惑度。
  • 实验结果：
    • 大多数模型在序列长度增加时验证困惑度上升。
    • xLSTM模型在较长序列上保持了较低的验证困惑度，显示了其在处理长序列时的优势。
• 语言基准测试
  • 数据集：300B个Token训练的SlimPajama
  • 比较内容：不同模型在下一词预测任务中的验证困惑度（Validation Perplexity）随参数数量变化的情况。
  • 图表说明：验证困惑度越低，表示模型的预测性能越好。
  • 实验结果：
    • 所有模型的验证困惑度随着参数数量的增加而下降，说明更大参数的模型在下一词预测任务上表现更好。
    • xLSTM模型（特别是xLSTM[7:1]和xLSTM[1:0]）在所有参数数量下都表现出色，验证困惑度较低，说明其在语言建模任务中的性能优越。
    • xLSTM模型比Mamba表现好，而Mamba比Llama表现好。这表明xLSTM在处理大规模语言建模任务时，具有明显的优势。

代码实现：https://github.com/AI-Guru/xlstm-resources?tab=readme-ov-file

创新点：

借助指数门控混合记忆和新内存结构，LSTM增强为sLSTM和mLSTM。二者的结合构成了xLSTM模块，进一步堆叠可以实现大模型化。

TTT

• 论文名称：Learning to (Learn at Test Time): RNNs with Expressive Hidden States，202407
• 论文地址：https://arxiv.org/abs/2407.04620
• 论文代码：https://github.com/test-time-training/ttt-lm-pytorch
• 论文作者：Yu Sun, Xinhao Li, et al.
• 关键词和摘要总结
  • 关键词：Test-Time Training (TTT), RNNs, Sequence Modeling
  • 摘要总结：提出TTT层，该层在保持线性复杂度的同时，通过将隐藏状态视为机器学习模型，并使用自监督学习更新规则，以提高长上下文建模的表达力。
• 研究设计和结论
  • 问题： 虽然RNN有着线性的复杂度，但RNN在面对长下文时处理起来会比较困难。【RNN层必须将上下文压缩成固定大小的隐藏状态，其作为一种压缩启发式方法，更新规则需要发现成千上万甚至数百万个token之间的潜在结构和关系】
  • 方法

    • TTT层：

      • 将固定大小的隐状态变量 ( $s_t$) 更换为小的 MLP 网络，提升了表达能力和自适应性。
      • 由传统的线性组合加激活函数更新，变为梯度下降法，实现边测试边学习。
      • 输出规则：$z_t=f\left(x_t;W_t\right)$
        • 大白话理解：输出token $z_t$基于由 f 使用更新后的权重$W_t$进行预测 $x_t$。
      • 更新规则：$W_t=W_{t-1}-\eta \nabla \ell \left(W_{t-1};x_t\right)$
        • 引入并行化：因为$W_t$在两个地方依赖于$W_{t-1}$无法并行化，一个是在减号前，一个是在$\nabla l$内部。
        • 问题核心：$\nabla l$内部包含了大部分计算，需要重点对$\nabla l$做并行化
        • 梯度下降GD：
        • 批量梯度下降：对所有这些变量相对于$W_0$进行计算，$G_t=\nabla \ell \left(W_0;x_t\right)$，但$W_t$实际上只比$W_0$多一步梯度步长，导致批量梯度下降的有效搜索空间会相对比较小。
        • 在线梯度下降：在线梯度下降，$W_t$距离$W_0$有t步之遥。
        • 小批量梯度下降：在批量梯度下降与在线梯度下降折中，即将批量大小设置为相对较小的b。$G_t=\nabla \ell \left(W_{t^{\prime }};x_t\right)$，其中$t^{\prime }=t-\mathrm{mod}(t,b)$。

    • 训练：内外环嵌套训练

      • 外环：RNN 大网络，负责整体模型参数优化。包含一个 Task 对象，定义了模型参数和损失函数 MSE。 forward方法接收输入序列并遍历每个标记，调用 Learner 对象进行训练和预测。在整个网络的训练过程中， Task 类的参数（${\theta }_K$ 、 ${\theta }_V$ 、${\theta }_Q$ ）会被优化，以提高模型的整体性能。借鉴了 QKV 模式，但本质上并不是自注意力机制。
      • 内环：MLP 小网络，负责具体训练和预测。在每个时间步内，Learner 对当前输入进行训练，计算损失函数的梯度并更新模型参数。 Learner 使用更新后的模型对当前输入进行预测，并返回预测结果。

    • 模型变体：

      • 参数化模型：尝试了简单的线性模型和多层 MLP（Y轴），优化器方面尝试了使用所有数据的批量梯度下降和小批量梯度下降法（X轴）。 TTT-linear 在使用批量梯度下降（Batch GD）时， TTT-Linear 的输出等价于线性注意力。
        • 对于TTT-Linear，$f_{lin}(x)=Wx$，其中W是方阵。
        • 对于TTT-MLP，$f_{MLP}$有两层，类似于Transformer中的MLP。
        • 为了在TTT期间获得更好的稳定性，f总是包含层归一化(LN)和残差连接。即，$f(x)=x+\mathrm{LN}\left(f_{\mathrm{res}}(x)\right)$，其中$f_{res}$可以是$f_{lin}$或$f_{MLP}$。
      • 非参数化模型：如 Nadaraya-Watson 估计器，理论上等价于自注意力机制。【唬人的，连固定的模型和参数都不要，就是一种
        非线性加权求和运算，根据训练数据直接进行预测。】
  • 实验结果：TTT-Linear 和 TTT-MLP 在长上下文处理和推理延迟方面表现优于 Transformer 和 Mamba 模型。
    • Pile 数据集：左图 2k 上下文长度上，蓝色的 TTT-Linear（M）和红色 Mamba 表现相当，橙色的 TTT-MLP（M）略差。右图 8k 上下文长度：蓝色 TTT-Linear（M）和橙色 TTT-MLP（M）表现优于红色 Mamba。
    • Books 数据集：
      • 左图2k 上下文长度：蓝色 TTT-Linear（M）和橙色 TTT-MLP（M）表现优于红色Mamba 和紫色 Transformer。 右图32k 上下文长度： TTT-Linear（M）和 TTT-MLP（M）在长上下文。
      • 1k 到 32k 上下文长度所有上下文长度： TTT-Linear 和 TTT-MLP 表现稳定，长上下文下优于 Transformer 和 Mamba。
    • TTT-MLP 的表现，橙色短上下文略差，越往右长上下文表现更好下表现更好
    • NVIDIA A100 GPU 上的延迟：
      • 前向延迟：蓝色 TTT-Linear 和红色 Mamba 延迟相同，绿色 Transformer 延迟较高。
      • 生成延迟：蓝色 TTT-Linear 和红色 Mamba 延迟低于绿色 Transformer 和橙色 TTT-MLP。
• 论文贡献
  • 创新点：
    • 引入了边测试边学习的框架，提升了自适应能力。
    • 结合内外环模型的思路，实现了模型参数的宏观和微观优化。
    • 隐状态的神经网络化，增强了模型的表达能力和动态调整能力。
  • 不足：
    • 新的学习模式带来了新的问题，如初始化、学习时长和边学边干中的冲突和矛盾需要进一步研究
    • 非参数化 TTT 模型尚未进行实际的实验验证。

📜 参考资料 :

https://zhuanlan.zhihu.com/p/716838305
https://space.bilibili.com/1921388479
https://blog.csdn.net/v_JULY_v/article/details/134923301
https://blog.csdn.net/v_JULY_v/article/details/132178447

📌 [ 笔者 ]   文艺倾年
📃 [ 更新 ]   2024.9.5
❌ [ 勘误 ]   /* 暂无 */
📜 [ 声明 ]   由于作者水平有限，本文有错误和不准确之处在所难免，
              本人也很想知道这些错误，恳望读者批评指正！