Locas：你的模型就是局部支持参数化记忆的最优初始化器 Locas: Your Models are Principled Initializers of Locally-Supported Parametric Memories
用模型自身参数初始化旁路FFN记忆，实现高效测试时训练
前置知识
测试时训练 (Test-Time Training, TTT)
TTT是一类在模型推理阶段利用测试输入的自监督信号来更新模型参数的方法。与传统的离线微调不同，TTT在推理时执行在线参数更新，使模型能够动态适应当前输入的分布。典型做法是对测试数据计算辅助自监督损失，然后通过梯度下降更新模型参数。TTT的关键优势在于能够在不访问原始训练数据的情况下，让模型适应新的域分布或任务。
本文提出的Locas本质上是一种高效的TTT方法，通过旁路参数化记忆替代传统的参数空间更新，理解TTT的基本框架是理解本文动机的基础。
前馈网络 (FFN) 作为键值记忆
Transformer中的FFN层可以被重新解释为软查找表或内部注意力机制。具体来说，一个两层FFN $\text{FFN}(A) = V^\top \phi(K^\top A)$ 中，矩阵 $K$ 的每一列 $k_j$ 可视为"键"向量，矩阵 $V$ 的每一行 $v_j$ 可视为"值"向量。每个键值对 $(k_j, v_j)$ 定义了一个记忆槽位：键向量通过与输入的内积决定何时激活该槽位，值向量决定检索什么内容。这种视角最早由Geva等人在2021年提出，为理解Transformer的知识存储机制提供了重要洞察。
本文的核心创新是基于这种FFN-键值记忆视角来设计旁路记忆模块，将FFN视为可扩展的参数化记忆是整个方法的理论基础。
LoRA (Low-Rank Adaptation)
LoRA是一种参数高效微调方法，它通过低秩分解将权重更新分解为两个小矩阵的乘积 $W = W_0 + BA$，其中 $B \in \mathbb{R}^{d_{out} \times r}$, $A \in \mathbb{R}^{r \times d_{in}}$，$r$ 远小于原始维度。这种方法显著减少了可训练参数数量，同时通过冻结预训练权重 $W_0$ 并仅更新低秩增量 $BA$ 来保持模型的原始知识。LoRA在各种下游任务中已被广泛采用。
TempLoRA（基于LoRA的TTT方法）是本文的主要对比基线，理解LoRA的工作机制对于理解本文提出的方法与已有方法的本质区别至关重要。
灾难性遗忘 (Catastrophic Forgetting)
灾难性遗忘是指神经网络在学习新任务或新知识时，会突然且大幅度地丧失之前学习到的能力。在测试时训练的场景中，当模型通过梯度更新适应新输入时，可能会干扰预训练阶段学习到的通用知识，导致模型在原始任务上的性能下降。这是TTT方法面临的核心挑战之一。
本文的一个重要贡献是证明旁路FFN架构（不修改原始权重）能够有效减少灾难性遗忘，MMLU实验结果显示遗忘率仅为0.1-0.2%。
GLU-FFN (Gated Linear Unit FFN)
GLU-FFN是现代LLM（如LLaMA、Qwen、Mistral）中广泛采用的FFN变体，其计算形式为 $\text{GLU-FFN}(A) = V^\top \cdot (\sigma(G^\top A) \odot K^\top A)$，其中 $G$ 是门控矩阵，$K$ 是上投影矩阵，$V$ 是下投影矩阵，$\sigma$ 是SiLU激活函数，$\odot$ 表示逐元素乘法。相比传统两层MLP，GLU通过门控机制实现了更精细的特征选择。
本文提出的Locas-GLU变体与现代LLM的GLU-FFN结构对齐，使得该方法能够无缝集成到现有的主流模型中，是方法实用性的关键。
研究动机
当前测试时训练（TTT）方法面临参数效率和计算效率的双重瓶颈。以TempLoRA为代表的方法需要在推理时对所有线性投影（注意力Q、K、V、O以及FFN的gate、up、down）引入低秩适配器，导致额外参数量巨大——在Qwen3-1.7B上需要73.4M额外参数，并且需要4.7倍的计算时间。更严重的是，这些方法通过修改现有权重矩阵来实现适应，直接干扰了预训练表征，导致显著的灾难性遗忘——在MMLU评估中，TempLoRA在记忆一本PG-19书籍后会造成0.6%的性能下降，在使用更大秩（r=512）时遗忘率高达1.2%。此外，随机初始化的低秩适配器需要大量梯度步数才能收敛，进一步增加了推理延迟。这些问题使得TTT方法在实际部署中面临效率和稳定性的双重挑战。
本文的目标是本文旨在通过引入一种全新的参数化记忆框架Locas，从根本上提升测试时训练的参数效率和计算效率。具体目标包括三个方面：第一，通过利用骨干模型自身的参数、激活和梯度来指导记忆模块的初始化，实现快速收敛，从而大幅减少所需的梯度更新步数；第二，通过设计旁路FFN架构（sideway FFN），在不修改原始模型权重的前提下实现真正的模型容量扩展，从而最小化灾难性遗忘；第三，使该方法能够无缝集成到现代基于GLU-FFN架构的大语言模型中，实现即插即用的高效持续学习能力。
与已有工作不同的是，本文的独特切入角度在于重新诠释FFN层的键值记忆本质，并将其转化为一个可操作的设计原则。与现有TTT方法（如TempLoRA）采用随机初始化然后通过大量梯度更新来学习不同，本文提出了一种"原则性初始化"（principled initialization）的范式：对于Locas-MLP变体，直接使用骨干模型的隐藏激活 $A_{it}$ 作为键向量初始化，使用梯度信号 $G_{it}$ 作为值向量初始化；对于Locas-GLU变体，通过分析模型原生FFN在记忆过程中的激活模式，选择激活重要性最高的Top-K维度进行参数克隆。这种"利用模型自身来初始化新记忆"的思想，本质上是在模型的激活空间中进行非线性主成分分析（nonlinear PCA），从而将任务相关信息压缩到少量关键维度中。这种方法从根本上区别于所有依赖随机初始化的TTT方法，是本文最核心的创新。
核心方法
Locas的核心思想可以用一个直觉来理解：如果把Transformer的FFN层看作一个软查找表，那么向模型添加新知识就相当于向这个查找表写入新的键值对。关键洞察是，这些新键值对的初始化不应该随机，而应该利用模型在处理新信息时产生的内部信号——激活值和梯度——来指导初始化。就像一个有经验的图书管理员不会随机摆放新书，而是根据书的内容和读者的查询模式来决定最佳位置，Locas利用模型的"内部视角"来确定新记忆的最佳初始状态。技术路线上，Locas首先将FFN重新解释为软查找表记忆（每个中间维度 $j$ 对应一个键值对 $(k_j, v_j)$），然后引入一个与骨干FFN并行的旁路记忆模块，该模块共享相同的FFN设计但拥有独立的可学习参数。记忆模块的输出通过缩放因子 $\tau$ 加到骨干FFN的输出上，实现真正的容量扩展而不修改原始权重。
Locas与已有方法（如TempLoRA）的本质区别在于两点。第一，架构层面：TempLoRA通过低秩更新 $W = W_0 + BA$ 直接修改原始权重矩阵，这会干扰预训练表征；而Locas引入一个与原始FFN并行的旁路模块 $H_{out} = \text{FFN}(A) + \tau \cdot \text{Locas}(A)$，原始权重完全不被修改，从而实现了"真正的容量扩展"而非"参数干扰"。第二，初始化层面：TempLoRA采用随机初始化，需要大量梯度步数才能收敛；而Locas-MLP通过复用模型的激活和梯度信号来初始化，键向量 $k_i^{new} \leftarrow \text{Normalize}(A_{it})$ 直接使用当前token的激活，值向量 $v_i^{new} \leftarrow \epsilon \cdot \text{GlobalNormalize}(G_{it})$ 使用反向传播的梯度信号；Locas-GLU则通过分析骨干FFN的激活模式，选择最活跃的维度进行参数克隆。这种原则性初始化使得记忆模块从一开始就与模型的内部表征对齐，实现了"热启动"而非"冷启动"。
方法步骤详情
Locas-GLU的具体执行流程包括以下步骤：（1）记忆初始化阶段，给定输入文本块 $x_{<T}$，首先进行一次前向传播，在每一层 $i$ 计算GLU-FFN的中间激活 $M_{it} = \sigma(W^{G,i} A_{it}) \odot (W^{K,i} A_{it})$；（2）计算每个中间维度 $j$ 的激活重要性 $\alpha_i^j = \frac{1}{T} \sum_{t=1}^{T} |M_{it,j}|$，按重要性降序排列；（3）选择Top-K最重要的维度构建记忆，键矩阵和门控矩阵通过克隆骨干FFN对应行初始化 $K_i \leftarrow \text{Normalize}([W^K]_{j \in S_i^r})$，值矩阵零初始化 $V_i \leftarrow 0$；（4）设置输出缩放因子 $\tau = \frac{1}{r} \cdot \frac{1}{m} \sum_{j=1}^{m} \|W^{down}[j,:]\|_2$；（5）推理阶段，模型逐token处理流式上下文，记忆模块通过标准反向传播更新；（6）应用权重范数裁剪，对每个向量 $w$ 强制执行 $w \leftarrow \frac{w}{\max(\|w\|_2, 1)}$，限制记忆每步的影响范围；（7）记忆模块的输出通过缩放因子加到骨干FFN输出上，实现知识的逐步积累。
技术新颖性
Locas的技术新颖性体现在多个维度。首先，"原则性初始化"是一个全新的概念，区别于所有TTT方法的随机初始化范式——对于Locas-MLP，本文证明了该初始化在给定假设下是"逐步最优"的（step-wise optimal），即在每个时间步和梯度更新步都是最优的；对于Locas-GLU，本文证明了激活引导的参数克隆将记忆初始化锚定在原始FFN参数的主子空间中。其次，旁路FFN架构（sideway FFN）的设计非常巧妙——它不是简单地添加一个并行模块，而是通过精心设计的输出缩放 $\tau$ 和权重裁剪机制，实现了"双重保护"防止灾难性遗忘：零初始化确保记忆模块初始时不影响模型行为，权重裁剪确保每步影响有界，输出缩放确保记忆贡献与骨干FFN的典型贡献尺度匹配。第三，将FFN的键值记忆视角与PCA的理论框架连接——Top-K激活选择本质上是在模型激活空间中进行非线性PCA，选择最重要的主成分方向来近似当前上下文的支持流形，这为方法的有效性提供了理论解释。
实验结果
实验结果全面验证了Locas在多个维度上的优势。在PG-19全书语言建模任务上，Locas-GLU在Qwen3-1.7B-Base上以仅11.0M额外参数（TempLoRA的15%）和1.8倍相对时间（TempLoRA的38%）达到了19.04的困惑度（200K上下文），而TempLoRA需要73.4M参数和4.7倍时间才达到19.13的困惑度。在记忆宽度实验中，Locas-GLU即使在极低维度 $r=16$（仅2.8M参数）也能达到19.14的困惑度，与TempLoRA在 $r=64$（73.4M参数）时的19.13相当，实现了26倍的参数压缩。MMLU遗忘评估显示，Locas-GLU在记忆整本PG-19书籍后仅造成0.2%的MMLU性能下降（60.4%到60.2%），而TempLoRA造成0.6%下降（60.4%到59.8%）；当扩大记忆容量到 $r=512$ 时，Locas-GLU的遗忘率几乎不变（0.1%），而TempLoRA的遗忘率翻倍（1.2%）。在LoCoMo长对话问答任务上，Locas-GLU在Qwen3-1.7B-Base的单跳问题上达到41.6% F1，相比全注意力基线的37.3%提升11.5%，相比TempLoRA的37.7%提升10.3%；在时序推理问题上达到34.1% F1，相比TempLoRA的29.1%提升17.2%。值得注意的是，在无上下文评估中，Locas-GLU在多跳问题上达到8.4% F1（TempLoRA仅4.5%），证明其更有效地将对话事实内化为持久的参数化记忆。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| PG-19全书语言建模（Qwen3-1.7B-Base，200K上下文） | 困惑度 (PPL) | 19.04 | TempLoRA: 19.13 | PPL降低0.09，同时额外参数减少85%（11.0M vs 73.4M），计算时间减少62%（1.8x vs 4.7x） |
| PG-19记忆宽度对比（Qwen3-1.7B-Base，r=16 vs r=64，200K上下文） | 困惑度 (PPL) | Locas-GLU r=16: 19.14, r=64: 19.04 | TempLoRA r=64: 19.13 | Locas-GLU在r=16（2.8M参数）时已接近TempLoRA在r=64（73.4M参数）的性能，参数效率提升26倍 |
| MMLU遗忘评估（Qwen3-1.7B-Base，记忆PG-19全书后） | MMLU准确率 (%) | 60.2%（下降0.2%） | TempLoRA: 59.8%（下降0.6%） | 遗忘率降低67%（0.2% vs 0.6%）；在r=512时遗忘率仅0.1%，TempLoRA为1.2% |
| LoCoMo单跳问答（Qwen3-1.7B-Base，全注意力上下文） | F1分数 (%) | 41.6% | 全注意力基线: 37.3%, TempLoRA: 37.7% | 相比全注意力提升11.5%，相比TempLoRA提升10.3% |
| LoCoMo时序推理（Qwen3-1.7B-Base，全注意力上下文） | F1分数 (%) | 34.1% | 全注意力基线: 33.5%, TempLoRA: 29.1% | 相比TempLoRA提升17.2%，表明旁路架构更好地保留了事件的时序结构 |
| LoCoMo单跳问答（Qwen3-4B-Base，全注意力上下文） | F1分数 (%) | 47.6% | 全注意力基线: 41.2%, TempLoRA: 41.5% | 相比全注意力提升15.5%，相比TempLoRA提升14.7%，展示了更大模型的有效性 |
局限与改进
尽管Locas展现了显著的优势，但仍存在若干局限性。首先，Locas-MLP变体（具有更清晰理论保证的版本）与使用GLU-FFN的现代LLM存在兼容性问题——论文指出，对于非两层MLP FFN训练的模型，MLP输入处的逐段线性可分性通常较差，这限制了该变体的适用范围。其次，论文提出的非线性SVD（NL-SVD）压缩算法虽然理论上优雅，但在实践中表现不佳：在200K上下文长度下，NL-SVD的困惑度为17.89，而标准反向传播为17.90，几乎没有差异，但NL-SVD需要12.7倍的相对计算时间（反向传播仅1.7倍）。这主要是因为NL-SVD需要至少float32精度（无法使用bfloat16混合精度），且SVD操作在GPU上的优化程度远不如标准梯度计算。第三，Locas-GLU的激活重要性计算需要一次额外的前向传播来分析激活模式，这在冷启动场景下可能引入额外延迟。第四，实验主要在相对较小的模型（0.6B-4B参数量）上进行验证，对于更大规模模型（如70B+）的有效性尚未得到验证。第五，论文未充分讨论记忆模块的容量上限——当需要记忆的信息量远超记忆模块的参数容量时，性能如何退化尚不清楚。
独立分析的弱点
从独立分析的角度，Locas存在几个值得关注的弱点。第一，激活重要性 $\alpha_i^j$ 的计算采用简单的绝对值平均，这种度量可能无法准确捕捉对下游任务真正重要的维度——例如，某些维度可能激活值较低但对特定推理任务至关重要。改进方向是引入任务感知的重要性度量，如梯度加权的激活重要性。第二，权重范数裁剪的阈值硬编码为1.0，这个阈值可能不是所有场景下的最优选择——在需要更强适应能力的场景中（如域偏移较大的情况），可能需要更大的阈值；在需要更保守适应的场景中，可能需要更小的阈值。改进方向是设计自适应阈值策略，根据当前输入的分布特征动态调整。第三，输出缩放因子 $\tau$ 的计算依赖于骨干FFN的权重范数统计量，但在记忆模块学习过程中，其输出分布可能会发生变化，静态的 $\tau$ 可能不是最优的。改进方向是将 $\tau$ 也作为可学习参数，或根据记忆模块的当前状态动态调整。第四，论文未讨论记忆模块的"遗忘"机制——当处理极长上下文时，早期记忆可能被后续记忆覆盖或干扰，缺乏显式的记忆管理策略。
未来方向
论文作者提出了几个有前景的未来研究方向。第一，动态记忆分配策略——当前方法使用固定的记忆宽度 $r$，但不同层、不同上下文可能需要不同容量的记忆。未来可以探索根据输入的复杂度和重要性动态调整每层的记忆维度。第二，层次化记忆架构——当前方法使用单层记忆模块，但人类记忆具有多尺度时间依赖性（工作记忆、短期记忆、长期记忆）。未来可以设计多层记忆层次结构，不同层次以不同时间尺度更新和遗忘。第三，与检索增强生成（RAG）系统的集成——Locas的参数化记忆可以与非参数化的外部检索系统形成互补，构建混合记忆架构，其中高频、核心的知识存储在参数化记忆中，低频、详细的信息存储在外部向量数据库中。第四，基于本文的PCA解释框架，未来可以探索更高效的子空间学习方法，如在线增量PCA或稀疏PCA，以进一步提升记忆初始化的质量。第五，将Locas应用于其他模态（如视觉、多模态）的测试时训练场景，验证其通用性。
复现评估
从复现评估的角度，本文具有较好的可复现性但存在一定门槛。代码层面，论文基于Qwen3系列模型和标准的HuggingFace Transformers库，实验设置相对清晰——使用SlimPajama数据集预训练的0.6B MLP-FFN模型、Qwen3-0.6B-Base/1.7B-Base/1.7B等公开模型。数据方面，PG-19和LoCoMo都是公开可用的基准数据集，不涉及私有数据。算力需求方面，论文在Qwen3-1.7B上进行实验，单张消费级GPU（如RTX 4090）即可复现主要结果；但Qwen3-4B的实验可能需要更大显存的GPU。复现难度中等——Locas-GLU的核心实现相对简单（激活分析、Top-K选择、参数克隆、标准反向传播），但需要注意输出缩放因子 $\tau$ 的计算细节和权重裁剪的实现。论文未明确说明是否开源代码，这可能会影响精确复现的难度。建议复现者重点关注三个关键细节：Top-K激活选择的具体实现、权重范数裁剪的向量化实现、以及输出缩放因子与骨干FFN统计量的对齐方式。
论文图表