Samba：用于高效无限上下文语言建模的简单混合状态空间模型
计算与语言
2025-03-03 v3 机器学习
摘要
高效建模具有无限上下文长度的序列一直是一个具有挑战性的问题。以前的方法要么 suffer from 计算复杂度呈二次增长，要么在长度泛化方面受限。本文中，我们提出了 Samba，一种简单的混合架构，将选择性状态空间模型（SSM）Mamba 与滑动窗口注意力（SWA）逐层组合。Samba 通过压缩给定序列为循环隐藏状态，同时仍能通过注意力机制精确地记住最近的记忆。我们将 Samba 扩展至 38 亿参数，训练了 3.2 万亿 token，并在多个基准测试中显著优于最新模型。在 4K 长度的序列上预训练后，Samba 在上下文长度为 100 万的零样本情境下显示出改进的困惑度。当在 4K 长度的序列上微调后，Samba 能够高效地外推至 256K 的上下文长度，在 Passkey 检索任务中实现完美的记忆回顾，并在具有挑战性的 Phonebook 任务上显示出优于全注意力模型的检索外推能力。作为一种线性时间序列模型，Samba 相对于 Transformer 使用分组查询注意力处理 128K 长度用户提示时，吞吐量提高 3.73 倍；在生成 64K token 时实现 3.64 倍的加速。我们的代码已公开于 https://github.com/microsoft/Samba。
引用
@article{arxiv.2406.07522,
title = {Samba: Simple Hybrid State Space Models for Efficient Unlimited Context Language Modeling},
author = {Liliang Ren and Yang Liu and Yadong Lu and Yelong Shen and Chen Liang and Weizhu Chen},
journal= {arXiv preprint arXiv:2406.07522},
year = {2025}
}
备注
Accepted by ICLR 2025. Camera-ready Version