Latest Exact Match Attention

📄 arXiv: 2609.25802v1 📥 PDF

作者: Moritz Brösamle

分类: cs.LG, cs.CC, cs.CL

发布日期: 2026-09-22

🔗 代码/项目: GITHUB

💡 一句话要点

提出最新精确匹配注意力机制以提升Transformer性能

🎯 匹配领域: 支柱九：具身大模型 (Embodied Foundation Models)

关键词: 注意力机制 Transformer 精确匹配 字级随机存取 模型训练 自然语言处理 长文本生成

📋 核心要点

  1. 现有的Transformer模型在处理长序列时面临计算和内存效率的挑战，尤其是在精确匹配的场景中。
  2. 本文提出的LEMA机制通过二值化查询和键，确保每个查询仅关注最新的完全匹配键，从而提高了效率。
  3. 实验结果表明，LEMA模型在合成任务中超越了GDN，并在生成速度上与GDN持平，显示出其在长距离记忆检索中的优势。

📝 摘要（中文）

本文介绍了一种新的注意力变体——最新精确匹配注意力（LEMA），该方法对Transformer模型中的查询和键进行了二值化处理，使每个查询仅关注最新的完全匹配键。我们证明了LEMA Transformer能够模拟字级随机存取存储器（word-RAM），并且相较于以往的硬注意力变体，LEMA的精确匹配限制使得其在计算和内存方面与word-RAM之间建立了密切的对应关系。通过提出一种训练方法，LEMA模型在合成关联回忆任务中表现优异，超越了固定状态大小的Gated DeltaNet（GDN）。

🔬 方法详解

问题定义：本文旨在解决现有Transformer模型在处理长序列时的计算和内存效率问题，尤其是在需要精确匹配的场景中，传统方法的灵活性不足，导致性能瓶颈。

核心思路：LEMA机制通过将查询和键进行二值化处理，使得每个查询仅关注最新的完全匹配键。这种设计不仅提高了计算效率，还简化了模型的内存管理。

技术框架：LEMA模型的整体架构包括二值化的查询和键生成模块、精确匹配注意力计算模块以及基于直通估计器的训练方法。模型通过这些模块实现高效的注意力机制和训练过程。

关键创新：LEMA的最大创新在于其对精确匹配的严格限制，使得word-RAM能够以独立于上下文长度的成本模拟LEMA Transformer，从而在计算和内存方面建立了更紧密的联系。

关键设计：在训练过程中，采用了直通估计器来处理二值化操作，并使用软注意力替代品逐步调整至LEMA。此外，模型的参数设置和损失函数设计也经过精心调整，以确保在合成任务中获得最佳性能。

🖼️ 关键图片

fig_0
img_1
img_2

📊 实验亮点

在合成关联回忆任务中，LEMA模型的表现超越了固定状态大小的GDN，显示出其在处理长距离记忆检索方面的优势。此外，LEMA模型在生成速度上与GDN持平，尽管其状态大小不断增长，展现了良好的效率。

🎯 应用场景

LEMA机制的潜在应用场景包括自然语言处理中的长文本生成、信息检索和对话系统等领域。其高效的记忆管理和计算能力使其在需要处理大量信息的任务中具有实际价值，未来可能推动更复杂模型的开发与应用。

📄 摘要（原文）

We introduce latest exact match attention (LEMA), an attention variant for transformers where queries and keys are binarized and each query attends only to the latest exactly matching key. We prove that LEMA transformers with chain of thought can simulate word-RAMs, as was recently shown for the less restrictive rightmost hard attention. In contrast to prior hard attention variants, the restriction to exact matches enables an efficient converse direction: word-RAMs can simulate LEMA transformers at a cost per token independent of the context length. Together, these results yield a close correspondence between the two computational models in terms of both compute and memory. Beyond the theory, we propose a training method for LEMA transformers that handles their non-differentiable operations with a straight-through estimator for the binarization and a soft attention surrogate annealed towards LEMA. On a synthetic associative recall task, LEMA models trained this way use their growing state to store and recall a large number of associations, outperforming gated DeltaNet (GDN) with its fixed state size. As a first scaling test, we train LEMA language models with up to 834 million parameters. They match softmax transformers of around half their size in loss and, on repeated rare phrases and a needle-retrieval task, remain behind softmax transformers but recall across longer distances than GDN models of comparable size. Finally, we implement dictionary-based inference for LEMA transformers and show constant generation speed comparable to GDN despite their growing state, with the dictionaries residing in main memory rather than VRAM. Code is available at https://github.com/moritzbroe/latest_exact_match_attention.