Linear Memory SE(2) Invariant Attention

TL;DR

This paper introduces a linear memory SE(2) invariant attention mechanism for processing spatial data, enabling scalable and efficient modeling of relative poses in autonomous driving tasks.

Contribution

It proposes a novel SE(2) invariant scaled dot-product attention that reduces memory complexity from quadratic to linear, improving scalability in spatial data processing.

Findings

Achieves linear memory scaling with scene size.

Demonstrates improved performance over non-invariant architectures.

Practical implementation in autonomous driving scenarios.

Abstract

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.