点击名片
关注并星标
#TSer#
本文来聊聊 xLSTM 具体做了哪些改进创新。
【论文标题】xLSTM: Extended Long Short-Term Memory
【论文地址】https://arxiv.org/abs/2405.04517
【开源代码】https://github.com/NX-AI/xlstm
LSTM回顾与局限
在20世纪90年代，长短期记忆网络（Long Short-Term Memory, LSTM）的核心思想——恒定误差旋转木马和门控机制被引入，用以解决循环神经网络中的梯度消失问题。自那时以来，LSTM 已经经历了时间的考验，并为众多深度学习的成功案例做出了贡献，特别是它们构成了第一批大型语言模型（LLMs）。
原始 LSTM 的计算涉及几个关键组件：输入门、遗忘门、输出门和单元状态。
图1：LSTM 单元的标准计算步骤
LSTM 在处理最近邻搜索问题时难以修订其存储决策，这限制了其在动态更新信息方面的能力。 LSTM 的存储容量有限，这导致它在处理稀有标记预测任务时表现不佳，因为信息必须被压缩进标量单元状态中。 由于隐藏状态之间的连接，LSTM 缺乏并行处理的能力，这影响了其在大规模数据处理中的效率。
xLSTM的创新与架构
xLSTM 的关键在于“扩展”这个词，它不仅仅是对 LSTM 的简单改进。xLSTM 将原始 LSTM 扩展为 LSTM 的多种变体，例如 sLSTM 和 mLSTM，每种变体都针对特定的性能和功能进行了优化，以处理各种复杂的序列数据问题。
总体而言，xLSTM 主要通过两方面改进来增强LSTM的能力：（1）引入指数门控；（2）引入新型记忆结构（sLSTM 和 mLSTM）
图2：xLSTM家族
01
指数门控
这是对 LSTM 中传统门控机制的一种改进，允许模型更有效地更新其内部状态。指数门控通过引入适当的归一化和稳定技术，使得 LSTM 能够更好地处理信息流，特别是在需要修订存储决策的场景中。
02
新型记忆结构
（1）sLSTM
sLSTM 在 LSTM 的基础上添加了标量更新机制。该设计通过对内部存储单元进行细粒度控制来优化门控机制，使其更适合处理具有细微时间变化的序列。
图3：sLSTM的前向传播方程
sLSTM 具备了跨单元的内存混合能力，允许不同存储单元之间进行信息交互。sLSTM 还支持多头部结构，每个头部可以独立处理信息，但头部之间并不直接混合内存，而是在每个头部内部的单元之间进行内存混合。这种设计为模型提供了一种新的内存混合方式，有助于捕捉更复杂的序列特征。
（2）mLSTM
mLSTM 通过将原始 LSTM 中的向量运算扩展为矩阵运算，显著增强了模型的内存容量和并行处理能力。
为了提高计算效率并实现并行处理，mLSTM 特别放弃了传统的隐藏状态间的递归连接，即内存混合。对于 mLSTM 而言，多头结构和多单元结构在功能上是等价的，这意味着它可以灵活地根据需要选择使用哪一种结构来优化性能和计算效率。
图4：mLSTM的前向传播方程
03
xLSTM架构
图5：xLSTM blocks
xLSTM 块应在高维空间中对过去进行非线性总结，以便更好地分离不同的历史或上下文。分离历史是正确预测下一个序列元素（如下一个 token）的先决条件。研究者在此采用了 Cover 定理。该定理指出，在高维空间中，非线性嵌入模式比在原始空间中更有可能被线性分离。
post up-projection 的残差块（如Transformer），它在原始空间中非线性地总结过去，然后线性映射到高维空间，应用非线性激活函数，并线性映射回原始空间。（见图5的左面板和图2的第三列）
pre up-projection 的残差块（如状态空间模型），它线性映射到高维空间，在高维空间中非线性地总结过去，然后线性映射回原始空间。对于包含 sLSTM 的 xLSTM 块，主要使用 post up-projection 块。对于包含 mLSTM 的 xLSTM 块，使用 pre up-projection 块，因为在高维空间中内存容量变得更大。
与 Transformer 不同，xLSTM 网络在计算上具有线性复杂度，在内存上具有与序列长度相关的恒定复杂度。由于 xLSTM 内存具有压缩性，因此非常适合工业应用和边缘部署。
xLSTM性能分析
研究者的实验结果表明，xLSTM 在多个任务上展现出了卓越的性能，包括在合成任务、长距离竞技场（Long Range Arena）以及大规模语言建模实验中的表现。这表明 xLSTM 能够有效地处理长序列和大型上下文，并且在语言建模方面具有强大的能力。
xLSTM 通过引入指数门控和矩阵内存结构，显著提高了模型的扩展性和内存容量。在多查询关联记忆（Multi-Query Associative Recall）任务中，xLSTM 表现出了优秀的记忆能力，尤其是在需要记忆大量键值对时。
图6：xLSTM的指数门控与记忆混合测试
xLSTM 在与 Transformers 和 State Space Models（状态空间模型）等现有技术的比较中表现出色。在 150 亿个 token 的训练数据上，xLSTM 实现了最低的验证集困惑度（perplexity），并且在 300 亿个 token 的训练数据上，xLSTM 在不同模型尺寸下仍然保持了优异的性能。据此研究人员认为 xLSTM 未来有望能够与 Transformer 进行 PK。
表1：在SlimPajama数据集上训练的不同语言模型的比较
总结
扫下方二维码，加入时序人学术星球
星球专注于时间序列领域的知识整理，前沿追踪
提供论文合集、视频课程、问答服务等资源
220+篇专栏笔记，已有210+小伙伴加入
价格随着内容丰富而上涨，早入早享优惠哦~
时间序列学术前沿系列持续更新中 ⛳️
后台回复"讨论"，加入讨论组一起交流学习 🏃
往期推荐阅读
觉得不错，那就点个在看和赞吧
