Orthros: Phase-Aware Heterogeneous Attention for Efficient Transformers¶

会议: NeurIPS 2026
arXiv: 2610.13735
OpenReview: https://openreview.net/forum?id=Ggocm0VXS9
领域: LLM 效率
关键词: 异构注意力, Prefilling-Decoding 解耦, 线性复杂度, 长文本检索, 阶段感知注意力

一句话总结¶

Orthros 提出了一种阶段感知异构注意力机制，将效率分配从层间混合转向推理阶段解耦：在全序列执行线性复杂度的滑动窗口注意力构建 KV 缓存，仅在负责预测新 token 的位置并行执行全量交叉注意力，在长上下文下实现高达 21.9 倍的 Prefill 加速且无损召回能力。

研究背景与动机¶

在现代长上下文场景（如高分辨率多模态理解、代码仓库分析以及长程智能体工作流）中，模型需要处理数十万至数百万 token 的交互序列。标准 Transformer 中自注意力的二次方计算复杂度 $\mathcal{O}(T^2)$ 迅速成为显存和延迟的致命瓶颈。现有的线性注意力和滑动窗口注意力方案虽然能将复杂度降至线性 $\mathcal{O}(T)$，但其依赖的状态压缩或历史键值剔除机制会造成严重的上下文召回能力退化。为了平衡性能与效率，学术界与工业界广泛采用空间混合注意力（如在不同层交替堆叠线性和全注意力），但这使得模型整体的渐进复杂度依然被最繁重的全注意力层所束缚，最终仅能取得常数倍加速，无法根本解决长序列 Prefill 的二次方计算瓶颈。

更为激进的方案如“解码器-解码器”（Decoder-Decoder）架构通过跨层注意力共享单个 KV 缓存来跳过生成前的全局计算，但这不仅因跨半数网络共享单层缓存而损伤模型容量，还锁死了模型宏观拓扑结构，丧失了对既有开源基础模型的即插即用迁移能力。根本矛盾在于：既有方案均试图在“模型结构”维度（层间或宏观网络切分）分配计算开销，忽视了自回归生成中 Prefilling（计算受限、无生成任务的上下文消化）与 Decoding（显存带宽 I/O 受限、单步自回归预测）在硬件瓶颈与功能属性上的根本异质性。

本文从阶段解耦（Phase-Aware Decoupling）的视角出发，将计算复杂度的分配从模型宏观结构转移至序列处理的不同阶段。自回归生成的本质是：给定 Prompt 序列，并在其末尾自回归生成目标 token；在绝大多数长 Prompt 场景中，生成长度远小于输入长度。核心 idea：将上下文消化与输出生成解耦，统一层内对全序列使用滑动窗口注意力构建 KV 缓存，仅对担负下一个 token 预测的位置附加全量交叉注意力，从而在计算受限的 Prefilling 阶段达成纯线性渐进复杂度，并在 Decoding 阶段将双 token 计算完全掩盖于显存 I/O 访存中。

方法详解¶

整体框架¶

