B'MOJO：基于全息记忆与消失记忆的基石模型混合状态空间实现
机器学习
2024-07-10 v1 计算与语言
神经与进化计算
摘要
我们描述了一系列支持转导推理的架构，允许记忆扩展到有限但事先未知的上限，同时在推理时高效利用有限资源。当前架构使用这些资源来表示数据，要么通过在有限时间跨度内以全息方式存储（如 Transformer 中的“上下文”），要么在无限时间跨度内以消失方式存储（如状态空间模型 SSMs 中）。近期的混合架构已结合了全息记忆和消失记忆，但存在限制，无法让设计者或学习过程无缝调节这两种记忆，也无法扩展全息记忆的范围。我们借助随机实现理论的思想，发展出一类模型 B'MOJO，用于在元素可组合模块中无缝结合全息记忆和消失记忆。总体架构可用于实现能够通过异步更新的记忆进行检索的模型，从而实现对短期全息记忆（“in-context”）、永久结构记忆（“in-weights”）、消失记忆（“in-state”）以及长期全息记忆（“in-storage”）的访问。我们表明，Transformer、现有的 SSMs 如 Mamba，以及诸如 Jamba 等混合架构，都是 B'MOJO 的特例，并描述了一个基本实现，将在硬件上高效堆叠和扩展。我们在转导推理任务上测试 B'MOJO，例如联想记忆，其中它在现有 SSMs 和混合模型中表现更佳；作为基准，我们测试普通语言建模，其中 B'MOJO 在参数规模与 Transformer 和 SSMs 相当时，实现的困惑度与之相当，训练速度可快达 10%。最后，我们表明 B'MOJO 在调节全息记忆和消失记忆方面的能力，使其在最长序列长度达训练时见到的四倍（32K token）的推理中表现更好。
引用
@article{arxiv.2407.06324,
title = {B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory},
author = {Luca Zancato and Arjun Seshadri and Yonatan Dukler and Aditya Golatkar and Yantao Shen and Benjamin Bowman and Matthew Trager and Alessandro Achille and Stefano Soatto},
journal= {arXiv preprint arXiv:2407.06324},
year = {2024}
}