Graph Machine: 通过边实现更好的预训练

TL;DR Graph Machine 用可微分的边指针实现 O(n) 状态下的稀疏动态路由，替换 Qwen3-0.6B 中 75% 稠密层后，仅每头检索 4/4096 token 即可维持甚至略优预训练 loss。

问题

问题背景

序列建模的核心权衡在于状态大小、访问模式与动态寻址三者的组合。当前主流模型各自偏向一端：RNN/SSM 保持常数级状态导致历史信息压缩，Transformer 保留 O(n) 状态却承受 O(n²) 全注意力计算，滑动窗口/稀疏注意力 将访问降到 O(1) 但路由静态，无法根据当前 token 内容动态选择历史位置。

现有方法局限

• 固定大小状态（如 RNN、Mamba）强制将全部历史压缩进固定向量，长程依赖容易丢失细节；
• 全注意力 的 O(n²) 计算与显存开销在长序列上不可接受，推理成本随长度平方增长；
• 静态稀疏路由（如 Sliding Window、BigBird）预先固定可访问的位置集合，与当前 token 内容无关。即便某些历史 token 对当前预测至关重要，只要落在窗口外便无法被检索，模型只能依赖多层堆叠间接传递信息，效率低下且上限受限。

为什么难且重要

要实现动态稀疏访问，每个 token 需要输出 O(log n) 位的地址信息，并直接索引 O(n) 状态中的特定条目。这在传统深度学习框架中不可微分，因为地址选择是离散操作。Graph Machine 提出的可微指针追逐机制模拟了“引用-跳转”过程，使整个路由可端到端训练。业界关注该问题，因为大模型推理成本与上下文长度直接挂钩，若能以 2~4 个 token 检索替代全注意力，长文档、代码库、多轮对话等场景的吞吐量将大幅提升。

行业类比

类似推荐系统中的记忆增强网络：不是遍历全部用户历史，而是根据当前 query 动态检索最相关的少量记录，在保持质量的同时将计算复杂度降到近似常数。

核心洞察

• **动态指针路由**：GM 将稀疏注意力中的路由决策参数化为可学习的边（整数索引），每个 token 提供 Θ(log n) 位地址，从而在 O(n) 状态上实现内容相关的 O(1) 访问。这与滑动窗口等静态稀疏方法（可访问位置预先固定，无法根据 token 内容改变）形成根本差异，GM 的边通过可微的 `referral mechanism`（类似指针追踪）更新，使路由模式可端到端训练。对工程而言，这意味着模型能学会为每个 token 选择最相关的历史状态，在保持线性复杂度的同时可能逼近 dense attention 的信息聚合能力。
• **图结构状态建模**：GM 将序列状态显式分解为节点特征（浮点隐状态）和边索引（整数指针），用图的角度统一 RNN/SSM 的压缩状态和 Transformer 的全状态。这种分解允许状态大小 Θ(n) 且访问 Θ(1)，避免了以往方法在状态容量和访问复杂度之间的两难。在 Qwen3-0.6B 上替换 75% dense 层并用 15.7B tokens 预训练，每 KV 头仅检索 4 个 token（共 4096）时 loss 反而略优于 dense baseline，表明大规模预训练中动态稀疏路由的潜力，但推理时额外地址存储和路由计算开销需进一步权衡。

方法

输入与状态

Graph Machine 层接收序列的 hidden_states（节点特征）与已存储的 edge_indices（整数指针）。状态规模为 O(n)，其中 n 为当前序列长度；每步仅新增一个 token，因此状态增量恒定。

关键模块

• Sparse Edge Referral (SER)：通过可微的指针追逐机制更新边目标。每条边指向历史 token 的索引，SER 根据当前 token 内容动态调整少量边的指向，相当于可学习的路由决策。
• Sparse Edge Attention (SEA)：基于 SER 更新后的边索引，仅从 n 个 token 的状态中检索固定数量（如 2 或 4 个）的 KV 对，执行稀疏注意力混合。
• Sparsify / Mix 操作：Sparsify 将稠密隐藏状态压缩为路由信号；Mix 将检索到的上下文合并回主表示。

输出

每层输出更新后的节点特征与边索引，供下一层复用。在 Qwen3-0.6B 中，75% 的稠密 Transformer 层被替换为 GM 稀疏层，每 KV 头仅检索 2/4096 token 时损失退化轻微，检索 4 token 时略有改善。

与同类方法差异

