NVIDIA NeMo 通过混合状态空间模型集成增强LLM能力

根据NVIDIA技术博客的报道，NVIDIA宣布在其NeMo框架中集成混合状态空间模型（SSM），这是人工智能领域的重要举措。此项发展有望提升大语言模型（LLM）的效率和能力。

基于Transformer模型的进展

自2017年引入transformer模型架构以来，AI计算性能迅速发展，使得更大更强的LLM得以创建。这些模型在智能聊天机器人、计算机代码生成甚至芯片设计中找到了应用。

为了支持这些先进LLM的训练，NVIDIA NeMo提供了一个用于构建、定制和部署LLM的端到端平台。NeMo内集成了Megatron-Core，这是一个基于PyTorch的库，提供了训练LLM所需的基本组件和优化。

状态空间模型的引入

NVIDIA的最新公告包括对状态空间模型（SSM）预训练和微调的支持。此外，NeMo现在还支持基于Google DeepMind描述的Griffin架构的模型训练。

替代模型架构的优势

尽管transformer模型通过注意机制在捕捉长距离依赖方面表现出色，但其计算复杂度随着序列长度的增加呈二次方增长，导致训练时间和成本增加。然而，SSM通过克服与注意机制相关的多种限制，提供了一种具有吸引力的替代方案。

SSM以其计算和内存方面的线性复杂度著称，使其在建模长距离依赖关系上更加高效。它们还提供了与transformer模型相当的高质量和准确性，并且在推理过程中需要更少的内存。

SSM在长序列训练中的效率

由于其高效处理序列建模任务的特点，SSM在深度学习社区中越来越受欢迎。例如，当序列长度增加到256K时，SSM的一个变体Mamba-2层比transformer层快18倍。

Mamba-2采用了结构化状态空间对偶（SSD）层，将SSM计算重新表述为矩阵乘法，从而利用NVIDIA Tensor Cores的性能。这使得Mamba-2在保持与transformer相竞争的质量和准确性的同时，训练速度更快。

增强性能的混合模型

通过结合SSM、SSD、RNN和transformer的混合模型可以利用每种架构的优点，并减轻各自的缺点。NVIDIA研究人员最近的一篇论文描述了混合Mamba-Transformer模型，这些模型在标准任务上的表现超过了纯transformer模型，预计推理时可快8倍。

这些混合模型还显示出更高的计算效率。随着序列长度的增加，训练混合模型所需的计算量的增长率远低于纯transformer模型。

未来展望

NVIDIA NeMo对SSM和混合模型的支持标志着实现新AI智能水平的重要一步。初期功能包括对Mamba-2、Griffin架构、混合模型组合以及各种模型微调的支持。未来的版本预计将包括更多的模型架构、性能优化并支持FP8训练。

欲了解更多详细信息，请访问NVIDIA技术博客。