DeepSeek 最新论文 《Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models》 代码解读

仓库连接如下：deepseek-ai/Engram

DeepSeek 新作：Engram —— 为大模型开启 $O(1)$ 查找的“条件存储”新纪元

在大型语言模型（LLM）的演进中，混合专家模型（MoE）通过计算稀疏性（Conditional Computation）成功扩展了模型容量。然而，DeepSeek 的研究团队发现，Transformer 在处理知识检索时效率并不高，它们往往在用昂贵的计算去“模拟”简单的查找过程。

为此，DeepSeek 提出了 Engram：一种现代化的 N-gram 查找模块。它引入了存储稀疏性（Memory Sparsity），通过 $O(1)$ 复杂度的查找直接获取知识。本文将结合代码实现，深入剖析 Engram 的核心设计。

一、核心思路：从计算稀疏到存储稀疏

传统的 MoE 是在激活层做文章，而 Engram 是在“记忆”上做文章。

N-gram 索引：不再只依赖当前 Token，而是利用前 $N$ 个 Token 的组合（N-gram）作为 Key。
$O(1)$ 查找：通过哈希函数将 N-gram 映射到超大规模的嵌入空间（Engram Vocab）。
条件门控：检索到的记忆不是直接相加，而是根据当前隐藏状态（Hidden States）进行“条件合并”。

二、代码实现详解

代码实现主要分为四个部分：分词器预处理、N-gram 哈希映射、多头嵌入检索 以及 Engram 核心逻辑。

1. 分词器压缩 (CompressedTokenizer)

为了让 Engram 更好地泛化（例如让 “Apple” 和 “apple” 指向同一个记忆点），DeepSeek 引入了规范化和压缩机制。

class CompressedTokenizer:
    """
    对原始 Token 进行规范化处理（如转小写、去除重音、NFKC 规范化），
    并将原始词表映射到一个更紧凑的‘规范化词表’中。
    """
    def __init__(self, tokenizer_name_or_path):
        self.tokenizer = AutoTokenizer.from_pretrained(tokenizer_name_or_path)
        # 规范化流程：转小写、处理空格、规范化字符
        self.normalizer = normalizers.Sequence([
            normalizers.NFKC(),
            normalizers.Lowercase(),
            normalizers.Replace(Regex(r"[ \t\r\n]+"), " "),
            normalizers.Strip(),
        ])
        self.lookup_table, self.num_new_token = self._build_lookup_table()

    def _build_lookup_table(self):
        # 遍历原始词表，将解码后的文本规范化后作为 Key
        # 从而实现多个相似 Token ID 对应同一个 Compressed ID
        # ... (构建映射表的逻辑)
        return lookup, len(new_tokens)

2. N-gram 哈希映射 (NgramHashMapping)

Engram 的灵魂在于如何快速定位 N-gram。由于 N-gram 的组合爆炸（例如 3-gram 的组合数极大），DeepSeek 使用了哈希分桶。

class NgramHashMapping:
    """
    计算输入序列的 2-gram, 3-gram 哈希值。
    使用位运算（XOR）和质数取模来保证哈希分布的均匀性。
    """
    def _get_ngram_hashes(self, input_ids: np.ndarray, layer_id: int) -> np.ndarray:
        # 1. 准备滑动窗口（Shift k）
        # tokens[0] 是当前 token, tokens[1] 是前一个, tokens[2] 是前前一个
        base_shifts = [shift_k(k) for k in range(self.max_ngram_size)]
        
        all_hashes = []
        for n in range(2, self.max_ngram_size + 1):
            # 2. 对 N-gram 进行哈希混合
            # mix = (token_0 * mult_0) ^ (token_1 * mult_1) ...
            tokens = base_shifts[:n]
            mix = tokens[0] * multipliers[0]
            for k in range(1, n):
                mix = np.bitwise_xor(mix, tokens[k] * multipliers[k])
            
            # 3. 多头映射：每个 N-gram 对应多个哈希头，减小冲突影响
            for j in range(self.n_head_per_ngram):
                mod = self.vocab_size_across_layers[layer_id][n-2][j]
                all_hashes.append(mix % mod) # 取模到各自的质数桶中
        return np.stack(all_hashes, axis=2)

3. 多头嵌入与条件门控 (Engram Module)

这是 Engram 模块的主体。它不仅查找记忆，还通过当前模型的 hidden_states 计算一个“门控信号”（Gate），实现条件查找。

class Engram(nn.Module):
    def forward(self, hidden_states, input_ids):
        """
        hidden_states: [B, L, D] (来自 Backbone 的特征)
        input_ids: [B, L] (原始 Token ID)
        """
        # 1. 计算所有 N-gram 的哈希 ID
        hash_input_ids = self.hash_mapping.hash(input_ids)[self.layer_id]
        
        # 2. 查找大规模 Embedding 空间
        # embeddings 形状: [B, L, num_heads * head_dim]
        embeddings = self.multi_head_embedding(hash_input_ids).flatten(start_dim=-2)
        
        # 3. 条件门控计算 (Conditional Gating)
        gates = []
        for hc_idx in range(backbone_config.hc_mult):
            # 将检索到的记忆投影到 Key 空间
            key = self.key_projs[hc_idx](embeddings)
            normed_key = self.norm1[hc_idx](key)
            
            # 将当前层特征作为 Query
            query = hidden_states[:, :, hc_idx, :]
            normed_query = self.norm2[hc_idx](query)
            
            # 计算相似度作为门控信号，决定读取多少记忆
            gate = (normed_key * normed_query).sum(dim=-1) / math.sqrt(D)
            gate = gate.sigmoid().unsqueeze(-1)
            gates.append(gate)
            
        # 4. 合并记忆并应用短卷积（ShortConv）增加局部上下文
        value = torch.stack(gates, dim=2) * self.value_proj(embeddings).unsqueeze(2)
        output = value + self.short_conv(value)
        return output

4. 短卷积 (ShortConv)

为了弥补 N-gram 静态查找可能丢失的局部时序信息，Engram 在最后加入了一个轻量级的深度可分离卷积（Depth-wise Conv）。

Python

class ShortConv(nn.Module):
    def __init__(self, hidden_size, kernel_size=4, dilation=1):
        # 使用分组卷积实现深度可分离卷积，极其轻量
        self.conv = nn.Conv1d(
            in_channels=total_channels,
            out_channels=total_channels,
            kernel_size=kernel_size,
            groups=total_channels, # 每个通道独立卷积
            padding=(kernel_size - 1) * dilation,
            dilation=dilation,
        )

三、总结：为什么要关注 Engram？

DeepSeek 的这篇论文提出了一个非常有前景的方向：Sparsity Allocation（稀疏性分配）。

U 型比例定律：论文发现计算稀疏（MoE）和存储稀疏（Engram）之间存在最优比例。盲目增加模型参数（计算量）不如分配一部分给“静态记忆”。
性能飞跃：在同等 FLOPs 的情况下，加入 Engram 的模型在 MMLU 和推理任务上表现显著优于纯 MoE 模型。
未来的轴：这为 LLM 的扩展提供了一个新维度——不再只是堆叠层数或专家数，而是构建一个可无限扩展的底层“知识库”。

Engram 的代码实现展示了 DeepSeek 如何巧妙地将古老的 N-gram 思想与现代的高性能计算（哈希、门控、稀疏性）结合。这种“返璞归真”的工程创新，或许正是通往更智能模型的一条捷径。

DeepSeek 新作：Engram —— 为大模型开启 $O(1)$ 查找的“条件存储”新纪元

一、 核心思路：从计算稀疏到存储稀疏

二、 代码实现详解