MoLe-VLA | AAAI 2026 | Paper Reading

MoLe-VLA: Dynamic Layer-skipping Vision Language Action Model via Mixture-of-Layers for Efficient Robot Manipulation

作者提出的模型主要用于解决机器人系统在部署时面临挑战，通过作者设计的时空感知路由器和自知识蒸馏方法，基于层混合的方式构建了全新的VLA.

工作类型(首次/改进)	技术路线	创新点	日期
改进	VLA	Efficient	2026-03-14

1.What？

本文提出一种混合层视觉-语言-动作模型，其在大型语言模型输入阶段引入新型层选择路由器以实现结构稀疏性。受混合专家模型（MoE）中路由机制的启发——该机制可在单一LLM层内实现横向专家级激活——我们将此概念纵向扩展以实现层级激活。

作者将每个LLM层视为一个专家，利用仿生路由器管理层跳跃，模拟大脑皮层-皮层下回路的选择性激活

MoLe在整体处理输入特征时动态选择最相关层

作者提出的STAR通过模态特定处理显式保留空间视觉范式与任务上下文动态性，并实现单步自适应跨模态融合。

跳过某些层不可避免地会降低模型整合视觉与语言信息、推理其关系、做出上下文感知决策或预测以指导动作的认知表达能力。针对此问题，我们提出认知自知识蒸馏（CogKD），以全层骨干网络为教师模型、MoLe跳跃层模型为学生模型

2.Why？

作者在中对OpenVLA在RLBench仿真器上的分析显示，相邻层输出之间的余弦相似度超过90%，但首层与末层的特征差异显著。这表明跳过相邻层以减少计算量的潜力，但也凸显了早期退出策略的局限性。舍弃深层网络可能丢失关键语义信息。此外，神经科学中浅层脑假说的最新突破性进展也提出，生物认知通过平衡深度层级处理与并行皮层-皮层下回路的捷径作用来实现高效推理，而非完全依赖逐层的信号传递。

作者通过两个原则实现了SBH(浅层脑假说)：

➊ 层级化语义编码：任务关键特征分布在各个层级，需要动态选择而非固定深度处理；

➋ 神经启发的稀疏化：模拟大脑的深度-并行性权衡，我们仅在层编码的语义与任务相关时激活该层。

3.How？

混合专家

MoE范式通过条件计算增强了模型容量，同时保持了计算效率。对于输入x∈Rd，标准MoE定义为：
MoE(x) = Σ_{i=1}^{Ne} G_i(x) · E_i(x)， (1)
其中Ne是专家数量，E_i：Rd→Rd代表第i个专家网络，G(x) = {G_1(x), …, G_{Ne}(x)}是满足Σ_{i=1}^{Ne} G_i(x)=1的门控函数。门控权重计算如下：
G(x) = Softmax(W_g · x + b_g)， (2)
其中W_g∈R^{Ne×d}且b_g∈R^{Ne}是可学习参数。为提升效率，应用了带top-k选择的稀疏门控。为解决输入被路由到少数专家导致的负载不均衡问题，引入负载均衡损失L_lb：
L_lb = (1/Ne) Σ_{i=1}^{Ne} ( ( Σ_{n=1}^N v_i(x_n) ) / ( Σ_{n=1}^N v_i(x_n) + ε ) )^2， (3)
其中若第i个专家被topk门控为输入x_n选中，则v_i(x_n)=1，否则为0。该损失促进了专家均衡并提升了效率。

混合层：MoLe-VMA

视觉语言动作模型。面对长度为L的语言指令l，机器人从传感器（如来自摄像头的RGB图像）获取观察o_t，以预测具有7自由度（DoF）的夹爪动作空间并执行：
a*_t = [Δx, Δy, Δz, Δφ, Δθ, Δψ, g]， (4)
其中Δx、Δy和Δz是末端执行器的相对平移偏移量，Δφ、Δθ、Δψ表示旋转变化，g∈{0,1}指示夹爪打开/关闭状态。我们的基础VLA模型主要由视觉编码器E、MLLM π和动作模块A组成。视觉编码器E包括DINO-v2（Oquab et al. 2023）和Siglip（Zhai et al. 2023），它将输入图像o_t编码为一系列信息记号v_t。对于多模态融合，在视觉编码器E生成的视觉表示之上建立MLLM，作为有效的多模态特征提取器π，形式化表示为：
f_t = π(l, E(o_t))， (5)
其中输出f_t代表MLLM最后一层在时间步t的隐藏状态序列，对应认知记号。这为后续动作模块解释和推导期望动作提供了条件。遵循CogAct（Li et al. 2024a），我们的动作模块A将从输出特征f_t中提取的认知特征作为输入，预测最终动作a*_t。我们的视觉、语言和动作模块通过最小化动作模块预测的噪声与真实噪声之间的均方误差进行端到端训练。以扩散头为例，损失函数定义为：

Ltask = Eε∼N (0,1),i ||εˆi − ε|| (6)

其中εˆi是第i步去噪过程中对含噪动作a_t^的预测噪声，ε为相应的真实噪声。

基于MoLe路由器的层跳过机制。 我们提出MoLeVLA以提升LLM在机器人任务中的效率，由于机器人任务推理需求较简单，许多Transformer层未被充分利用。MoLe采用轻量级路由器，在推理过程中自适应跳过非必要的Transformer层，在降低计算成本的同时保持性能。如图2所示，对于给定的K层MLLM π，MoLe路由器处理输入嵌入x_k ∈ R^{b×n×d}，生成二值门控向量G_{mol}(x) = {G_k}{k=1}^K，其中G_k ∈ [0,1]。为确保效率，仅将G{mol}(x)中top-k值设为1*，决定哪些层π_k以隐藏特征h_k执行，其余层则跳过：

