Birdie:通过奖励驱动目标与课程 advancing 状态空间模型
计算与语言
2025-02-25 v5 人工智能
机器学习
摘要
高效的状态空间模型(SSMs),如线性循环神经网络和线性注意力变体,相比 Transformer 具有计算优势,但在需要长程上下文检索的任务(如文本复制、联想回忆和长上下文问答)上表现不佳。此前为应对这些挑战的努力多集中于架构修改,往往重新引入计算效率问题。本文提出一种新颖的训练过程 Birdie,在不改变架构的前提下显著增强 SSM 的上下文检索能力。我们的方法结合双向输入处理与通过强化学习优化的动态混合专用预训练目标。我们引入一种新的双向 SSM 架构,可无缝从双向上下文处理过渡到因果生成。实验评估表明,Birdie 在多号码电话簿查询、长段落问答和填空等检索密集型任务上显著提升性能,缩小了与 Transformer 的性能差距,同时保留了计算效率。我们的发现强调了训练过程在利用 SSM 固定状态容量方面的重要性,为提升其能力开辟了新方向。所有代码和预训练模型可在 https://www.github.com/samblouir/birdie 获取,支持 JAX 和 PyTorch。
引用
@article{arxiv.2411.01030,
title = {Birdie: Advancing State Space Models with Reward-Driven Objectives and Curricula},
author = {Sam Blouir and Jimmy T. H. Smith and Antonios Anastasopoulos and Amarda Shehu},
journal= {arXiv preprint arXiv:2411.01030},
year = {2025}
}
备注
Accepted to EMNLP 2024 (Main Conference)