Titan-Llama给预训练模型外挂神经记忆：冻结主干，长文问答准确率从23.87%升到40.51%
核心概要
Titan-Llama把预训练Llama-3.2-1B-Instruct前四层改成块状注意力，并在每个KV头上外挂一个固定大小的神经记忆模块，用测试时训练把历史键值压进模块，主干参数保持冻结。在QuALITY上，相比仅用分段注意力窗口，准确率从23.87%升到40.51%；在LongHealth上从14.50%升到40.75%。评测配置下，中等上下文长度预填充快20%以上、解码快37%，稳态解码显存低24%。
深度剖析
预训练注意力模型不必从头重训，也能获得固定大小的长上下文记忆：Titan-Llama只后训练约占总量1%的外部权重，就让被分段注意力截断的历史信息重新可用。 此前的做法要么保留完整KV缓存、让长上下文推理越来越贵，要么改用线性注意力或状态空间模型，而这通常需要重新训练一个新模型，无法复用已有的注意力主干。 在Llama-3.2-1B-Instruct上，与仅使用分段注意力窗口的基线相比，QuALITY准确率从23.87%升到40.51%，LongHealth从14.50%升到40.75%。
历史注意力被拆成两部分：最近窗口仍走精确注意力，更早的键值由神经记忆模块读出，再按一个逐通道学习的门控权重与局部输出相加。 神经记忆模块是一个三层网络，每个KV头一个，隐藏维64，用GELU激活，新增可训练参数共17.23M；它通过测试时训练更新，即推理时对模块权重做梯度下降。 每处理完一个固定长度的词块，模块就用键值重构的平方误差损失更新一次，使用学习到的动量、衰减和步长；评测在单张80GB NVIDIA A100上进行，训练序列最长16K token。
启示与展望
演示范围是Llama-3.2-1B-Instruct：前四层改为块状注意力并各加神经记忆模块，其余层保留全注意力，训练序列最长16K token，推理测量在单张80GB A100、批大小为1下完成。
键值重构只是有用记忆的替代目标，并不保证精确复现历史注意力；在32K等超过16K训练长度的设置上性能下降，多值检索在每个长度都低于全注意力，块大小对RULER的影响也随任务和长度而异。