h_k = G_k · π_k(h_{k-1}) + (1 − G_k) · h_{k-1}。 (7)

传统MoE路由器将Token分配给专家，而MoLe路由器则跳过整层，避免冗余计算。这提升了推理效率与响应速度，使MoLe特别适用于如操作和导航等需要轻量级自适应处理的实时机器人任务。

时空感知路由器

我们提出STAR路由机制，通过模态特定处理与自适应跨模态融合保留空间视觉模式与任务感知动态，优化VLA任务的层选择。给定视觉特征v_t ∈ R^{b×n_{img}×d}与文本特征l ∈ R^{b×n_{text}×d}，两种模态通过可学习矩阵W_p ∈ R^{d×d_1}投影到共享潜在空间：

h_{img} = v_t · W_p, h_{text} = l · W_p。 (8)

我们从h_{img}计算空间路由权重S ∈ R^{b×N_e}，以处理图像块并捕获几何关系与物体布局：

S = W_s^{(2)} · φ(W_s^{(1)} · h_{img} + b_s^{(1)})， (9)

其中φ表示GELU，空间池化保持视觉推理的平移不变性。同时，从h_{text}通过Transformer模块推导出时序路由权重T ∈ R^{b×N_e}，以提取指令感知的时序动态，这里的“时序”对应语言指令隐含的任务执行阶段（例如“先抓取后放置”）：

T = W_t · Φ(Transformer(h_{text}))。 (10)

Transformer建模文本命令中的步骤依赖性，平均池化Φ则提取阶段转换信号。可学习的锐度控制器α ∈ [0,1]根据指令复杂度自适应调整路由焦点：

α = σ(W_τ^⊤ · h_{text}^{[CLS]} + b_τ)， (11)

其中σ为Sigmoid函数，α较高时强调空间专家以精确估计抓取姿态，α较低时激活几何-文本对齐专家以满足任务条件适配。最后，最终门控权重G ∈ R^{b×N_e}通过动态锐化融合结合两种模态S与T：

G = S + T。 (12)

通过整合空间与时序信息，我们的方法使路由器能够选择LLM层，自适应优化VLA任务性能。该方法高效，每样本FLOPs仅为O(N_e(d^2 + N_t^2{ext}))，而标准MoE框架为O(N_e d)，其中d ≫ N_text, d^2。该设计确保了高适应性与计算效率。

认知自知识蒸馏

为缓解层跳过导致的认知损失，我们采用自蒸馏，以原始模型为教师，MoLe为学生。遵循Token级蒸馏方法，学生从教师处重构每个Token：给定f^{(t)} ∈ R^{n×d}和f^{(s)} ∈ R^{n×d_s}，Token重构如下：

L_{mimic} = 1/N ||f^{(t)} − μ(f^{(s)})||2^2， (13)
其中μ(·)为投影层。然而，式(13)对所有Token均衡蒸馏，对于某些视觉Token与文本相关性更高的复杂操作任务可能次优。因此，我们利用认知Token e_tc ∈ R^{1×d}（即LLM输出序列的最后一个Token）进行自适应蒸馏。该Token自然封装了多模态输入下语言指令与视觉信息融合形成的全局任务语义。每个模型拥有各自的认知Token（教师为e{t,c}^{(t)}，学生为e_{c,t}^{(s)}）。我们通过计算蒸馏过程中各Token与认知Token的相似度，识别出感兴趣Token（ToIs）M：
M^{(i)} = η(e_c^{(i)} f^{(s)})，i ∈ {s, t}， (14)
其中η表示Sigmoid函数。接着，利用教师与学生认知Token生成的ToIs交集决定各Token的蒸馏程度，即M = M^{(t)} ⊙ M^{(s)}，因为**蒸馏Token应包含对教师和学生均重要的Token。因此，式(13)可更新为：
L_{cog-mimik} = 1/N ||M ⊙ f^{(t)} − μ(M ⊙ f^{(s)})||*2^2， (15)
其中⊙表示逐元素矩阵乘法。此外，我们引入与认知Token配对的Reverse-KL，优先覆盖教师分布的模式，同时抑制学生的虚假模式。*原因在于，对于VLA蒸馏，反向KL损失可防止对分布外视觉-语言对的过度自信预测，同时保留语义多样性，比前向KL具有更好的泛化性。因此，我们按照前述方式得到L{cog-reversekl}以增强分布约束：
L_{cog-reversekl} = (M ⊙ f^{(s)}) log((M ⊙ f^{(s)})/(M ⊙ f^{(t)}))。 (16)
最终，我们的CogKD损失可表示为：
L_{cog} = (1 − λ_1)L_{cog-mimik} + λ_1L_{cog-reversekl}， (17)
其中λ_1设为0.5以平衡损失。

优化目标

对于教师模型的更新，我们使用预训练参数初始化两个模型，并通过指数移动平均（EMA）以更新权重α=0.999更新教师模型。最终训练目标由L_task、L_cog与L_lb组合得到：

L_{MoLe} = L_{task} + λ_2L_{cog} + λ_3L_{lb}， (18)

其中λ_2与λ_3为两个超参数，默认分别设为0.5和0.1。

4.Takeaways:

我觉得论文的核心主要是跳过机制，和我的想法核心思想类似，但是作者设计更为细致，我可能就只想着跳过整个LLM前向传播了。

局限性&未来工作：

论文中未说