🤖 AI

Mask-Based Priors Are More Persistent than Query-Key Initializations

本文表明，直接利用任务级交互结构来初始化加法注意力掩码，比查询-键（query-key）初始化能提供更持久且有效的归纳偏置，使 Transformer 能够在布尔推理任务上实现近乎完美的外推，并在低数据量算术场景中提升性能。

原作者： Mingze Ma, Hemanth Saratchandran, Cameron Gordon, Simon Lucey

发布于 2026-08-04

📖 1 分钟阅读☕ 轻松阅读

原作者： Mingze Ma, Hemanth Saratchandran, Cameron Gordon, Simon Lucey

原始论文采用 CC BY 4.0 许可（http://creativecommons.org/licenses/by/4.0/）。 ✨ 这是对下方论文的AI生成解释。它不是由作者撰写或认可的。如需技术准确性，请参阅原始论文。阅读完整免责声明

想象一下，你正在教一个机器人玩解谜游戏。你向它展示了的一百万个游戏范例，它因此变得非常擅长。但随后，你让它玩一个稍微有些不同的新版本游戏，而这个版本是它从未见过的。机器人并没有去理解新的规则，而是找到了一个“捷径”——一个在旧范例中有效、但在新规则下完全错误的简单技巧。这是人工智能领域，特别是针对一种类脑计算机模型——Transformer——中一个著名的难题。这些模型是许多现代 AI 工具（如编写故事、解决数学问题和识别面部）背后的引擎。科学家们发现，虽然这些机器人非常擅长记忆它们所见过的东西，但它们往往无法进行“外推”（extrapolate），即无法将知识应用于新情境，因为它们总是会被困在最简单的答案而非“真正”的答案上。

研究人员提出的核心问题是：我们如何阻止机器人走捷径？其中一个想法是在机器人开始学习之前，通过先教给它规则来给它一个“领先优势”。在 Transformer 的世界里，这被称为归纳偏置（inductive bias）。这就像是在侦探开始破案前先给他一张地图。如果地图是正确的，侦探就能快速找到罪犯；如果地图是错误的，或者侦探忽略了地图，他就会迷失方向。长期以来，科学家们试图通过调整机器人的内部“传感器”（称为查询与键投影，Query and Key projections）来绘制这些地图，使它们能自然地观察到正确的线索。但事实证明，机器人由于太渴望从数据中学习，很快就会忘记你给它的地图，并回到走捷径的习惯上。

这篇题为《基于掩码的先验比查询-键初始化更具持久性》（Mask-Based Priors Are More Persistent than Query-Key Initializations）的论文深入探讨了这个问题。来自澳大利亚机器学习研究所的 Mingze Ma 及其团队提出了一个简单而深刻的问题：有没有一种方法，能给机器人一张它“无法忘记”的地图？他们测试了两种绘图方法。第一种是旧方法：调整传感器。第二种是更简单的新方法：使用一种特殊的掩码（mask），直接将地图绘制在机器人的“注意力”机制上。

以下是他们的发现。当使用旧方法（调整传感器）时，机器人起初拥有正确的地图，但一旦训练开始，机器人就会抹掉这张地图，并退回到错误的捷径答案。这就像是在白板上写字，而机器人立刻将其擦除。然而，当使用新的基于掩码的方法时，地图却留了下来。他们没有改变传感器，而是直接向注意力分数中添加了一个永久的、可学习的偏差——一个“推动力”。这个推动力是如此持久，以至于经过数千步训练后，机器人仍然记得规则。

结果非常显著。在机器人通常会失败的复杂逻辑谜题（称为布尔任务）中，新方法让机器人在未见过的题目上达到了接近 100% 的正确率，而旧方法仍停留在 75%。这仿佛机器人终于学会了像人类一样思考，而不是像一个只会背小抄的记忆者。这种成功不仅限于逻辑谜题；团队还在数学问题（如加法和乘法运算大数字）以及图像识别方面进行了测试。在每种情况下，机器人学习得更快，且犯错更少，尤其是在数据量较少的情况下。

该论文表明，秘诀不仅在于你教给机器人“什么”，还在于你把知识教在“哪里”。通过将结构化规则直接放入注意力掩码中，而不是隐藏在复杂的传感器计算里，机器人就能将规则留在脑海中。作者展示了这种简单的技巧如何像一个坚固的锚一样，即使在数据浪潮试图将其冲离航道时，也能稳住机器人的学习方向。虽然这并不能解决 AI 的所有问题（团队也承认他们仍需预先知道规则才能绘制地图），但它为构建更聪明、更可靠、能够真正从经验中学习而非仅仅靠记忆的 AI 提供了一个强大的新工具。

技术摘要：基于掩码的先验比查询-键初始化更具持久性

问题陈述

Transformer 在处理未见数据（GOTU）时表现出系统性的失效模式，尤其是在布尔外推任务上。虽然这些模型可以完美拟合观测到的训练数据，但它们通常倾向于根据“最小度插值器”（MDI）进行外推——这是一种更简单的、低阶的函数，它虽然匹配训练域，但在处理未见输入时会失败——而非真正的、高阶的目标函数。这种行为表明存在一种特定的归纳偏置，即 Transformer 偏好低阶结构的解。

现有的改进尝试通过注入 查询-键（QK）初始化（例如，设置 $W_Q$ 和 $W_K$ 使 $QK^\top$ 近似于期望的注意力模式）来注入结构先验，但已被证明是不充分的。本文认为，虽然基于 QK 的方法可以在初始化时产生期望的注意力模式，但这种结构是间接且依赖于内容的。因此，注入的先验在训练过程中会被迅速覆盖或“冲刷掉”，因为模型会重新利用查询和键的投影来学习任务相关的表示，从而无法改变最终的外推规则。

