Kalman Delta Networks: 不确定性感知的联想记忆

TL;DR 将线性注意力重构为线性高斯状态空间模型，用卡尔曼滤波显式跟踪记忆不确定性，提出扫描兼容的 KDN 近似，在语言建模上超越 DeltaNet 等线性注意力基线。

问题

线性注意力（Linear Attention）在前沿语言模型中用于高效长上下文推理，核心是固定大小的循环联想记忆，每个token需在线决定写入内容和覆盖强度。

现有Delta-rule模型（如DeltaNet、Gated DeltaNet）从当前token嵌入学习写入强度，但不追踪记忆估计的置信度。这导致写入无法适应累积证据：例如某记忆位置已被多次可靠更新，系统仍以固定强度覆盖，可能丢失关键信息。技术上，它们用各向同性代理预测协方差，忽略协方差追踪，使更新退化为固定增益Kalman滤波器特例。

该问题难在不确定性感知与GPU并行的平衡。最优估计要求Kalman滤波，但精确的Riccati递归是密集且状态依赖的，不适合线性注意力扫描。业界需要既保持常数内存解码，又提升记忆质量，因此对不确定性建模是重要方向。

类比：类似推荐系统中根据用户历史交互置信度动态调整新行为权重，而非每次同等覆盖旧兴趣。

核心洞察

• 将线性注意力的联想记忆显式建模为带不确定性的隐状态，并用卡尔曼滤波执行最优递归估计。不同于 Delta 规则只从当前 token 学习一个固定写入强度，KDA 跟踪记忆协方差，使每次写入根据累积证据和观测可靠性自适应加权，避免过早覆盖重要关联，从而提升记忆的长期稳定性。
• 推导出扫描兼容的 KDN 近似：通过对角变分推断将后验协方差投影为对角高斯，使不确定性递归成为 Mobius 映射，支持 associative scan 的对数并行深度。原始卡尔曼 Riccati 递归是稠密状态依赖的，不适用于 GPU 并行扫描；Diagonal KDN 和 Isotropic KDN 在保留关键不确定性信息的同时实现硬件高效实现，化解了贝叶斯记忆与线性注意力效率之间的矛盾。
• 揭示 Delta 规则模型本质上是固定增益卡尔曼滤波器，其写入强度相当于用一个 token-wise 各向同性替代预测协方差。这一形式化将 Delta 规则统一到卡尔曼滤波框架下，说明 Delta 模型实际采用了简化的不确定性假设；而 KDN 通过显式协方差跟踪提供了更通用的写入策略，为设计更精细的记忆更新机制奠定了理论基础。

方法

输入与状态建模

Kalman Delta Networks (KDNs) 接收线性注意力中经过投影的 query / key / value 表示，并将递归联想记忆重新形式化为 线性-高斯状态空间模型。记忆状态作为隐变量，每个 token 的写入视为带噪声的观测，从而为显式跟踪不确定性提供概率基础。

关键模块：卡尔曼滤波更新与扫描近似

1. 预测：状态转移同时传播记忆状态和协方差，预测当前 token 处的先验分布。
2. 更新：使用 Kalman gain 对残差写入加权，增益由累积证据和观测可靠性自适应决定——相比之下，Delta-rule 模型相当于采用固定增益的特例。
3. 高效近似：精确追踪需要状态相关的 Riccati 递归，不利于 GPU 并行扫描。本文提出 Diagonal KDN（在线平均场变分推断，将后验投影到对角高斯）和 Isotropic KDN（每头单标量不确定性）。两者的不确定性递归为 Möbius 映射，可结合对数并行深度的关联扫描与硬件高效的 chunkwise 实现。

输出与差异点

输出为更新后的记忆状态及其不确定性估计，可直接用于后续线性注意力查询。与现有线性注意力基线（如 Gated DeltaNet / KDA）相比，KDNs 首次在递归联想记忆中保留预测协方差的精确估计，而非用 token-wise 各向同性代理替代，从而在 750M 和 1.3B 参数规模下稳定提升 perplexity 与平均下游准确率。

实验

实验设计

论文在 750M 和 1.3B 参数规模上进行受控预训练，评估语言建模困惑度、常识推理平均准确率以及上下文检索能力。基线包括 DeltaNet、Gated DeltaNet、Mamba 等 state-of-the-art 线性注意力模型。实验还通过消融研究考察信息尺度、观测噪声与过程噪声的影响，并对比了吞吐量。

