Causal analysis improves VLA efficiency

Causal analysis of the VLM–action-expert interface revealed a single-layer design. It preserves task performance while removing 97% of cross-attention parameters.
AE-only FLOP reduction: 93.5% · VLM + AE FLOP reduction: 40.2%

Aug 2, 2026. We apply activation patching and layer ablation to the VLM–action-expert interface of MolmoBot, π0.5, and MolmoAct2, and find that each action expert reads a relatively narrow band of middle-to-late layers. We use these insights to train more efficient versions of MolmoBot and π0.5.

The idea in one animation. A vision-language model reads the camera image and the instruction; an action expert turns its features into robot motion. At the start the expert reads every VLM layer. Our causal analysis finds that one mid-to-late layer carries the instruction, so the other reads can be removed, and the VLM layers above that point are then never used.

Abstract

Most vision-language-action policies have two parts: a vision-language model (VLM) that interprets the scene and instruction, and an action expert that produces robot motion. In MolmoBot, the action expert receives features from all 36 VLM layers via cross attention. We ask whether all 36 connections are needed. We use activation patching to test one connection at a time. We swap the features sent from one VLM layer while keeping the image and robot state fixed. Layer 24 produced by far the largest change in the robot's instructed reach, while most other layers had little effect. This result suggested a simpler model in which the action expert reads only layer 24.

We tested that design in roughly 1,000 held-out simulation rollouts per variant. The layer-24-only model reached 60.8% success when pretrained in that architecture and 59.9% after a 50K-step retrofit. The original all-layer model reached 59.5%. Removing the other 35 cross-attention reads reduced cross-attention parameters by 97% and action-expert parameters by 27%. Under an assumed 928-token context, the static action-expert FLOP count fell by 93.5%. We found similar concentration in the middle-to-late layers of π0.5 and MolmoAct2, although the relevant layers and inputs differed by model. A separate 50-trial robot pilot compared the original π0.5 policy with a version restricted to layers 14–16. These experiments found no statistically significant different in task success rate or time to completion.

1 Introduction

Many VLA policies combine a pretrained vision-language model with an action expert that denoises a short action sequence. The expert reads the VLM through cross-attention in MolmoBot and MolmoAct2, or through a shared prefix cache in π0.5. This creates a practical question: does the expert decide what to do, or execute a decision already encoded by the VLM? The answer determines where to monitor or steer the policy. It also identifies which parts of the interface may be unnecessary.

We study three policies: MolmoBot-DROID (36-layer Molmo2-4B VLM and 36-block DiT action expert), π0.5 (18-layer PaliGemma backbone and Gemma-300m action expert), and MolmoAct2-DROID (36 backbone/action-expert layer pairs). The causal analyses use fixed-noise, inference-time interventions on simulated or recorded scenes. Section 6 tests the resulting architectural prediction with closed-loop rollouts. Figure 1 gives the cross-model overview.

2 Method: activation patching at the interface

We compare two runs that differ only in the instruction ("Pick up the tomato" vs. "Pick up the plate"), holding the image and robot state fixed, and move one activation from one run into the other. Because nothing else differs, the moved activation isolates the instruction-dependent variable, and both runs stay on the model's input distribution. This is the interchange intervention of Geiger et al. [5]. The activation we move sits at a specific point in the interface: action-expert block k reads VLM layer k, so patching the cross-attention context supplied to one block isolates one point in the VLM-to-action-expert handoff. We sweep all 36 layers one at a time, then patch a six-layer window jointly, because two layers that each carry part of the signal can both look weak in a single-site sweep. To score what each patch does to the behavior, we report the fraction of the reach gap recovered, where the reach gap is the difference between the reaches produced under the two instructions. We subtract the base reach from the patched reach, project that difference onto the counterfactual-minus-base action direction along the tomato–plate axis, and normalize it: 0 means no effect, and 1 means the patch transports the behavior completely. Because the score is signed and specific to this contrast, it separates movement toward the counterfactual from an arbitrary perturbation. We found no established activation-patching methods reference for continuous-action policies, so the metric follows first principles and language-model best practices. One further source of variation has to be averaged out: the action expert denoises from sampled noise, so every measurement averages 64 paired flow-noise draws, with common random numbers for the clean and patched runs.

