Efficient Equivariant Transformer for Self-Driving Agent Modeling

Scott Xu; Dian Chen; Kelvin Wong; Chris Zhang; Kion Fallah; Raquel Urtasun

cs.ROApr 1, 2026

Efficient Equivariant Transformer for Self-Driving Agent Modeling

Authors: Scott Xu, Dian Chen, Kelvin Wong, Chris Zhang, Kion Fallah, Raquel Urtasun

Organizations: Waabi · University of Toronto

Abstract

Accurately modeling agent behaviors is an important task in self-driving. It is also a task with many symmetries, such as equivariance to the order of agents and objects in the scene or equivariance to arbitrary roto-translations of the entire scene as a whole; i.e., SE(2)-equivariance. The transformer architecture is a ubiquitous tool for modeling these symmetries. While standard self-attention is inherently permutation equivariant, explicit pairwise relative positional encodings have been the standard for introducing SE(2)-equivariance. However, this approach introduces an additional cost that is quadratic in the number of agents, limiting its scalability to larger scenes and batch sizes. In this work, we propose DriveGATr, a novel transformer-based architecture for agent modeling that achieves SE(2)-equivariance without the computational cost of existing methods. Inspired by recent advances in geometric deep learning, DriveGATr encodes scene elements as multivectors in the 2D projective geometric algebra $\mathbb{R}^*_{2,0,1}$ and processes them with a stack of equivariant transformer blocks. Crucially, DriveGATr models geometric relationships using standard attention between multivectors, eliminating the need for costly explicit pairwise relative positional encodings. Experiments on the Waymo Open Motion Dataset demonstrate that DriveGATr is comparable to the state-of-the-art in traffic simulation and establishes a superior Pareto front for performance vs computational cost.

Figures & tables

Panel of Figure 1 — Figure 1: The DriveGATr architecture. Figure 1(a) provides an overview. The poses and features of $N_{actor}$ agents and $N_{map}$ nodes in each scene are encoded as multivectors in ${\mathbb{R}^{*}_{2,0,1}}$ and scalars. These tensors are processed by N transformer blocks, each consisting of agent and map cross attention, temporal causal self-attention, equivariant MLPs and invariant adapters. Each of these modules contain skip connections, closely mimicking a standard transformer. Figure 1(b) describes the attention block. For all attention blocks, the query inputs are agent of interests. The key and value inputs are agents and map nodes for the cross attentions, attending across elements per-timestep. For self-attention, attention spans the temporal axis per-agent. Figure 1(c) describes the equivariant MLP block. Table 1 provides details on ${\mathbb{R}^{*}_{2,0,1}}$ encodings. Section 3.4 provides details on each of the equivariant primitive layers. Section 3.5 describes the invariant adapter block.

Geometry / transform	Representation	Transform inverse
Line $ax+by+c=0$	$ae_{1}+be_{2}+ce_{0}$	-
Point $(x,y)$	$xe_{20}+ye_{01}+e_{12}$	-
Translation by $(a,b)$	$1-\frac{a}{2}e_{01}+\frac{b}{2}e_{20}$	$1+\frac{a}{2}e_{01}-\frac{b}{2}e_{20}$
Counterclockwise rotation of angle $\theta$	$\cos\frac{\theta}{2}-\sin\frac{\theta}{2}e_{12}$	$\cos\frac{\theta}{2}+\sin\frac{\theta}{2}e_{12}$

Table 1: 2D encodings of geometries and transforms.

Method	# Params	Finetuned	RMM $\uparrow$	Kinematic $\uparrow$	Interactive $\uparrow$	Map-based $\uparrow$	minADE $\downarrow$
BehaviorGPT	3M	-	0.7438	0.4254	0.7233	0.7976	1.3804
SMART-7M	7M	-	0.7678	0.4894	0.7306	0.8163	1.3532
SMART-7M + CAT-K	7M	✓	0.7709	0.4937	0.7333	0.8185	1.2953
Transformer	3M	-	0.7257	0.4759	0.7058	0.7654	1.6720
Transformer + RPE	3M	-	0.7251	0.4708	0.6953	0.7808	1.7486
Transformer + DRoPE	3M	-	0.7206	0.4629	0.7017	0.7604	1.9193

Table 2: Results on WOSAC 2024 validation set. RMM is the Realism Meta Metric used for ranking. DriveGATr achieves the best RMM among its baselines. Compared to publicly available state-of-the-art methods, DriveGATr achieves comparable realism without top-

k

sampling and closed-loop fine-tuning.

Panel of Figure 2 — Figure 2: Training curves. We compare the envelope of minimal training loss per FLOP ( left ), compute efficiency as the number of agents scale ( middle ), and sample efficiency as the training dataset grows ( right ). Compared to the baselines, DriveGATr establishes a superior Pareto front in performance and computation cost thanks to its compute and sample efficiency.

Panel of Figure 3 — Figure 3: Robustness to roto-translations. We compare Transformer ( left ), Transformer + DRoPE ( middle ), and DriveGATr ( right )’s robustness to roto-translations. In each figure, we overlay rollouts from the original coordinate frame vs one rotated by 90°and translated by 100m forward. Blue trajectories visualize model predictions in the original input, and red visualize predictions in the transformed scene. DriveGATr produces consistent trajectories despite closed-loop execution, demonstrating its robustness to roto-translations.

