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

AI summary

Overview

• Research area: Interpretability and inference efficiency of vision-language-action (VLA) models for autonomous driving, studied on the ORION driving model evaluated in the CARLA-based Bench2Drive benchmark.
• Technical level: Advanced. The paper assumes familiarity with transformer decoder stacks, linear probing, logit-lens style readouts, and open-loop trajectory metrics (ST-P3 ADE).
• Scope in one sentence: The paper probes the single planning token of a 32-layer driving VLA at every decoder depth, shows that navigation intent is linearly decodable after the first decoder block while compatibility with the frozen trajectory planner improves only late, and uses that diagnostic to prune 8 of 32 decoder layers within an approximately 5% relative open-loop error increase for a measured 1.33x decoder speedup.

What This Paper Is About

Driving VLA models route decisions through a deep language-model decoder, and it is unclear how much of that depth the actual action requires. The authors take a representative ORION-based driving VLA whose entire plan is compressed into one planning token that a frozen generative planner decodes into waypoints, and they ask when, along the 32 decoder layers, the information the planner needs actually becomes usable. Their goal is both diagnostic (understanding the depth-wise distribution of planner-relevant computation) and practical (testing whether redundant layers can be removed to cut latency without meaningfully hurting trajectory quality).

Key Contributions

1. A trajectory-space native-head lens. Instead of training a separate readout, the authors decode the planning token from all 33 hidden states (the embedding output and each of the 32 layers) through the model's own frozen planner, measuring how compatible each intermediate representation is with the deployed planner. This is described as an analogue of the logit lens applied to trajectories rather than vocabulary logits, paired with a linear probe of navigation intent.

2. A depth-wise analysis across five Bench2Drive ability categories. The study stratifies results by merging, overtaking, emergency braking, giving way, and traffic-sign compliance, showing early linear command decodability where the command varies and a consistent progression of planner compatibility across the evaluated ORION checkpoint.

3. A pruning study using planning-token angular deviation as the removal criterion. The authors show that planning-token input-output cosine is a better guide than isolated intermediate decode quality: 8 of 32 layers can be removed within an approximately 5% relative open-loop error increase, yielding a measured 1.33x decoder speedup.

4. A learned-readout control. Small trained heads mapping layer-1 representations to the final-layer token space show that early layers already contain planning information that the frozen planner simply cannot consume in its native format.

Main Findings

• Command intent is linearly decodable almost immediately. The command probe is at chance (16.7%) at hidden state 0, where the planning token is still context-free, rises to 97.7% after the first decoder block, and peaks at 98.6%. The navigation command is never provided to the model as text, so this structure must be generated internally.

• Planner compatibility develops late. Error under the frozen-planner readout falls slowly through the middle layers and reaches its minimum only at the final hidden state, where Avg-L2 is 2.11 m. More than half of the total reduction over the embedding baseline is reached only by hidden state 15. The normalized figure spans an absolute maximum of 6.23 m at hidden state 0 and an absolute minimum of 2.11 m at hidden state 32.

• Early intent does not imply early planner-readiness. At hidden state 1 the command probe is already near its ceiling, yet decoding the same token gives an Avg-L2 worse than the embedding baseline. The early layers hold the discrete intent in a form a linear classifier can read, but not yet in the format the frozen planner turns into a low-error trajectory.

• A late non-monotonic spike appears. The per-layer trajectory error shows a spike at hidden-state indices 29 to 31 before the final hidden state recovers. The authors treat this as an observed directional excursion rather than a resolved mechanism.

• The pattern holds across driving abilities. In four of five ability categories, half of the total error reduction is reached by hidden state 15 (hidden state 17 for give way), with the minimum at the final hidden state in all five. Where the navigation command has sufficient variation (merging and traffic-sign compliance), it is linearly readable from the first hidden state at 95 to 98% at hidden state 1. In the emergency-brake, give-way, and overtaking families the command distribution is near-constant or strongly imbalanced, so the command probe is uninformative there.

