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

DriveMamba proposes a task-centric, scalable state space model for efficient end-to-end autonomous driving that replaces the sequential Transformer-based paradigm with a unified Mamba decoder featuring linear-complexity operators and bidirectional trajectory-guided scanning to overcome information loss, cumulative errors, and computational inefficiencies in handling spatiotemporal inputs.

The Gist

Imagine you are teaching a brand new driver how to navigate a busy city.

The Old Way (The "Assembly Line" Problem)
Most current self-driving systems work like a rigid factory assembly line.

1. Station 1 (Perception): A robot looks out the window and says, "I see a red car."
2. Station 2 (Prediction): A second robot takes that note and says, "The red car is moving left."
3. Station 3 (Planning): A third robot takes that info and says, "Okay, I'll turn right."

The problem? If the first robot makes a tiny mistake (like misjudging the red car's speed), that error gets passed down the line, gets magnified, and the final decision might be dangerous. Also, this "assembly line" is slow because every station has to wait for the previous one to finish. It's like trying to run a marathon while waiting for a slow friend to tie their shoe at every mile marker.

The New Way: DriveMamba (The "Super-Organized Brain")
The authors of this paper, DriveMamba, propose a completely different approach. Instead of an assembly line, imagine a super-organized brain that looks at everything at once and makes decisions instantly.

Here is how it works, using simple analogies:

1. The "Token" Menu (No More Heavy Lifting)

Old systems try to build a massive, 3D map of the entire world (called a "Dense BEV") before making a decision. It's like trying to draw the entire city of New York in perfect detail before deciding where to turn. It takes too much time and memory.

DriveMamba is smarter. It treats every piece of information (a car, a pedestrian, a lane line) as a simple "token" or a sticky note.

• Instead of drawing the whole city, it just writes down: "Car at 5 o'clock," "Lane at 12 o'clock," "Myself at center."
• It organizes these sticky notes into a neat list based on where they are in space and time. This makes the system incredibly fast and lightweight.

2. The "Mamba" (The Efficient Reader)

The core of this system is a new type of AI engine called Mamba.

• The Old Engine (Transformer): Imagine reading a book where, to understand the last sentence, you have to re-read every single previous sentence to see how they connect. As the book gets longer, this gets exhausting and slow. This is how older self-driving cars work; they get bogged down as the scene gets complex.
• The Mamba Engine: Imagine a reader who has a magical memory. They can read a long story, remember the important parts, and instantly understand the context without re-reading the whole thing. They process information in a straight line (linear complexity).
• The Result: DriveMamba can "read" a long, complex driving scene (like a 20-second video of traffic) just as fast as a short one. It doesn't get tired or slow down.

3. The "Trajectory-Guided" Scan (The Driver's Gaze)

This is the paper's "secret sauce."

• Old systems look at the whole world equally. They stare at a parked car on the other side of the street just as hard as the car directly in front of them.
• DriveMamba mimics a human driver's eyes. It uses a "Local-to-Global" scan.
  • First, it focuses intensely on the immediate path (the "Local" part).
  • Then, it expands its view to the surroundings (the "Global" part).
  • Crucially, it follows a predicted path (a "trajectory"). If the car plans to turn left, the system prioritizes looking at the left lane and the cars there, ignoring the irrelevant stuff on the right. It's like a spotlight following your intended path.

4. The "Unified" Brain (No More Silos)

In the old assembly line, the "Perception" team and the "Planning" team rarely talk.
In DriveMamba, everything happens in one single room.

• The system learns how a pedestrian's movement (Prediction) directly affects where the car should steer (Planning) simultaneously.
• It realizes, "Oh, that pedestrian is stepping out, so I need to slow down now," without waiting for a separate step to finish. This creates a much smoother, safer, and more reactive driving style.

Why Does This Matter?

• Speed: It runs 10 times faster than some of the best previous systems. It can make decisions in milliseconds, which is crucial for avoiding accidents.
• Efficiency: It uses much less computer power (memory), meaning it could eventually run on the standard computer inside your car, not just a supercomputer.
• Safety: By looking at the whole picture at once and prioritizing what matters (the path ahead), it makes fewer mistakes and handles complex, chaotic traffic better.

In a nutshell:
If previous self-driving cars were like a slow, clunky assembly line that got confused by long traffic jams, DriveMamba is like a highly focused, lightning-fast driver who knows exactly where to look, remembers everything important, and makes split-second decisions without getting overwhelmed. It's a giant leap toward making self-driving cars that are as safe and efficient as a human expert driver.

Technical Summary

1. Problem Statement

End-to-End Autonomous Driving (E2E-AD) aims to unify perception, prediction, and planning into a single differentiable framework. However, current state-of-the-art methods face three critical limitations:

• Sequential Information Loss: Most approaches (e.g., UniAD, VAD) follow a sequential "perception-prediction-planning" paradigm using separate Transformer decoders. This manual ordering causes inevitable information loss and cumulative errors across modules, restricting flexible task synergy.
• Inefficiency and Scalability: Existing methods often rely on dense Bird's Eye View (BEV) feature generation or quadratic-complexity attention mechanisms (Transformers). These approaches are computationally expensive, struggle with long-term temporal modeling, and face memory bottlenecks when scaling to high-resolution inputs or large models.
• Lack of Ego-Centric Interaction Modeling: Current models often treat task relations statically or globally, failing to capture the dynamic, ego-centric interaction order (e.g., prioritizing objects directly in the vehicle's path) necessary for safe planning.

2. Methodology: DriveMamba

The authors propose DriveMamba, a Task-Centric Scalable paradigm that replaces the sequential Transformer architecture with a Unified Mamba Decoder based on State Space Models (SSMs).

Core Architecture

• Sparse Tokenization: Instead of generating dense BEV features, DriveMamba converts multi-view images and task queries into sparse token sequences.
  • Image Tokens: Encoded via visual backbones (ResNet, ViT, or VMamba) and converted to tokens.
  • Task Tokens: Pre-defined queries for Agents (dynamic objects), Maps (static lanes), and Ego (planning trajectory).
  • Initialization: Tokens are initialized with semantic embeddings and Task-Centric Spatial-Temporal Positional Embeddings (SPE, TPE, Task PE). Crucially, 3D positions are estimated via a depth prediction branch to enable precise spatial sorting.
• Unified Mamba Decoder: A single decoder stack processes all tokens simultaneously using Bidirectional Mamba (B-Mamba) layers. This allows for:
  • View Correspondence Learning: Implicitly learning the relationship between raw sensor inputs and task outputs without dense BEV construction.
  • Task Relation Modeling: Capturing dynamic inter-dependencies between different tasks (e.g., how map data influences agent prediction) in parallel.
  • Long-term Temporal Fusion: Utilizing a FIFO memory queue to store history task queries (rather than heavy BEV features) for efficient streaming temporal modeling.

Key Innovation: Hybrid Spatiotemporal Scan

To adapt the 1D Mamba sequence model to 3D spatiotemporal data, the authors design a Hybrid Spatiotemporal Scan method:

1. Bidirectional Scanning: Forward and backward SSMs are used to preserve spatial locality.
2. Trajectory-Guided "Local-to-Global" (L2G): A novel scan order where tokens are sorted based on their proximity to the ego-vehicle's interpolated future trajectory. This ensures the model prioritizes "in-path" objects, mimicking human driving attention.
3. Layer-wise Alternation: The scan strategy alternates between "Horizontal-First" and "Vertical-First" across decoder layers to balance semantic perception and planning efficiency.

Training Strategy

• Single-Stage End-to-End: The model is trained jointly on detection, mapping, motion prediction, and planning losses without sequential dependencies.
• Iterative Optimization: Task reference points and embeddings are updated iteratively within the decoder blocks.

3. Key Contributions

1. Task-Centric Scalable Paradigm: The first E2E-AD framework to integrate view correspondence, task relations, and temporal fusion into a Unified Mamba Decoder, eliminating the need for sequential modular designs and dense BEV features.
2. Hybrid Spatiotemporal Scan: A novel bidirectional scanning mechanism combining "Local-to-Global" (trajectory-guided) and "Spatial-to-Temporal" strategies. This preserves spatial locality for perception while optimizing interaction order for planning.
3. Linear Complexity & Efficiency: By leveraging the linear-complexity Selective State Space Model (SSM), DriveMamba achieves significant reductions in GPU memory usage and inference latency compared to quadratic-complexity Transformers, enabling scalable long-horizon modeling.
4. State-of-the-Art Performance: Demonstrates superior performance across open-loop (nuScenes) and closed-loop (Bench2Drive) benchmarks, achieving high Driving Scores with significantly lower latency.

4. Experimental Results

The model was evaluated on nuScenes (open-loop) and Bench2Drive (closed-loop).

• Efficiency:
  • DriveMamba-Tiny achieves 17.9 FPS (55.8ms latency), which is 10x faster than UniAD and 4x faster than VAD.
  • It reduces GPU memory consumption by 68.8% compared to Transformer-based baselines when scaling input resolution.
• Open-Loop Planning (nuScenes):
  • DriveMamba-Tiny reduces average L2 error by 42.1% compared to UniAD and 38.9% compared to VAD.
  • Collision rates are significantly lower (e.g., 0.15% vs. 0.71% for UniAD).
• Closed-Loop Planning (Bench2Drive):
  • DriveMamba-Base achieves a Driving Score of 65.50, outperforming DriveTransformer-Large (63.46) despite having fewer parameters and higher efficiency.
  • It demonstrates robust interactive planning, successfully handling complex scenarios like "HazardAtSideLane" by dynamically adjusting trajectories based on traffic conditions.
• Scalability: Performance scales effectively with deeper decoders and stronger backbones (e.g., ViT-L/16), showing consistent improvements in both perception and planning metrics.

5. Significance

DriveMamba represents a paradigm shift in autonomous driving architecture. By moving away from the computationally heavy, sequential Transformer designs toward a parallel, linear-complexity State Space Model, it addresses the critical bottlenecks of scalability and efficiency.

• Practical Deployment: Its high inference speed (up to 17.9 FPS on a Tiny model) makes it highly suitable for deployment on resource-constrained automotive chips.
• Future Research: It establishes a new baseline for "Task-Centric" learning, proving that dynamic task relations and ego-centric attention can be modeled more effectively without explicit sequential ordering or dense feature maps. This opens the door for larger, more capable E2E-AD models capable of handling long-term, high-resolution spatiotemporal data.