arXiv:2501.12352cs.LGcs.AI2025-01被引 65

用测试时回归统一解释序列模型设计原理

Test-time regression: a unifying framework for designing sequence models with associative memory

  • 将记忆与召回建模为回归问题,通过测试时优化实现关联记忆
  • 揭示线性注意力无法捕捉令牌间相关性,解释软最大注意力归一化的有效性
  • 提出新注意力变体,适合研究序列建模理论与架构设计的学者

序列模型是现代深度学习的核心。然而快速发展的各类架构(如Transformer与递归模型)看似无关。本文提出一个统一框架,基于关联回忆(即检索上下文相关标记)的实证重要性,将记忆与召回形式化为两步过程:记忆(建模为回归问题)与检索。结合这两步的层通过输入标记的“测试时回归”实现关联回忆。包括线性注意力、状态空间模型、快速权重编程器、在线学习器和软最大注意力在内的主流结构,均作为三个设计选择的特例:回归权重、回归函数类与测试时优化算法。该方法解释了线性注意力在捕捉令牌间相关性上的局限,并为软最大注意力中的查询-键归一化提供了数学依据。此外,揭示了设计空间中未探索区域,并由此推导出软最大注意力的高阶推广形式。本工作不仅实现统一,还将序列建模与经典回归方法连接,为构建更强大且理论严谨的架构开辟道路。

原文摘要 · Abstract (English)

Sequence models lie at the heart of modern deep learning. However, rapid advancements have produced a diversity of seemingly unrelated architectures, such as Transformers and recurrent alternatives. In this paper, we introduce a unifying framework to understand and derive these sequence models, inspired by the empirical importance of associative recall, the capability to retrieve contextually relevant tokens. We formalize associative recall as a two-step process, memorization and retrieval, casting memorization as a regression problem. Layers that combine these two steps perform associative recall via ``test-time regression'' over its input tokens. Prominent layers, including linear attention, state-space models, fast-weight programmers, online learners, and softmax attention, arise as special cases defined by three design choices: the regression weights, the regressor function class, and the test-time optimization algorithm. Our approach clarifies how linear attention fails to capture inter-token correlations and offers a mathematical justification for the empirical effectiveness of query-key normalization in softmax attention. Further, it illuminates unexplored regions within the design space, which we use to derive novel higher-order generalizations of softmax attention. Beyond unification, our work bridges sequence modeling with classic regression methods, a field with extensive literature, paving the way for developing more powerful and theoretically principled architectures.

序列模型注意力机制回归框架架构统一

Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。