人工海马体网络用于高效长序列建模
计算与语言
2025-12-18 v2 人工智能
机器学习
摘要
长序列建模面临在类RNN模型中压缩固定大小记忆的效率与注意力机制Transformer中无损增长记忆的保真性之间的根本性权衡。受认知科学中的多存储模型启发，我们引入一种人工神经网络记忆框架。该方法维持Transformer KV缓存的滑动窗口作为无损短期记忆，同时通过称为人工海马体网络（Artificial Hippocampus Network, AHN）的可学习模块反复压缩窗口外信息到固定大小的紧凑长期记忆中。为验证该框架，我们使用现代类RNN架构实例化AHN，包括Mamba2、DeltaNet和GatedDeltaNet，用于增强开源大语言模型。我们还提出一种高效自蒸馏训练方法，其中冻结基础模型的所有参数，仅优化AHN的参数。对于推理，本方法默认设置注意力的滑动窗口大小为32k，AHN仅在序列长度超过32k窗口时激活，从而解决该尺度下注意力二次复杂度问题。大量实验在长上下文基准测试LV-Eval和InfiniteBench上表明，增强AHN的模型在滑动窗口基线上 consistently 超越，并在性能上与全注意力模型相当甚至更好，同时大幅降低计算和内存需求。例如，使用AHN增强Qwen2.5-3B-Instruct模型，可将推理FLOPs降低40.5%，内存缓存降低74.0%，在LV-Eval（128k序列长度）上的平均得分从4.41提升至5.88。代码已公开：https://github.com/ByteDance-Seed/AHN。
引用
@article{arxiv.2510.07318,
title = {Artificial Hippocampus Networks for Efficient Long-Context Modeling},
author = {Yunhao Fang and Weihao Yu and Shu Zhong and Qinghao Ye and Xuehan Xiong and Lai Wei},
journal= {arXiv preprint arXiv:2510.07318},
year = {2025}
}
备注
Code: https://github.com/ByteDance-Seed/AHN