Depth-Wise Probing and Pruning of the Planning Token in a Driving Vision-Language-Action Model
This study introduces depth-wise probing and pruning of the planning token in a driving VLA model, revealing early linear decodability of semantic intent and enabling 1.33× speedup with minimal performance loss.
Key Findings
Methodology
This paper proposes a trajectory-space native-head probing approach, extracting the planning token from each decoder layer and decoding it using the model’s frozen planner. By comparing the decoded trajectories with ground truth and measuring trajectory error (Avg-L2), as well as training linear classifiers to predict navigation commands, the study analyzes how task-relevant information is distributed across layers. The layers are ranked based on the angular deviation they induce in the planning token, guiding layer pruning. This methodology combines internal representation decoding with task-specific metrics, providing a comprehensive view of information flow within the model.
Key Results
- Semantic intent in the planning token is linearly decodable with 97.7% accuracy after the first decoder layer, far above the 16.7% chance level, indicating early presence of meaningful command information. The trajectory compatibility with the frozen planner improves gradually, reaching a minimum Avg-L2 error of 2.11 meters only at the final layer. Learning readouts from the first layer recover much of this gap, showing early semantic information exists but is not yet in a format usable by the planner. Pruning the top 8 layers (out of 32) based on angular deviation induces only about 5% increase in open-loop error, while achieving a 1.33× speedup, with no statistically significant degradation across five driving ability categories.
- Across five ability categories (merging, overtaking, emergency braking, giving way, traffic sign compliance), the progression of planner compatibility is consistent. Command decodability is high from the first layer, while trajectory accuracy improves steadily, confirming the generality of the information distribution pattern.
- Layer ranking based on cosine similarity of input-output representations effectively identifies redundant layers. Removing up to 8 layers maintains performance within acceptable bounds, significantly reducing inference latency, and demonstrating the potential for efficient model compression without sacrificing task performance.
Significance
This work provides a profound understanding of how deep models encode task-relevant information, revealing that early layers already contain sufficient semantic intent, while later layers refine trajectory compatibility. Such insights enable targeted model pruning, reducing computational costs crucial for real-time autonomous driving. The methodology bridges the gap between interpretability and efficiency, offering a pathway to deploy large-scale models in safety-critical systems. The approach also opens avenues for explainability in other multi-modal transformer applications, fostering trust and robustness in AI systems. Overall, this research advances both theoretical understanding and practical deployment of deep models in autonomous vehicles.
Technical Contribution
The key technical innovation lies in the trajectory-space native-head probing, which directly assesses the model’s internal representations against the native planner’s output. Unlike traditional feature attribution methods, this approach evaluates the task-specific compatibility of intermediate tokens by decoding trajectories and measuring their errors. The angular deviation-based layer ranking introduces a novel criterion for layer removal, balancing information preservation and computational efficiency. Combining linear command decoding with trajectory error analysis provides a comprehensive picture of how semantic and geometric information evolve across layers. These contributions offer a new toolkit for interpretability and model compression in large transformer-based systems.
Novelty
This study is the first to systematically analyze the distribution of planning-related information across decoder layers in a driving VLA model using trajectory decoding as a diagnostic. The innovative use of the frozen planner as a decoding head, combined with angular deviation ranking for layer pruning, distinguishes it from prior feature attribution or pruning methods. Unlike existing approaches that focus solely on feature importance or model size reduction, this work emphasizes task-specific compatibility and information redundancy, providing actionable insights for model acceleration without performance loss. It bridges the gap between interpretability and efficiency in complex multi-modal models.
Limitations
- The experiments are limited to the ORION model and Bench2Drive dataset; different architectures or data distributions may exhibit different information distribution patterns, limiting direct generalization.
- The pruning strategy relies on static angular deviation metrics, which may not account for dynamic scene complexities or rare scenarios, potentially affecting robustness.
- While inference speed improves significantly, the impact on model robustness, safety, and generalization under diverse real-world conditions remains to be thoroughly validated. Future work should incorporate robustness evaluations and adaptive pruning strategies.
Future Work
Future research will explore the applicability of the proposed probing and pruning techniques across different model architectures and datasets, aiming for a universal framework for model interpretability and acceleration. Combining dynamic, scene-aware pruning with fine-tuning could further enhance robustness and generalization. Extending the analysis to multi-task settings and integrating causal inference methods may deepen understanding of internal information flow. Additionally, efforts to optimize hardware-specific implementations will facilitate deployment in real-world autonomous systems, ultimately advancing the development of safe, efficient, and explainable AI for autonomous driving.
AI Executive Summary
The rapid evolution of autonomous driving systems has seen a surge in the adoption of large-scale deep learning models, particularly transformer-based architectures that integrate vision, language, and action modalities. These models leverage extensive semantic knowledge to handle complex driving scenarios, but their deep structures introduce significant computational latency, hindering real-time deployment. Existing optimization strategies, such as knowledge distillation and early exit mechanisms, often lack a detailed understanding of how task-relevant information is distributed across model layers.
This study addresses this gap by investigating the internal representations of a representative vision-language-action (VLA) model, ORION, used for autonomous driving. The core idea is to analyze how the model encodes the planning information necessary for trajectory generation. By extracting the planning token—a compressed representation of the entire driving plan—from each decoder layer and decoding it using the model’s native, frozen planner, the authors measure the trajectory error and command decodability at each depth. This approach, termed trajectory-space native-head probing, provides a direct window into the evolution of task-specific information within the network.
The empirical findings are striking. Early in the network, the semantic command—such as 'merge' or 'turn left'—can be decoded with 97.7% accuracy, indicating that the model’s internal representations already contain rich, linearly accessible intent information. Meanwhile, the compatibility of the planning token with the frozen planner, reflected in trajectory accuracy, improves gradually, reaching its minimum error only at the final layer (2.11 meters). This dissociation reveals that semantic intent is present early but not immediately actionable, while geometric consistency develops later.
Building on this insight, the authors propose a layer pruning strategy based on the angular deviation induced by each layer on the planning token. By removing the layers with minimal impact on the token’s direction, they achieve a 1.33× speedup with less than 5% increase in open-loop trajectory error. Importantly, this pruning does not significantly degrade performance across different driving scenarios, demonstrating the method’s robustness and practical utility.
These results have profound implications. They suggest that large transformer models for autonomous driving contain substantial redundancies, which can be exploited to accelerate inference without sacrificing safety or accuracy. The methodology also enhances interpretability, providing a clearer understanding of how information flows and transforms within deep networks. Looking ahead, this work opens pathways to more efficient, explainable, and adaptable autonomous systems, capable of operating reliably in real-world environments with reduced computational costs.
Overall, this research bridges the gap between deep model interpretability and practical deployment, offering a scalable framework for analyzing and optimizing complex multi-modal AI systems in autonomous driving and beyond.
Deep Dive
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.
References (20)
Eliciting Latent Predictions from Transformers with the Tuned Lens
Nora Belrose, Zach Furman, Logan Smith et al.
Orion: A Holistic End-To-End Autonomous Driving Framework by Vision-Language Instructed Action Generation
Haoyu Fu, Diankun Zhang, Zongchuang Zhao et al.
Senna: Bridging Large Vision-Language Models and End-to-End Autonomous Driving
Bo Jiang, Shaoyu Chen, Bencheng Liao et al.
Perceiver: General Perception with Iterative Attention
Andrew Jaegle, Felix Gimeno, Andrew Brock et al.
Multi-Modal Fusion Transformer for End-to-End Autonomous Driving
Aditya Prakash, Kashyap Chitta, Andreas Geiger
BranchyNet: Fast inference via early exiting from deep neural networks
Surat Teerapittayanon, Bradley McDanel, H. T. Kung
OpenVLA: An Open-Source Vision-Language-Action Model
Moo Jin Kim, Karl Pertsch, Siddharth Karamcheti et al.
GPT-Driver: Learning to Drive with GPT
Jiageng Mao, Yuxi Qian, Hang Zhao et al.
Reason2Drive: Towards Interpretable and Chain-based Reasoning for Autonomous Driving
Ming Nie, Renyuan Peng, Chunwei Wang et al.
Generalized Predictive Model for Autonomous Driving
Jiazhi Yang, Shenyuan Gao, Yihang Qiu et al.
VAD: Vectorized Scene Representation for Efficient Autonomous Driving
Bo Jiang, Shaoyu Chen, Qing Xu et al.
Bench2Drive: Towards Multi-Ability Benchmarking of Closed-Loop End-To-End Autonomous Driving
Xiaosong Jia, Zhenjie Yang, Qifeng Li et al.
Planning-oriented Autonomous Driving
Yi Hu, Jiazhi Yang, Li Chen et al.
DriveGPT4: Interpretable End-to-End Autonomous Driving Via Large Language Model
Zhenhua Xu, Yujia Zhang, Enze Xie et al.
LMDrive: Closed-Loop End-to-End Driving with Large Language Models
Hao Shao, Yuxuan Hu, Letian Wang et al.
CARLA: An Open Urban Driving Simulator
Alexey Dosovitskiy, Germán Ros, Felipe Codevilla et al.
LayerSkip: Enabling Early Exit Inference and Self-Speculative Decoding
Mostafa Elhoushi, Akshat Shrivastava, Diana Liskovich et al.
Reducing Transformer Depth on Demand with Structured Dropout
Angela Fan, Edouard Grave, Armand Joulin
LLM-Pruner: On the Structural Pruning of Large Language Models
Xinyin Ma, Gongfan Fang, Xinchao Wang
ShortGPT: Layers in Large Language Models are More Redundant Than You Expect
Xin Men, Mingyu Xu, Qingyu Zhang et al.