3 Results: single-layer localization in MolmoBot

The core experiment fixes the image and robot state. MolmoBot sees a tomato and a plate. The only thing that changes between two cases in the instruction: "Pick up the tomato" or "Pick up the plate." The action-expert block k reads VLM layer k, so patching one layer isolates one point in the causal chain. The outcome is the predicted 16-step reach along the tomato–plate axis, averaged over 64 paired flow-noise draws. With no intervention the tomato instruction moves +0.084 m and the plate instruction moves −0.096 m, and because every other input is identical, that difference isolates the instruction-dependent causal signal.

The signal could have appeared in early layers, spread across depth, or traveled through the image patches, and the results rule out all three possibilities. It is localized: patching layer 24 of 36 recovers 79% of the reach gap, flipping the plate-instructed reach from −0.096 m to +0.046 m, and a six-layer window (L21–L26) recovers the full effect. We also find that it travels through the instruction text rather than the image: patching the object token alone has no effect, however the surrounding instruction positions flip the reach to +0.079 m. The image-patch positions carry no signal at layer 24. 

We preform two further experiments to verify the importance of the L21-L26 causal band. Firstly replacing the band's activations with their average across scenes wipes out the decision, leaving a reach of −0.005 m instead of the instructed −0.096 m, while doing the same to unrelated layers changes nothing. Secondly, we preform steering experiments. By adding a single direction at layer 24 (the difference of means vector between the tomato and mean plate activations), shifts the reach further toward the tomato the harder we push, and a random direction of the same size never does. The pattern holds under three fresh flow-noise seeds: each one peaks at layer 24, with recovered gaps of 0.7658, 0.7711, and 0.7924.

4 Results: the same concentrated band in π0.5 and MolmoAct2, at different layers

Two additional policies show the same broad pattern as MolmoBot: a concentrated causal band between the VLM and the AE. The band's location and its carrier, meaning which input positions transmit the signal (vision patches, instruction tokens, or proprioceptive state tokens), vary depending on the model.

π0.5

π0.5 connects its action expert to an image-and-language prefix cache. Across 400 held-out DROID steps, layer-wise ablation identifies a mid-to-late band at layers 13–16, with a peak at layer 14. Modality-restricted ablations place the mid-to-late band in the vision positions and show an additional proprioceptive-state contribution at layer 14. Instruction positions contribute almost nothing. Patching the full prefix from another scene redirects the action, while swapping only the instruction barely moves it. On these scenes, π0.5's action expert mainly reads vision and state.

MolmoAct2

The same protocol on 200 DROID steps places MolmoAct2's causal band at action-expert blocks 17–20. The 10-token proprioceptive-state span carries most of the effect, not the image patches. Masking the state span and template positions reproduces the full-interface ablation (0.78 vs. 0.75). Masking image positions shifts the action by 0.04; masking instruction positions shifts it by 0.005.

5 Results: the reach target forms in layers 12–22, before it is read at 24

The task variable encoded within the MolmoBot residual stream partially explains the observed causal relationship between the VLM and the AE. Across held-out objects and positions, the target object position becomes decodable between layers 12 and 22 and peaks at r = 0.93 at layer 22. The representation's composition changes with depth. Position and distance-to-base peak near layer 22; object identity and size peak in the earliest layers.

6 Results: rollouts with cross-attention at only one layer

The causal result predicts that an action expert reading only VLM layer 24 should preserve performance. We test that prediction by retaining cross-attention in one block and removing it from the other 35.