• Angular deviation is the right pruning criterion. Ranking layers by the angular deviation they induce in the planning token (highest input-output cosine first) keeps error within about 5% of the full model up to k = 8 removed layers, and within 1% up to k = 4 (Avg-L2 2.06 to 2.17 m). The removed set at k = 8, using zero-based decoder-layer indices, is {9, 10, 13, 17, 18, 19, 27, 30}.

• The competing ranking is a poor guide. Ranking layers by their own decoded error targets the late spike at hidden-state indices 29 to 31, which decode badly in isolation yet are needed by later layers, and it raises error by 27% at the first layer removed. Contiguous late-block and random orders sit between the two criteria, but only the cosine order keeps error flat out to eight layers.

• Speedup is real but modest. Removing the eight layers reduces decoder latency from 497.52 to 373.22 ms, a measured 1.33x speedup at a layers-executed fraction of 0.75. Larger removals trade more error for more speed, up to 4x at k = 24. The identity-patch proxy and a physically sliced decoder produce the same planning token (cosine 1.0 at every k).

• No family-specific degradation is statistically resolved. On the same 160 cold-start frames per family, comparisons include traffic sign 1.22 to 1.24 m, emergency brake 2.10 to 2.19 m, and merging 3.08 to 2.90 m. The authors state that overlapping marginal normal-approximation intervals are not an equivalence test and cannot exclude small effects.

• Learned readouts recover much of the early-layer gap. A residual MLP adapter (4096 to 1024 to 4096) trained on 3,536 disjoint training scenes lowers layer-1 Avg-L2 from 8.51 to 3.47 m, and a cross-attention resampler with 8 latent queries and a 1024-dimensional bottleneck (about 40M parameters) reaches 2.98 m against a final-layer baseline of 2.17 m. This confirms that the native-planner curve measures representation compatibility rather than the first presence of planning information.

• Vision tokens carry most of the resampler's usable signal. In test-time masking, retaining only the 529 vision tokens stays within 0.15 m of the unablated 599-token sequence, while retaining only the 70 text tokens or the single planning token degrades toward the raw layer-1 baseline.

• Streaming preserves the trend. On the same 5,291 route-ordered frames, streaming lowers the 1, 2, and 3 s errors by 0.008, 0.035, and 0.073 m while preserving the depth-wise trend.

Methodology in Plain English

The authors study ORION, a driving model that processes six surround-view camera streams through a vision backbone and perception heads, projects the resulting query tokens into a language-model embedding space, concatenates them with a fixed text prompt, and runs a 32-layer LLaMA-style decoder. One special token is appended, and its 4096-dimensional hidden state after the final layer is the planning token. A compact generative planner, trained jointly with the base model and kept frozen throughout the study, maps that token to six waypoints over a 3 s horizon.

Their central move is to intercept the planning token at every depth rather than only at the end. At each of the 33 hidden states they apply the model's final normalization and pass the token through the same frozen planner, producing a trajectory that can be scored against ground truth with the ST-P3 Average Displacement Error. Because the planner is fixed, every layer is scored in an identical metric space, answering the counterfactual: if the model had to commit to a plan using this layer's representation, how good would the plan be? They also fit a regularized multinomial logistic regression classifier at each hidden state to predict the six Bench2Drive commands, which the model never receives as input text.

For the pruning experiment, they exploit the fact that each LLaMA layer adds a pure residual, so patching a layer to the identity is equivalent to deleting it. They rank layers by the cosine similarity between a layer's input and output hidden states at the planning-token position, remove the highest-cosine layers first, and compare against a decode-marginal order, a contiguous late-block order, and a random order. Latency is measured as the wall-clock of one non-autoregressive FP32 decoder forward pass at batch size 1 on a single A100 80 GB GPU after warm-up, excluding language generation and planner decoding. Four evaluation sets are used: a 2,000-frame ability-balanced probing set, an 800-frame ability-balanced pruning set, a 500-frame command-balanced reshaping set, and a 5,291-frame route-ordered streaming set.

Why This Matters

