Lngram v2: 具有可解释离散表示的潜在 N-Gram 记忆

# Lngram v2：具备可解释离散表示的潜在N元语法记忆
来源：https://arxiv.org/html/2609.03426

###### 摘要
Transformer缺乏原生查找机制，需要重复的密集计算来识别和重用局部静态模式。Lngram v1通过离散潜在N元语法寻址引入了与分词器无关的条件记忆，但其记忆容量与骨干网络宽度耦合，导致参数和激活成本高昂，限制了可扩展性。我们提出了Lngram v2，它解耦了路由数量、记忆维度和骨干网络宽度，并引入了上下文感知的分组查询注意力读出，以独立扩展记忆容量。零值接收槽和反事实替代梯度进一步提高了读出选择性和路由可训练性，同时保留了硬离散寻址。在不同规模的视觉语言模型（VLM）上的实验显示了一致的改进，包括成功扩展到300亿参数模型。与Lngram v1相比，Lngram v2在保持或提高语言建模性能的同时，大幅减少了总参数量和每标记的激活参数。进一步分析表明，其离散ID保留了连续隐状态的大量语义结构，使得仅从ID就能恢复语义，并在不同数据集上保持稳定的ID-语义关联。这些结果确立了Lngram v2作为一种高效且可扩展的潜在条件记忆机制，其离散地址也为分析内部模型表示提供了结构化接口。

††脚注文本：1 北京邮电大学。2 快手科技。* 项目负责人。†‡ 通讯作者：王小杰和温斌。

