基于门控关联检索的通用三重潜在压缩

# 通用三阶隐状态压缩与门控关联检索
来源: https://arxiv.org/html/2606.05175

###### 摘要

压缩后的隐状态能否恢复高阶的token交互，而无需显式枚举token三元组或依赖特定基准的解析？我们研究了一个小型通用三阶隐状态序列模型家族，它维护一个运行中的token状态以及一条压缩后的成对记忆通路。由此产生的基础架构是循环的、无类型的，并且仅从原始token序列中训练。我们首先评估了三种变体——密集成对记忆（三阶隐状态）、槽压缩成对记忆（三阶槽）、成对记忆加局部卷积（三阶混合）——与标准Transformer基线在字节级WikiText-2、联想召回和吞吐量扩展上的比较。在等宽设置下，所有三种隐状态变体在验证字节每比特上都优于Transformer基线，其中三阶混合达到4.766，而基线为5.124。在近似参数匹配的语言模型跟进实验中，三阶混合仍然优于Transformer，在169,424个参数下达到5.067，而Transformer在174,848个参数下为5.124。有针对性的接近预算扫描发现了一个更强的优胜者：三阶混合，d=64，状态维度88，在177,752个参数下达到4.827，并在三个随机种子上平均为4.806，而Transformer基线为5.143。纯隐状态模型在重联想召回任务上并未改进，因此我们添加了一个面向召回的混合模型：独立的键值记忆、前一个token查询、最后一层检索以及门控后期融合路径。这个门控混合模型在d=64时，在三个随机种子上平均联想召回率达到41.9%，而Transformer为25.0%，普通三阶混合为11.6%，并在最佳种子上达到100.0%。我们还在一个基于分词器的MiniMind语言模型基准上，在共享的三种子、120步协议下进行了测试。在那里，Transformer基线平均token损失为7.317，transformer_plus_triple为7.239，三阶混合为7.035，而面向召回的門控混合模型为6.766。增益对种子敏感，并且在当前的Python参考实现中，这些模型仍然远慢于注意力机制。后续的预热加课程重训练尝试将门控模型回调到12.7%的平均召回率。我们将这些结果解释为：隐状态压缩有助于通用语言建模（包括在token化预训练基准上），而精确检索则受益于独立的门控记忆路径，而非仅靠单个压缩状态。代码和实验脚本可在https://github.com/xiaol/Autoresearch_ideas.git 获取。

## 1 引言

