PaperlayerView exact revision on arXiv

Source-linked AI summary

ROCKET: Residual-Oriented Multi-Layer Alignment for Spatially-Aware Vision-Language-Action Models

Guoheng Sun, Tingting Du, Kaixi Feng, Chenxiang Luo, Xingguo Ding, Zheyu Shen, Ziyao Wang, Yexiao He, Ang Li

arXiv:2602.17951v1cs.CVcs.AI

TL;DR

VLA backbones predominantly pretrained on 2D imagery can lack stable 3D structure, and single-layer alignment misses complementary cross-layer signals. ROCKET aligns residual streams across layers using a shared projector, with Matryoshka-style sparse activation to balance alignment losses across depths. ROCKET consistently improves spatial grounding and generalization across datasets and VLA backbones, reaching SOTA performance with ∼4% of the training compute required by prior SOTA methods.

  • Problem

    VLA backbones predominantly pretrained on 2D imagery can lack stable 3D structure, and single-layer alignment misses complementary cross-layer signals.

    Source: §1 · p. 1, §3.2 · p. 3

  • Method

    ROCKET aligns residual streams across layers using a shared projector, with Matryoshka-style sparse activation to balance alignment losses across depths.

    Source: §1 · p. 2

  • Results

    ROCKET consistently improves spatial grounding and generalization across datasets and VLA backbones, reaching SOTA performance with ∼4% of the training compute required by prior SOTA methods.

    Source: §1 · p. 2, §7 · p. 8

  • Takeaways & Limitations

    ROCKET offers a residual-oriented alignment framework that the authors describe as a simple and scalable route to more reliable 3D-aware robotic manipulation.

    Source: §7 · p. 8

    The framework assumes access to a frozen 3D-aware vision foundation model that provides geometry-grounded patch-token representations.

    Source: §3.1 · p. 3

Abstract

from arXiv · show

Vision-Language-Action (VLA) models enable instruction-following robotic manipulation, but they are typically pretrained on 2D data and lack 3D spatial understanding. An effective approach is representation alignment, where a strong vision foundation model is used to guide a 2D VLA model. However, existing methods usually apply supervision at only a single layer, failing to fully exploit the rich information distributed across depth; meanwhile, na\"ive multi-layer alignment can cause gradient interference. We introduce ROCKET, a residual-oriented multi-layer representation alignment framework that formulates multi-layer alignment as aligning one residual stream to another. Concretely, ROCKET employs a shared projector to align multiple layers of the VLA backbone with multiple layers of a powerful 3D vision foundation model via a layer-invariant mapping, which reduces gradient conflicts. We provide both theoretical justification and empirical analyses showing that a shared projector is sufficient and outperforms prior designs, and further propose a Matryoshka-style sparse activation scheme for the shared projector to balance multiple alignment losses. Our experiments show that, combined with a training-free layer selection strategy, ROCKET requires only about 4% of the compute budget while achieving 98.5% state-of-the-art success rate on LIBERO. We further demonstrate the superior performance of ROCKET across LIBERO-Plus and RoboTwin, as well as multiple VLA models. The code and model weights can be found at https://github.com/CASE-Lab-UMD/ROCKET-VLA.

1. Introduction

VLA models built on predominantly 2D-pretrained backbones can lack stable 3D structure, while single-layer alignment depends on a difficult post-hoc layer choice and naive multi-layer alignment can cause gradient interference. ROCKET uses a shared projector and sparse activation to align across depth, reporting state-of-the-art performance with about 4% of prior SOTA training compute.

1. Introduction

  • 2D-pretrained VLA backbones may lack stable 3D structure, limiting generalization when manipulation depends on geometry, viewpoint changes, or fine-grained spatial relations.Source: §1 · p. 1
  • Implicit representation alignment is typically more inference-efficient and easier to scale than adding 3D inputs or using external depth estimators.Source: §1 · p. 1
  • Single-layer alignment is sensitive to layer choice, while independent projectors for multi-layer alignment can learn inconsistent mappings and cause gradient interference and performance collapse.Source: §1 · pp. 1–2
  • ROCKET aligns residual streams across multiple layers with a shared projector and Matryoshka-style sparse activation to reduce gradient conflicts and balance alignment losses across depths.Source: §1 · p. 2
  • ∼4% of the training compute required by prior SOTA methods is enough for ROCKET to reach SOTA performance.Source: §1 · p. 2