Orthros Attention 将自回归序列处理显式拆分为两类功能向量：承载上下文并物化逐层 KV 缓存的基础 token 序列 $\{ \bm{x}_i \}_{i=1}^T$，以及仅在需要预测后续 token 的位置实例化的输出专属 token 序列 $\{ \bm{x}'_i \}_{i=T_{\text{in}}+1}^T$（其中 $T_{\text{in}}$ 为纯上下文 Prompt 长度，$T_{\text{out}} = T - T_{\text{in}}$ 为输出预测位置数）。

整个前向传播高度统一：基础 token 通过轻量局部滑动窗口注意力（SWA）以 $\mathcal{O}(T)$ 复杂度消化局部上下文并物化完整的逐层键值（KV）缓存；输出专属 token 则携带独立的 Query，对已经物化的完整历史 KV 缓存执行全量交叉注意力（Full Cross-Attention），以获取无损的全局信息。这两组计算在单层注意力模块内通过沿序列维度拼接实现并行执行，随后直接送入标准的 FFN 与残差连接。在宏观层面上，Prefilling 阶段仅对 Prompt 最后一个 token 实例化单个输出 token（$T_{\text{out}}=1$），使繁重的长上下文 Prefill 降至纯线性；而在自回归 Decoding 阶段，每一步仅处理一对解耦 token（1 个基础 token 增量更新滑动窗口 KV，1 个输出 token 交叉注意历史全量 KV 预测下一词）。

%%{init: {'flowchart': {'rankSpacing': 24, 'nodeSpacing': 28, 'padding': 6, 'wrappingWidth': 400}}}%%
flowchart TD
    A["输入序列 X"] --> B["输入张量拼接与分组<br/>基础序列 X 与输出序列 Xcross"]
    B --> C["1. 阶段感知异构注意力解耦机制<br/>双流分支并行计算"]
    C --> D["2. 共享投影与轻量仿射变换 (SO)<br/>共享权重 + 独立缩放偏移向量"]
    D --> E["3. 统一注意力层级算子实现<br/>Flash-Attn 局部 SWA 与全局 CA"]
    E --> F["4. 四重流水线统一范式<br/>PT / SFT / Prefill / Decoding"]
    F --> G["最终隐状态提取与 Logits 预测"]

关键设计¶

1. 阶段感知异构注意力解耦机制：将计算复杂度从结构分配转移至阶段解耦

传统自注意力强制让每一个 token 都对此前所有 token 施加因果全注意力，导致 Prefilling 阶段面对数十万 token 时产生不可承受的二次方 FLOPs。Orthros 重新审视自回归生成的非对称需求：位于上下文内部的中间 token 仅需聚合局部语境以提供高质量表征，唯有真正执行“预测下一个词”的位置才需要跨越全历史的全局信息检索。因此，Orthros 将全长为 $T$ 的输入序列划分为前 $T_{\text{in}}$ 个纯上下文 token 与后 $T_{\text{out}}$ 个生成目标 token。对于基础 token，采用窗口大小为 $W$ 的滑动窗口注意力生成输出 $\bm{o}_t$： $$\bm{o}_t = \frac{\sum_{j=\max(1, t-W+1)}^{t} \exp\left(\frac{\bm{q}_t \bm{k}_j^\top}{\sqrt{d}}\right) \bm{v}_j}{\sum_{m=\max(1, t-W+1)}^{t} \exp\left(\frac{\bm{q}_t \bm{k}_m^\top}{\sqrt{d}}\right)}$$ 而对于后 $T_{\text{out}}$ 个位置（$t \in [T_{\text{in}}+1, T]$），附加生成全局全量交叉注意力输出 $\bm{o}'_t$： $$\bm{o}'_t = \frac{\sum_{j=1}^{t} \exp\left(\frac{\bm{q}'_t \bm{k}_j^\top}{\sqrt{d}}\right) \bm{v}_j}{\sum_{m=1}^{t} \exp\left(\frac{\bm{q}'_t \bm{k}_m^\top}{\sqrt{d}}\right)}$$ 在长上下文 Prefilling 时，$T_{\text{in}} = T - 1, T_{\text{out}} = 1$，仅末尾单个 token 进行一次全局检索，前序 $T-1$ 个 token 严格运行 $\mathcal{O}(T \cdot W)$ 的线性计算，从渐进上打破了 Prefilling 的二次方墙；而在 Decoding 阶段，$T=1, T_{\text{out}}=1$，每次迭代仅计算 2 个 token 的向量运算，其显存加载权重的 I/O 开销彻底淹没了这微小的双 token 算力，实现了阶段收益的最大化。

2. 共享投影与轻量仿射变换 (SO)：无额外参数开销的高效表征分化

如果为基础 token 和输出专属 token 分别配置独立的 $W_Q, W_K, W_V$ 投影矩阵，不仅会大幅增加参数量，更会在显存带宽受限的 Decoding 阶段成倍增加权重加载延迟。为了兼顾表征专一性与参数复用，Orthros 采用统一的线性投影层，并在 Query 侧引入轻量级缩放与偏移仿射变换（Scale-Offset, SO）： $$\bm{q}_i = \bm{\alpha} \odot (\bm{x}_i \bm{W}_Q) + \bm{\beta}, \quad \forall i \in [1, T]$$ $$\bm{q}'_i = \bm{\alpha}' \odot (\bm{x}'_i \bm{W}_Q) + \bm{\beta}', \quad \forall i \in [T_{\text{in}}+1, T]$$ 其中 $\bm{\alpha}, \bm{\beta}, \bm{\alpha}', \bm{\beta}' \in \mathbb{R}^{1 \times d}$ 为可学习的低维向量，开销相比庞大的投影矩阵几乎可以忽略。该设计允许模型利用同一套基底投影权重，同时使全局检索 Query 与局部聚合 Query 在特征空间中解耦，消融实验证实去除 SO 会使模型在下游任务的平均分出现显著退步。

3. 统一注意力层级算子实现：单层自包含与开源基座即插即用转换

与破坏网络拓扑结构的跨层共享（Decoder-Decoder）或宏观交替架构不同，Orthros 将异构计算内聚封装在每一个独立的注意力层中。在骨干网络中，对于输入的 $T$ 个 token，直接将末尾 $T_{\text{out}}$ 个位置的隐状态向量复制并拼接在序列尾部，形成 $(T + T_{\text{out}}) \times d_{\text{model}}$ 的扩展输入张量。经过层归一化后，前 $T$ 个 token 生成全局共享的 KV 缓存，后 $T_{\text{out}}$ 个 token 生成全局 Query。两个分支分别调用底层高效的 FlashAttention 内核（基础分支传入固定局部窗口大小 $W$，交叉分支传入无界窗口 $\infty$），随后将输出张量重新拼接并通过统一的输出投影矩阵 $\bm{W}_O$ 映射。非注意力层（如 MLP/FFN）由于具有按 token 独立运算的天然特性，完全无需修改即可天然透明兼容。这种自包含设计使得任何现有的开源 Transformer（如 Llama 3.1、Qwen 3）均可无缝将注意力算子热替换为 Orthros，仅需极轻量的数据持续预训练即可完成转换。

4. 四重流水线统一范式：严格保证训练与推理的算法一致性

很多高效架构在推理加速时破坏了自注意力因果掩码，导致训练-推理存在算法差异。Orthros 通过统一参数化 $(T_{\text{in}}, T_{\text{out}})$ 实现了四种场景下的无缝泛化： - 预训练（Pre-training, PT）：由于每个 token 均需计算交叉熵损失，$T_{\text{in}} = 0, T_{\text{out}} = T$。全序列复制，每个位置都计算局部 SWA 与全局 Cross-Attention。 - 指令微调与强化学习（SFT / RL）：Prompt 部分无需计算梯度，$T_{\text{out}}$ 精确等于 Target Response 的 token 长度，无需对 Prompt 的中间位置执行全注意力，大幅节省微调算力。 - 长文本预填充（Prefilling）：$T_{\text{in}} = T - 1, T_{\text{out}} = 1$。极低延迟消化巨量上下文，仅在最后一个 token 执行一次全局检索。 - 自回归解码（Decoding）：$T = 1, T_{\text{out}} = 1$。输入单一新生成 token，基础流更新滑动窗口 KV，交叉流基于全局历史 KV 预测下一 token。

损失函数 / 训练策略¶

Orthros 遵循标准自回归语言模型的交叉熵损失函数。对于输出预测序列 $t \in [T_{\text{in}}+1, T]$，模型仅从骨干网络第 $L$ 层输出提取末尾 $T_{\text{out}}$ 个输出专属 token 的隐状态，经最终的 RMSNorm 与 LM Head 投影计算预测概率分布： $$\mathcal{L} = - \sum_{t=T_{\text{in}}+1}^{T} \log P(x_{t+1} \mid \bm{x}'_t)$$ 训练采用 AdamW 优化器，结合余弦退火学习率调度策略。在基座模型转换实验中，使用 80B token 的高质量数据（Nemotron-CC、数学与代码混合语料），在序列长度 32,768 下执行 10,000 步持续预训练；长文本扩展阶段则在 ProLong 数据集上继续微调 20B token，将有效上下文窗口无损拓展至 128K。

实验关键数据¶

主实验¶

在 1.3B 规模下，全量使用 FineWeb-Edu 100B token 从头预训练，Orthros 在常识推理任务上取得了超越基线 Transformer 及多种主流线性注意力架构的优异表现。

架构 / 模型	ARC-C	ARC-E	HellaSwag	OpenBookQA	PIQA	Social-IQa	WinoGrande	平均得分
Transformer	37.80	70.58	56.21	37.80	72.91	41.10	57.46	53.41
Orthros (Ours)	39.08	71.09	57.19	40.20	71.82	42.12	56.35	53.98
GatedDeltaNet (GDN)	38.82	70.66	56.28	40.20	71.44	40.63	55.01	53.29
Gated Linear Attention (GLA)	36.86	67.97	53.29	38.20	70.73	40.07	56.20	51.90
RetNet	35.84	69.82	54.33	37.00	71.93	39.76	55.01	51.96
SWA (Window=512)	36.95	70.16	56.11	38.80	72.42	40.94	57.14	53.22

在大模型转换方面，将现成预训练基座模型（Llama 3.1-8B 与 Qwen 3-8B）转换为 Orthros 并以 80B token 持续预训练（CPT），模型表现不仅没有退化，甚至超越了原架构继续训练的 baseline：

基座模型	评估配置	ARC-C	ARC-E	HellaSwag	OpenBookQA	PIQA	Social-IQa	WinoGrande	MMLU	平均得分
Llama 3.1-8B	原始基座 (Base)	53.50	81.57	78.99	45.20	79.76	46.67	73.09	65.00	65.47
Llama 3.1-8B	Orthros 转换	57.08	84.13	79.10	45.60	80.25	48.93	74.35	68.65	67.26
Llama 3.1-8B	原架构 CPT	57.51	83.80	78.93	45.40	80.25	47.80	73.40	68.89	67.00
Qwen 3-8B	原始基座 (Base)	57.34	81.73	78.49	42.20	79.27	54.61	72.53	76.82	67.87
Qwen 3-8B	Orthros 转换	59.39	84.81	78.06	44.80	79.27	50.51	72.38	77.76	68.37
Qwen 3-8B	原架构 CPT	57.59	83.29	77.59	43.20	78.78	50.00	71.98	77.96	67.55

在上下文信息召回基准测试（DROP、NarrativeQA、SWDE、TriviaQA、FDA、SQuAD）中，Orthros 1.3B 平均分达到 34.4，全面超越基线 Transformer（33.9）以及 GDN（28.4）、GLA（27.1）、SWA（28.8）等线性模型，彻底破除了高效模型召回率暴跌的魔咒。此外，8B 转换模型在 128K 长度的“大海捞针”（Needle In A Haystack）测试中取得了全绿的近乎完美召回率。

消融实验¶

针对 Scale-Offset（SO）仿射变换机制与模型跨尺寸可扩展性进行了详细消融实验（10B token 训练）：

参数规模	模型配置	ARC-C	ARC-E	HellaSwag	OpenBookQA	PIQA	WinoGrande	平均得分
160M	Orthros w/o SO	25.86	47.64	33.14	32.20	62.62	50.28	41.96
160M	Orthros (w/ SO)	26.19	46.68	34.26	30.80	63.33	51.22	42.08
160M	Transformer 基线	25.77	46.25	33.64	30.60	62.08	52.64	41.83
1.3B	Orthros w/o SO	30.03	55.05	44.83	34.40	69.31	51.85	47.58
1.3B	Orthros (w/ SO)	31.48	56.19	46.64	34.80	69.04	54.30	48.74
1.3B	Transformer 基线	31.06	57.07	45.30	33.60	68.99	53.91	48.32
6.8B	Orthros w/o SO	32.59	60.44	49.39	36.00	70.02	52.96	50.23
6.8B	Orthros (w/ SO)	35.84	65.57	53.43	36.60	72.69	54.46	53.10
6.8B	Transformer 基线	33.45	60.90	49.84	34.80	70.35	52.49	50.31

关键发现¶

大模型尺度优势扩大：从 160M 扩展到 6.8B，Orthros 相对 Transformer 的优势不断扩大。在 6.8B 下，Orthros（53.10）大幅领先标准 Transformer（50.31），领先达 2.79 个百分点，展现出优秀的 Scaling 特性。
SO 变换不可或缺：在 6.8B 规模下，去除 SO 仿射变换导致平均得分从 53.10 骤降至 50.23（下降 2.87 分），证明共享投影权重下利用轻量向量实现局部与全局 Query 角色分化的极高有效性。
超长文本推理极致提速：在 8 张 GPU 评测下，当输入序列扩展至 100 万 token 时，Orthros 实现了相比 Qwen 3 高达 21.9 倍的 Prefill 加速；在端到端 SGLang 评测中（128K 输入 / 128 输出），首字延迟 TTFT 从 Transformer 的 77.25s 骤降至 17.02s，吞吐量提升 3.5 倍。

亮点与洞察¶

将宏观结构混合重构为微观阶段解耦：跳出了既往在 Transformer 深度方向拼凑混合注意力（Hybrid）或粗暴砍掉半个网络缓存（Decoder-Decoder）的传统思路，敏锐抓住自回归中 Prefill 与 Decode 不同的硬件瓶颈，用单一解耦层达成理论与硬件层面的最优对齐。
Decoding 额外算力的零成本隐藏：理论上解码阶段每步需要计算 2 个 token，但在长文本生成中，Decoding 受到显卡高带宽显存（HBM）搬运模型权重的严格限制；双 token 带来的极小算力增加被完完全全地遮盖在内存等待时间中，换来了 Prefilling 阶段数量级的爆发式加速。
基座模型数据高效转换迁移性：仅凭 80B 持续预训练 token 即可无痛将 Llama 3.1 / Qwen 3 改造为全线性 Prefill 的高效模型，无需从头烧录数万卡月的大规模预训练，工程落地可行性极高。

局限与展望¶

预训练阶段的 FLOPs 开销加倍：为了严格维护训练与推理的算法自洽，全序列预训练时必须把所有 token 视作预测目标（$T_{\text{out}}=T$），导致前向反向计算开销相比普通因果语言模型翻倍。
KV 缓存物理显存占用未减少：Orthros 保证了逐层完备的 KV 缓存用于高精度检索，虽然计算复杂度降为线性，但 KV 缓存本身的显存驻留大小与原版 Transformer 一致，在极端长上下文下的显存占用问题仍需借助 GQA 或量化解决。
专家混合模型 (MoE) 适配挑战：在 MoE 架构中，解耦的双 token 对可能被路由到不同的物理 Expert，这会破坏计算局部性并加剧显存 I/O 负担，未来需要专门针对解耦 token 设计绑定路由策略。

评分¶

新颖性: ⭐⭐⭐⭐⭐ 阶段感知解耦思想打破了传统的空间混合结构定势，概念干净且硬件适配性极强。
实验充分度: ⭐⭐⭐⭐⭐ 覆盖从头预训练、160M-6.8B 扩展消融、基座转换、长文本大海捞针到百万级 Prefill 与真实端到端推理测速。
写作质量: ⭐⭐⭐⭐⭐ 逻辑严密，算法伪代码与宏观四重范式解释透彻。
价值: ⭐⭐⭐⭐⭐ 兼顾线性 Prefill 加速与 100% 全注意力检索保真度，极具工业级替代潜力。