The paper reframes a common efficiency question: rather than assuming that layers which look unimportant in isolation can be deleted, it shows that what matters is how much a layer rotates the representation the downstream module has to consume. This matters for interpretability research because it extends lens-style analysis from discrete vocabulary logits to continuous actions, and it matters for deployment because latency in safety-critical driving stacks is a real constraint.

Real-world applications:

• Latency reduction on embedded automotive compute. Removing 8 of 32 decoder layers cuts decoder latency from 497.52 to 373.22 ms in the measured setting, a directly relevant scaling knob for on-vehicle inference.
• Model diagnostics during development. The per-layer compatibility curve gives engineers a way to see where a VLA's plan becomes usable, which can guide where to invest capacity or where a frozen downstream module is mismatched to upstream representations.
• Informing redesign of planner interfaces. The layer-1 readout experiments show that early representations carry planning information in a format the frozen planner cannot use, which motivates retraining or adapting planner heads rather than only shrinking the decoder.
• Risk-aware compression planning. The per-family breakdown, even where underpowered, sketches how to check whether compression behaves differently across driving abilities such as merging, overtaking, emergency braking, giving way, and traffic-sign compliance.

Industry relevance: the work comes from Robert Bosch GmbH with academic collaborators at the Karlsruhe Institute of Technology, and it targets a production-relevant question. The authors note that the decoder accounts for roughly half of per-frame latency while perception components dominate the remaining budget, so an Amdahl-style projection based on the measured decoder share gives an end-to-end speedup of approximately 1.13x, a value they did not directly time end to end. That framing is a candid statement of the limits of post-hoc decoder compression in this architecture.

Future Directions

• Closed-loop evaluation. All pruning results here are open-loop and do not establish closed-loop safety or deployment readiness, so measuring closed-loop behavior under the pruned decoder is the most direct next step.
• Beyond a single checkpoint and benchmark. The findings are limited to one ORION checkpoint and the Bench2Drive setup, and the authors describe them as checkpoint-specific, so testing whether the late planner-compatibility schedule and the cosine pruning criterion transfer to other VLA checkpoints is open.
• Closing the early-layer format gap. The learned readouts recover much of the layer-1-to-final gap but do not reach the final layer, raising the question of what transformation would make early planning-token representations fully consumable by the frozen planner, and whether planner retraining would help.
• Better statistics for per-family claims. The per-family pruning intervals are underpowered for equivalence claims, so larger per-family sample sizes are needed before any statement about family-specific degradation can be made.

Target Audience

Readers who will benefit most are researchers and engineers working on vision-language-action models, autonomous driving stacks, and model efficiency, particularly those interested in interpretability tools that operate through a model's own downstream modules. Practitioners concerned with inference latency on embedded automotive hardware will find the pruning frontier directly relevant, while interpretability researchers will find the trajectory-space lens and the layer-1 versus final-layer comparison the most distinctive parts. The paper is written for an audience already comfortable with transformer internals, linear probes, and trajectory error metrics.

Authors’ abstract

Vision-language-action (VLA) models route driving decisions through a deep language model, but it is unclear how much of that depth the action itself requires. We study a representative driving VLA whose entire plan is carried by a single planning token that a generative planner decodes into a trajectory. Borrowing the planner as a trajectory-space logit lens, we decode the planning token from every one of the 32 decoder layers and measure two signals: the linear decodability of the navigation command and trajectory compatibility with the frozen native planner. Our diagnostic shows that semantic intent is linearly decodable early: command-probe accuracy reaches 97.7\% after the first decoder layer, compared with 16.7\% chance. In contrast, compatibility with the frozen native planner improves gradually across depth, with open-loop Avg-L2 reaching its minimum of 2.11\,m only at the final layer. Learned readouts from the first layer recover much of this gap, indicating that planning information is already present early but is not yet represented in the format expected by the deployed planner. Ranking decoder layers by the angular deviation they induce in the planning token permits removal of 8 of 32 layers within an approximately 5\% relative open-loop error increase and yields a measured 1.33$\times$ decoder speedup. At the evaluated sample size, no family-specific degradation is statistically resolved. These findings are limited to the evaluated ORION checkpoint and Bench2Drive setup.