2. Related Work

Prior work improves VLA spatial grounding with explicit 3D inputs, depth-aware encodings, or representation alignment from pretrained 3D foundation models. Representation-alignment methods usually supervise one layer, while extending distillation across layers introduces projector-related optimization challenges.

2. Related Work

  • 3D information has been added to VLA models through point-cloud inputs or depth-aware encodings, while 3D foundation models offer spatial representations as supervision for 2D vision backbones.Source: §2 · p. 2
  • Representation-alignment methods commonly match student features from one layer to a fixed teacher layer, and the preferred layer varies across methods and empirical settings.Table 1 compares methods for spatial supervision; its caption identifies OpenVLA’s choice as layer 24 of a 32-layer backbone.Source: §2 · p. 2, Table 1 · p. 2
  • Multi-layer distillation can support knowledge transfer, but layer-wise projectors used to reconcile feature spaces can introduce gradient interference in VLA models.Source: §2 · pp. 2–3

3. Theoretical Framework: Residual Dynamics and Gradient Coherence

The framework treats a pretrained VLA policy as a student supervised by a frozen 3D-aware vision teacher. It argues that independent projectors can produce conflicting cross-layer gradients, whereas a shared projector structures those interactions and improves gradient coherence.

3.1. Background

  • A pretrained VLA policy is treated as the student, trained on action demonstrations and applicable to different action parameterizations.The student receives an observation and instruction; the teacher is a frozen 3D-aware vision model that provides geometry-grounded patch-token representations.Source: §3.1 · p. 3
  • Single-layer alignment matches student visual tokens at a selected depth to teacher representations through a lightweight projector, using a similarity-based loss within a joint task-and-alignment objective.Source: §3.1 · p. 3

3.2. The Gradient-Conflict of Multi-Layer Alignment

  • Single-layer alignment requires post-hoc layer selection and misses complementary signals across depth, motivating simultaneous alignment of multiple layers.Source: §3.2 · p. 3
  • Residual representations evolve through layer updates, and hidden states within each model tend to converge into a shared feature space, motivating a consistent student-to-teacher residual-stream mapping.Under naive alignment, learned projectors diverge and become nearly orthogonal in parameter space.Source: §3.2 · p. 3
  • In a Pre-LN residual stream, early-layer gradients receive an additive mixture of signals from aligned future layers, so whether alignment helps depends on whether those signals add constructively or cancel.Source: §3.2 · p. 4
  • With independent projectors, cross terms determine constructive or destructive gradient interference and can remain weak, sign-unstable, or destructive across layers.Source: §3.2 · p. 4
  • Independent projectors can leave cross-layer interactions unconstrained, and observed gradients remain poorly aligned throughout training; this design also hurts downstream VLA manipulation performance.The figure caption describes cosine similarity between gradients in shallow layers; red indicates coherence and blue indicates interference.Source: §3.2 · pp. 3–4, Figure 1 · p. 4

3.3. Motivation

  • A shared projector gives cross-layer interactions a structured decomposition and biases gradients toward constructive interference, with higher cross-layer cosine similarity and downstream manipulation gains reported in experiments.The theoretical account describes a common PSD reference plus a controlled deviation, with a signal-aligned lower bound on a specified error-signal subspace.Source: §3.3 · p. 4

4. Methodology

ROCKET aligns multiple VLA layers with a 3D vision teacher using one shared mapping, then varies that projector’s activated capacity across depth to balance alignment losses.

4.1. Overview of ROCKET

  • ROCKET aligns multiple layers and uses a shared projector to learn a common mapping between residual streams, addressing gradient conflicts across layers.The overview also notes that a simple layer-selection strategy lets ROCKET outperform baselines.Source: §4.1 · pp. 4–5
  • Matryoshka-style activation progressively increases the fraction of active projector parameters from shallow to deep layers to prevent shallow layers from dominating.The nested projectors share parameters while exposing increasing effective capacity at deeper layers.Source: §4.1 · p. 5, §4.3 · pp. 5–6

4.2. Multi-layer Alignment with Shared Projector

  • A single shared projector replaces layer-specific projectors to learn a student-to-teacher mapping that remains consistent across aligned depths.ROCKET applies the shared mapping to the selected student and teacher layer pairs.Source: §4.2 · p. 5

