xLSTM模型学习笔记

笔记来源：bilibili

LSTM 回顾

原始的 LSTM 是为了解决 RNN 时序反向传播中梯度消失和爆炸问题而提出的。

其所谓的门控机制，其实就是一种时序上的注意力机制，相当于把不同时间进行"掺和"，是对时序信息的一种选择性控制。从这个视角看，它与 Transformer 和 mamba 都异曲同工之妙。核心思想都是选择性控制信息流动，更好地处理时序数据或序列信息。门控机制通过固定的结构和参数来控制信息流，而注意力机制通过动态计算权重来控制信息流。因此，门控机制可以看作是一种特定形式的时序注意力机制，对不同时间步的信息进行选择性控制和"掺和"。可以认为是一种约束版或者简化版的注意力机制。

其缺点

LSTM 有三个主要局限性：

在处理长序列时效率低
记忆容易有限
不可以并行处理数据

而 transformer 借助网格模块堆叠，参数规模扩充和 GPU 并行处理拼算力，在一定的程序上解决了以上的问题，所以 transformer 实现了赶超。

初级版：sLSTM 改进注意力机制

改进的点：

输入门和遗忘门的激活函数从 sigmoid 改成了指数函数（红色部分）
引入了一个归一化状态 $n_t$ ，相应的隐层 $h_t$ 的计算方式亦了，改成了 $c_t/n_t$ （公式 10）
引入了一个额外状态 $m_t$ 来进一步稳定门控

改进的原因如下：

指数函数相比于 sigmoid 函数，具有更大的输出范围和更大的梯度，可以减轻梯度消失的问题，使得梯度在反向传播过程中不会迅速减小，从而使得模型在训练时能够更有效地更新权重。
指数函数的增长速度比 sigmoid 函数快，对输入变化更加敏感，因此，可以更迅速地强烈的调整输入和遗忘门的输出，使得模型能够更快地捕捉到输入信息的变化，更加选择性地记住或忘记信息，从而提高模型的记忆和遗忘能力。
强烈的选择性可以让模型能够更准确地保留重要信息和丢弃不重要的。在特定任务（如长序列的最近邻搜索或稀有事件预测）中表现得尤为显著，能够显著提升模型的性能。

引入归一化和状态 $m_t$ 都是为了稳定，因为指数激活函数可能导致数值过大而溢出，前者相当于搞了一个大分母。后者通过下面的公式进行：

第一个式子中，使用了 log 的作用就是防止输入门和遗忘门都不要太猛。然后根据 $m_t$ 来调整输入了门与遗忘门，这样就实现了对输入门与遗忘门的调节。

在原论文中，还证明了在瘦身传播中使用 $f'_t$ 和 $i'_t$ 替换 $f_t$ 和 $i_t$ 不会改变整个网络的输出，也不会改变参数损失的导数。

增加了这些公式相当于增加了新的记忆单元，它们之间通过连接从长短期记忆状态，借助门控（阀门）i, f, 0 进行记忆混合。门控就是选择，也是一种时序注意力机制的体现。

中级版：mLSTM 改进内存处理

解决了敏感度，某种程序上也是长序列处理效率问题，为了增强 LSTM 的存储能力，文章将 LSTM 的记忆单元从一个标量 c 增加到矩阵 C。而且在这里引入了 transformer 键值对的概念，更新规则如下：
$C_t=C_{t-1}+v_tk_t^T$
在将输入投影到键和值之前，mLSTM 进行层归一化，使得均值为零。同时，将协方差更新规则，也就是优化器整合到 LSTM 架构中，遗忘门对应于衰减率，输入门对应于学习率，而输出门则缩放检索到的向量，最终形成了下面的迭代公式：

与之前的 sLSTM 对比，最大的区别之一就是状态和权重参数都变成了矩阵的形式，对应的运算变成了向量矩阵简洁和哈达玛积。区别之二是增加了 $q_t$ ， $k_t$ , $v_t$ 这种键值对的计算公式，优化了自注意力机制，多了好几个权重矩阵增强了模型表达能力。其他的公式基本没变。相当于，记忆单元没变，只是每个单元扩容了记忆的容量。

此外，这种框架可以使用多头模型，头与头之间没有记忆混合，因此可以充分并行，可以提升并行能力。

高级版：xLSTM 大模型

Cover 定理

Cover 定理：它及衍生的高维空间中非线性映射理论是现代大模型设计的重要理论依据之一。尤其是在深度学习和大规模神经网络的设计中，直到了关键的作用。

大模型中，激活函数通过非线性变换将数据映射到高维空间，使得模型可以捕捉复杂的模式和特征，增强模型的表达能力。深度网络在权重矩阵和激活函数共同作用，将输入数据逐步映射到越来越高的维度。这使得在低维空间难以分离的模式在高维空间国变得线性可分。Transformer 模型就是通过多头注意力机制在高维空间中进行并行处理，使得不同位置的特征可以相互影响和结合，从而提高了模型的性能。

Cover 定理为这些设计提供了理论支持，解释了为什么通过高维空间国的非线性映射可以提高模型的性能。