高阶token交互处于序列建模的核心，但捕获它们的标准方式——全自注意力——带来二次方序列交互，并且高度依赖高度优化的内核(Vaswani et al., 2017 (https://arxiv.org/html/2606.05175#bib.bib1))。循环和线性时间替代方案旨在用压缩状态取代该成本，但往往在此过程中损失表达能力。

本文研究一个狭窄的问题：通用的隐状态能否将类似三元组的token交互压缩得足够好，从而在标准语言模型基准上有所帮助？这个问题由两个观察启发。首先，矩阵状态和混合循环模型表明，更丰富的状态更新可以恢复更多通常归因于注意力的行为(Mishra and others, 2026 (https://arxiv.org/html/2606.05175#bib.bib12); Peng et al., 2024 (https://arxiv.org/html/2606.05175#bib.bib13))。其次，针对特定机制的合成基准虽然有用，但不足以证明架构有助于通用序列建模。

因此，我们聚焦于刻意简单的基准栈：字节级WikiText-2语言建模(Merity et al., 2017 (https://arxiv.org/html/2606.05175#bib.bib17))、基于分词器的MiniMind预训练基准、合成联想召回压力测试，以及Apple MPS上的序列长度吞吐量分析。我们的模型家族是刻意的通用型：它只接收token ID，不解析有类型的角色，也不使用任何任务特定的符号头。

本文现在报告三个相互关联的故事。首先，基础三阶隐状态家族已经在通用字节级语言建模上优于小型Transformer。其次，这些增益迁移到了token化的MiniMind基准，包括一个transformer_plus_triple兼容模型，它保留标准自注意力查找，仅将三阶交互作为后期读取分支添加。第三，原始的纯隐状态模型在联想召回上失败，最佳修复方案并非将精确检索强制放入同一个压缩状态中。相反，一个面向召回的混合模型添加了独立的键值路径和门控后期融合。

#### 贡献。

- • 我们定义了一个通用的三阶隐状态架构家族，结合了运行中的token状态与压缩的成对记忆通路。
- • 我们扩展了这个家族，增加了一个面向召回的混合模型，保留三阶隐状态路径用于压缩，同时添加独立的键值记忆、前一个token查询和门控后期融合用于检索。
- • 我们在仓库中提供了一个专用的架构图和一个独立的基准驱动程序。
- • 我们展示了所有三种等宽三阶隐状态变体在小型字节级WikiText-2上均优于Transformer基线，并且最强的三阶混合变体在近似参数匹配后仍然改进。
- • 我们添加了一个基于分词器的MiniMind基准，并展示了所有测试的三阶增强模型在平均token损失上均优于Transformer；transformer_plus_triple提供了更干净的兼容路径，而门控混合模型在这个小型基准中是LM质量最强的模型。
- • 我们展示了纯隐状态家族在联想召回上失败，但门控混合检索路径在三个种子上的平均表现可以优于Transformer在该基准上的表现，尽管存在显著的种子敏感性。
- • 我们清楚地展示了相应的系统注意事项：在当前参考实现中，这些模型均未实现时钟时间吞吐量的胜利。

#### 代码可用性。

## 2 相关工作

第一个背景是日益增长的关于二次注意力高效替代方案的文献。深度循环架构如Universal Transformers(Dehghani et al., 2019 (https://arxiv.org/html/2606.05175#bib.bib2))、恢复循环更新的线性注意力公式(Katharopoulos et al., 2020 (https://arxiv.org/html/2606.05175#bib.bib3))、留存和状态空间替代方案如RetNet和Mamba(Sun et al., 2023 (https://arxiv.org/html/2606.05175#bib.bib7); Gu and Dao, 2023 (https://arxiv.org/html/2606.05175#bib.bib8))，以及最近的循环复兴如RWKV、Gated DeltaNet、xLSTM和矩阵状态语言模型(Peng et al., 2023 (https://arxiv.org/html/2606.05175#bib.bib16); Zhu and others, 2026 (https://arxiv.org/html/2606.05175#bib.bib14); Beck et al., 2024 (https://arxiv.org/html/2606.05175#bib.bib11); Mishra and others, 2026 (https://arxiv.org/html/2606.05175#bib.bib12); Peng et al., 2024 (https://arxiv.org/html/2606.05175#bib.bib13))——它们都研究用压缩状态和线性时间序列交互能恢复多少注意力的行为。我们的基础三阶隐状态家族属于这一路线，但明确地将运行中的token状态与面向高阶交互的压缩成对记忆路径分离开。

第二个背景是显式记忆和检索增强。乘积键记忆、最近邻语言模型和记忆Transformer表明，精确或近似联想查找可以补充参数化序列建模(Lample et al., 2019 (https://arxiv.org/html/2606.05175#bib.bib9); Khandelwal et al., 2019 (https://arxiv.org/html/2606.05175#bib.bib10); Wu et al., 2022 (https://arxiv.org/html/2606.05175#bib.bib6))。我们的面向召回的混合模型在精神上最接近这个家族和最近的混合联想记忆公式(Lufkin et al., 2026 (https://arxiv.org/html/2606.05175#bib.bib24))，但它是刻意的最小化：检索仅在输出层添加，有损的循环压缩路径保持不变。

最后，最近面向生产的混合模型通常结合多种序列机制，而不是承诺单一原语。压缩Transformer、Jamba、Samba和Hymba展示了混合局部压缩、注意力、循环或状态空间组件的不同方式(Rae et al., 2019 (https://arxiv.org/html/2606.05175#bib.bib5); Lieber et al., 2024 (https://arxiv.org/html/2606.05175#bib.bib25); Ren et al., 2024 (https://arxiv.org/html/2606.05175#bib.bib26); Dong et al., 2024 (https://arxiv.org/html/2606.05175#bib.bib27))。我们这里的目标更小规模、更具诊断性：使用受控的基准栈来分离隐压缩帮助什么与精确检索在单个通用架构中帮助什么。

## 3 方法

### 3.1 通用三阶隐状态层

每层接收token特征xt∈Rdx_t ∈ R^d，并产生四组投影：

at, bt, qtl, qtr ∈ RH×S, (1)

连同用于运行状态sts_t和压缩成对记忆通路的可学习衰减门。状态更新为

st = λt ⊙ s_{t-1} + (1−λt) ⊙ at, (2)

其中λt是一个可学习的每token保留门。然后密集成对记忆存储前一个状态与当前写入向量之间的交互：

Pt = γt ⊙ P_{t-1} + (1−γt) ⊙ (s_{t-1} ⊗ bt), (3)

其中Pt ∈ RH×S×S。读取首先用左查询探查压缩成对记忆，然后将结果与右查询组合：

ot = Wo vec( (Pt qtl) ⊙ qtr ). (4)

层输出将这个混合读取添加回残差流，然后应用FFN块。图2 (https://arxiv.org/html/2606.05175#S3.F2) 对比了这条高阶通路与普通自注意力：自注意力一次对一个查询和一个键进行评分，而三阶隐状态路径在压缩之前显式地从token三元组形成一个联合特征。

### 3.2 变体

我们研究三个通用变体。

- • 三阶隐状态：密集成对记忆Pt。
- • 三阶槽：用可学习槽替换密集成对记忆，存储低秩左因子和右因子。
- • 三阶混合：保持密集成对记忆，并在循环栈之前添加局部卷积路径。

与基准感知的类型隐状态模型不同，这些变体不暴露特殊角色、标签或符号任务头。

### 3.3 Transformer+三元组兼容基线

为了区分“三阶交互有帮助”与“完全替换注意力”，我们还评估了一个兼容基线，它保持标准Transformer栈，只在输出添加一个三阶隐状态层。设h_t是最终Transformer隐藏状态，e_t表示token嵌入流。我们运行一个单层三阶隐状态层，使用token源左键，没有关联或直接查找分支，以产生高阶读取c_t。最终logits为

lt = W_vocab h_t + W_vocab ( LN(c_t) ). (5)

这个transformer_plus_triple模型保留普通自注意力键值查找，并将三阶通路视为附加的高阶残差头。我们仅在MiniMind和联想召回跟进中将其用作诊断性兼容基线。

### 3.4 门控混合联想检索

基础三阶隐状态家族能很好地压缩上下文，但不提供明确的内容可寻址查找。为了测试精确检索是否应该单独处理，我们添加一个面向召回的混合模型，它保留三阶隐状态栈，仅在最顶层增强一个显式键值记忆。

给定token特征e_t，检索路径将键值对(k_{t-1}, v_t)写入只追加的记忆，使用可学习投影

kt = Wk e_t, vt = Wv e_t. (6)

在检索时，我们使用前一个token表示查询记忆，

qt = Wq e_{t-1}, (7)

并计算与存储的归一化键的top-k相似度得分：

α_{t,i} ∝ exp( q̂t^T k̂i / √da ). (8)

检索到的记忆向量为

mt = ∑_i α_{t,i} v_i. (9)

然后我们应用可学习读取门，并只在输出层融合：

m̃t = σ(Wg qt) ⊙ mt, lt = W_vocab h_t + W_vocab ( LN(m̃t) ). (10)

这个设计使得三阶隐状态状态不受检索写入影响，并将精确查找与有损循环压缩分离。图1 (https://arxiv.org/html/2606.05175#S3.F1) 总结了这种分离：三阶交互首先计算，然后通过压缩隐状态路径存储，而精确检索作为下游可选分支保持独立。

参见图注图1：通用三阶隐状态家族的简单直观图。三个token通过显式的高阶混合步骤共同交互，而不仅仅通过成对注意力分数。产生的交互然后被写入一个压缩的循环隐状态路径，而可选的确切键值分支可以保持独立，仅在输出处融合。参见图注图2：自注意力计算成对查询-令牌兼容性并用加权和聚合值。三阶注意力则在交互被压缩成隐状态之前，从token三元组显式形成一个高阶特征。

## 4 实验设置

#### 语言建模。

我们在WikiText-2原始文本上训练字节级自回归模型(Merity et al., 2017 (https://arxiv.org/html/2606.05175#bib.bib17))。输入是UTF-8字节，词汇表大小为259。等宽设置使用d=64，3层，4头，序列长度128，批量大小16，在Apple MPS上运行80个优化器步骤。

#### MiniMind基于分词器的LM。

为了测试字节级增益是否迁移到更标准的token化设置，我们还在仓库中提供的MiniMind预训练语料库和分词器上进行基准测试。我们遵循现有MiniMind工具包使用的相同小型标准协议：隐藏大小96，2层，4头，序列长度96，批量大小4，120个优化器步骤，32个评估批次，学习率3e-4，确定性1-in-20留出，种子{1,2,3}。比较的模型包括Transformer、三阶隐状态、三阶混合、transformer_plus_triple和面向召回的門控混合模型。

#### 联想召回。

我们还在一个合成键值检索任务上训练每个模型，任务包含四个键值对、随机填充和一个最终查询。模型运行200步，批量大小32，序列长度128。准确性仅在最终答案token上测量。

#### 面向召回的混合跟进。

由于原始的纯隐状态家族和早期无门控双记忆变体在联想召回上仍接近随机猜测，我们运行了第二个专门的实验来测试图1中的门控混合设计。