We evaluate each variant on held-out FrankaPickDroidBench scenes with roughly 1,000 closed-loop rollouts. A model is either pretrained with the single-layer architecture or retrofitted from the released all-layer model with a 50K-step fine-tune.

Layer 24 preserves performance

The pretrained layer-24 model reaches 60.8% success, and the 50K-step retrofit reaches 59.9%. The all-layer baseline reaches 59.5%. An existing model can therefore adopt the single-layer architecture without repeating the expensive pretraining stage.

Layer 12 needs full pretraining

The 50K-step layer-12 retrofit reaches nearly the same flow-matching loss as layer 24 (0.0063 vs. 0.0062), but only 47.5% rollout success. Full pretraining raises layer-12 success to 57.0%. That is +9.5 points above the retrofit (p<0.001) and 3.9 points below layer 24 (p=0.077). The retrofit works at layer 24, but the shallower interface needs pretraining. This further validates our causal result.

The single-layer model is smaller and cheaper to run

Reading only layer 24 removes 97% of cross-attention parameters and 27% of action-expert parameters, reducing total model parameters by 2.9%. It also removes the long-context projection and cross-attention work from 35 of 36 blocks. At the assumed 928-token context and 16-token action horizon, the static linear-and-attention count falls from 110.9 to 7.18 GMAC (billions of multiply-accumulate operations) per action-expert call, a 93.5% reduction. The context is an assumed inference shape, not a value read from the model code. Across plausible context lengths, the reduction ranges from about 91% to 95%. Flow matching calls the expert ten times, so one action generation falls from 1.109 to 0.0718 TMAC (trillions of multiply-accumulate operations) of action-expert work. These are static operation counts, not measured latency.

The single-layer design also permits VLM truncation because layers 25–35 are never consumed. Combining that truncation with the cheaper action expert reduces the static VLM-plus-action-expert count from 7.230 to 4.323 TMAC per action chunk, a 40.2% reduction.

7 Results: a real-robot pilot with 25 trials per policy

All MolmoBot results above come from simulation. A separate study tests the design principle on physical hardware.

A separate π0.5 study uses a physical Franka Emika with a simple pick and place task. The vanilla policy is a π0.5 checkpoint we trained ourselves on the full DROID dataset with absolute joint positions as the action output. The pruned policy is that same checkpoint, fine-tuned with the action expert reading only VLM layers 14–16, this model's own causal band rather than MolmoBot's single layer 24. Both arms succeed on 22 of 25 trials (88%, Wilson interval 70.04–95.83; Fisher exact p = 1.0). We find no statistically significant difference in success rate between the policies (Pi05 22/25, Pi05 ablated 22/25).

8 Limitations and open questions

• All MolmoBot causal results and the architecture comparison are simulation (MolmoSpaces); the only real-robot evidence is the 25-trial-per-arm π0.5 pilot in Section 7.
• The causal-localization results are single-chunk (open-loop): they characterize the policy's first action chunk under intervention, not closed-loop behavior. Closed-loop performance enters only through the architecture comparison in Section 6.
• π0.5's "instruction largely bypassed" evidence rests mostly on paraphrase-style counterfactuals.
• The π0.5 real-robot test is a pilot: 25 trials per arm, with a timing comparison which does not reach statistical significance.

9 Related work

Mechanistic interpretability of VLAs

Several recent studies analyze VLA internals. Häon et al. [6] locate interpretable directions in π0 and OpenVLA, then steer behavior zero-shot on a physical UR5. VLA-Trace [7] uses representation similarity and attention-knockout interventions on π0.5 and OpenVLA. It traces modality-specific control pathways, reports weak fine-grained semantic following, and leaves the architecture unchanged. Swann et al. [8] train sparse autoencoders on π0.5 and OpenVLA. Most features correspond to memorized demonstration episodes; a minority encode general motion and semantic primitives. They validate steering in LIBERO simulation. Zhang et al. [9] identify a linguistic-blindness failure on contradictory instructions and correct it with inference-time attention recalibration on π0, π0.5, and OpenVLA-OFT, including a real Franka arm. Linear probes also recover symbolic state from OpenVLA's hidden layers [10]. This work localizes, probes, or steers at inference time, but does not use localization to change the interface.

