PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Rongyu Zhang, Menghang Dong, Yuan Zhang, Liang Heng, Xiaowei Chi, Gaole Dai, Li Du, Yuan Du, Shanghang Zhang

arXiv:2503.20384v2cs.ROcs.AI

TL;DR

MLLM-based robotic systems face high computational demands, while layer redundancy coexists with important semantic differences across layers. MoLe-VLA dynamically selects LLM layers using STAR and restores cognitive information with CogKD, achieving up to 8% higher model performance while reducing computational costs by ×5.6.

  • Problem

    MLLM deployment in real-world robotic systems is constrained by high computational demands, and existing layer-sparsification methods can overlook semantic information in critical final layers.

  • Method

    MoLe-VLA treats LLM layers as experts, uses the Spatial-Temporal Aware Router to select relevant layers dynamically, and applies Cognition Self-Knowledge Distillation to compensate for information lost through layer-skipping.

  • Results

    MoLe-VLA improves model performance by up to 8% while reducing computational costs by ×5.6 across real-world and RLBench environments.

  • Takeaways & Limitations

    MoLe-VLA enables efficient and adaptable robotic systems by dynamically activating key LLM layers while preserving essential information.

Abstract

from arXiv · show

Multimodal Large Language Models (MLLMs) excel in understanding complex language and visual data, enabling generalist robotic systems to interpret instructions and perform embodied tasks. Nevertheless, their real-world deployment is hindered by substantial computational and storage demands. Recent insights into the homogeneous patterns in the LLM layer have inspired sparsification techniques to address these challenges, such as early exit and token pruning. However, these methods often neglect the critical role of the final layers that encode the semantic information most relevant to downstream robotic tasks. Aligning with the recent breakthrough of the Shallow Brain Hypothesis (SBH) in neuroscience and the mixture of experts in model sparsification, we conceptualize each LLM layer as an expert and propose a Mixture-of-Layers Vision-Language-Action model (MoLe-VLA, or simply MoLe) architecture for dynamic LLM layer activation. We introduce a Spatial-Temporal Aware Router (STAR) for MoLe to selectively activate only parts of the layers based on the robot's current state, mimicking the brain's distinct signal pathways specialized for cognition and causal reasoning. Additionally, to compensate for the cognitive ability of LLMs lost in MoLe, we devise a Cognition Self-Knowledge Distillation (CogKD) framework. CogKD enhances the understanding of task demands and improves the generation of task-relevant action sequences by leveraging cognitive features. Extensive experiments conducted in both RLBench simulation and real-world environments demonstrate the superiority of MoLe-VLA in both efficiency and performance. Specifically, MoLe-VLA achieves an 8% improvement in the mean success rate across ten tasks while reducing computational costs by up to x5.6 compared to standard LLMs.

1. Introduction

MoLe-VLA addresses the computational burden of robotic MLLMs by dynamically selecting LLM layers according to spatial-temporal task information. STAR performs layer routing, while CogKD preserves cognitive information lost through layer skipping, yielding improved efficiency and performance.

  • Robotic MLLM deployment is limited by high computational demands despite strong multimodal understanding and generalization abilities.
  • Adjacent LLM layers exhibit substantial redundancy, but first and last layers differ significantly, limiting simple early-exit strategies.OpenVLA analysis found cosine similarity between consecutive layer outputs above 90%.
  • MoLe-VLA treats each LLM layer as an independent expert and dynamically activates layer combinations through an input-stage router.The design extends mixture-of-experts routing vertically from expert-wise activation to layer-wise activation.
  • STAR processes visual spatial features and textual temporal dependencies, then selects the top-k layers using softmax probabilities.This routing uses spatial-temporal information from the robot’s current environment.
  • CogKD uses the full-layer model as teacher and the layer-skipping model as student to preserve task comprehension and action generation.A learnable cognition token integrates visual tokens and language guidance.
  • MoLe reduces computational costs by ×5.6 while improving model performance by up to 8% in RLBench and real-world experiments.
  • The framework combines Shallow Brain Hypothesis-inspired dynamic activation, STAR routing, and CogKD to improve efficiency and recover cognitive information.

2. Related works

Related work improves VLA efficiency through architectural design, compression, and dynamic networks, while sparse mixture-of-experts models reduce computation by activating only part of the model.

  • VLA models integrate visual and linguistic inputs to interpret environments and generate executable actions.
  • Existing VLA efficiency methods include efficient architectures, model compression, and dynamic networks.
  • Sparse MoE architectures reduce inference computation by activating only a small portion of the computation graph.
  • Prior dynamic routing methods choose between computational paths, whereas MoLe routes among standard block computations at the layer level.

3. Methods

