7D-AI系列：DeepSeek Engram 架构核心数学公式

Engram的github仓库地址：https://github.com/deepseek-ai/Engram

本文档详细说明了 Engram 架构中使用的核心数学公式及其解释。

一、N-gram 哈希计算公式

1.1 公式说明

N-gram 哈希是 Engram 架构的核心，用于将 N-gram 序列映射到哈希索引，实现 O(1) 的查找。

1.2 数学表达式

对于 N-gram 序列[t₀, t₁, ..., tₙ₋₁]：

步骤 1：混合计算（XOR 混合）

mix = (t₀ × m₀) ⊕ (t₁ × m₁) ⊕ ... ⊕ (tₙ₋₁ × mₙ₋₁)

参数说明：

• tᵢ：第 i 个 token 的 ID（压缩后的 token ID）
• mᵢ：该层第 i 个位置的随机乘数（奇数，通过r × 2 + 1生成）
• ⊕：按位异或（XOR）运算

代码实现：

mix=(tokens[0]*multipliers[0])forkinrange(1,n):mix=np.bitwise_xor(mix,tokens[k]*multipliers[k])

步骤 2：哈希索引计算

hash = mix mod p

参数说明：

• mix：混合后的值
• p：质数模数（每个头使用不同的质数，减少哈希冲突）
• mod：取模运算

代码实现：

head_hash=mix%mod# mod 是质数

1.3 完整公式

对于 N-gram 大小为 n，第 j 个头的哈希索引：

hash_{n,j} = ((t₀ × m₀) ⊕ (t₁ × m₁) ⊕ ... ⊕ (tₙ₋₁ × mₙ₋₁)) mod p_{n,j}

其中p_{n,j}是第 n 个 N-gram 的第 j 个头使用的质数模数。

1.4 设计要点

1. XOR 混合的优势：
   • 快速计算
   • 良好的分布特性
   • 可逆性（在已知乘数的情况下）
2. 质数模数的优势：
   • 减少哈希冲突
   • 更好的分布均匀性
   • 数学上更稳定
3. 多层乘数：
   • 每层使用不同的随机乘数，确保不同层的哈希函数不同
   • 通过种子确保可复现性：seed_layer = seed + PRIME_1 × layer_id

二、门控机制计算公式

2.1 公式说明

门控机制用于动态融合静态 N-gram 记忆与动态隐藏状态，决定 N-gram 记忆的融合强度。

2.2 数学表达式

步骤 1：相似度计算（点积）

s = (K_norm · Q_norm) / √d

参数说明：

• K_norm：归一化后的键（从 N-gram 嵌入生成）
  • K = Linear_key(E_ngram)
  • K_norm = RMSNorm(K)
• Q_norm：归一化后的查询（从隐藏状态生成）
  • Q = hidden_states[:, :, hc_idx, :]
  • Q_norm = RMSNorm(Q)
• d：隐藏层维度（hidden_size）
• ·：点积（内积）运算

代码实现：

key=self.key_projs[hc_idx](embeddings)normed_key=self.norm1[hc_idx](key)query=hidden_states[:,:,hc_idx,:]normed_query=self.norm2[hc_idx](query)gate=(normed_key*normed_query).sum(dim=-1)/math.sqrt(backbone_config.hidden_size)

步骤 2：门控激活函数（特殊设计）

gate = σ(√(|s| + ε) × sign(s))

参数说明：

• s：相似度值
• ε = 1e-6：防止除零的小常数
• sign(s)：符号函数（保留原始符号）
• σ：Sigmoid 函数，σ(x) = 1 / (1 + e^(-x))

代码实现：

gate=gate.abs().clamp_min(1e-6).sqrt()*gate.sign()gate=gate.sigmoid()

激活函数设计说明：

1. 先取绝对值并限制最小值：|s| + ε
2. 开方：√(|s| + ε)
3. 恢复符号：sign(s) × √(|s| + ε)
4. 最后应用 sigmoid：σ(...)

这种设计可以：

