Source-linked AI summary
DriveMamba: Task-Centric Scalable State Space Model for Efficient End-to-End Autonomous Driving
Haisheng Su, Wei Wu, Feixiang Song, Junjie Zhang, Zhenjie Yang, Junchi Yan
TL;DR
E2E-AD systems face information loss from manually ordered modules, costly dense BEV processing, and quadratic attention that limits spatiotemporal scalability. DriveMamba addresses these issues with sparse token representations and a unified Mamba decoder, while experiments demonstrate effectiveness, efficiency, and scalability across nuScenes and Bench2Drive.
Problem
E2E-AD systems suffer from information loss and restricted task relations under sequential modular designs, while dense BEV processing and quadratic attention hinder efficient spatiotemporal scaling.
Method
DriveMamba uses sparse image and task tokens in a unified Mamba decoder for dynamic task relations, implicit task-sensor correspondence, and long-term temporal fusion.
Results
Experiments on nuScenes and Bench2Drive demonstrate DriveMamba's effectiveness, efficiency, and scalability; its tiny model reaches 17.9 FPS in open-loop planning evaluation.
Takeaways & Limitations
DriveMamba provides a scalable unified framework for efficient end-to-end autonomous driving with dynamic task modeling, view correspondence learning, and long-term temporal fusion.
Takeaways & Limitations
The framework's excessive entanglement constrains flexible problem localization and debugging, motivating combinations of joint and divided modeling with parallel decoders.
Abstract
from arXiv · showhide
Recent advances towards End-to-End Autonomous Driving (E2E-AD) have been often devoted on integrating modular designs into a unified framework for joint optimization e.g. UniAD, which follow a sequential paradigm (i.e., perception-prediction-planning) based on separable Transformer decoders and rely on dense BEV features to encode scene representations. However, such manual ordering design can inevitably cause information loss and cumulative errors, lacking flexible and diverse relation modeling among different modules and sensors. Meanwhile, insufficient training of image backbone and quadratic-complexity of attention mechanism also hinder the scalability and efficiency of E2E-AD system to handle spatiotemporal input. To this end, we propose DriveMamba, a Task-Centric Scalable paradigm for efficient E2E-AD, which integrates dynamic task relation modeling, implicit view correspondence learning and long-term temporal fusion into a single-stage Unified Mamba decoder. Specifically, both extracted image features and expected task outputs are converted into token-level sparse representations in advance, which are then sorted by their instantiated positions in 3D space. The linear-complexity operator enables efficient long-context sequential token modeling to capture task-related inter-dependencies simultaneously. Additionally, a bidirectional trajectory-guided "local-to-global" scan method is designed to preserve spatial locality from ego-perspective, thus facilitating the ego-planning. Extensive experiments conducted on nuScenes and Bench2Drive datasets demonstrate the superiority, generalizability and great efficiency of DriveMamba.
1 INTRODUCTION
DriveMamba addresses limitations of sequential, dense-BEV and quadratic-attention E2E-AD designs with a task-centric unified Mamba decoder. It combines sparse token representations, dynamic task relations, view correspondence learning, and long-term temporal fusion for efficient spatiotemporal modeling.
- E2E-AD methods commonly sequence perception, prediction, and planning with Transformer modules, which can cause information loss and cumulative errors.
- Dense BEV methods are computationally expensive for long-range perception and impractical for long-term temporal fusion, while sparse query methods retain quadratic-attention scalability limits.
- DriveMamba replaces dense BEV construction with sparse image and task tokenization, then uses a unified decoder to model task-task relations and task-sensor correspondence in parallel.
- DriveMamba integrates view correspondence learning, task relation modeling, and long-term temporal fusion in a unified framework designed for efficiency and scalability.
- The framework encodes multi-view images as token sequences, initializes spatiotemporal task queries, and applies bidirectional serialization for view correspondence, task relations, and temporal fusion.
2 METHOD
DriveMamba formulates image features and task outputs as position-aware tokens, then uses unified bidirectional Mamba decoding to model view correspondence, task relations, and long-term temporal context. Its hybrid scans preserve spatial locality by matching scan organization to task and ego-trajectory structure.
- Preliminaries: Mamba maps input sequences to outputs through a hidden state, while its selective scanning mechanism adaptively derives parameters from the input.The bidirectional variant processes flattened visual sequences in forward and backward directions for spatially aware modeling.
- Token Formulation: DriveMamba adds task-specific spatial, temporal, and task positional embeddings to distinguish tokens across locations, timestamps, sensors, and task types.Sensor tokens come from multi-view image features, while agent, map, and ego queries represent dynamic objects, static map elements, and ego behavior.
- Token Formulation: Depth prediction assigns image tokens 3D positions using camera intrinsics, extrinsics, and predicted depth, enabling their spatial ordering with task queries.Motion-aware layer normalization compensates for object movements before unified decoding.
- Hybrid Spatiotemporal Scan: DriveMamba extends 2D bidirectional scanning to spatial-temporal scans, including horizontal-first, vertical-first, ego-centric, and trajectory-centric local-to-global orders.Trajectory-centric scanning weights task queries by distance to interpolated future ego waypoints, emphasizing in-path queries.
- Hybrid Spatiotemporal Scan: Different scan orders serve different functions: alternating horizontal and vertical scans support view correspondence, trajectory-centric scans support task relations and planning, and spatial-first scans arrange history memory.The method uses linear-complexity Mamba to handle high-resolution and long-term memory tokens efficiently.
- Discussion: The design prioritizes task-specific spatial locality: trajectory-guided local-to-global scanning models interaction order from the ego perspective for interactive planning.This task-oriented scan design differs from vision-oriented structured scan patterns and Transformer planners that neglect ego-perspective interaction order.
- Unified Mamba Decoder: The unified decoder contains B-Mamba layers for view correspondence learning, task relation modeling, and long-term temporal fusion within a joint decoding framework.Task queries directly extract semantics from raw sensor features, while prior timestamp queries are stored in FIFO queues and motion-compensated before fusion.
3 EXPERIMENTS
Experiments evaluate DriveMamba across open- and closed-loop planning, ablations, scalability, design choices, robustness, and trajectory-centric scan generalization. Results show strong planning performance, efficient Transformer comparisons, scalable task-centric decoding, and robustness to several initialization and depth conditions.
- Main Results: DriveMamba-Tiny reduces average L2 error by 42.1% versus UniAD and 38.9% versus VAD, while reducing average collision rate by 11.8% and 31.8%, respectively.Inference takes 55.8 ms, corresponding to 17.9 FPS.
- Main Results: DriveMamba-Base achieves superior closed-loop planning performance with fewer parameters and higher efficiency than DriveTransformer-Large.The comparison is conducted on Bench2Drive, where L2 error and closed-loop driving performance do not always align across methods.
- Modular Study: The unified decoder’s VCL, LTF, and TRM components support sensor-context extraction, streaming temporal modeling, and dynamic inter-task and intra-task relation learning.The ablation identifies VCL as especially important for randomly initialized queries to extract context from sensor tokens.
- Paradigm Study: DriveMamba’s linear-complexity SSM increases processing speed by 3.2× and reduces GPU memory consumption by 68.8% when scaling input resolution.The comparison concerns Transformer-based end-to-end planners and long sequential tokens.
- Scalability Study: Scaling the visual encoder mainly improves open-loop perception, whereas stacking decoder layers contributes most to ego-planning with a small perception decrease.The authors associate this pattern with planning-oriented rather than general perception learning.
- Planning-oriented Perception Learning: CIPO perception performance improves monotonically with additional unified-decoder layers, supporting planning-oriented task synergy and dynamic relation learning.CIPO contains objects and lanes within 5 m of each ground-truth future waypoint.
- Initialization of Trajectory Prior: Different trajectory-prior initializations cause only trivial changes in perception and planning, indicating training stability and robustness of the egocentric design.The evaluated initializations are Origin, Uniform, and Random waypoint offsets.
- Robustness Analysis: Ground-truth depth improves perception by 4.9 NDS, while planning changes by only -0.02% Collision Rate, indicating greater planning robustness to depth errors.Noise in camera extrinsics or predicted depth affects perception more substantially than planning.
4 CONCLUSION
DriveMamba is presented as a task-centric, scalable Mamba paradigm for efficient E2E-AD. It unifies task relation modeling, view correspondence learning, and long-term temporal fusion over sparse representations with hybrid spatiotemporal scanning.
- 4 CONCLUSION: DriveMamba integrates dynamic task relation modeling, implicit view correspondence learning, and long-term temporal fusion into a unified decoder that scales by stacking layers.The decoder operates on sparse representations rather than dense BEV features.
- 4 CONCLUSION: A hybrid spatiotemporal scan captures task-related inter-dependencies while preserving spatial locality for ego-planning.The scan is designed to avoid spatial locality loss during unified decoding.
A.1 END-TO-END PLANNING
End-to-end planning methods differ in whether they optimize planning implicitly or integrate tasks explicitly. The cited literature also motivates state-space modeling for long-range dependencies and Mamba’s linear-complexity selective scanning.
- End-to-End Planning: Implicit E2E planning directly optimizes planning but offers less interpretability and controllability, whereas explicit methods integrate perception, prediction, and planning in unified frameworks.The explicit methods described here commonly use sequential BEV-based designs, while other approaches introduce vectorization, parallel decoders, graphs, or multiple sensors.
- End-to-End Planning: Recent query-based methods use sparse representations, but the cited literature identifies continuing inefficiency, inflexibility, and limited dynamic task-relation modeling.These limitations motivate the task-centric design developed in DriveMamba.
- State Space Models: State Space Models model long-range sequence dependencies, while Mamba adds selective scanning and a linear-complexity operator for efficient sequential processing.The passage frames these properties as the basis for applying Mamba to E2E-AD token sequences.
B EVALUATION METRICS
The evaluation uses task-specific metrics for perception, motion prediction, and planning. These metrics include detection and mapping accuracy, trajectory displacement and miss rates, and planning error and collision measures.
- Perception: Object detection is evaluated with mAP, mATE, mASE, mAOE, mAVE, mAAE, and NDS, while online mapping uses mean AP across three map classes.The map classes are lane divider, pedestrian crossing, and road boundary.
- Motion Prediction: Motion prediction uses minADE, minFDE, and MR, computed within matched true positives with a 1.0 m matching threshold and a 2.0 m miss-FDE threshold.These metrics follow standard motion-prediction protocols described in the passage.
- Planning: Planning is evaluated using L2 error and collision rate, with L2 evaluation aligned to VAD and collision evaluation addressing drawbacks in earlier implementations.The cited drawbacks include occupancy-map discretization that can produce false collisions.
C IMPLEMENTATION DETAILS
DriveMamba uses a fixed perception range and ResNet-50 image backbone, and predicts a three-second ego trajectory from two seconds of history.
- The default implementation uses a 60m×30m longitudinally and laterally perception range with ResNet-50 image encoding.
- DriveMamba predicts a 3-second future ego-vehicle trajectory at 2Hz using 2 seconds of historical input.
- Trajectory interpolation uses T ′ set to 30.
D MORE DETAILS
The method combines token initialization, spatially aligned augmentation, trajectory-guided sorting, and hybrid scans to model task relations and temporal information.
- Spatial scanning: Ego-centric Local2Global Scan indexes N × N locations by their layer distance from the border and assigns position-dependent sequence indices.
- Data augmentation: Training augmentation jointly transforms agent, map, and trajectory targets while adjusting camera extrinsics to preserve spatial alignment.
- Trajectory prior: Ego queries use semantic and positional embeddings, with tokenized future waypoints initialized at the ego origin and iteratively refined by predicted offsets.
- Hybrid scan: Hybrid scanning alternates Horizontal-First and Vertical-First scans for view correspondence, uses Trajectory-Centric Local2Global scanning for planning, and applies spatial-first temporal scanning for long-term fusion.
E MORE ABLATION STUDY
Ablations show benefits from end-to-end training, an intermediate BEV size, trajectory-guided sorting, and the unified decoder’s middle-task performance and efficiency.
- Training strategy: End-to-end training outperforms divided training, while dense supervision supports convergence and model scaling.
- BEV size: A 50×50 BEV size provides the best planning performance, whereas smaller sizes blur query importance and larger sizes disrupt spatial continuity.
- Trajectory-guided scan: Using ground-truth future trajectories yields 0.26m average L2 error and 0.07% average collision rate, supporting trajectory guidance for query sorting.
- Middle-task performance: DriveMamba surpasses existing BEV-Centric and Query-Centric methods on nuScenes middle tasks.
- Efficiency: Long-term temporal fusion and memory propagation account for 52.7% of module runtime, while linear-complexity attention and parallel decoding support efficiency and scalability.
F DISCUSSION
DriveMamba replaces sequential Transformer processing with a unified, parallel Mamba decoder and reports higher efficiency and planning performance across comparisons.
- Motivation: Sequential Transformer planners can incur information loss, cumulative errors, and restricted task relations through manually ordered modules.
- Efficiency: DriveMamba increases processing speed 3.2× and reduces GPU memory usage 68.8% for long-time, high-resolution image sequences.
- Unified design: The unified decoder integrates task-centric spatiotemporal positional embeddings, trajectory-centric Local-to-Global scanning, and query-centric linear attention for parallel modeling.
- View correspondence: View correspondence learning aligns 2D sensor tokens with 3D task tokens through point-level depth estimation and hybrid 3D token sorting before linear-complexity SSM processing.
- Transformer comparison: DriveMamba-Tiny improves Driving Score by 6.07 over DriveTransformer-Small while reducing latency by 40.5% to 55.8ms versus 93.8ms.
- Related comparison: The paper contrasts DriveMamba’s pure Mamba decoder with DRAMA’s combined Mamba-Transformer decoder for camera- and LiDAR-based BEV features.
- Qualitative results: Qualitative closed-loop results illustrate DriveMamba’s perception, motion, and planning trajectories in a HazardAtSideLane interactive scene.
G LIMITATIONS AND SOCIAL IMPACT
DriveMamba’s highly entangled framework limits problem localization and debugging, while the paper presents deployment across constrained car-chip resources and illustrates closed-loop hazard responses.
- Limitations: Excessive framework entanglement constrains flexible problem localization and debugging.The authors identify combining joint and divided modeling with parallel decoders as future work.
- Social Impact: DriveMamba is described as deployable on mass-produced car chips with varying computing-resource limitations.The paper positions it as plug-and-play software for assisting human drivers with decision-making and safe driving.
- Closed-loop behavior: In a closed-loop HazardAtSideLane scenario, the ego vehicle brakes when a slow-moving hazard blocks the lane and maneuvers beside traffic when the side lane is available.The behavior is visualized across different timestamps on Bench2Drive test routes.