让注意力机制提前感知未来信息，提升语言建模效果
Causal Attention with Lookahead Keys
- 动态更新关键词，让早期词元提前融入后续上下文
- 在多个模型规模下降低验证困惑度，提升下游任务表现
- 保持自回归性的同时实现并行训练，适合大模型应用
标准因果注意力中，每个词元的查询、键、值（QKV）是静态的，仅编码先前上下文。我们提出一种名为CASTLE（Causal Attention with Lookahead Keys）的新注意力机制，使每个词元的键在上下文展开过程中持续更新。这些更新后的键被称为前瞻键，虽属于较早位置，却整合了相对该位置更晚出现的词元信息，同时严格保持自回归特性。尽管机制看似顺序执行，我们推导出其数学等价形式，避免显式生成每位置的前瞻键，从而支持高效并行训练。在语言建模基准测试中，CASTLE在不同模型规模下均优于标准因果注意力，显著降低验证困惑度，并提升多项下游任务性能。
原文摘要 · Abstract (English)
In standard causal attention, each token's query, key, and value (QKV) are static and encode only preceding context. We introduce CAuSal aTtention with Lookahead kEys (CASTLE), an attention mechanism that continually updates each token's keys as the context unfolds. We term these updated keys lookahead keys because they belong to earlier positions yet integrate information from tokens that appear later relative to those positions, while strictly preserving the autoregressive property. Although the mechanism appears sequential, we derive a mathematical equivalence that avoids explicitly materializing lookahead keys at each position and enables efficient parallel training. On language modeling benchmarks, CASTLE consistently outperforms standard causal attention across model scales, reducing validation perplexity and improving performance on a range of downstream tasks.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品；中文卡片由大模型生成，请以原文为准。