Efficiently Scanning and Resampling Spatio-Temporal Tasks with Irregular Observations

TL;DR

This paper introduces a novel algorithm for efficient sequence modeling of irregular spatio-temporal data, combining cross-attention and resampling to handle variable observation sizes, with applications in multi-agent tasks.

Contribution

The paper proposes a new method that alternates cross-attention and discounted resampling to improve efficiency in modeling irregular spatio-temporal sequences.

Findings

Achieves comparable accuracy with fewer parameters

Faster training and inference than existing methods

Effective in multi-agent intention tasks

Abstract

Various works have aimed at combining the inference efficiency of recurrent models and training parallelism of multi-head attention for sequence modeling. However, most of these works focus on tasks with fixed-dimension observation spaces, such as individual tokens in language modeling or pixels in image completion. To handle an observation space of varying size, we propose a novel algorithm that alternates between cross-attention between a 2D latent state and observation, and a discounted cumulative sum over the sequence dimension to efficiently accumulate historical information. We find this resampling cycle is critical for performance. To evaluate efficient sequence modeling in this domain, we introduce two multi-agent intention tasks: simulated agents chasing bouncing particles and micromanagement analysis in professional StarCraft II games. Our algorithm achieves comparable accuracy with a lower parameter count, faster training and inference compared to existing methods.