# Linear Memory SE(2) Invariant Attention

arXiv: https://arxiv.org/abs/2507.18597 · HTML: https://mlsanity.com/paper/2507.18597/

## Key facts

- No citing papers in the map yet (data through 14 Oct 2026).
- Builds on 21 of its 26 references that are in the map.
- First posted on arXiv 24 Jul 2025.

## Builds on (most cited first, 21 of 21)

- [Attention Is All You Need](https://mlsanity.com/paper/1706.03762/) (2017)
- [An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale](https://mlsanity.com/paper/2010.11929/) (2020)
- [LLaMA: Open and Efficient Foundation Language Models](https://mlsanity.com/paper/2302.13971/) (2023)
- [RoFormer: Enhanced Transformer with Rotary Position Embedding](https://mlsanity.com/paper/2104.09864/) (2021)
- [FlashAttention: Fast and Memory-Efficient Exact Attention with IO-Awareness](https://mlsanity.com/paper/2205.14135/) (2022)
- [Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains](https://mlsanity.com/paper/2006.10739/) (2020)
- [Mixed Precision Training](https://mlsanity.com/paper/1710.03740/) (2017)
- [Wayformer: Motion Forecasting via Simple & Efficient Attention Networks](https://mlsanity.com/paper/2207.05844/) (2022)
- [Parting with Misconceptions about Learning-based Vehicle Motion Planning](https://mlsanity.com/paper/2306.07962/) (2023)
- [Rotary Position Embedding for Vision Transformer](https://mlsanity.com/paper/2403.13298/) (2024)
- [Waymax: An Accelerated, Data-Driven Simulator for Large-Scale Autonomous Driving Research](https://mlsanity.com/paper/2310.08710/) (2023)
- [Language Conditioned Traffic Generation](https://mlsanity.com/paper/2307.07947/) (2023)
- [The Waymo Open Sim Agents Challenge](https://mlsanity.com/paper/2305.12032/) (2023)
- [GTA: A Geometry-Aware Attention Mechanism for Multi-View Transformers](https://mlsanity.com/paper/2310.10375/) (2023)
- [GoRela: Go Relative for Viewpoint-Invariant Motion Forecasting](https://mlsanity.com/paper/2211.02545/) (2022)
- [BehaviorGPT: Smart Agent Simulation for Autonomous Driving with Next-Patch Prediction](https://mlsanity.com/paper/2405.17372/) (2024)
- [SceneDiffuser: Efficient and Controllable Driving Simulation Initialization and Rollout](https://mlsanity.com/paper/2412.12129/) (2024)
- [Scenario Diffusion: Controllable Driving Scenario Generation With Diffusion](https://mlsanity.com/paper/2311.02738/) (2023)
- [Closed-Loop Supervised Fine-Tuning of Tokenized Traffic Models](https://mlsanity.com/paper/2412.05334/) (2024)
- [Learning Realistic Traffic Agents in Closed-loop](https://mlsanity.com/paper/2311.01394/) (2023)
- [Learning the RoPEs: Better 2D and 3D Position Encodings with STRING](https://mlsanity.com/paper/2502.02562/) (2025)

## Abstract (arXiv)

Processing spatial data is a key component in many learning tasks for autonomous driving such as motion forecasting, multi-agent simulation, and planning. Prior works have demonstrated the value in using SE(2) invariant network architectures that consider only the relative poses between objects (e.g. other agents, scene features such as traffic lanes). However, these methods compute the relative poses for all pairs of objects explicitly, requiring quadratic memory. In this work, we propose a mechanism for SE(2) invariant scaled dot-product attention that requires linear memory relative to the number of objects in the scene. Our SE(2) invariant transformer architecture enjoys the same scaling properties that have benefited large language models in recent years. We demonstrate experimentally that our approach is practical to implement and improves performance compared to comparable non-invariant architectures.