层级库(LoL):在运行时自适应排序 Transformer 层 Library of Layers (LoL): Adaptively Ordering Transformer Layers at Runtime
路由器在推理时从共享层库中动态挑选下一个 Transformer 层,统一参数复用与稀疏激活
前置知识
混合专家(Mixture of Experts, MoE)
MoE 把 Transformer 的 FFN 替换为多个并行『专家』,由小路由器按 token 隐藏状态为每个 token 选 top-$k$ 个专家计算,其余闲置。总参数量大而每 token 实际执行的 FLOPs 少,实现条件计算;负载常用辅助损失或无损均衡维持。
LoL 在『大库』极限下与稀疏 MoE 直接对比(每层 3 专家取 2、约 17 ExaFLOPs),理解 MoE 的稀疏激活与负载均衡机制,才能读懂该对比实验的设计与结果解读。
循环/复用 Transformer(looped transformers)
不堆叠 $N$ 个不同层,而是把同一组 Transformer 块循环应用 $D$ 次(如 Universal Transformer、Huginn),参数在深度间共享:总参数少而有效深度大,FLOPs/参数比高,理论上可达到图灵完备,特别适合需要迭代精化的算法型推理任务。
LoL 在『小库』极限下等价于 looped 模型,论文用 Huginn 架构作为主要对比基线;只在中间层做复用、前后加固定层的设计也直接借鉴该领域的实践经验。
因果注意力与 Q-K 错配
自回归模型要求每个位置只注意之前的 token(因果掩码)。若相邻 token 在同一深度被分到不同注意力层,query 与 key 由不同投影矩阵生成,两组向量空间不一致,点积相似度失去意义;作者发现这种逐 token 路由会导致训练不稳、性能变差。
这正是 LoL 采用『窗口级路由 + 跨深度错位窗口注意力(strided window attention)』的根本原因,是理解整个方法设计动机的核心概念。
Top-1 路由与无损负载均衡(loss-free balancing)
路由器对每个输入只选得分最高的专家(top-1 硬选择)。为让梯度流向离散选择且保持负载均衡,loss-free 方法(Wang et al., 2024)为每个专家维护仅用于 argmax 的偏置 $b_m$,并把被选专家输出按 softmax 概率 $\alpha$ 缩放,无需辅助损失。
LoL 的路由器训练完全复用该机制:被选专家的残差更新按 $\alpha_k$ 缩放使梯度流入路由器,附录消融还证明 loss-free 优于传统辅助损失。
FiLM 条件调制
FiLM 用条件向量生成逐特征的缩放 $\gamma$ 与偏置 $\beta$,对表示做仿射变换 $x \mapsto \gamma \odot x + \beta$。LoL 在每个块的每个 Norm 后施加深度依赖的 FiLM,使同一专家在不同路由深度产生不同行为。
这是让『同一个块被多次使用但每次行为不同』成为可能的关键组件;没有它,复用层在不同深度只能执行完全相同的函数,表达力会受限。
线性探针与 AUROC
线性探针在冻结的模型内部表示上训练一个线性分类器(如 ridge 回归),检验某种信息是否以线性可读出的形式存在于表示中;AUROC(ROC 曲线下面积)衡量二分类的排序质量,取值 0.5 为随机猜测、1.0 为完美区分,对类别不平衡较为稳健。
安全性实验正是用 ridge 探针读取路由概率中的有害性信息,并以 AUROC 报告结果(窗口级 0.92、序列级 0.95),这是评判『路由器编码了语义信息』这一主张的关键量化指标。
研究动机
标准 Transformer 的前向传播顺序在架构设计时就固定死了:无论输入是一段代码、一句法语还是一道数独题,第 1 层到第 $N$ 层永远按同一顺序执行。这与软件工程的实践形成鲜明对比——论文开篇举例:数据库查询引擎会根据查询内容把过滤、连接、排序、聚合等原语按不同序列组合,模块化编程让同一批组件在不同程序中动态拼装。这种刚性带来双重浪费:一方面,有用的『计算原语』可能需要在网络的不同深度重复表示,浪费参数;另一方面,每个 token 都要执行全部层的计算,FLOPs 效率受限。已有方案各有局限:MoE(Shazeer et al. 2017;Switch Transformer)只在每个固定深度上稀疏化选择 MLP,块的执行顺序依然固定;looped 模型(Universal Transformer、Huginn)把同几层重复循环 $D$ 次,所有输入共享同一条循环路径;layer skipping(Layerskip、Mixture-of-Recursions 等)只能动态跳过层,不能重排层。换言之,所有这些方法都没有把『层的执行顺序』本身当作可学习的、依赖输入的自由度。
本文的目标是本文要构建一个『层的顺序由模型自己决定』的架构:一个小型路由器在推理时为每段上下文从全局层库中挑选下一个要执行的 Transformer 块,端到端训练、无需任何关于顺序的显式监督。具体目标有四个层面:(1)用单一超参数——库大小 $L$——在同一架构家族内连续覆盖两个极端:小库迫使层被反复复用(类似 looped 模型的参数高效 regime),大库允许很多层在一次前向中闲置(类似 MoE 的稀疏 regime);(2)在与 dense、looped、MoE 三类基线参数/FLOPs 匹配的预训练中保持竞争性(360M 与 1B 两个规模、最多约 20B tokens);(3)证明路由路径是可解释的:不同语言、不同解题算法应映射到不同的层执行序列;(4)证明路径可用于控制:推理时强制某条路径即可让模型切换解题策略,为模型行为的安全与控制提供新接口。
与已有工作不同的是,已有工作要么在固定顺序下做稀疏(MoE),要么在固定顺序下做复用(looped)。与本文最接近的 Gong et al.(2024)也是先预训练再追加层复用、最终模型并不逐层动态选择,而且采用 top-$k>1$ 的软路由、通过平均所选专家的注意力原始权重来回避 attention 空间错配问题。LoL 的独特切入是:把『下一步执行哪个块』作为首要建模对象,从零开始端到端训练 top-1 逐层选择的全局层库。为此必须直面两个此前工作绕开的难题:其一,若逐 token 选层,不同位置的 query 和 key 会由不同的注意力层产生,Q-K 空间错配导致训练不稳——LoL 用跨深度错位的窗口注意力(strided window attention)让同窗口 token 共享一个专家来化解;其二,路由决策只能使用生成阶段同样可得的信息才能保持因果性——LoL 让每个窗口只依据更早窗口的隐藏态均值做决策,且窗口边界逐深度平移以逐步扩大信息视野。
核心方法
直觉上,LoL 把前向传播想象成图书管理员反复从书架取书:书架(层库)固定,但每次取哪本、某本取多少次,由当前内容决定。技术上模型由三段组成:固定顺序、全序列因果注意力的 prelude 与 coda,中间夹一个 routed core。核心在 $D$ 个路由深度上循环执行:每个深度把序列切成若干长 $W$ 的窗口,窗口边界按步幅 $S$ 逐深度平移,偏移量 $\delta_\ell = (\ell-1)S \bmod W$,位置 $i$ 属于窗口 $k$ 当 $w_i^{(\ell)} = \lfloor (i+\delta_\ell)/W \rfloor$——这就是 strided window attention:窗口内部做因果注意力,窗口间靠边界错位逐深度传递信息(类似滑窗叠加后感受野随深度扩大)。每个窗口由线性路由器 top-1 地从 $L$ 个专家块中选一个执行;路由输入是之前所有窗口平均隐藏态的均值,首窗口例外,用可学习深度向量 $e^{(\ell)}$ 与首 token 隐藏态的平均。被选专家的残差更新按路由概率 $\alpha_k$ 缩放(loss-free 负载均衡),让梯度流向路由器;块内每个 Norm 后加 FiLM 深度调制。训练用 growing:核心深度从浅逐渐增长($1-\sqrt{\,\cdot\,}$ 调度、前 75% 训练),先便宜后深。360M 模型为 4 prelude + 4 coda + 16 路由深度,1B 模型为 4+4+8。
核心创新在于:『层的执行顺序』第一次成为端到端学习的输入条件函数。与 MoE 的本质区别:MoE 在网络中每个固定深度上为每个 token 选择一个小型 FFN,骨架顺序不变;LoL 选择的是完整 Transformer 块(注意力+FFN),且同一块可以出现在路由路径的任意深度、任意多次,顺序与复用次数完全由数据决定。与 looped 模型的本质区别:looped 模型所有输入走同一条固定循环路径,而 LoL 的每个窗口可以走不同路径,且 $L$ 是一个连续旋钮——$L$ 小时退化为高度复用的循环模型,$L$ 大时变为高稀疏模型,一个家族统一了 dense、depth-recurrent、MoE 三种范式。与 Gong et al. 的区别:top-1 硬选择而非加权平均多个专家的注意力权重,保留了『路径』这个离散、可解释、可干预的对象。配套的 strided window attention 是让这一切在因果语言建模下可行的关键机制:它同时解决了 Q-K 错配(同窗口共享专家)、因果性(路由只用历史窗口)与跨窗口信息流(边界错位)三个约束。
方法步骤详情
完整流程对应附录 Algorithm 1(自然语言实验取 $W=128$、$S=64$、上下文 1024)。第一步,输入 $x \in V^n$ 过 prelude 得 $s^{(0)}=P(x)$。第二步,对每个路由深度 $\ell=1,\dots,D$ 计算偏移 $\delta_\ell=(\ell-1)S \bmod W$,把位置分入窗口 $I_k$,共 $K_\ell$ 个窗口。第三步,对每个窗口求平均隐藏态 $\mu_k$。第四步,构造因果路由输入:$k=0$ 时 $r_k=\tfrac{1}{2}(e^{(\ell)}+s^{(\ell-1)}_0)$;$k>0$ 时 $r_k=\tfrac{1}{k}\sum_{j<k}\mu_j$,只用更早窗口。第五步,路由 logits 加高斯噪声鼓励探索后做 top-1 选择 $\rho_k=\arg\max_m(a_{k,m}+b_m)$($b_m$ 为 loss-free 均衡偏置),权重 $\alpha_k=a_{k,\rho_k}$。第六步,窗口内 token 用专家 $E_{\rho_k}$ 做窗口内因果注意力与 FFN,块内每个 Norm 后施加 FiLM 调制,输出 $y_k$。第七步,残差更新 $s^{(\ell)}_{I_k}=s^{(\ell-1)}_{I_k}+\alpha_k(y_k-s^{(\ell-1)}_{I_k})$,$\alpha_k$ 保证梯度经概率路径流入路由器。第八步,重复 $D$ 次后由 coda 做全因果注意力并 softmax 得输出。训练用 Muon+AdamW、文档间注意力隔离与 growing 调度;360M 模型(GPT-2 Medium 形状加 RoPE/RMSNorm/SwiGLU)为 4+4+16,1B 模型(Llama-3.2-1B 形状)为 4+4+8。
技术新颖性
技术新颖性体现在四个机制的组合。其一,全局共享的块级库加深度无关的 top-1 路由:专家是完整 Transformer 块而非 FFN,库被所有路由深度共享,区别于逐层 MoE,也区别于 Huang et al.(2026b)仍在固定深度放置 MoE 层的全局专家池设计。其二,strided window attention 是新的注意力掩码:逐深度把窗口边界平移 $S$ 使信息场随深度扩大,同时严格保持因果性与 Q-K 一致性;作者发现逐 token 路由的 Q-K 错配会直接导致训练失败,此前工作(如 Gong et al. 用权重平均)是刻意回避而非解决该问题。其三,FiLM 深度调制让同一块在不同深度『变身』,用 $O(1)$ 额外参数为每个专家扩展出 $D$ 个行为变体,缓解块复用的表达瓶颈。其四,growing 深度渐增借鉴 looped 模型的训练技巧但首次用于动态路由架构,在早期训练大幅节省 FLOPs(图 7 显示只带来极小的最终损失上升)。消融进一步支撑设计:loss-free 均衡优于辅助损失(图 8),路由权重跨深度共享与否影响甚微(图 9)。
实验结果
实验分两条线:可解释性/可控性与规模化竞争性。(1)多语言:360M 模型在 FineWeb-2 的 10 种语言等比混合上训练 10B tokens,把每个窗口各深度的专家选择拼成 one-hot 向量做 PCA,语言间路径明显分离,共享拉丁字母书写的语言(德/西/法/土)聚簇(图 3 左,图 11 显示分化是逐渐学出来的)。(2)安全性:用 FineWeb 有害性分数(0 分安全、4-5 分不安全)等比混合训练 10B tokens、训练不给标签,之后在路由概率上训练 ridge 探针:窗口级测试 AUROC 0.92,序列级 0.95;开发集交叉验证显示单专家概率(0.8244)与专家频率对数似然比(0.8300)也能恢复大部分信号,打乱标签后跌回 0.5014(表 2),排除泄漏。(3)数独:5 种策略、20M 序列约 10B tokens,训练时只有路由器能看到策略标识 $[S_x]$;推理时用中性占位符替换标识并强制 mode path,策略遵循率 S1–S4 为 100%、S5 为 99.5%,答案正确率 98.5%–100%(图 3 中、图 14)。(4)迷宫:BFS/DFS/A*/双向搜索/死端填充 5 种算法共享同一输出语法、只在访问格子顺序上不同;强制路径后模型以 99.79%–100% 复现指定算法的探索顺序(图 3 右、图 18),说明层序控制了算法本身。(5)1B 数学:16 层模型(4+4+8),Nemotron+OpenMath 数据约 90 ExaFLOPs(约合 20B tokens),训练损失与参数匹配 dense 基线持平(图 4 左),五项基准均在 2 个标准误内与最优基线相当,ARC-Easy 上 LoL 50.3% 显著高于全注意力 vanilla 的 40.8%(表 1)。(6)库大小消融:$L=2/4/8$ 匹配 Huginn $2\times8/4\times4/8\times2$(360M、10B FineWeb tokens),评估损失相当但 LoL 种子间方差更大(图 5);对 3 选 2 MoE 在约 17 ExaFLOPs 下 LoL(24 专家×16 深度)训练损失高约 0.05(2.8 vs 2.75,图 6)。专家利用率热图(图 4 右、图 6 右)显示专家呈深度特化:有的只被特定深度使用,有的跨深度通用。
查看结构化数据
| 任务 | 指标 | 本文 | 基线 | 提升 |
|---|---|---|---|---|
| 1B 数学模型 · ARC-Easy | 准确率 (%) | 50.3 ± 1.0 | Vanilla 40.8 ± 1.0;Vanilla 滑窗 49.5 ± 1.0 | +9.5 vs 全注意力 vanilla;与最强基线持平 |
| 1B 数学模型 · ARC-Challenge | 准确率 (%) | 29.1 ± 1.3 | Vanilla 28.2 ± 1.3 | +0.9(2 个标准误内持平) |
| 1B 数学模型 · HellaSwag | 准确率 (%) | 31.2 ± 0.5 | Vanilla 31.4 ± 0.5 | −0.2(持平) |
| 1B 数学模型 · PIQA | 准确率 (%) | 59.7 ± 1.1 | Vanilla 56.6 ± 1.2 | +3.1 |
| 1B 数学模型 · GSM8K | 准确率 (%) | 27.9 ± 1.2 | Vanilla 29.7 ± 1.3 | −1.8(2 个标准误内持平) |
| 安全性判别(路由概率线性探针,序列级) | AUROC | 0.95(窗口级 0.92) | 打乱标签对照 0.5014 | 远超随机的 0.5,证明路由线性编码有害性信息 |
| 数独策略遵循(强制 mode path) | 策略遵循率 (%) | S1–S4 100%,S5 99.5% | 随机猜策略 20%(1/5) | 近满分复现指定解题策略 |
| 迷宫策略遵循(强制 mode path) | 策略遵循率 (%) | 99.79%–100% | 随机 20%(1/5) | 共享输出语法下仍精确复现指定搜索算法 |
| 360M vs 稀疏 MoE(每层 3 专家取 2,17 ExaFLOPs) | 训练交叉熵 | 2.8 | MoE 2.75 | −0.05(MoE 略优,LoL 具竞争性) |
| 360M vs 循环 Huginn(L=2/4/8,10B FineWeb tokens) | 评估交叉熵 | 与对应配置持平 | Huginn 2×8 / 4×4 / 8×2 | 持平(LoL 种子间方差略大) |
局限与改进
作者坦承这是对动态层排序的初步探索。具体局限:路由按窗口而非 token 粒度进行,一个 128-token 窗口内所有 token 绑定到同一专家,梯度信号因窗口聚合而变粗——换来稳定性但牺牲 token 级自适应;与参数匹配的稀疏 MoE 相比,LoL 在真实文本上训练损失仍高约 0.05;$L$ 较小时最终损失对随机种子更敏感(图 5 误差条明显更大),作者归因于路由初始化决定训练中探索的路径;实验止步于 1B 参数、约 20B tokens、1024 上下文,且 1B 数学数据是合成的(图 4 的损失绝对值不可与 FineWeb 实验直接比较);mode path 干预需先在带策略标签的验证集上统计路径,无法零样本指定。我的补充观察:窗口边界逐深度移动意味着同一 token 在不同深度的 K/V 属于不同窗口,KV cache 难以跨深度复用,推理系统优化空间受限;窗口内所有 token 强制同路径,对『同一窗口内混合语言/任务』的场景表达能力存疑;安全探针的标签本身来自另一个分类器(Maini et al.),0.95 的 AUROC 衡量的是与噪声标签的一致性而非真实安全性;此外论文未系统消融窗口大小 $W$ 与步幅 $S$ 的选择。
独立分析的弱点
(1)窗口级 top-1 路由粒度粗:长窗口内任务或语言混杂时无法分支处理。改进:两级路由(窗口选主专家、token 级微调残差权重),或按内容自适应调整窗口大小。(2)核心深度 $D$ 推理时固定:不能按问题难度分配计算量。改进:结合提前退出,或允许不同窗口执行不同有效深度(部分窗口提前进入 coda),实现输入自适应的 test-time compute。(3)路由对初始化敏感、小库时种子间方差大(图 5)。改进:对路由路径做熵正则、多路由头集成,或用固定路径的 looped 模型热启动后再放开路由自由度。(4)Q-K 一致性依赖窗口内共享专家,代价是注意力感受野被限制在窗口内($W=128$),长程依赖只能靠 prelude/coda 与跨深度错位缓慢传播。改进:窗口内引入少量全局锚点 token(类似 attention sink)携带长程信息。(5)KV cache 与推理效率:窗口边界逐深度移动,K/V 无法直接复用,连续批处理优化受限。改进:设计固定分区方案,或定量研究共享块的权重驻留(cache 友好)能否弥补。(6)只在 1B 参数与 20B tokens 量级验证,scaling 行为未知,growing 与路由在大模型上的稳定性需在 7B+ 规模检验。
未来方向
作者明确提出的方向:其一,可控行为——既然路由特征能线性区分安全/不安全文本与语言,能否在 chat 场景中干预路由以抑制说谎、图谋(scheming)等行为(引用 Park et al. 的评估框架)是开放问题;其二,测试时扩展——LoL 原则上可以在推理时使用比训练时更大的路由深度 $D$,为 latent reasoning 提供按需加深计算图的 test-time scaling 机制;其三,结合测试时学习——把路由对象从固定 Transformer 块换成不断增长的测试时可学习 adapter 库,让参数化记忆对所有层可用。基于本文成果可延伸的方向:把 mode path 提取做成自动化的『算法发现』工具,在无标签数据上聚类路由路径并解释其语义;将层库与 MoE 组合(库中每个块本身是 MoE)以同时获得执行顺序稀疏与 FFN 稀疏;理论层面刻画 top-1 路由路径与可计算算法类的对应关系(可借助 looped transformer 的图灵完备性结果);以及多模态路由——视觉 token 与文本 token 是否天然需要不同的层执行顺序。
复现评估
复现条件较好。代码已开源(github.com/arahmani1/library-of-layers);所有训练数据公开:FineWeb/FineWeb-2(通用与多语言)、Nemotron+OpenMath(数学)、Maini et al. 的 FineWeb 有害性分数(安全);数独/迷宫为程序生成且论文给出完整规则(数独 20M 序列约 10B tokens、保留 30–36 个提示数字;迷宫 8×8 无环、7.2M 迷宫 36M 序列)及格式示例(表 3、表 4、图 15、图 19)。架构与训练超参完整:GPT-2 Medium 形状改造(RoPE/RMSNorm/SwiGLU/无 bias 注意力)、$W=128$、$S=64$、上下文 1024、Muon+AdamW、loss-free 均衡、growing $1-\sqrt{\,\cdot\,}$ 调度,前向传播伪代码(Algorithm 1)逐行可查。算力为学术实验室量级:360M×10B tokens 约 10–17 ExaFLOPs,1B 约 90 ExaFLOPs,约合数十块 A100/H100 训练数天到数周。主要难点在于自定义注意力掩码(窗口边界逐深度循环平移)与路由的工程实现,需要修改 attention kernel;但伪代码详尽、设计选择均有消融支撑,综合难度中等;探针与 PCA 分析门槛很低,普通研究者也能快速验证。
论文图表
dense 固定深度模型中全注意力、滑窗 128、错位窗口(shifted window 128, stride 64)三种注意力模式的训练损失对比,全注意力最低。
基线设定依据:说明 LoL 采用窗口注意力隐含一定代价,使与滑窗基线的公平对比成为必要。
迷宫强制路由下的五指标柱状图:策略遵循率 99.975%/100%/99.9%/99.79%/99.935%,trace 为所请求策略的比例 99.994%–100%,前缀符合率 99.963%–99.985%,答案正确率 99.925%–99.985%,padding 移除 99.965%–99.985%。
迷宫实验的完整定量结果:在共享输出语法的苛刻设定下仍近满分复现算法。
前向传播伪代码:输入 $x$、深度 $D$、窗口 $W$、步幅 $S$;依次执行 prelude、逐深度的窗口划分($\delta_\ell$)、路由输入构造(首窗口用 $e^{(\ell)}$、后续用历史均值)、top-1 选择、窗口内专家执行与 $\alpha_k$ 缩放残差更新,最后 coda 加 softmax 输出。
方法的精确形式化规范,逐行对应实现,是复现工作的最关键参考。