Depth-Wise Probing and Pruning of the Planning Token in a Driving Vision-Language-Action Model

Technical Summary: Depth-Wise Probing and Pruning of the Planning Token in a Driving VLA

Problem Statement
Vision-language-action (VLA) models for autonomous driving increasingly rely on deep transformer architectures to route perception inputs into driving decisions. While these models leverage broad world knowledge to handle complex scenarios, their structural depth introduces significant computational latency, which is problematic for safety-critical controllers. It remains unclear how task-relevant computation is distributed across the decoder layers. Specifically, it is unknown whether all layers are necessary for generating a valid trajectory or if specific layers can be removed to reduce latency without substantial performance degradation. This study investigates the spatial distribution of planning information within a representative driving VLA (ORION) to determine the optimal point for pruning.

Methodology
The authors analyze the ORION model, which processes surround-view camera streams and a fixed text prompt to generate a single "planning token." This token is decoded by a frozen generative planner into a 3-second trajectory (6 waypoints). The study employs two primary diagnostic tools:

1. Trajectory-Space Logit Lens: Instead of training external classifiers to interpret intermediate layers, the authors decode the planning token from every one of the 32 decoder layers using the model's own frozen generative planner. This measures the "compatibility" of each intermediate representation with the deployed planner. The resulting trajectories are scored using the ST-P3 Average Displacement Error (ADE).
2. Linear Probing: A linear classifier is trained on the planning token at each layer to predict the discrete navigation command (e.g., "merge," "overtake"). This measures the linear decodability of semantic intent.
3. Pruning Strategy: To identify removable layers, the authors rank decoder layers based on the angular deviation (cosine similarity) they induce in the planning token (input vs. output). Layers that minimally rotate the token are pruned first. This "representational redundancy" metric is compared against pruning based on isolated decode quality, contiguous late-block removal, and random removal.

Key Findings
The study reveals a distinct divergence in how different types of information evolve through the network depth:

• Early Semantic Decodability: Semantic intent (the navigation command) becomes linearly decodable very early. Command-probe accuracy reaches 97.7% after the first decoder layer (compared to 16.7% chance), despite the command never being provided as textual input.
• Late Planner Compatibility: In contrast, compatibility with the frozen native planner improves gradually. The trajectory error (Avg-L2) decreases slowly through the middle layers and reaches its minimum (2.11 m) only at the final hidden state.
• Representation vs. Presence: The gap between early decodability and late planner compatibility suggests that while geometric and planning information is present early, it is not yet in the specific format required by the frozen planner. Learned readouts (MLP adapters and cross-attention resamplers) applied to the first layer can recover much of this gap, confirming that the information exists but is representationally incompatible with the downstream module.
• Pruning Efficacy: Ranking layers by the angular deviation they induce in the planning token is a superior pruning criterion compared to removing layers based on their isolated decode quality.
  • Removing 8 of the 32 layers (specifically layers {9, 10, 13, 17, 18, 19, 27, 30}) results in an approximate 5% relative increase in open-loop error.
  • This removal yields a measured 1.33× decoder speedup (reducing latency from ~497ms to ~373ms).
  • Pruning based on isolated decode quality (targeting layers with high error in isolation) is ineffective, as these layers are often critical for subsequent computations.

Results and Limitations
The pruning study was evaluated across five Bench2Drive ability categories (merging, overtaking, emergency braking, giving way, and traffic-sign compliance). The progression of planner compatibility was consistent across all categories. At the evaluated sample size (160 frames per family for pruning analysis), no family-specific degradation was statistically resolved, though the authors note that overlapping confidence intervals do not constitute an equivalence test.

The authors explicitly state that these findings are limited to the evaluated ORION checkpoint and the Bench2Drive simulation setup. The results are open-loop; they do not establish closed-loop safety or deployment readiness. Furthermore, because the decoder accounts for only roughly half of the total per-frame latency (with perception dominating the rest), the projected end-to-end speedup is approximately 1.13×.

Significance
The paper provides a diagnostic framework for understanding the internal mechanics of driving VLAs. It demonstrates that high-level semantic intent is established early in the network, while the refinement of that intent into a planner-compatible format occurs progressively through the depth. The study validates that "redundancy" in terms of directional rotation (cosine similarity) is a better indicator for layer removal than the raw error of an intermediate layer when decoded in isolation. This offers a principled approach to optimizing the computational efficiency of VLA-based driving systems without retraining the core model or the planner.