用共享乘积键记忆替代FFN：23M参数下与MoE打平

在Transformer架构中，前馈网络（FFN）承担了大部分参数和计算量。最近，Reddit社区分享的一项实验尝试用共享乘积键记忆（shared product-key memory）取代FFN，在2300万参数规模下与混合专家（MoE）模型打成平手，但作者将其定性为「负面结果」。这项实验连同开源的推理基准一起，为高效模型设计提供了有价值的参考。

实验背景：FFN的替代方案探索

Transformer中的FFN层通常占据模型参数的三分之二以上。为了降低计算成本，研究者提出过多种替代方案，包括MoE（混合专家）和基于记忆的模块。乘积键记忆（product-key memory）是一种通过键值对存储和检索信息的结构，其核心思想是利用可学习的键向量与查询向量的内积来加权聚合值向量。

这项实验的出发点是：能否用一个共享的乘积键记忆模块完全替代FFN，从而在保持性能的同时减少参数或计算？作者在2300万参数的小规模模型上进行了对比，基线包括标准Transformer（使用FFN）和同等参数量的MoE模型。

乘积键记忆（Product-Key Memory） 是一种基于键值检索的神经网络模块，最早由 Lample 等人于 2019 年提出。它将记忆组织为两组键向量的笛卡尔积：每个记忆槽由一对键（k₁, k₂）的乘积索引，查询向量分别与两组键计算内积得到权重，再聚合对应的值向量。与传统稠密记忆相比，乘积键结构将记忆容量从 O(N) 扩展到 O(N²) 而只需 O(N) 的存储和计算，因此在同等开销下可容纳更多记忆槽。在 Transformer 中，它通常被插入到注意力层之后，作为 FFN 的替代或补充，用于存储事实性知识或模式。其核心超参数包括键的维度、每组键的数量以及值向量的维度，这些选择直接影响检索精度与计算开销的平衡。

实验结果：与MoE持平，但未超越

根据Reddit帖子披露的信息，使用共享乘积键记忆的Transformer在语言建模任务上达到了与MoE相近的困惑度（perplexity），参数总量为2300万。然而，与标准FFN Transformer相比，该方案并没有展现出明显优势。作者将这一结果称为「负面结果」，意味着假设未得到验证——共享记忆模块并未带来预期的效率或性能提升。

这一发现与部分已有研究相呼应：在较小规模下，记忆模块的收益可能被其自身的开销抵消。乘积键记忆需要计算查询与所有键的内积，当键的数量较多时，计算复杂度会显著上升。此外，共享记忆意味着所有层使用同一组键值，可能限制了模型在不同抽象层次上学习差异化表示的能力。

困惑度（Perplexity） 是语言建模中最常用的内在评估指标，定义为模型对测试集每个 token 预测概率的几何平均的倒数，即 PPL = exp(-(1/N)Σ log p(xᵢ))。困惑度越低，说明模型对文本的预测越准确。它等价于模型在每个 token 上的平均分支因子：困惑度为 20 意味着模型平均在 20 个候选词中犹豫。在相同词表和数据集下，困惑度可直接比较不同架构的建模能力，但它对分词方式和词表大小敏感，跨设置对比时需谨慎。本实验中，共享乘积键记忆与 MoE 达到相近困惑度，说明两者在语言建模的似然拟合上能力相当，但困惑度并不直接反映下游任务表现或推理效率，这也是作者强调需结合其他指标综合判断的原因。

为什么共享乘积键记忆未能胜出？

从架构设计角度看，FFN在每个位置独立作用，通过非线性变换增加模型容量；而记忆模块依赖于检索机制，其有效性高度依赖于键的质量和查询的分布。在训练初期，键向量随机初始化，检索到的信息可能噪声较大，导致梯度信号不稳定。共享机制进一步加剧了这一问题，因为不同层对记忆的访问模式可能相互冲突。

另一个可能的原因是参数效率的衡量方式。虽然记忆模块可能减少了FFN的参数量，但检索操作引入了额外的计算图和内存访问开销。在2300万参数的小模型上，这些开销的相对占比更高，可能掩盖了参数减少带来的好处。

混合专家（MoE, Mixture of Experts） 是一种条件计算架构，由多个并行的“专家”子网络和一个门控网络组成。门控网络为每个输入 token 计算各专家的权重，通常只激活 top-k 个专家（如 top-2），从而在总参数量很大的情况下保持每个 token 的计算量恒定。MoE 的优势在于参数规模可随专家数量线性增长而计算量近似不变，但代价是路由不均衡、训练不稳定以及显存占用高。本实验将 MoE 作为基线之一，是因为它同样旨在替代或增强 FFN 的容量效率：与共享乘积键记忆的“检索式”思路不同，MoE 走的是“稀疏激活”路线。两者在 23M 参数下打平，说明在该规模下稀疏激活与键值检索带来的收益相近，但都未能超越稠密 FFN 的性价比。

开源推理基准：附带的价值

除了负面结果本身，作者还开源了一个推理基准（open reasoning benchmark）。该基准旨在评估模型在逻辑推理任务上的表现，可能包含多步推理、数学问题或常识推理等题型。虽然帖子未详细说明基准的具体内容，但开源基准对于社区复现和比较不同架构的性能不能忽视。

负面结果和开源基准的结合，体现了开放研究的精神：即使假设未成立，实验过程和工具仍能推动领域进步。其他研究者可以基于此基准测试自己的记忆增强模型，避免重复踩坑。

对高效Transformer设计的启示

这项实验提醒我们，在小规模模型上，架构创新需要谨慎评估。FFN之所以成为标配，是因为它在参数效率和训练稳定性之间取得了良好平衡。任何替代方案都需要在多个维度上证明自己：参数更少、计算更低、性能不降。共享乘积键记忆目前尚未展示出全面优势，但这并不意味着记忆机制没有前景。

未来的方向可能包括：分层共享（不同层使用不同的记忆子集）、稀疏检索（只激活部分键）、以及与MoE结合使用。例如，可以用记忆模块处理高频模式，用FFN处理长尾知识。此外，在更大规模（如1亿参数以上）上重复该实验，或许能揭示记忆模块的缩放规律。

总结

Reddit上分享的这项实验用共享乘积键记忆替代FFN，在2300万参数下与MoE性能相当，但未超越标准Transformer。作者诚实标记为负面结果，并开源了推理基准。这一工作为高效模型架构的探索提供了实证数据，也提醒社区：小规模下的负面结果同样有价值，它们划定了当前方法的边界，指引后续研究避开低效路径。

对于关注AI编程和模型架构的开发者，可以关注该开源基准，并在自己的项目中测试记忆模块的变体。