摘要： 由LSTM奠基者Sepp Hochreiter领衔的团队，近日推出了经过重大优化的xLSTM 7B模型。该模型在推理速度上超越了同等规模的Mamba模型高达50%，并实现了权重和代码的全面开源，为大型语言模型（LLM）的推理效率带来了新的突破。
近年来，大型语言模型（LLM）凭借其强大的问题解决能力，在各个领域取得了显著进展。然而，随着模型规模的不断扩大，推理速度成为了制约LLM应用的关键瓶颈。传统的Transformer架构虽然占据了主流地位，但在处理长序列输入时，计算量会呈平方级增长，严重影响了推理效率。
为了解决这一问题，上世纪90年代兴起的LSTM（长短期记忆网络）重新受到了关注。作为LSTM的提出者和奠基者，Sepp Hochreiter 及其团队于去年推出了xLSTM，旨在通过扩展LSTM的参数规模，使其成为Transformer的有力替代品。xLSTM具有与序列长度线性相关的计算扩展能力和稳定的内存占用，但在扩展至更大参数规模时，其推理速度和效率并未得到充分评估。
近日，Sepp Hochreiter等来自NXAI、JKU的研究者再次对xLSTM进行了优化，成功将其扩展到了70亿参数。这项研究成果不仅突破了xLSTM的参数规模限制，更在推理速度上取得了显著提升。
xLSTM 7B：性能与效率的双重提升
具体而言，xLSTM 7B模型基于DCLM数据集，使用128块H100 GPU，在8192上下文长度下训练了2.3万亿token。研究者对原始xLSTM架构进行了改进，确保训练效率和稳定性，同时保持任务性能。新架构依靠mLSTM单元和并行训练模式，实现高性能的同时最大化速度。
- 架构优化： 通过修改模块架构，研究者优化了吞吐量，在低维空间运行 mLSTM 并添加前馈 MLP 层，同时去除了不必要的组件以提高 GPU 利用率。优化后的架构在保持相似性能的同时，将 token 吞吐量提高了 2 到 4 倍。
- 训练稳定性： 研究者还优化了训练稳定性，特别是 mLSTM 单元的门控机制，有效解决了梯度问题。
在各类任务评估中，xLSTM 7B 与同规模 Transformer 和 Mamba 模型表现相当。通过架构优化，该模型在推理效率测试中实现了最高的预填充和生成吞吐量，同时保持最低的 GPU 内存占用。
论文作者之一 Günter Klambauer 表示，xLSTM 7B 成为了最快、最高效的 7B 语言模型！
技术细节：mLSTM单元的优化
xLSTM 7B 架构的核心是 mLSTM 单元，它的循环和并行模式可以实现高效的训练和推理。为了充分发挥该单元的潜力，研究者重新审视了相邻块结构的设计。
研究者通过解决以下四个限制来优化 mLSTM 块以实现最大效率：
- 在模型的嵌入维数而不是更高维数的空间中操作 mLSTM 单元，并在每个 mLSTM 层之后放置位置前馈 MLP 层。
- 放弃通道卷积和可学习的残差连接等操作，并用密集线性层替换块查询、键和值投影。
- 确保每个 head 的门预激活都是独立计算的。
这些优化显著提升了模型的计算效率和GPU利用率，为xLSTM扩展到更大规模奠定了基础。
开源赋能：加速LLM推理研究
更重要的是，研究团队将xLSTM 7B模型的权重和代码全部开源，为LLM推理加速领域的研究者提供了宝贵的资源。这一举措有望加速相关技术的发展，推动LLM在更多领域的应用。
结论与展望
xLSTM 7B模型的问世，标志着LSTM在LLM领域迎来了新的发展机遇。通过架构优化和技术创新，xLSTM不仅突破了参数规模的限制，更在推理速度上实现了显著提升。模型的开源发布，将进一步促进LLM推理加速领域的研究和应用。
随着LLM技术的不断发展，推理效率将成为其应用普及的关键因素。xLSTM 7B模型的成功，为我们提供了一个新的方向，也为未来的LLM架构设计带来了新的启示。
参考文献：
- xLSTM 7B: A Recurrent LLM for Fast and Efficient Inference: https://arxiv.org/pdf/2503.13427
- 代码地址：https://github.com/NX-AI/xlstm
- Hugging Face 地址：https://huggingface.co/NX-AI/xLSTM-7b
Views: 0