FactorEngram 用共享字典分解 n-gram 记忆,在 340M 与 1B 骨干上把 8K 长上下文检索准确率最高提升 43.3 个百分点
核心概要
该工作提出 FactorEngram,把查找式 n-gram 记忆从「每个模式一个整体嵌入」改为「在跨模式共享的基向量字典上检索稀疏系数」,并用同一字典做基级别的上下文门控,在 340M 与 1B 参数 Transformer 骨干上改善了语言建模、下游任务与长上下文检索表现。
深度剖析
FactorEngram 把局部 token 模式的记忆表示为共享字典上的稀疏正则化系数,而不是每个 n-gram 独占一个稠密嵌入,使相关模式可以通过共同基向量共享参数,而非仅靠哈希碰撞共享。 相对 Engram 把每个检索到的 n-gram 嵌入当作整体单元、存放在各自哈希槽中并用单一标量门控的做法,这里把记忆的最小单元从「n-gram 嵌入向量」改为「字典系数」,并沿用稀疏编码思路对系数施加惩罚。 论文在 340M 与 1B 两个骨干规模上从零联合训练,340M 用 30B token、1B 用 120B token(FineWeb-Edu,上下文 8192),并与纯 Transformer 及作者复现的 Engram 对比;消融显示完整分解设计相对 Engram(含 unigram)把 WikiText 困惑度从 22.32 降到 21.29、LAMBADA 从 24.42 降到 21.69,平均下游准确率提高 0.92 个百分点。
基级别上下文门控让上下文可以逐个调节每个记忆分量:骨干隐状态被投影为查询并与每个基向量打分,得到的权重在重建前缩放对应系数。 Engram 的标量门只能整体放大或抑制嵌入,无法在「the bank」这类多义模式中保留与上下文相关的分量、抑制无关分量;这里把用于判断上下文相关性的向量与用于重建输出的向量统一为同一套字典。 消融中把基级别门控换成标量门控(保留系数-字典表示)后,所有指标下降,平均下游准确率从 49.82 降到 48.15,最难的 NIAH-3 从 58.3 降到 9.4。
FactorEngram 同时覆盖单个 token 与多 token n-gram,并系统研究了记忆分支的插入位置,发现插在中间层注意力子层之前是有效配置。 论文指出 STEM 只检索单 token 嵌入、Engram 只检索 2-gram 与 3-gram,且既有方法把记忆分支插在固定位置;这里把模式覆盖与插入位置都作为受控变量研究。 在 24 层 340M 骨干上先扫单层插入深度(第 12 层在各项评估间较均衡),再固定第 12 层扫第二层,10 & 12 层组合取得最高 NIAH-3;层内比较显示「注意力之前」在全部准确率指标上最好,而「FFN 之前」的 NIAH-3 仅 18.8,「FFN 内部」各项最差。
在两个骨干规模上,FactorEngram 相对 Transformer 与 Engram 整体改善评估表现,长上下文检索增益尤其突出。 相对 Engram,340M 上 NIAH-2 与 NIAH-3 分别提升 41.5 与 43.3 个百分点;1B 上检索增益在 9.3 到 27.6 个百分点之间,同时 WikiText 困惑度与 Transformer 相当、下游准确率略低于 Engram。 结果来自论文表 1 与附录表 7 的完整数值,覆盖 WikiText、LAMBADA 困惑度,PIQA、HellaSwag、WinoGrande、ARC-Easy、ARC-Challenge、SocialIQA、BoolQ 七项下游准确率,以及三个 8K 上下文 NIAH 变体。
启示与展望
这项工作面向以查找式记忆扩展 LLM 参数的研究与工程场景:记忆作为辅助分支加入,注意力与前馈模块保持不变,因此可直接叠加在既有 Transformer 骨干上。论文给出的默认配置是 24 层 340M 骨干、记忆插在第 10 与 12 层、位于注意力子层之前,覆盖 unigram、bigram、trigram,稀疏强度取中等值。作者指出 1B 参数模型对边缘计算与端侧部署仍有价值,因此该结果对资源受限的部署场景具有直接意义;同时作者把「这些收益能否延续到更大规模」列为后续工作。
论文自述的局限是仅在 340M 与 1B 参数上评估,更大规模是否保持收益尚待验证。消融显示稀疏强度存在最优点:从 0 增大到中等强度时整体表现改善,但继续增大后所有评估指标都差于无正则化模型,因此稀疏强度的选择仍是需要按规模与数据重新确认的超参数。1B 规模上 WikiText 困惑度与 Transformer 相当、下游平均准确率略低于 Engram,说明不同指标上的收益并不一致。此外,基向量没有预定义语义标签,字典与系数是与骨干联合学到的,其可解释性并未在本文中展开。
