昨天，DeepSeek发布了一篇新工作[1]，推出了一个叫Engram的模块。
这个工作和之前发布的 mHC，性质差不多，都是在模型算法层面上的改进，而不是大的模型版本更迭。
本文来看看这个 Engram 的模块是个什么东西。
动机
这篇文章的标题是：Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
翻译一下就是：大语言模型稀疏建模的新方向：条件化记忆与可扩展查找
看标题就知道，这篇文章的目的就是稀疏建模，和前面DSA之类的目的一样，那就是省钱。
昨天，我写了一篇文章，回顾了一下生成式语言模型的机理。
会发现，当前GPT架构的模型，都很暴力，不管输入的问题是长还是短，是简单还是复杂，都需要进行注意力计算。
举个例子，如果问模型：1+2=？
模型会先用分词器划分token，划分成1、+、2、=、?，这几个token，然后embedding成特征向量后，再变成Q、K、V三部分，计算Attention。
如果我们被问到这样的简单计算题，人脑中是否会经历这个计算过程呢？
如果问999+345=?，这种略复杂的计算题，脑子还需要转一下去计算，但1+2=3这种问题，未必要计算，因为小学课本里见过无数次这样的计算题，已经潜移默化成了静态知识的一部分，不需要再去算，所以就能“脱口而出”。
生活中，还有很多这样的问题，比如问模型一些历史题，模型如果有这方面的静态知识储备，直接输出来就行了，而不需要专门去推理，浪费算力。
这就是 Engram 的动机。这个出发点还是挺符合直觉的。
Engram具体架构
Engram具体是怎么做的呢？下图是该模块的架构图。
图中，当模型处理到Great这个Token时，它和前一个token会组成一个2-Gramthe Great，和前两个token会组成一个3-GramAlexander the Great。
这两个分别取计算哈希值，然后查表，获取到历史已经存在Embedding结果。
相当于之前这两个或三个Token放在一起的组合，如果在记忆表里出现过，那就可以直接拿出它们之前计算得到的中间结果，参与下一轮的计算。
这个原理有点类似于KV Cache，本质上还是空间换时间，只不过KV Cache对顺序有严格要求，Engram只要求满足局部顺序2/3个token的顺序。
这样做，会存在两个问题：
- 问题一：不同的词汇虽然语义相同(比如Apple和apple)，但变成Token之后，可能会非常离散，因为分词器会严格给不同的token分配不同ID，这一点如果没看懂，可以看我上一篇文章，里面有详细说明分词器的原理。
- 问题二：这张词汇表的内容是巨大的，因为要存不少token和其组合的中间结果，如果用一张哈希表来进行计算，容易产生哈希冲突：即不同的token组合却映射成了相同结果。
对于问题一，Engram给出的解决方案是用后处理的方式，比如全部转小写、去除特殊符号，让它们尽可能表示一致。
对于问题二，Engram给出的解决方案是用多张哈希表进行计算。如果一张哈希表的冲突概率是p，K张哈希表同时冲突的概率就是p K p^KpK，这样冲突的可能性就极低了。哈希表的查询，计算复杂度是O(1)，这样所带来的额外计算量微乎其微。
为什么Engram在这里用2/3-Gram，而不用4-Gram。这就涉及到一个复杂度的问题。
Engram并不是作为一个孤立模块存在的，在整个大模型中，有多个关键的部分都插入了Engram，这就意味着记忆表需要记录这些组合的多个静态中间状态，4-Gram会导致组合的种类呈指数级上升，从而使表的体积巨大。
然而，只用2/3-Gram会导致一个明显的问题是，模型太关注局部表现了。
一段话中，三个单词最多只能表示一种局部信息，这三个词放到不同的上下文语境中，结果有可能不一样。
如果静态查找的内容不符合上下文情景，反而会给模型带来“噪音”。
为了解决这个问题，Engram又引入了一种门控机制。
上图中，2-Gram Embedding 和 3-Gram Embedding 得到的结果进行拼接，然后通过两个线性层(Linear)映射成K和V。
整段内容的中间状态(input hidden)，也直接输入进去，作为Q，参考transformer的计算公式，来计算相关度。
如果相关度高于门控阈值，当前全局语境与查到的记忆值具有一致性，那么记忆值就可以被直接利用；
如果相关度低于门控阈值，意味着查询到的记忆值是噪音，会被丢弃，不去干扰后面的计算。
系统效率优化
理解完Engram的原理后，不难发现，这个模块的核心，是要维护一张记忆表，这非常占用空间，如果把这张表都放进显存中，会严重挤占部署模型所需要的空间。
为此，DeepSeek又做了如下图所示的优化。
在训练的时候，记忆表的内容需要高频变化，所以该内容还是需要驻留在显存中，以保证训练速度。
但是在推理时，记忆表的内容已经固定，只需要读取就行了，它可以直接放在内存里，利用前面层的计算作为缓冲区，来进行内容读取，从而降低读取速度的影响。
并且，作者还提到，可以根据内容出现的频率进行单独优化，比如频繁访问的内容，将其缓存在更快的存储层（例如 GPU HBM 或 DRAM），低频出现的长尾分布内容再放到内存，这样对延迟的影响最小。
资源分配的讨论
作者进一步从资源分配的角度，讨论了该如何合理分配计算资源。
假设：
- 𝑃 t o t 𝑃_{tot}Ptot表示模型的总参数量。
- 𝑃 a c t 𝑃_{act}Pact表示模型的激活的参数量。
- 𝑃 s p a r s e 𝑃_{sparse}Psparse(𝑃 t o t 𝑃_{tot}Ptot-𝑃 a c t 𝑃_{act}Pact)表示模型剩余可分配的参数量。
在MoE模型中，𝑃 s p a r s e 𝑃_{sparse}Psparse全部用来扩容专家(未激活的专家)，但Engram提出，可以把这部分参数进行重新分配。
作者给出了一个比例因子𝜌，
- 𝜌=1，表示这部分参数全部给MoE。
- 𝜌<1，MoE的部分分到𝜌，Engram分到1-𝜌。
作者发现，在总参数量和激活参数量固定的情况下，只调节𝜌，模型在验证集上的损失会出现一条U型曲线。
这意味着，纯MoE和纯Engram都不是最优解，最优解是混合两者，即将𝜌的取值限定为 0.75 到 0.8范围。
实验
为了验证Engram的有效性，作者基于Dense、MoE、Engram分别构建了几个模型，并在通用的Benchmark上进行测试，结果如下表所示。
这里需要注意，Engram的模型里面是包含MoE结构的，例如，Engram-27B中，路由专家从原本的72个消减成55个，节省出来的参数量用来构建Engram。
结果发现，在相同参数量27B的情况下，Engram果然比MoE效果要好，并且作者还在40B的参数量上，验证了Engram的scaling能力，发现是有效的，依然能继续涨点。
没准他们现在已经在进行上百B的Engram实验了，期待后续进一步scaling的效果。
分析
Engram有效的原理是什么，为了弄清这一点，作者又进行了可视化分析。
如下图所示，图a是通过LogitLens，来计算不同层的输出与最终结果的KL散度。
可以发现，Engram能够比纯MoE整体的KL散度下降的更快，这就意味着它的浅层就能获得MoE的更深层的信息。
图b和图c是通过 CKA 来分析Engram和纯MoE不同层之间的相似性。图中，白线在对角线以上，这就意味着Engram的浅层就和MoE的深层有很大的相似性。
这种理解非常直观，Engram就是给模型单独开了一条快速通道，让它能直接追上MoE算半天之后的结果。
看到这里，想到了Engram很像是听播客，听完行业大佬的分享，一下子获取到了原本可能需要十多年后才有的认知，这何尝不是一种 Engram 的具象化呢。
总结
看完文章，其实会发现DeepSeek做的很“克制”，它不像OpenAI那样，直接扩个十倍(GPT-2->GPT-3)看看再说，而是先考虑参数不变，看看方法的有效性。
总之，Engram 会比 mHC 更有普适性，作为即插即用的模块，估计又可以催生出一些科研成果了。
作为MoE架构的改进，如果它成为一种新的范式，那么，内存又该涨价了。
参考
[1] https://github.com/deepseek-ai/Engram/blob/main/Engram_paper.pdf