## 1 引言
Transformer已成为大型语言和多模态模型的主流骨干架构（Vaswani等人，2017 (https://arxiv.org/html/2609.03426#bib.bib1)；Devlin等人，2019 (https://arxiv.org/html/2609.03426#bib.bib31)；Brown等人，2020 (https://arxiv.org/html/2609.03426#bib.bib32)；Dosovitskiy等人，2021 (https://arxiv.org/html/2609.03426#bib.bib33)；Radford等人，2021 (https://arxiv.org/html/2609.03426#bib.bib34)；Alayrac等人，2022 (https://arxiv.org/html/2609.03426#bib.bib35)）。然而，注意力和前馈网络必须同时处理上下文相关的推理和更接近条件查找的操作，例如识别常见的多标记实体和重复的局部模式。由于标准Transformer缺乏原生查找机制，这些模式仍然必须通过密集计算来近似，消耗了本可用于支持组合推理的模型容量（Cheng等人，2026 (https://arxiv.org/html/2609.03426#bib.bib36)）。

为了将局部模式匹配与骨干网络分离，Cheng等人（2026）(https://arxiv.org/html/2609.03426#bib.bib36)引入了基于标记ID N元语法的条件记忆。然而，其地址依赖于分词器ID和固定哈希，使得该机制对分词边界、哈希冲突敏感，并且难以统一扩展到文本之外。Lngram v1则将隐状态离散化为可学习符号，并在潜在空间中构建精确的N元语法地址，将经典N元语法建模（Brants等人，2007 (https://arxiv.org/html/2609.03426#bib.bib37)）扩展到神经表示，实现了对文本、视觉和跨模态输入的条件记忆。尽管如此，Lngram v1仍有两个局限。首先，其路由数量与骨干网络宽度耦合，导致记忆大小和读出成本随模型维度增加，限制了可扩展性。其次，虽然离散Lngram ID直接决定记忆地址，但尚不清楚这些ID是否保留了其来源连续表示的语义结构。现有的表示分析方法通常需要连续激活以及探测器或学习到的字典（Alain和Bengio，2016 (https://arxiv.org/html/2609.03426#bib.bib2)；Hewitt和Manning，2019 (https://arxiv.org/html/2609.03426#bib.bib3)；Tenney等人，2019 (https://arxiv.org/html/2609.03426#bib.bib4)；Bau等人，2017 (https://arxiv.org/html/2609.03426#bib.bib5)；Huben等人，2024 (https://arxiv.org/html/2609.03426#bib.bib6)）。如果Lngram ID本身保留了语义信息，它们可能为记忆寻址和表示分析提供一个有用的接口。

因此，我们提出了Lngram v2，它解耦了路由数量、记忆维度和骨干网络宽度，使记忆容量可以独立扩展。检索到的不同阶N元语法记忆被组织为记忆标记，当前隐状态通过分组查询注意力（GQA）（Ainslie等人，2023 (https://arxiv.org/html/2609.03426#bib.bib7)）执行上下文相关的读出。固定的零值接收槽允许模型抑制不可靠的检索，而按路由流式处理在大规模下控制了中间激活。为了训练硬离散路由，我们使用反事实查找替代梯度，该梯度提供依赖于记忆的优化信号，同时在训练和推理期间保持前向寻址路径一致，遵循离散决策替代优化的一般原则（Bengio等人，2013 (https://arxiv.org/html/2609.03426#bib.bib8)）。我们进一步使用仅接受多路由离散签名的ID-only语义读取器研究Lngram离散表示中编码的语义。该读取器保留了连续隐状态超额语义读出能力的65.77%–84.27%。在保留数据上，单个路由代码表现出与人物、动物、车辆和常见物体等概念的可重现关联，而大部分语义能力来源于路由间的分布式组合。这些结果表明，Lngram ID不仅作为记忆地址，还作为内部表示的可统计分析的离散接口。

我们在多个规模的视觉语言模型上评估了Lngram v2。在匹配的训练数据、计算预算和骨干网络优化设置下，它持续提升了Keye2B和Keye30B的性能，并成功扩展到300亿参数模型。参数匹配的稀疏FFN和乘积键记忆（PKM）基线表明，这些增益不能仅由额外参数或稀疏容量来解释。与Lngram v1相比，Lngram v2将模块总参数减少了82.6%，每标记激活参数减少了95.2%，同时保持或改善了验证损失。总体而言，Lngram v2提供了一个可扩展的潜在空间条件记忆机制，以及用于分析模型表示的结构化离散接口。

## 2 架构
### 2.1 概述
如图1 (https://arxiv.org/html/2609.03426#S2.F1)所示，Lngram v2是一个插入到Transformer解码器层中的条件记忆分支，用于将局部模式匹配和存储从骨干网络中卸载。忽略批量维度，第ℓ层的输入为 H(l) = [h1, ..., hT]ᵀ ∈ ℝᵀˣᵈ。

Lngram v2将每个隐状态映射到多个离散符号，为精确记忆查找构建局部N元语法地址，并将检索到的条目组织为按路由的记忆标记。当前隐状态随后通过分组查询注意力（GQA）对这些标记执行上下文相关的读出，并辅以固定的零值接收槽来抑制不可靠的检索。输出进一步由短程因果卷积处理，并作为残差注入骨干网络。前向传递始终使用硬离散寻址，而离散化投影则使用反事实替代梯度进行训练。对于大量路由，读出可以沿路由维度流式处理以减少激活内存。

### 2.2 多路由离散寻址
Lngram v2通过独立投影将隐状态离散化：
U = RMSNorm(H), Z = UW_q, W_q ∈ ℝᵈˣᴿᴹ,
其中R是路由数量，M是每个路由的位数。对于位置t、路由 r ∈ {0, ..., R-1} 和位 j ∈ {0, ..., M-1}，b_{t,r,j} = I[z_{t,r,j} > 0]，这些位被打包成离散符号 a_{t,r} = ∑_{j=0}^{M-1} b_{t,r,j} 2^j, a_{t,r} ∈ {0, ..., K-1}, K = 2^M。
每个隐状态因此被映射到R个并行的离散符号。由于地址源自潜在表示而非分词器ID，相同机制适用于跨模态。

### 2.3 精确N元语法检索
对于每个阶 n ∈ N，Lngram v2维护一个记忆表 E⁽ⁿ⁾ ∈ ℝᴿᴷⁿˣᵈᵐ，其中 dₘ 是记忆条目维度。对于路由 r ∈ {0, ..., R-1}，在位置t结束的N元语法被编码为 g_{t,r}⁽ⁿ⁾ = rKⁿ + ∑_{i=0}^{n-1} a_{t-n+1+i,r} K^i，并通过 m_{t,r}⁽ⁿ⁾ = E⁽ⁿ⁾[g_{t,r}⁽ⁿ⁾] 精确检索。不同路由占据不相交的地址范围，为每个路由-N元语法组合产生唯一条目，无需近似搜索。无效前缀和跨越打包序列边界的N元语法被掩码。不同阶的检索条目被连接成 x_{t,r} = Concat_{n∈N} m_{t,r}⁽ⁿ⁾ ∈ ℝ^{|N|dₘ}，每个位置产生R个记忆标记。

### 2.4 上下文感知读出
Lngram v2使用当前隐状态在检索到的路由标记之间进行上下文选择。令 H_q 和 H_kv 表示查询和KV头的数量，头维度 d_h：H_q d_h = d, G = H_q / H_kv。
查询和记忆键/值为 Q_t = Reshape(Norm(h_t) W_Q) ∈ ℝ^{H_q × d_h}, [K_{t,r}; V_{t,r}] = Reshape(Norm(x_{t,r}) W_{KV}), K_{t,r}, V_{t,r} ∈ ℝ^{H_kv × d_h}。
对于查询头 a ∈ {0, ..., H_q-1}，定义 κ(a) = ⌊a / G⌋，并计算 s_{t,a,r} = (Q_{t,a}ᵀ K_{t,r,κ(a)}) / √d_h。
为了避免将概率质量强加给不可靠的记忆，我们引入了一个固定logit b₀的零值接收槽：
α_{t,a,r} = exp(s_{t,a,r}) / (exp(b₀) + ∑_{r'=0}^{R-1} exp(s_{t,a,r'})),
o_{t,a} = ∑_{r=0}^{R-1} α_{t,a,r} V_{t,r,κ(a)}。
我们设置 b₀ = log(R(1-μ₀)/μ₀)，其中μ₀指定分配给记忆路由的初始总概率质量。
头输出被投影并通过深度可分离因果卷积精炼：
ṽ_t = Concat_{a=0}^{H_q-1} o_{t,a} W_O, 且 Ṽ = [ṽ_1, ..., Ṽ_T]ᵀ。
然后序列表示被精炼为 Y = Ṽ + SiLU(DWConv1D(Norm(Ṽ)))，核大小为4。
结果作为前注意力残差注入：H⁽ℓ⁾ ← H⁽ℓ⁾ + Y。
对于大的R，路由维度可以分块处理，同时在线累积softmax归一化和值加权和，避免同时存储所有路由激活。

### 2.5 通过离散寻址的反向传播
前向传递始终使用硬决策 b_{t,r,j} = I[z_{t,r,j} > 0]，这是不可微的。因此，我们使用反事实记忆查找来训练离散化投影，同时保持硬前向路径。对于N元语法窗口中的一个符号，令 z = (z_0, ..., z_{M-1}), p_j = σ(τ z_j)。
保持其他符号固定，将该符号替换为 c ∈ {0, ..., K-1} 得到反事实检索 E⁽ᶜ⁾。其概率为 P(c | z) = ∏_{j=0}^{M-1} p_j^{β_j(c)} (1-p_j)^{1-β_j(c)}，给出条件期望 ̃E(z) = ∑_{c=0}^{K-1} P(c | z) E⁽ᶜ⁾。
此期望仅用于路由梯度；实际前向输出仍为硬查找 E_hard。
在大规模实验中，我们使用一位近似。对于位j，令 E_j⁽⁰⁾ 和 E_j⁽¹⁾ 表示将该位固定为0或1同时固定所有其他位时获得的反事实检索。给定上游梯度g，∂L/∂z_j ≈ λ τ p_j (1-p_j) ⟨g, E_j⁽¹⁾ - E_j⁽⁰⁾⟩。
仅替换离散化投影的反向梯度：通过硬查找选择的记忆条目接收标准梯度，而反事实条目仅用于估计路由梯度。因此，训练和推理共享完全相同的硬寻址和精确查找前向路径。进一步的推导和实现细节见附录A (https://arxiv.org/html/2609.03426#A1)。

## 3 分析
Lngram v2将连续隐状态映射为离散ID用于记忆寻址。我们研究这些ID是否保留原始表示的语义，这种语义是否可以在不访问连续隐状态的情况下恢复，以及它们是局限于单个路由代码还是分布于多个路由。

### 3.1 实验设置
我们分析了一个2B参数的Keye2B+Lngram v2模型中的两个Lngram层，该模型训练了大约200亿额外标记（第4节 (https://arxiv.org/html/2609.03426#S4)）。每层包含128个路由，每个路由16个代码，为每个标记产生离散签名 C = (c_1, ..., c_128)。我们使用10,000张COCO 2017图像（Lin等人，2014 (https://arxiv.org/html/2609.03426#bib.bib9)），分为7,000张参考图像和3,000张保留图像。所有统计标准和读取器仅在参考集上确定，并在保留集上独立评估。

我们考虑三个语义目标：标题提及、视觉存在和局部存在，分别对应于文本描述中表达的概念、完整图像和局部图像区域。我们首先验证离散空间没有坍塌。在分析的两层中，每个路由平均使用约八个有效代码，最频繁的代码约占出现次数的30%，并且所有4,096个理论可能的代码在每层中的相对频率介于10⁻⁵和0.3之间。这些代码形成一个非平凡的离散空间。