Efficiently Scanning and Resampling Spatio-Temporal Tasks with Irregular Observations
This addresses a bottleneck in spatio-temporal tasks with irregular observations, such as multi-agent intention tasks, though it appears incremental in combining recurrent and attention-based approaches.
The paper tackles the problem of sequence modeling with varying-size observation spaces by proposing a novel algorithm that alternates between cross-attention and discounted cumulative sum, achieving comparable accuracy with lower parameters and faster training/inference compared to existing methods.
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.