Titan-Llama：用外部神经记忆模块改造预训练大语言模型
Titan-Llama: Retrofitting Pre-trained LLMs with External Neural Memory Modules
AI总结：
本研究提出Titan-Llama，通过在测试时更新参数的分层记忆架构改造预训练Transformer，提升了长上下文任务的准确率与推理效率，相关代码已在GitHub开源。
摘要：
Transformer因注意力机制允许跨序列位置并行计算，可高效进行大规模预训练，但长上下文推理需要越来越大的KV缓存。高效的循环替代方案降低了推理的渐近复杂度，但其序列依赖性限制了训练期间的逐词并行性，导致更难扩展。此外，当前扩展长上下文任务的范式在精确保留信息与保持长上下文推理效率之间存在矛盾。为解决这一差距，我们提出Titan-Llama，一种用于改造基于注意力的语言模型的后训练方法，采用在测试时更新参数的分层记忆架构。Titan-Llama应用分块注意力掩码，对近期标记保留全保真度注意力，同时将历史上下文压缩为循环的固定大小神经记忆模块（NMMs）。我们的方案允许替换预训练的注意力层，无需从头训练新的语言模型。在Llama-3.2-1B-Instruct上，Titan-Llama在QuALITY数据集上将分段注意力窗口的准确率从23.87%提升至40.51%，在LongHealth数据集上从14.50%提升至40.75%。推理基准测试显示，在评估配置中，预填充速度提升超过20%，中等上下文长度下解码速度提升37%，稳态解码内存使用量降低24%。这些结果共同表明，结合利用测试时训练的记忆架构机制可有效调节预训练模型的长上下文行为，为适配预训练Transformer实现高效长上下文推理提供了实用途径。我们的代码已在GitHub上发布。
英文摘要：
Transformers are efficient to pre-train at scale because attention permits parallel computation across sequence positions, but long-context inference requires increasingly large KV caches. Efficient recurrent alternatives reduce the asymptotic complexity of inference, but their sequential dependence limits token-wise parallelism during training, making them harder to scale. Further, the current paradigm of extending long context tasks creates competing notions of retaining information precisely and keeping long-context inference efficient. To address this gap, we introduce Titan-Llama, a post-training method for retrofitting attention-based language models with a hierarchical memory architecture that updates parameters at test-time. Titan-Llama applies a blockwise attention mask and preserves full-fidelity attention over recent tokens while compressing historical context into recurrent, fixed-size neural memory modules (NMMs). Our formulation allows pre-trained attention layers to be replaced without training a new language model from scratch. On Llama-3.2-1B-Instruct, Titan-Llama raises accuracy over segmented attention windows from 23.87\% to 40.51\% on QuALITY and from 14.50\% to 40.75\% on LongHealth. Inference benchmarks show more than 20\% faster prefill, 37\% faster decode at medium context lengths, and 24\% lower steady-state decode memory usage in the evaluated configurations. Together, these results demonstrate that incorporating memory architecture mechanisms that utilize test-time training effectively modulates long-context behavior of pre-trained models, offering a practical route to adapting pre-trained transformers for efficient long-context inference. Our code is released on Github \href{ this https URL }{GitHub}
机构、总结、中文摘要均由AI分析而来，可能存在错误，请以论文正文为准。
