[Paper] RMAAT：星形胶质细胞启发的记忆压缩与重放用于高效长上下文Transformer

发布: 2周前 (2026年1月2日 GMT+8 02:34)

7 min read

原文: arXiv

Source: arXiv - 2601.00426v1

请提供您希望翻译的具体文本内容，我将按照要求将其翻译为简体中文并保留原有的格式。

Overview

本文介绍了 RMAAT，一种借鉴星形胶质细胞（大脑的支持细胞）理念的 Transformer 变体，用以解决长序列自注意力的著名二次成本问题。通过嵌入轻量级的“记忆压缩”循环和基于回放的训练方案，作者在长上下文基准测试中实现了竞争性的准确率，同时大幅削减计算量和 GPU 内存使用。

Astrocyte‑inspired memory tokens 在分段输入之间持久存在，充当先前上下文的压缩摘要。
Retention factor 源自模拟的长期可塑性（LTP），可在运行时自适应地压缩或扩展记忆令牌。
Linear‑complexity intra‑segment attention 受短期可塑性（STP）启发，在每个块内部消除二次复杂度的增长。
Astrocytic Memory Replay Backpropagation (AMRB)，一种利用已存储记忆状态来降低反向传播内存占用的训练算法。
Empirical validation on the Long Range Arena (LRA) 表明 RMAAT 在准确率上匹配或超越最先进水平，同时 FLOPs 下降约 40 %，GPU 内存占用降低约 30 %。

分段处理 – 将输入序列拆分为固定大小的块。每个块由标准的 Transformer 块处理，但不是在块结束后丢弃隐藏状态，而是更新并携带一小组 memory tokens（记忆标记）。
通过保留因子进行记忆压缩 – 每个段结束后，记忆标记会通过一个学习得到的门控机制，该机制模拟星形胶质细胞的 LTP：重要信息被保留（高保留），冗余信息被压缩（低保留）。这使得记忆大小与总序列长度无关，保持恒定。
块内线性注意力 – 在一个段内部，注意力使用基于核函数的线性 Transformer（例如 Performer‑style）近似，反映星形胶质细胞的 STP 动力学，使每块的计算成本为 O(N) 而非 O(N²)。
AMRB 训练 – 在反向传播时，算法回放存储的记忆状态，而不是为所有先前块保留完整的计算图。此回放模拟生物学的记忆巩固，显著降低长序列所需的激活内存。

整体流程可视化为一个循环：块 → 线性注意力 → 记忆更新（压缩） → 传递到下一个块，其中 AMRB 负责梯度流动。

基准	准确率 (RMAAT)	基线 (例如 Longformer)	FLOPs ↓	GPU 内存 ↓
ListOps	71.2 %	70.8 %	~38 %	~32 %
Text (Char)	84.5 %	84.1 %	~42 %	~30 %
Retrieval	88.9 %	88.3 %	~35 %	~28 %

准确率：RMAAT 在所有 LRA 任务上与现有高效 Transformer 持平或略有提升。
计算与内存：线性注意力的分段机制以及压缩记忆同时降低了 FLOPs 和峰值 GPU 内存，使得在单块 16 GB GPU 上能够处理最长达 8 K token 的序列（普通 Transformer 会出现 OOM）。
消融实验：去除保持因子或 AMRB 重放会导致准确率下降 5‑10 %，并显著增加内存消耗，验证了两种星形胶质细胞机制的重要性。

长文档处理 – RMAAT 可以直接嵌入法律合同、代码库或科学论文等需要超过几千个 token 的上下文的流水线，而无需多 GPU 分片。
边缘与移动推理 – 恒定大小的内存占用和线性注意力使其能够在 RAM 受限的设备上运行，为设备端摘要或转录打开了可能。
训练效率 – AMRB 的 replay 策略降低了梯度检查点（gradient checkpointing）的内存开销，允许在预训练时使用更大的批量或更长的序列，从而加快收敛速度并降低云端成本。
类神经设计 – 该工作展示了一种将生物可塑性概念转化为软件原语的具体方法，鼓励进一步探索基于大脑的 AI 模型优化。

总体而言，RMAAT 将神经科学灵感与工程实用性相结合，指向了一类面向日益增长上下文窗口的记忆高效 Transformer 的新方向。