4.3. Matryoshka-style Sparse Activation

  • The training objective combines action-prediction loss with the multi-layer alignment regularizer, using λ = 0.5 by default.The alignment loss uses 1 − cos(a, b) by default.Source: §4.3 · p. 6

5. Main Result

ROCKET matches Spatial Forcing on LIBERO using about 4% of its compute budget, and is also evaluated with different VLA models and on robustness and bimanual benchmarks.

5.1. Experimental Setup

  • Unless otherwise specified, pretrained OPENVLA-7B weights are fine-tuned on LIBERO and LIBERO-PLUS; LoRA uses 50,000 steps (batch size 32), while LIBERO full fine-tuning uses 30,000 steps (batch size 64).RoboTwin 2.0 follows the official training protocol, using LoRA fine-tuning for 30,000 steps with batch size 16.Source: §5.1 · p. 6

5.2. Performance of ROCKET Across Different Settings

  • 81.7% average success on LIBERO-Plus exceeds the baseline’s 80.0%, with the largest gains under Robot and Layout shifts.The benchmark reports seven perturbations and averages results across Spatial, Object, Goal, and Long task groups.Source: Figure 5 · p. 7, §5.2 · p. 7
  • At about 4% of Spatial Forcing’s compute budget, ROCKET matches its LIBERO performance and substantially outperforms prior models.The comparison includes standard 2D VLA models, explicit-3D methods, and single-layer representation-alignment baselines.Source: §5.2 · p. 7
  • ROCKET improves PI0.5 performance by 2.3% under the same full-fine-tuning setting.The comparison is against direct fine-tuning and Spatial Forcing.Source: Table 3 · p. 6, §5.2 · p. 7
  • On RoboTwin 2.0 bimanual tasks, ROCKET has a clear advantage in Easy and is only slightly behind the strongest fine-tuned baseline in Hard.Evaluation uses ALOHA assets and the Easy and Hard task sets.Source: §5.2 · p. 7

6. Discussion

ROCKET’s sparse activation reaches 98.5% after the shared projector reaches 98.2%; the study also examines layer count, layer-selection strategies, and training efficiency. It finds that roughly 8–16 aligned layer pairs offer the best trade-off and that reasonable selection strategies provide gains without expensive post-hoc search.

6.1. Optimal Number of Alignment Layers

  • Roughly 8–16 aligned layer pairs yield the best trade-off; performance saturates or degrades beyond this range, so ROCKET uses N = 10 by default.The study fixes projector capacity while varying the number of aligned pairs.Source: §6.1 · p. 7, Table 4 · p. 7

6.2. Criteria for Layer Selection

  • Most reasonable layer-selection strategies provide consistent gains without expensive post-hoc search, and E2M-Last1 is the default.Strategies include uniform, stage-based, E2M-Last1, and similarity-based selection.Source: §6.2 · pp. 7–8, Table 5 · p. 8

6.3. Data and Computational Efficiency

  • ROCKET performs strongly with 1%, 5%, and 10% of LIBERO data, is about 10% ahead of baseline around 10k steps, and reaches SOTA at 24× lower cost than the previous best method.The limited-data subsets are each trained for 30k steps.Source: §6.3 · p. 8, Table 6 · p. 8, Table 7 · p. 8

6.4. Ablation Study

  • 98.5% performance follows adding Matryoshka-style sparse activation to the shared projector, which reaches 98.2%.The ablation reports independent projectors at 80.0%.Source: §6.4 · p. 8, Table 8 · p. 8

7. Conclusion

ROCKET transfers 3D spatial priors into 2D-pretrained VLA models through residual-oriented multi-layer alignment and a shared projector. Experiments report improved spatial grounding and generalization across datasets and VLA backbones, with near-SOTA LIBERO performance at a small fraction of the compute.

7. Conclusion

  • ROCKET transfers 3D spatial priors from a vision foundation model into 2D-pretrained VLA models using residual-oriented multi-layer alignment.A layer-invariant shared projector mitigates gradient interference, while Matryoshka-style sparse activation balances supervision across depths.Source: §7 · p. 8
  • Experiments show improved spatial grounding and generalization across datasets and VLA backbones, with near-SOTA LIBERO performance at a small fraction of the compute.Source: §7 · p. 8
  • ROCKET offers a simple and scalable route to more reliable 3D-aware robotic manipulation.Source: §7 · p. 8