MoLe-VLA applies conditional computation to vision-language-action models by routing inputs to selected layers, while its VLA pipeline maps visual-language features to robot actions. STAR dynamically selects layers using spatial and temporal information, and CogKD supports efficient layer skipping during inference.

  • 3.1. Preliminary: Mixture-of-Experts: Conditional computation combines expert outputs through learned gating, with sparse top-k selection improving efficiency and load balancing encouraging balanced expert utilization.The MoE formulation uses gating weights over expert networks and introduces a load-balance loss for selected experts.
  • VLA model: The VLA model encodes camera observations and language with a vision encoder and MLLM, then uses a cognition feature to predict actions.The action module is trained end-to-end using mean squared error between predicted and ground-truth diffusion noises.
  • 3.2. Mixture-of-Layers: MoLe-VLA: MoLe-VLA comprises the MoLe architecture, STAR, and CogKD for dynamically selecting LLM layers in vision-language-action models.The framework figure identifies STAR and CogKD as the two principal components of MoLe-VLA.
  • 3.2. Mixture-of-Layers: MoLe-VLA: MoLe adaptively skips non-essential transformer layers by setting only the top-k router values to one and carrying hidden features through skipped layers unchanged.This layer-level routing avoids token-wise expert allocation and targets reduced inference cost for robotic tasks.
  • 3.3. Spatial-Temporal Aware Router: STAR projects visual and textual features into a shared latent space, computes spatial weights from visual features, and derives temporal weights from text features.The router uses the spatial structure of visual inputs and temporal dependencies in language inputs to select layers dynamically.
  • Algorithm 1: Algorithm 1 takes observations, language instructions, ground-truth noise, and the total layer count to produce the MoLe student loss while computing skip indices with STAR.The procedure explicitly separates skip-index computation from the student-model loss output.

11 Step 4: Compute skip indices via STAR router

The method dynamically routes inputs to selected LLM layers using spatial-temporal information, then uses cognition-aware self-distillation to preserve task-relevant representations while reducing computation.

  • 11 Step 4: Compute skip indices via STAR router: STAR combines spatial visual features and temporal textual dependencies into routing information for dynamically selecting relevant LLM layers.The router produces final gating weights through temperature-scaled Gumbel-Softmax selection.
  • 11 Step 4: Compute skip indices via STAR router: The router's computation costs O(Ne(d2 + Ntext2)) FLOPs per sample, compared with O(Ned) for standard MoE when d ≫ Ntext, d2.This design is intended to provide adaptive layer selection with computational efficiency.
  • 3.4. Cognition self-Knowledge Distillation: CogKD uses the original full-layer model as teacher and the layer-skipping MoLe model as student to compensate for reduced cognitive expressiveness.The framework is designed to preserve grasping ability while mitigating cognitive collapse.
  • 3.4. Cognition self-Knowledge Distillation: A learnable cognition token integrates visual tokens with language guidance and identifies tokens of interest for task-relevant distillation.The teacher and student each have a cognition token, and token-interest masks are based on cognition-token similarity.
  • 3.4. Cognition self-Knowledge Distillation: CogKD combines cognition-feature mimicry and Reverse-KL losses, with λ1 set to 0.5 to balance them.The final cognition loss is a weighted combination of Lcog-mimic and Lcog-reversekl.
  • 3.5. Optimization Objective: The overall MoLe objective combines task, cognition, and load-balancing losses, using λ2 = 0.5 and λ3 = 0.1 by default.These terms jointly define the final training objective.

4. Experiments

MoLe-VLA is evaluated in RLBench and real-world settings for manipulation performance, efficiency, scalability, quantization, and component effectiveness. Across these evaluations, it selectively skips layers while retaining or improving task performance.

  • Experimental setup: RLBench evaluates ten tabletop manipulation tasks, while real-world experiments use an FR3 robot and three object-interaction tasks.RLBench uses a Franka Panda with front-view RGB input; real-world deployment uses an FR3 platform.
  • Simulation results: 60.8% mean success rate is achieved by MoLe-CogAct across ten RLBench tasks using only half of the LLM layers.This exceeds DeeR at 59.2% and MoD at 56.4%.
  • Efficiency analysis: 19% of backbone FLOPs and 2× faster inference are achieved while maintaining similar success rates to the full-layer backbone.MoLe requires 0.309 seconds per iteration while retaining the highest reported mean success rate of 60.8%.
  • Quantization: 58.8% success rate at 15.7 Hz is achieved after 8-bit quantization while using 55% of CogAct’s GPU memory.The comparison is against FP16 CogAct on an NVIDIA 4090D.
  • Scalability: +2.7%, +3.6%, and +1.5% mean-success-rate improvements are reported for Small, Base, and Large models over full-layer CogAct.MoLe-Large reaches a 71.5% mean success rate.
  • Ablation study: 60.8% mean success rate is obtained with STAR, cognition tokens, MSE, and Reserve KL losses, a +3.6% gain over baseline CogAct.The baseline CogAct mean success rate is 57.2%.
  • Real-world evaluation: 80% success rate is achieved on the real-world pour water task while reducing LLM computational cost by 50%.The task requires precise 3D position and rotation predictions.
  • Real-world evaluation: MoLe-VLA predicts continuous 7-DoF end-effector poses that support precise execution along planned trajectories.The pour water demonstration includes grasping, lifting, positioning, and gripper rotation.

