LACE: 用于跨线程探索的格子注意力机制

# LACE：用于跨线程探索的晶格注意力机制

来源：https://arxiv.org/html/2604.15529

###### 摘要

当前的大语言模型独立地进行推理。虽然通常会并行采样多个推理路径，但这些轨迹互不交互，且常常以相同的冗余方式失败。我们介绍LACE，一个框架，将推理从独立试验的集合转变为协调的并行过程。通过重新利用模型架构以支持跨线程注意力，LACE允许并发推理路径在推理过程中分享中间见解并相互纠正。一个核心挑战是缺乏展现此类协作行为的自然训练数据。我们通过合成数据管道解决这一问题，该管道显式教导模型在线程间通信和错误纠正。实验表明，这种统一的探索显著优于标准并行搜索，推理准确度提升超过7个百分点。我们的结果表明，当允许并行推理路径交互时，大语言模型可以更有效。机器学习、ICML

参见图1 图1：标准隔离采样与我们的LACE框架的比较。蓝线代表常见的因果注意力，紫线表示我们的晶格注意力。顶部：传统LLM以隔离方式执行推理，这些独立线程缺乏通信，通常导致以相同方式失败。底部：LACE引入晶格注意力，支持跨线程交互。这种架构允许并发路径分享中间见解并进行即时纠正，引发协同推理，确保多样化探索和发现最优解的更高概率。

## 1 引言

推理很少是一条直线。当人类解决困难的证明或复杂的计划时，我们不会在真空中坚持单一路径（Kahneman，2011）。我们同时探索多个假设，允许一个线程的失败在另一个线程中转向搜索，在认知（Johnson-Laird，2010）和神经生理学（Kay等，2020）上都有体现。我们通过这种持续的内部对话应对不确定性，这是一个附带思维的过程。相比之下，大语言模型被限制在严格的顺序生成中（Vaswani等，2017），其中每个线程相互独立，生成k个不同的解决方案以选择最佳的（Cobbe等，2021；Zheng等，2025；Wu等，2025）。然而，这种方法本质上效率低下，因为采样时存在隔离。就像k个代理在分离的房间中解谜一样，线程无法分享突破或警告死路，如图1所示。因此，计算被浪费在冗余而非协同上；线程经常陷入相关的错误（Kim等，2025），因为隔离损害了多样性，减少了发现最优答案的可能性。此外，对隔离样本的事后验证不仅效率低下，而且容易出现内部偏见，因为模型可能在所有独立线程中始终倾向于某些有缺陷的推理模式（Madaan等，2023）。这激发了我们的核心问题：推理线程能否在生成过程中通信，以分享原位见解并协作解决问题？

在本文中，我们介绍LACE（用于跨线程探索的晶格注意力），一个重新配置变换器以支持附带思维的框架。我们将标准的1-D因果注意力推广为2-D晶格注意力结构，引入一个宽度维度，允许信息不仅在时间（令牌）上流动，也在线程上流动。这种架构在单个前向传递中将推理从一组隔离的独立事件转变为统一的协作探索。通过连接这些线程，LACE积极地多样化搜索策略以防止冗余失败（即不再犯同样的错误）并启用原位评估以即时识别和选择最优解决方案。如图1所示，这是隔离采样器按设计无法实现的能力。

实现这一能力需要克服重大的数据稀缺性挑战。标准预训练语料库缺乏多线程推理。与常规单线程数据不同，LACE需要相关但逻辑多样的并行线程，这是社区主要未探索的领域。通过高温采样（Toshniwal等，2024b）、少样本提示（Toshniwal等，2024a；Zheng等，2025）或多模型采样（Muennighoff等，2025）生成多个解决方案的数据集产生表面上的改述而非逻辑多样的推理路径。此外，它们是自成一体、相互独立的轨迹。同时，大多数并行化数据集（Wu等，2025）仅包含交错的短并行段在否则顺序的轨迹内，严重限制了持续跨线程协作的范围。基于这些，我们呈现一个合成数据管理管道，构造具有显式交互点的并行轨迹。与朴素的高温采样相比（产生表面改述），我们的管道生成具有本质上不同解决路径的线程。我们利用这些数据进行连续预训练，其次是监督微调（SFT）和强化学习（RL）。具体来说，我们在SFT期间引入随机抖动以强制跨线程信息共享，并设计专门的RL奖励以激励线程多样化探索并自我评估其解决方案以最大化准确性。可视化和数值结果表明我们的方法涌现地学会了在线程间分享见解并即时选择最佳解决方案，在AIME 25（Zhang & Math-AI，2025）、AIME 24（Zhang & Math-AI，2024）和LiveBench（White等，2025）等具有挑战的推理基准上优于独立采样基线。值得注意的是，这些能力涌现虽然模型仅在我们的合成数据上训练，展示了对现实世界任务的强大泛化。

我们的贡献总结如下：

- •我们提出**晶格注意力**，一个有原则的1-D因果注意力向2-D的推广，使多个推理线程能够在生成过程中即时通信，将隔离的并行采样转变为统一的协作探索。

