带前瞻键的因果注意力机制
计算与语言
2025-09-30 v2 机器学习
摘要
在标准因果注意力中，每个 token 的查询、键和值（QKV）是静态的，且仅编码前文上下文。我们提出了一种带前瞻键的因果注意力机制（CASTLE），该注意力机制随着上下文的展开持续更新每个 token 的键。我们将这些更新后的键称为前瞻键，因为它们属于较早的位置，却整合了相对于这些位置出现在更晚的 token 的信息，同时严格保持了自回归性质。尽管该机制看似是顺序的，我们推导出一个数学等价形式，避免了在每个位置显式地具体化前瞻键，从而实现了高效的并行训练。在语言建模基准测试中，CASTLE 在不同模型规模上始终优于标准因果注意力，降低了验证困惑度，并提升了一系列下游任务的性能。
引用
@article{arxiv.2509.07301,
title = {Causal Attention with Lookahead Keys},
author = {Zhuoqing Song and Peng Sun and Huizhuo Yuan and Quanquan Gu},
journal= {arXiv preprint arXiv:2509.07301},
year = {2025}
}