与滑动窗口等静态稀疏路由不同，GM 的边是动态、可微更新的指针，能以 O(1) 实际检索成本访问远距离历史，同时保持 O(n) 状态规模，无需固定窗口或压缩历史。

实验

实验设计

Graph Machine 在 Qwen3-0.6B 上将 75% 的 dense Transformer layers 替换为 GM sparse layers,从头预训练 15.7B tokens。每个 sparse layer 的每个 KV head 仅从 4096 个 token 中检索 2 或 4 个 token,对比不同检索宽度对 loss 的影响。基线为原始 dense Qwen3-0.6B。

关键发现

• 检索宽度为 2 时(仅 0.05% state 访问),loss 仅轻微下降,说明动态路由能在极小访问量下捕获关键上下文。
• 检索宽度为 4 时,最佳模型 loss 略微优于 dense 基线,表明 sparse dynamic routing 不仅压缩计算,还可能带来轻微正则或更好的注意力聚焦。
• 这验证了信息论论证:每个 token 提供 Θ(log n) 比特动态地址即可实现常数访问,同时保持 O(n) 状态。

与基线/同类方法对比解读

与 sliding-window attention 等静态稀疏方法相比,GM 的 edge indices 是内容相关的动态指针,不局限于固定邻域,因此能用更少的访问覆盖长程依赖;与 RNN/SSM 的固定大小状态相比,GM 保留完整 O(n) 节点特征,仅索引稀疏,避免历史压缩瓶颈。工程上,该架构在长序列预训练中可将注意力计算量降低约 1000 倍(4096→4),同时保持或提升 loss,为高效 Transformer 替代方案提供了新证据。

行业影响

落地场景

Graph Machine (GM) 的 O(n) 状态 + O(1) 稀疏动态访问 适合需要长上下文但资源受限的场景。例如：

• 电商推荐：用户行为序列可达数万条，GM 用可微指针动态检索相关历史 item，替代 dense attention，在千万级候选上保持实时响应。
• 医疗记录分析：电子病历长文本中，GM 可稀疏路由到关键病史节点，辅助诊断而不丢失全量上下文。

商业价值

• 降本：显存与计算从 O(n^2) 降到 O(n)，长序列推理成本大幅下降，同等硬件可支持更长上下文窗口。
• 增收：完整用户历史能提升推荐点击率和转化，例如电商场景 A/B 测试可验证 GM 替代 attention 后的 GMV 增量。
• 体验：流式对话或实时决策延迟降低，长文档问答响应更快。

跟现有产品/工作流的接口

GM 的 sparse layer 设计为 Transformer 的 drop-in 替代，论文用 Qwen3-0.6B 替换 75% dense layers 预训练，说明可接入现有 LLM 训练栈。集成时需：

1. 实现 sparse edge referral 和 sparse edge attention 算子，可基于 PyTorch 扩展或定制 CUDA kernel。
2. 推理引擎需支持动态索引检索，类似稀疏图计算或内存增强网络，可参考 vLLM 的分页注意力机制扩展。
3. 对现有模型继续预训练或微调，无需重头训练，降低迁移成本。

GM 与 滑动窗口注意力 的核心差异在于动态路由：窗口方法静态排除远距离 token，GM 每步根据 content 选择 Θ(1) 个相关状态，信息密度更高，适合 streaming 与超长序列场景。

局限

• **实验规模有限**：当前仅在 **Qwen3-0.6B** 上替换 75% 密集层并预训练 **15.7B tokens**，参数量和训练数据均远小于主流大模型（通常数百亿参数、数万亿 tokens）。GM 的 **边索引** 与 **动态路由** 在更大规模下是否稳定、能否保持损失优势尚未验证，直接扩展到生产级模型存在不确定性。
• **对比基线不足**：论文仅与原始 dense Transformer 比较损失，未直接对比现有稀疏注意力或线性注意力方法（如 **Routing Transformer**、**Reformer**、**BigBird**、**Performer**、**Mamba** 等）。因此难以判断 GM 在 **效率-性能权衡** 上的相对优势，也无法证明其路由机制比其他稀疏方案更有效。
• **工程实现挑战**：动态指针追踪需要维护离散的整数边索引，导致 **非连续内存访问** 与 **负载不均衡**，难以直接利用 **FlashAttention** 等高度优化的 dense attention kernel。边的可微更新（**referral mechanism**）涉及离散决策，训练稳定性和并行效率可能需要额外工程技巧，实际推理吞吐可能不理想。