IA	DA	RMM $\uparrow$	minADE $\downarrow$
✓		0.7408	1.6176
	✓	0.7483	1.6053
✓	✓	0.7478	1.5798

Table 3: Ablation studies on the WOSAC 2% validation split. On the top , we ablate the importance of the invariant adapter (IA) and distance-awareness (DA). On the bottom , we ablate the number of map tokens each agent attends to in agent-to-map attention.

Panel of Figure 4 — Figure 4: Robustness to roto-translations. Additional visualizations showing Transformer ( left ), Transformer + DRoPE ( middle ), and DriveGATr ( right )’s robustness to roto-translations. In each figure, we overlay rollouts from the original coordinate frame vs one rotated by 90°and translated by 100m forward. Blue trajectories visualize model predictions in the original input, and red visualize predictions in the transformed scene.

Panel of Figure 5 — Figure 5: Visualizations of multi-modal DriveGATr rollouts. In each figure, we overlay trajectories from multiple ( $k=8$ ) rollouts. Different actors of the same color corresponds to the same rollout.

Method	Training latency	Training peak memory	Inference latency	SE (2) equivariant?
Transformer	196.86 ms	3.28 GB	42.64 ms	$\color[rgb]{1,0,0}\times$
Transformer + RPE	328.70 ms	15.19 GB	82.48 ms	$\color[rgb]{0,1,0}\checkmark$
Transformer + DRoPE	195.51 ms	5.08 GB	46.03 ms	$\color[rgb]{1,0,0}\times$
DriveGATr -3M	264.34 ms	6.35 GB	58.01 ms	$\color[rgb]{0,1,0}\checkmark$

Table 4: Peak memory and latency comparisons of the baselines and DriveGATr . All methods in the same group use the same context length. The training latency times the combined forward and backward passes. The inference latency are averaged over the entire 8s rollout. Both training and inference use batch size = 8. All entries are evaluated a single A5000 with fixed random seed.

Explore similar work

Drive3R: Block-wise Streaming Geometric Foundation Models for Autonomous Driving
Oct 10, 2026Tianyi Yang, Haojun Yang, Jipeng Xiao +43D Foundation Models Streaming 3D Reconstruction
AutoLab, School of Artificial Intelligence, Shanghai Jiao Tong University · Voyager Research, Didi Chuxing · The Chinese University of Hong Kong, Shenzhen
DVGT-2: Vision-Geometry-Action Model for Autonomous Driving at Scale
Apr 1, 2026Sicheng Zuo, Zixun Xie, Wenzhao Zheng +6End-to-End Autonomous Driving Autonomous Driving Planning
Tsinghua University · Xiaomi EV · Peking University +1
Pictura: Perspective-View Self-Play at Scale for Driving
Jul 28, 2026Yuan Yin, Elias Ramzi, Marc Lafon +8Self-Play RL RL for Autonomous Driving
valeo.ai, Paris, France · Valeo Brain, Créteil, France · Sorbonne Université, CNRS, ISIR, F-75005 Paris, France

Oct 10, 2026J/K move · Enter open · S save

Drive3R: Block-wise Streaming Geometric Foundation Models for Autonomous Driving

Tianyi Yang, Haojun Yang, Jipeng Xiao, Xuesong Chen+3

AutoLab, School of Artificial Intelligence, Shanghai Jiao Tong University · Voyager Research, Didi Chuxing · The Chinese University of Hong Kong, Shenzhen

3D Foundation Models Streaming 3D Reconstruction Feed-Forward 3D Reconstruction

Open paper AskPDF arXiv

High-fidelity closed-loop simulation for autonomous driving increasingly relies on dense multi-view 3D reconstruction. While recent feed-forward geometric foundation models have revolutionized per-scene optimization, applying them to continuous, long-horizon driving sequences is computationally intractable due to GPU memory limits and the quadratic complexity of global attention. Existing scaling strategies are fundamentally limited, as causal streaming suffers from compounding scale drift and trajectory deviation, while rigid "chunk-then-align" heuristics fragment global context and produce trajectory discontinuities. To address these limitations, we introduce Drive3R, a block-wise streaming geometric foundation model. Its block-wise causal attention resolves the scalability bottleneck of fully bidirectional attention while avoiding the loss of local future context in purely causal streaming. Its confidence-aware re-anchoring further mitigates long-range drift by detecting unreliable streaming states, re-initializing the KV cache, and stitching refreshed segments through SE(3) alignment. Extensive experiments demonstrate that Drive3R establishes a new state-of-the-art on nuScenes, KITTI, and Waymo. Specifically, on nuScenes, it improves over the recent DVGT-2 baseline from 0.1806 to 0.1081 in AbsRel, from 0.6750 to 0.5344 in point-cloud Accuracy, and from 1.5812 to 0.9092 in ATE.

Abstract

Explore similar work

Drive3R: Block-wise Streaming Geometric Foundation Models for Autonomous Driving

DVGT-2: Vision-Geometry-Action Model for Autonomous Driving at Scale

Pictura: Perspective-View Self-Play at Scale for Driving