- •我们设计了一个**多线程后训练框架**，配备专门的**合成数据管道**，生成相关但逻辑多样的推理线程，解决关键的数据稀缺性挑战并激励协作探索。

- •我们演示了LACE在具有挑战的基准上显著优于独立采样基线，具有涌现的跨线程通信和即时解决方案自选择能力，泛化超出合成训练数据。

## 2 相关工作

#### 外部测试时搜索和缩放

大语言模型的范式正从参数缩放转向推理时计算缩放（Snell等，2024；Wu等，2024），最近的研究表明额外的推理时计算（通常框架为系统2思维）可以产生与缩放模型大小相当的收益（Li等，2025）。现有方法主要通过外部搜索或编排实现这种额外计算，包括顺序推理策略如思维链及其变体（Wei等，2022b；Kojima等，2022；Zhou等，2022）、结构化搜索方法如思维树（Yao等，2023）和基于自一致性的并行采样管道（Wang等，2022）、自举最优-N聚合（Rakhsha等，2025）、学习的验证器或奖励模型（Stiennon等，2020；Cobbe等，2021）、RL基础的并行推理框架如并行R1和原生并行推理器（Zheng等，2025；Wu等，2025）以及宽度导向的外部平行思维如ParaThinker（Wen等，2025）。尽管存在差异，这些方法主要在模型原生令牌级生成过程之外协调推理，并仍然容易受到独立采样轨迹间的偏见或相关失败的影响（Ichihara等，2025；Huang等，2023；Madaan等，2023）。

#### 生成时并行推理

更近期的工作开始将并行推理移入生成过程本身。ParaDecode OneSeq（Yu等，2025b）将多个分支打包到单一序列中以用于效率导向解码。Hogwild!（Rodionov等，2025）通过共享跨线程KV状态启用显式并发注意力，而GroupThink（Hsu等，2025）通过直接跨线程交互研究并发推理代理之间的令牌级协作。虽然LACE与这些方法共享生成时跨线程交互的目标，但其机制不同。它不是通过共享KV状态和非标准掩码将每个线程暴露于其他线程的令牌历史，而是通过标准注意力衍生表示上的轻量级门控侧路径引入*隐式*跨线程交互。这保留了标准因果注意力骨干，同时允许线程在生成过程中相互影响，将独立采样转变为显式针对冗余探索和相关错误的协作过程（Kim等，2025；Hong & Page，2004）。

## 3 方法

我们在三个部分介绍LACE方法：第3.1节详细说明晶格注意力机制，第3.2节概述训练框架，第3.3节描述数据管理管道。整个训练管道的概述如图2所示。

参见图2 图2：晶格注意力层概述。t表示线程索引，l表示令牌位置。晶格注意力通过在对齐的令牌位置跨不同线程注意标准注意力的输出，实现跨线程通信。生成的跨线程上下文通过门控残差连接融合回主路径。这种设计允许每个线程受益于其同伴的推理进展，同时保留令牌轴上的原始因果结构。

### 3.1 晶格注意力

#### 符号

令B表示批处理大小，N表示线程数，L表示上下文长度，D表示隐藏维度。我们将头维度表示为d，并对晶格注意力组件使用下标p。

#### 概述

为了支持不同线程间的信息流，我们提出晶格注意力，一个沿线程维度运行的附加注意力机制，垂直于令牌轴。朴素的直接在输入嵌入上应用跨线程注意力的设计面临计算和参数挑战。跨线程相同位置的令牌由于异步推理进展可能不强相关，而扩展注意力上下文会产生二次计算开销。此外，从零开始训练额外层容易扰动学好的因果注意力层，尤其是数据有限的情况下。这些挑战驱动我们提出一个**上下文感知**和**参数高效**的架构。

我们不对原始嵌入进行注意力操作，而是在标准缩放点积注意力（SDPA）的输出上操作，A_std∈R^((BN)×L×D_a)，该输出已通过因果注意力编码了丰富的上下文信息。这种设计选择使LACE能够继承标准注意力的有效上下文长度，同时避免冗余计算。

为实现参数效率，我们采用三个关键策略。首先，SDPA通过轻量级下采样投影到较低维空间，在跨线程注意力之前，显著降低计算成本。其次，受ControlNet启发（Zhang等），我们有选择地将晶格注意力层插入到基础模型的中间到最后层，其中跨线程通信对复杂推理最有益（Yang等，2024）。第三，我们采用带学习门控的残差连接以调节晶格注意力的贡献，允许模型动态平衡线程独立和线程感知处理。这些策略共同将额外参数限制在原始模型的少于11%，同时支持有效的跨线程信息交换，对预训练因果层的干扰最小。

#### 跨线程注意力

我们投影标准注意力输出并计算并行QKV：

Z = A_std W_down,p  (1)

Q_p, K_p, V_p = Proj(Z)  (2)

为了编码令牌和线程索引，我们应用3D RoPE（Su等，2024；Ma等，2025），其分割头维度：前d_t维编码令牌位置t，其余d_b=d-d_t维编码块索引n：

Q̃_p, K̃_p = RoPE_3D(Q_p, K_p; t, n)  (3)

我们从(BN, L, ·)重塑为(B, NL, ·)以启用跨块注意力...