How action experts read their backbone

This literature is the closest context for our architecture result. GR00T N1 [11] reads a single intermediate layer: its diffusion transformer cross-attends to the 12th layer of its Eagle-2 backbone instead of the final layer. The authors report faster and more successful behavior than with final-layer features. SmolVLA [12] uses less of the interface by alternating cross-attention with self-attention across action-expert blocks. It conditions on only the first half of the VLM layers. Its conditioning-depth ablation reports average success of 75.0, 78.5, 79.5, and 80.3 for 8, 16, 24, and 32 layers. RDT-1B [13] alternates image and text conditioning across successive cross-attention layers, so neither modality enters every layer.

Many other architectures read the full depth or only the final layer. π0 [14] and π0.5 [15] couple the expert to the backbone with shared attention over a per-layer prefix K/V cache at every layer. MolmoAct2 [16] uses per-layer cross-attention over 36 backbone/expert pairs. MolmoBot [17] is the model analyzed here. CogACT [18] conditions its diffusion module on one pooled token from the final layer; OpenVLA [19] and OpenVLA-OFT [20] read final-layer hidden states. MolmoAct [21] has no separate action expert. Mid-layer and partial-depth reads already exist, but their published justifications are empirical: either a success-versus-latency tradeoff or a redundancy heuristic. None reduces the interface to cross-attention at a single block.

Removing and reusing computation

VLA efficiency methods remove or reuse computation without testing what the removed parts carry. EfficientVLA [22] prunes language layers, reduces visual tokens, and caches intermediate action-head features across denoising steps. VLA-Cache [23] reuses static visual tokens across consecutive timesteps. DeeR-VLA [24] and A1 [25] exit the backbone early at each timestep. General transformer methods remove depth with structured dropout [27] or similarity heuristics that delete the least-changing blocks [28, 29]. The closest work is "Finetuning Vision-Language-Action Models Requires Fewer Layers Than You Think" [26]. It uses centered kernel alignment to find redundant twin layers in the backbone and policy head, removes up to half the depth of π0 and GR00T-N1.5, and validates on ten real-world tasks across four embodiments. The selection rule differs: that work uses representation similarity, while ours tests whether behavior transfers through a layer.

This literature also limits the compute comparison. Cross-attention keys and values depend on the conditioning, not the noisy action sample, so they can be computed once per action chunk and reused across denoising steps. π0's prefix K/V cache already does this. Diffusion work shows that cross-attention output can converge early enough to cache or skip [30], and sharing K/V across layers is also established [31]. The measured MolmoBot implementation recomputes the VLM feature projection and cross-attention K/V at every flow step

References

Models and data

1. Allen Institute for AI. MolmoBot-DROID: Molmo2-4B vision-language backbone with a 36-block DiT flow-matching action expert, per-layer cross-attention coupling. Released model checkpoint; evaluated in the MolmoSpaces simulation suite (FrankaPickDroidBench).
2. Physical Intelligence. π0.5 (pi05_droid) and the openpi repository. PaliGemma-backbone VLM with a Gemma-300m action expert coupled by shared self-attention over a per-layer prefix K/V cache; PyTorch port of the released JAX checkpoint used here, with a numerical-parity gate against the original.
3. Allen Institute for AI. MolmoAct2-DROID (HuggingFace: allenai/MolmoAct2-DROID). 36 backbone/action-expert layer pairs with dedicated cross-attention.
4. DROID robot manipulation dataset (two-camera setup); recorded episodes used for the π0.5 and MolmoAct2 analyses.

Literature