方法论

作者提出了 基于掩码的结构化初始化（Mask-Based Structured Initialization） 作为一种更具持久性的替代方案。该方法不再通过依赖于内容的查询和键投影的相似度分数间接编码先验，而是将结构先验直接注入到注意力逻辑值（logits）中，作为一种 有限的可学习加性偏置。

核心机制

标准的自注意力机制被修改如下：
$\text{Attention}(X) = \text{softmax}\left( \frac{XW_Q W_K^\top X^\top}{\sqrt{d}} + \log \sigma(M) \right) XW_V$
其中：

$M \in \mathbb{R}^{N \times N}$ 是一个可学习的掩码矩阵，初始化自任务级的交互结构。
$\sigma(\cdot)$ 是 Sigmoid 函数，确保 $\log \sigma(M) \leq 0$ 。
项 $\log \sigma(M)$ 作为一个加性偏置，在保留标准注意力机制的同时，选择性地降低不理想的交互权重。

设计原则

掩码基于目标任务的函数分解进行初始化：

任务对齐交互： 将注意力限制在必须根据目标规则进行交互的 Token 上（例如，奇偶校验中的成对交互，或加法中的逐位对齐）。
抑制捷径： 抑制会导致 MDI 解的交互（例如，多数投票中的全局求和捷径）。
可学习性： 与硬性的因果掩码或填充掩码不同，这些掩码是有限且可学习的，允许模型在优化过程中精炼先验，同时保持核心结构约束。

理论依据

论文提供了梯度层面的分析（定理 1 和推论 1）来解释基于掩码的先验为何更具持久性：

QK 初始化： 对 $Q$ 和 $K$ 的梯度随范数 $\|q\|$ 和 $\|k\|$ 缩放。在实践中，较大的特征范数会引发巨大的更新，从而迅速覆盖初始化的结构。
掩码初始化： 对掩码参数 $m_j$ 的梯度受限于 $p_i \sigma(-m_j) \leq p_i$ 。由于掩码与依赖于内容的特征规模解耦，并受到 Sigmoid 因子的调节，其更新更加缓慢且稳定，从而在整个优化过程中保留了结构偏置。

核心贡献

识别了先验持久性失效问题： 论文证明了基于 QK 的结构化初始化虽然在初始化时有效，但在针对布尔 GOTU 任务进行训练时，无法保留注入的先验，导致模型陷入默认的 MDI 偏置。
提出基于掩码的初始化方法： 这是一种新颖的方法，通过将结构先验直接作为加性注意力逻辑值偏置注入，实现了结构约束与依赖于内容的相似度得分的分离。
梯度层面解释： 提供了理论证明，表明加性掩码将结构偏置与特征规模解耦，从而实现了比 QK 方法更稳定的更新和更高的持久性。
实证验证： 在布尔推理、算术和视觉基准测试中的全面评估表明，基于掩码的初始化成功地将外推行为从 MDI 解转向了目标函数。

实验结果

布尔外推 (GOTU)

在 2Parity、3-Bit Majority 和 Cyclic3 等 Transformer 通常会失败（由于 MDI 偏置，在 2Parity 上准确率仅为 ~75%）的任务中：

Vanilla 及 QK 初始化模型： 无法逃脱 MDI 偏置，在 2Parity 上维持约 75% 的准确率，在 3-Bit Majority 上维持约 87%。
掩码初始化模型： 实现了 近乎完美的外推（2Parity 和 3-Bit Majority 达到 100% 准确率，Cyclic3 达到 99.98%）。损失函数显著下降（例如，在 2Parity 上从 1.379 降至 0.0029）。

算术任务

在低数据比例下的 3 位数加法 和乘法任务中：

掩码初始化显著优于 Vanilla 和基于 QK 的方法（SVD 和 Opt）。
在仅有 0.4% 训练数据的加法任务中，掩码初始化实现的序列准确率为 77.49%，而默认模型仅为 1.19%。
在 3% 数据量的乘法任务中，掩码初始化达到了 69.49% 的序列准确率，而默认模型仅为 5.01%。

视觉基准测试

在使用 ViT-Tiny 处理 CIFAR10、CIFAR100 和 ImageNet100 时：

掩码初始化始终取得最高的 Top-1 准确率，优于默认初始化和先前的结构化方法（Mimetic 和 Impulse）。
例如，在 CIFAR10 上，掩码初始化达到 95.66%，而默认模型为 94.20%，Impulse 为 95.12%。

NLP 任务

在显式结构先验定义较弱的 WikiText-2 和 TinyStories 上：

掩码初始化（使用因果相对距离先验）在困惑度（perplexity）和收敛速度方面带来了虽小但一致的提升，展示了其在“真实”先验并非严格已知时的适应能力。

重要性与主张

论文声称：先验所处的位置决定了它是否能改变外推行为。通过将结构先验从依赖于内容的 $QK^\top$ 投影转移到加性的注意力逻辑值中，作者证明了归纳偏置可以被赋予持久性。

这项工作的意义在于证明了注意力掩码不仅是用于因果性或局部性的架构约束，还可以作为 编码持久归纳偏置的简单基质。这种方法允许 Transformer 在不修改架构或引入新操作的情况下，克服其在逻辑和组合任务中向最小度插值器倾斜的默认倾向。作者指出，尽管该方法假设可以获取有意义的交互先验（即“预知未来/Oracle”假设），但它证明了当这些先验通过掩码正确编码时，可以从根本上改变 Transformer 的学习动力学和泛化能力。