• 增强门控的区分度
• 更好地控制 N-gram 记忆的融合强度
• 避免梯度消失问题

步骤 3：值融合

output = gate ⊙ V_proj

参数说明：

• V_proj：投影后的 N-gram 嵌入值（V_proj = Linear_value(E_ngram)）
• gate：门控值（形状为[B, L, 1]）
• ⊙：逐元素乘法（Hadamard 积）

代码实现：

value=gates*self.value_proj(embeddings).unsqueeze(2)

2.3 完整公式

对于第hc_idx个 hyper-connection 通道：

K = Linear_key(E_ngram) Q = hidden_states[:, :, hc_idx, :] s = (RMSNorm(K) · RMSNorm(Q)) / √d gate = σ(√(|s| + ε) × sign(s)) V = Linear_value(E_ngram) output = gate ⊙ V

三、残差连接公式

3.1 公式说明

通过短卷积增强特征，并使用残差连接保持梯度流动。

3.2 数学表达式

output = V_gated + Conv(V_gated)

参数说明：

• V_gated：门控后的值（gate ⊙ V_proj）
• Conv：短卷积操作（深度可分离卷积）
• +：残差连接（逐元素相加）

代码实现：

output=value+self.short_conv(value)

3.3 设计要点

1. 残差连接的优势：
   • 保持梯度流动
   • 允许模型学习残差（增量改进）
   • 提高训练稳定性
2. 短卷积的作用：
   • 捕获局部模式
   • 增强序列依赖关系
   • 提高特征表达能力

四、时间偏移公式

4.1 公式说明

时间偏移用于构建 N-gram 序列，通过填充和截取实现不同位置的 token 组合。

4.2 数学表达式

对于位置i的 token，N-gram 序列为：

[tᵢ₋ₖ, tᵢ₋ₖ₊₁, ..., tᵢ]

参数说明：

• k：偏移量（0 ≤ k < n）
• tᵢ₋ₖ：当i - k < 0时，使用pad_id填充

代码实现：

defshift_k(k:int)->np.ndarray:ifk==0:returnx shifted=np.pad(x,((0,0),(k,0)),mode='constant',constant_values=self.pad_id)[:,:T]returnshifted

4.3 示例

对于序列[t₀, t₁, t₂, t₃]和 3-gram：

• shift_0:[t₀, t₁, t₂, t₃]（无偏移）
• shift_1:[pad, t₀, t₁, t₂]（偏移 1）
• shift_2:[pad, pad, t₀, t₁]（偏移 2）

在位置i=2的 3-gram 为：[t₀, t₁, t₂]（来自shift_0[2]、shift_1[2]、shift_2[2]）

五、公式总结

5.1 核心公式列表

5.2 数据流

输入文本 ↓ Token 压缩 ↓ N-gram 哈希: hash = XOR_mix(tokens) mod prime ↓ 嵌入查找: E_ngram = Embedding[hash] ↓ 门控计算: gate = σ(√(|K·Q|/√d + ε) × sign(K·Q)) ↓ 值融合: V = gate ⊙ Linear_value(E_ngram) ↓ 卷积增强: output = V + Conv(V) ↓ 输出（增强的隐藏状态）

六、关键设计思想

1. O(1) 查找：通过确定性哈希函数，相同 N-gram 总是映射到相同的嵌入表索引，实现 O(1) 的查找复杂度。
2. 条件记忆：通过门控机制动态决定 N-gram 记忆的融合强度，而不是固定融合，提高了模型的灵活性。
3. 稀疏激活：仅在特定层（如第 1 层和第 15 层）使用 Engram，实现稀疏激活，平衡计算效率和性能。
4. 多头设计：每个 N-gram 使用多个头（如 8 个头），每个头使用不同的质数模数，减少哈希冲突。

七、参考文献

• Engram 论文：Conditional Memory via Scalable Lookup: A New Axis of Sparsity for Large Language Models
• 代码实现：engram_demo_v1.py

需要我帮你提取文档中的核心公式代码片段并整理成可直接运行的 Python 示例吗？