5. Conclusion

The paper concludes that MoLe-VLA optimizes VLA models by dynamically activating key LLM layers with STAR while using CogKD to preserve cognitive capacity. Experiments in RLBench and real-world environments support efficient and adaptable robotic systems.

  • Conclusion: MoLe-VLA dynamically activates key LLM layers with STAR to reduce redundancy while preserving essential information.The framework is inspired by the Shallow Brain Hypothesis.
  • Conclusion: CogKD addresses performance loss from layer skipping by enhancing efficiency and cognitive capacity.The conclusion presents CogKD as a complement to sparse layer activation.
  • Conclusion: RLBench and real-world experiments show that MoLe reduces computational costs for efficient and adaptable robotic systems.The conclusion states the supported scope without specifying a single benchmark value.

Supplementary Material

The supplementary material documents real-world hardware, training settings, additional comparisons, and data-scalability experiments. These materials extend reproducibility and evaluate MoLe under reduced-task and deployment-oriented conditions.

  • Supplementary analyses: The supplementary material provides hyperparameters and training-scaling analyses to support reproducibility and further exploration.Appendices A–C cover robot deployment, experiment hyperparameters, and training scalability.
  • Real-world deployment: Real-world experiments use an FR3 arm with a 3D-printed UMI gripper and a right-positioned GoPro 9 RGB camera.The setup addresses limitations of the FR3 default gripper.
  • Training details: Training on RLBench uses a batch size of 64 × 8, AdamW, 2 × 10^-5 learning rates, and 2500 warmup steps.The cited training configuration is summarized in the supplementary hyperparameter tables.
  • Benchmark comparison: Table 8 compares existing VLA models across ten RLBench tasks, with efficiency methods operating using only 50% of LLM layers.Results are color-coded for first and second place, with row colors indicating baseline type.
  • Data scalability: 82.7% success rate is achieved on three RLBench tasks with MoLe using 50% of baseline computational resources.CogAct, Random-skip, and DeeR achieve 71.0%, 64.1%, and 78.6%, respectively.

D. Hyperparameter analysis

The hyperparameter analysis evaluates MoLe-CogAct across ten RLBench tasks under 50% layer skipping. Performance improves with greater emphasis on specific layers, while smaller α values improve optimization stability.

  • Parameter effects: Increasing λ3 consistently improves performance by emphasizing the role of specific layers.The analysis evaluates λ1, λ2, λ3, and α configurations.
  • Parameter effects: Smaller α values lead to better optimization stability across the tested configurations.The reported effect is based on the parameter ablation study.
  • Evaluation setting: The study evaluates MoLe-CogAct under a 50% layer-skip setting across ten RLBench tasks.These settings define the scope of the hyperparameter comparison.

E. Layer skip analysis

MoLe-CogAct remains robust as layers are skipped, with substantial performance degradation appearing only under the most aggressive skipping. Supplementary evaluations show successful spatially demanding manipulation, while real-world failures expose limits in control, rotation, and reachable pose prediction.

  • Robustness under layer skipping: Performance remains stable across most tasks as skipped layers increase, with only slight success-rate declines through 24 skipped layers.A significant drop appears only when skipping 30 layers, corresponding to an almost 95% reduction in FLOPs.
  • Robustness under layer skipping: An almost 95% FLOPs reduction from skipping 30 layers produces a significant performance drop.
  • RLBench qualitative results: MoLe-CogAct accurately predicts 7-DoF end-effector poses across seven supplementary RLBench tasks, enabling smooth task completion.The visualizations evaluate capabilities of the efficient layer-skipping architecture.
  • RLBench qualitative results: The supplementary RLBench results highlight effectiveness on tasks requiring both spatial understanding and precise control.
  • Real-world failure analysis: Real-world failures include loss of control, rotational prediction errors, and poses beyond the Franka arm’s physical or workspace limits.Examples include pull-drawer force or slipping failures, pour-water angle errors, and unreachable detach-charger poses.
Loading 2503.20384v2…