DriveMamba: Task-Centric Scalable State Space Model for Efficient End-to-End Autonomous Driving

TL;DR

DriveMamba introduces a task-centric, scalable end-to-end autonomous driving framework that efficiently models long-term temporal dependencies and task relations, surpassing existing methods in accuracy and efficiency.

Contribution

It proposes a novel unified decoder with linear complexity that integrates dynamic task relation modeling, implicit view correspondence, and long-term temporal fusion for autonomous driving.

Findings

Outperforms existing methods on nuScenes and Bench2Drive datasets.

Demonstrates superior efficiency and generalizability.

Effectively models long-term temporal dependencies.

Abstract

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.