用测试时回归统一解释序列模型设计原理
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 官方产品;中文卡片由大模型生成,请以原文为准。