A. Training Details

Training freezes the teacher throughout and fine-tunes the student VLA model. For standard OpenVLA-7B training, the learning rate changes after 10,000 steps; search-oriented runs use a smaller batch and delay that decay.

A. Training Details

  • The teacher remains frozen throughout training while only the student VLA model is fine-tuned.Source: A. Training Details · p. 13
  • OpenVLA-7B trains for 50,000 steps with batch size 32, using a learning rate of 5 × 10^-4 for 10,000 steps before decaying to 5 × 10^-5.Source: A. Training Details · p. 13
  • Search-oriented experiments use batch size 16 and retain the 5 × 10^-4 learning rate for 30,000 steps before decaying to 5 × 10^-5.This adjustment reduces compute while avoiding under-training.Source: A. Training Details · p. 13
  • PI0.5 experiments use batch size 64 in PyTorch, while PI0 RoboTwin 2.0 experiments use batch size 16 with JAX.Source: A. Training Details · p. 13
  • In single-task settings, the alignment-loss coefficient is reduced from 0.5 to 0.125 because action loss typically converges quickly.Source: A. Training Details · p. 13

B. Performance Trajectories on LIBERO Across Training

The appendix describes LIBERO training trajectories, benchmark suites, VLA architectures, and projector similarities. It characterizes LIBERO’s task and data structure, details LIBERO-Plus and RoboTwin 2.0, and reports that independently initialized projectors learn distinct mappings whose discrepancies can lead to gradient conflicts.

B. Performance Trajectories on LIBERO Across Training

C.1. LIBERO

  • LIBERO comprises four task suites across three robot embodiments, with suites spanning spatial, object, goal, and long-horizon generalization; its simulated demonstrations provide synchronized RGB and proprioceptive observations.The original release contains approximately 20,000 demonstration episodes.Source: §C.1 · pp. 13–14

C.2. LIBERO-Plus

  • LIBERO-Plus expands the benchmark to over 100 manipulation tasks and hundreds of thousands of simulated expert episodes, with synchronized multimodal observations.Source: §C.2 · p. 14

C.3. RoboTwin 2.0

  • RoboTwin 2.0 provides 50 dual-arm tasks across five embodiments and over 100,000 expert trajectories for bimanual manipulation evaluation.Its object library contains 731 instances spanning 147 categories.Source: §C.3 · p. 14

D.1. OpenVLA

  • OpenVLA combines DINOv2 and SigLIP visual features, a Llama 2 7B backbone, and discretized 7-dimensional robot actions.Source: §D.1 · p. 15

D.2. OpenVLA-OFT

  • OpenVLA-OFT uses continuous-action prediction with LoRA fine-tuning, while OFT+ adds language-conditioned FiLM modulation of visual features.Source: §D.2 · p. 15

D.3. PI0

  • PI0 uses flow matching to generate action sequences with separate vision-language and action experts.Source: §D.3 · p. 15

D.4. PI0.5

  • PI0.5 discretizes robot state into prefix tokens and replaces PI0’s MLP timestep fusion with adaptive RMS normalization.Source: §D.4 · p. 15

F. Why We Need a Nested Sparse-Activation Mechanism to Balance Multiple Alignment Losses

VGGT and OpenVLA-7B representations are more similar at shallow depths, where a shared projector aligns more easily but may not serve deeper layers equally well. ROCKET addresses this mismatch by increasing projector capacity activated with depth.

F. Why We Need a Nested Sparse-Activation Mechanism to Balance Multiple Alignment Losses

  • Shallow VGGT layers are more similar to OpenVLA-7B, but a shared projector that aligns shallow layers may not align deeper layers equally well.The reported similarity between VGGT and OpenVLA-7B generally decreases with depth.Source: F. Why We Need a Nested… · pp. 15–16
  • ROCKET’s nested sparse activation uses fewer projector parameters for shallow layers and progressively activates more, reaching full capacity at the deepest layers.Source: F. Why We Need a Nested… · p. 16
  • CKA compares the pairwise similarity structure induced by two representation spaces and is normalized from HSIC.Source: F. Why We Need a Nested… · p. 16
  • Linear CKA compares OpenVLA-7B and VGGT representations across different feature dimensions when they share the same observations.The implementation resamples VGGT features to OpenVLA’s token layout before computing CKA.Source: F. Why We Need a Nested… · p. 18

