arXiv:2411.01030v4 Announce Type: replace-cross 摘要：高效的状态空间模型（SSMs），如线性递归神经网络和线性注意力变体，在计算效率上优于Transformer，但在执行需要长距离上下文检索的任务，如文本复制、关联回忆和长文本上下文的问答时却表现不佳。之前为了解决这些问题，主要集中在架构上的修改，但往往会重新引入计算上的低效性。在本文中，我们提出了一种新的训练过程Birdie，该过程显著增强了SSMs的上下文检索能力，而不会改变其架构。我们的方法结合了双向输入处理与动态混合的专业预训练目标，并通过强化学习进行优化。我们介绍了一种新的双向SSM架构，该架构能够无缝地从双向上下文处理过渡到因果生成。实验评估表明，Birdie在多号码电话簿查找、长段落问答和填充等检索密集型任务上的性能显著提高。这缩小了与Transformer的性能差距，同时保持了计算效率。我们的研究结果突出了训练过程对利用SSMs固定状态容量的重要性，并提供了一种新的方向以进一步提升其能力。所有代码和预训练模型均可在https://www.github.com/samblouir/birdie获取，支持JAX和PyTorch。