5. Atticus Geiger, Hanson Lu, Thomas Icard, Christopher Potts. “Causal Abstractions of Neural Networks.” NeurIPS 2021. arXiv:2106.02997.
6. Bear Häon, Kaylene Caswell Stocking, Ian Chuang, Claire Tomlin. “Mechanistic Interpretability for Steering Vision-Language-Action Models.” CoRL 2025, PMLR v305, 2743–2762. arXiv:2509.00328.
7. Haoyuan Shi et al. “VLA-Trace: Diagnosing Vision-Language-Action Models through Representation and Behavior Tracing.” arXiv:2605.30117.
8. Aiden Swann et al. “Sparse Autoencoders Reveal Interpretable and Steerable Features in VLA Models.” arXiv:2603.19183.
9. Ninghao Zhang et al. “Restoring Linguistic Grounding in VLA Models via Train-Free Attention Recalibration.” arXiv:2603.06001.
10. “Probing a Vision-Language-Action Model for Symbolic States.” arXiv:2502.04558.
11. Johan Bjorck et al. “GR00T N1: An Open Foundation Model for Generalist Humanoid Robots.” arXiv:2503.14734.
12. Mustafa Shukor et al. “SmolVLA: A Vision-Language-Action Model for Affordable and Efficient Robotics.” arXiv:2506.01844.
13. Songming Liu et al. “RDT-1B: a Diffusion Foundation Model for Bimanual Manipulation.” ICLR 2025. arXiv:2410.07864.
14. Kevin Black et al. “π0: A Vision-Language-Action Flow Model for General Robot Control.” arXiv:2410.24164.
15. Physical Intelligence. “π0.5.” arXiv:2504.16054.
16. Haoquan Fang et al. “MolmoAct2: Action Reasoning Models for Real-world Deployment.” arXiv:2605.02881.
17. “MolmoB0T: Large-Scale Simulation Enables Zero-Shot Manipulation.” arXiv:2603.16861.
18. “CogACT.” arXiv:2411.19650.
19. Moo Jin Kim et al. “OpenVLA: An Open-Source Vision-Language-Action Model.” arXiv:2406.09246.
20. Moo Jin Kim et al. “Fine-Tuning Vision-Language-Action Models: Optimizing Speed and Success.” arXiv:2502.19645.
21. “MolmoAct.” arXiv:2508.07917.
22. “EfficientVLA: Training-Free Acceleration and Compression for Vision-Language-Action Models.” arXiv:2506.10100.
23. “VLA-Cache: Towards Efficient Vision-Language-Action Model via Adaptive Token Caching in Robotic Manipulation.” arXiv:2502.02175.
24. Yang Yue et al. “DeeR-VLA: Dynamic Inference of Multimodal Large Language Models for Efficient Robot Execution.” NeurIPS 2024. arXiv:2411.02359.
25. “A1: A Fully Transparent Open-Source, Adaptive and Efficient Truncated Vision-Language-Action Model.” arXiv:2604.05672.
26. “Finetuning Vision-Language-Action Models Requires Fewer Layers Than You Think.” arXiv:2606.20246.
27. Angela Fan, Edouard Grave, Armand Joulin. “Reducing Transformer Depth on Demand with Structured Dropout.” ICLR 2020. arXiv:1909.11556.
28. Andrey Gromov et al. “The Unreasonable Ineffectiveness of the Deeper Layers.” ICLR 2025. arXiv:2403.17887.
29. Xin Men et al. “ShortGPT: Layers in Large Language Models are More Redundant Than You Expect.” Findings of ACL 2025. arXiv:2403.03853.
30. “Faster Diffusion via Temporal Attention Decomposition.” TMLR 2025 (earlier titled “Cross-Attention Makes Inference Cumbersome in Text-to-Image Diffusion Models”). arXiv:2404.02747.
31. William Brandon et al. “Reducing Transformer Key-Value Cache Size with Cross-Layer Attention.” NeurIPS 2024. arXiv:2405.12981.