G. Per-task Results on LIBERO-Plus

This block provides per-task LIBERO-Plus results and develops a theoretical account of cross-layer alignment gradients. Under Pre-LN residual-smallness, early-layer gradients approximate sums of future local distillation gradients, while a shared projector can make their interactions more coherent under stated conditions.

G. Per-task Results on LIBERO-Plus

H.1. Takeaway: shared projector makes cross-layer gradients add constructively

  • The analysis argues that, under Pre-LN residual-smallness, a shared projector makes cross-layer gradients add constructively.Source: §H.1 · p. 19

H.2. Setup: Pre-LN residual stream, nonlinear projector, and distillation loss

  • The setup models a Pre-LN student on a residual stream, attaches nonlinear projectors at selected layers, and sums their teacher-alignment losses.The loss uses a fixed student-to-teacher layer mapping and may use uniform mean weighting.Source: §H.2 · p. 19

H.3. Jacobians needed for gradient transport and projector-induced coupling

  • The derivation specifies residual-block and projector Jacobians to track gradient transport and projector-induced coupling across layers.Source: §H.3 · pp. 19–20

H.4. Exact gradient at an early layer and its superposition approximation

  • An early-layer gradient is written exactly through residual-block back-propagation, then approximated by removing deep transport effects under residual-smallness.Source: §H.4 · p. 20

H.5. Residual-smallness implies near-identity transport

  • Under residual-smallness, back-propagation remains close to identity, so an early-layer gradient is a sum of future local distillation gradients plus a depth-dependent error.For small depth or residuals, the transport bound is approximately (l − i)ε.Source: §H.5 · p. 20

H.6. Shared projector makes cross-layer interference constructive

  • Cross-layer interference depends on whether local gradient terms reinforce or cancel; sharing one projector gives their interactions a common operator structure.Source: §H.6 · p. 21, §H.6.1 · p. 21

H.6.1. CASE 1 (SHARED PROJECTOR): A COMMON PSD OPERATOR GOVERNS CROSS-LAYER INTERFERENCE.

  • With positively correlated teacher errors and controlled deviations, shared-projector cross-layer interference is biased toward non-destructive, often constructive interactions.The theorem’s conditions include bounded residual increments, projector-Jacobian Lipschitzness, and near-isometry on the error-signal subspace.Source: §H.6.1 · p. 22

H.6.2. CASE 2 (MULTIPLE INDEPENDENT PROJECTORS): LACK OF A UNIFORM CONSTRUCTIVE LOWER BOUND.

Independent projectors lack a uniform guarantee of constructive cross-layer interference, and their gradients can remain poorly aligned during training. The paper also describes similarity-based layer selection and reports selected layer indices for several VLA configurations.

H.6.2. CASE 2 (MULTIPLE INDEPENDENT PROJECTORS): LACK OF A UNIFORM CONSTRUCTIVE LOWER BOUND.

  • Independent projectors lack a uniform constructive lower bound without additional coupling assumptions, and their cross terms can permit destructive interference.Source: §H.6.2 · pp. 22–23
  • Independent projectors tend to decouple as training proceeds, with low cross-layer gradient cosine similarity persisting, while ROCKET’s shared-projector gradients remain aligned.The passage reports this pattern from early training through the mid-to-late stage.Source: §H.6.2 · p. 23

I.1. Similarity-based Layer Selection

  • Layer selection compares ten layers chosen by highest or lowest similarity using either a ten-layer window criterion or a single-layer criterion.Lower similarity is described as potentially identifying layers with greater information transformation.Source: Figure 15 · p. 26, §I.1 · p. 26

I.2. Detailed Chosen indices for All Layer Selection Strategies

  • For OpenVLA, selected layers use indices 2–18 plus the last layer; for PI0 and PI0.5, they use indices 2–10 plus the last layer.The OpenVLA count includes the projector as a layer; PI0 and PI0.5 each have 18 layers.Source: Table 10 · p. 27

Related papers

Loading 2602.17951v1…