关键发现

KDN 变体（Diagonal KDN 与 Isotropic KDN）在预训练困惑度与下游任务平均准确率上一致优于现有线性注意力模型。Diagonal KDN 通过在线变分推断将协方差投影到对角高斯族，保持扫描兼容性；Isotropic KDN 引入每头单一不确定性标量，进一步降低计算量。两种变体的不确定性递归均为 Mobius 映射，可进行对数并行深度的关联扫描。

与基线对比解读

相比 DeltaNet 的固定增益更新，KDN 显式跟踪记忆估计不确定性，使 Kalman gain 能够依据累积证据与观测可靠性自适应加权残差写入。这解决了 Delta-rule 模型每次更新强度仅由当前 token 决定、无法适应未来查询需求的问题。硬件高效的 chunkwise 实现使得 KDN 在 GPU 并行扫描上可行，理论复杂度与 DeltaNet 相当，但带来一致性能提升。消融表明信息尺度与噪声设置对模型行为有重要影响，后续可进一步探索更灵活的协方差近似。

行业影响

落地场景

KDNs 适合需要长上下文高效推理的 LLM 产品，如实时对话系统、代码助手、企业文档问答。其不确定性感知的记忆更新在流式处理、长文档摘要、多轮对话状态跟踪中可提升稳定性。具体案例：电商客服机器人处理长对话历史时，KDN 的卡尔曼增益能自适应决定何时覆盖旧关联，减少上下文混淆；内容平台的自动标签与检索系统处理长文本时，常量内存特性支持边缘端或低成本推理。

商业价值

核心收益在降本与体验提升：常量内存解码和并行扫描降低长序列推理的显存占用与延迟，直接减少 GPU 成本；不确定性感知写入减少错误覆盖，提高长文本任务准确率，降低人工审核与纠错成本。对于按 token 计费的 API 服务，更稳定的长上下文表现可提升企业客户留存与客单价。

与现有产品/工作流的接口

KDN 可作为线性注意力层的 drop-in 替换，集成到现有 Transformer 架构，只需替换注意力模块并调整训练流程。推理侧需实现chunkwise 扫描与自定义 CUDA kernel 以发挥并行优势；现有框架如 vLLM、TensorRT-LLM 可通过插件支持。训练时需适配预训练数据，因卡尔曼递归对数值稳定性敏感，可能需要混合精度与梯度裁剪调整。可借助现有 SSM 推理库（如 Mamba 生态）进行优化。

局限

• **近似协方差限制表达力**：为满足 GPU 并行扫描，Diagonal KDN 与 Isotropic KDN 分别采用**对角协方差**和**单标量各向同性协方差**近似完整后验。这牺牲了记忆状态各维度之间的相关性建模能力，可能限制模型捕获复杂的联想模式。在需要精细不确定性建模的任务中，近似估计会偏离最优卡尔曼滤波，而精确的 **Riccati 递归**又因计算密集不可用。论文未系统比较近似解与精确解在小型设置下的差距，也未展示该近似对最终性能影响的边界，因此无法判断近似误差在实际应用中的容忍度。
• **实验规模与任务覆盖有限**：预训练仅包含 **750M** 与 **1.3B** 参数模型，且评估局限于语言建模困惑度、常识推理和上下文检索平均准确率。缺乏更大规模（如 7B+）验证，也未在长上下文推理、代码生成或需要强记忆保持的任务上测试。此外，未深入分析内存不确定性在长序列中的实际行为（例如是否出现不确定性消失或爆炸），这限制了方法在真实生产规模下的可信度，难以断言其在大模型时代的普适性。
• **计算开销与部署门槛**：虽然 **chunkwise** 实现降低了扫描复杂度，但不确定性递归引入了额外状态与计算，可能增加推理延迟和显存占用。论文在消融中报告了吞吐量，但未与同等吞吐预算下的其他高效注意力方法（如 **FlashLinearAttention** 或基于 **Mamba** 的混合架构）进行公平对比。实际部署中，额外的**不确定性参数**与 **Mobius 映射**计算需要专门优化，否则在资源受限场景难以推广，这削弱了其作为轻量级替换方案的竞争力。