结论:以边指针实现极稀疏动态路由,密集注意力高度可裁剪

Graph Machine的核心结论是,语言模型预训练中密集全局注意力并非不可或缺。该研究9月3日(上海时间)以预印本形式首次公开,提出将序列表示的O(n)状态与O(1)稀疏访问通过可学习的边指针解耦:以整数边索引承载地址、以浮点边权重承载梯度,并通过边转介在层间迭代更新图结构,再以边注意力沿边聚合信息。作者在与稠密基线完全一致的骨架、超参与随机种子下,证明将75%稠密层替换为稀疏层后,模型仍能在标准预训练设置下收敛并保持可比损失。

研究对象与方法

研究对象为Graph Language Machines(GLM),以Qwen3-0.6B为稠密对照。方法上,GM维护三类状态:节点特征、边索引与边权重,采用 s1稀疏坐标表示。每个稀疏层包含稀疏边转介子模块与稀疏边注意力子模块:转介通过对软稀疏邻接矩阵求ℓ次幂后稀疏化来组合边,混合时对k个存储边以节点特征预测的混合矩阵与温度缩放进行加权;注意力则将查询-键分数作为节点因子、混合边权重对数作为边因子,以专家乘积形式相加后softmax,仅在每KV head检索2或4个位置。输入边初始化为指向自身与前k-1个前驱,支持input-refresh与dense-refresh两种边刷新及基于注意力证据的存储边重对齐。

评估条件与数据范围

评估在单一控制条件下进行:所有模型从零预训练于随机采样的FineWeb-Edu子集,少于一轮、打包为4,096 token序列(文档均长1,035±1,909),单一H100 SXM训练,Qwen3需53小时、GLM约150–160小时。测试集为同一FineWeb-Edu的固定held-out随机子集,最终损失约2.60。作者系统对比了不同稀疏预算(2,4,4,2)与(4,4,4,4)、存储边数16–32、转介步数0–6及是否启用dense-refresh的12种配置,报告了参数量、估算FLOPs与海量序列下的KV访问比例。

主要结果

  • 在0.098%稠密KV访问(每头2 token)下,最佳Theia模型最终损失仅比Qwen3高约0.014,表明大部分Transformer层可被极稀疏层替换而质量基本不变;在0.195%访问(每头4 token)下,最佳Hyperion-K16-R3-S实现约0.003的损失降低,以约11%更多参数与约19%更少注意力计算取得最优。 转介从0增至3步带来约0.026提升,额外步数在部分配置下继续有效(如Hyperion-K16-R4较R3再降0.003),但非单调(Theia-K24-R4较R3回退0.001)。 增加边数与dense-refresh有增益:K24-R3较K16-R3改善0.003,Hyperion-K16-R3-S较不带S版本改善0.004,但最优并非参数或计算量最大者。

限制、评审状态与意义

该工作明确为预印本,未经同行评审,作者结论限于特定规模与指标。局限包括:仅600M级模型与15.7B token,远小于前沿万亿级训练;仅以测试损失为聚合指标,未展开下游任务与长上下文评估;原型核效率不足,相对吞吐高度依赖硬件与实现;RoPE在稀疏注意力中被省略,转介的可解释性与归纳偏置尚未在语言建模中系统检验。尽管如此,研究首次在标准预训练场景中展示了以对数级地址位与直接索引检索替代稠密打分的架构原语,为在效率与关系遍历之间权衡提供了新自由度,也为后续定制核、更大规模与下游验证指明方向。

论文作者指出,高度可裁剪不等于可完全移除密集注意力,但在当前设置与度量下,密集计算可与稀疏内存操作进行实质权衡,相关结果需在更广条件下独立验证。