Efficiently Scanning and Resampling Spatio-Temporal Tasks with Irregular Observations

Authors: Bryce Ferenczi, Michael Burke, Tom Drummond
Primary area: learning on time series and dynamical systems
Keywords: sequence modeling, efficient training, efficient inference, spatio-temporal, multi-agent task

Abstract:
Various works have aimed at combining the inference efficiency of recurrent models and training parallelism of MHA 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. Variably sized, irregular observation spaces are relatively under-represented, yet they occur frequently in multi-agent domains such as autonomous driving and human-robot interaction. 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.

Decision: reject
Reviewer scores: 6, 6, 6, 6

Reviews and discussion:
<thread>
creation date:
1730486424665

thread_type:
official_review

<thread_object>
creator: reviewer_7wzf

<content>
summary:
The paper tackles the issues of spatio-temporal modeling in the context of variable observation spaces, such as when there are variable objects across episodes / scenes, etc or when the object number varies within an episode (such as when entities enter and leave an arena). The authors introduce a few baseline tasks representative of this problem and cover various attention-based encoding and temporal aggregation schemes, some of which are advertised as being novel (such as some of the encoder methods). They perform various analysis on different methodologies, such as task performance, memory, speed, etc on these tasks which highlight some of the benefits of various methods (such as the scan-based one).

strengths:
Overall the paper is interesting and approaches a very important problem I think in a comprehensive way, covering some benchmarks that seem well motivated to me as well as some methodologies for dealing with the variable observation problem. The experiments I think are clear enough and the conclusions drawn I think are sound, demonstrating that some methods perform better than others.

weaknesses:
There are places where the writing is good, and there are places where the writing and motivation is not so good. The introduction is very vague and lacks citations on some critical claims in the paper that make it difficult for readers who are not very familiar with the domain to understand well. The benchmarks are introduced well enough, but it's unclear why we need multi-player problems to address the core problem of variable observation spaces. Critically though when the methods come around the organization of the paper suffers, as it becomes extremely difficult to understand the motivation of different design choices, such as the scan operation (e.g., why was the scheme in figure 4 motivated, other than post-hoc performance?). The methods in Figure 3 seem interesting, but I had difficulty finding their consequence in the paper later on.

(Many of these and questions below were addressed sufficiently in the rebuttal and revision)

questions:
1) Could you clarify the contributions of the paper in terms of the models introduced, notably wrt to SSMs. I read the related works, but they did not do a good job connecting to what was done in prior sections.
2) Why is it important that the benchmarks include PvP / PvE games?
3) Did you ablate over decoders when you decided on the final scheme to compare methods? Is it possible that some results are decoder dependent (ie there's a relationship between decoder and encoder results).
4) Could you explain in more detail the motivation of the scan methods? What do you think is going on that makes this method work? What does this have to do with variable observation spaces?
5) What is Lx in 4.1.1? What does the first subscript mean in L_xt?
6) Does the model ever output observation predictions? You mention "sample the current observation" but it's unclear what you mean.
7) What do you mean by "learned initialization" in 4.1.2?
8) What is the consequence of those PvE encoders from Figure 3 on the results?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730553301094

thread_type:
official_review

<thread_object>
creator: reviewer_rrfx

<content>
summary:
The paper proposes a method for sequential prediction tasks that can handle observations of varying sizes at different times. It also simultaneously introduces two intention-prediction tasks in which to test this method. Experiments are performed against a few baselines and components of the method are ablated.

strengths:
- The contribution of the task based on StarCraft can be an interesting domain for future work to study
- The writing is in general fairly clear to follow

weaknesses:
I should start off by pointing out that I'm not precisely acquainted with this literature, so that will affect my judgment.

## Major points
1. One large issue I see with the paper is that I'm not totally certain why these particular tasks are good ones in order to explore modeling arbitrary-sized observation spaces. I think StarCraft as a domain makes sense, but there has been much prior work in StarCraft (which is not cited in the paper) that has used it in a different way, and I don't understand why the decision was made to couple "intention-prediction" tasks with the method. It seems that some domains like sequential decision making in healthcare, or some kind of variable-sized text prediction task would be equally fine as a benchmark domain. This is further complicated by the relatively unclear definition of the exact input and output for each task in Section 3. Given the motivation of these tasks is that it shows up often in interactions, it would seem that a space like Shared Autonomy might provide some more common benchmarks.

2. Another big issue that I see is that the paper does not make a strong enough case for the choice of baselines. I would expect there to be accompanying work in the domain previously with which to compare against, but because 4.1.2 does not cite any prior work, it's unclear to me that these choices are well-founded. Certainly I'm familiar with Mamba2, but the rest is unclear. Also the choice of encoders for the RNN methods feels somewhat unjustified.

3. The last major issues I see is that the results in Figure 5 and Table 1 feel somewhat inconclusive. Because nothing is controlled (either parameter count, memory, or training time) between methods, and there is no clear winner across metrics, it's difficult to see the benefit. When coupling these mixed results with the choice of new bespoke domains, it's even harder to understand the longer-term contribution. I don't think it's the case that a new method needs to be always better than prior work, but the fact that it does not appear to have particular advantages in computation that was claimed as a contribution ("efficiently address sequence modeling" at the end of Section 1), makes it hard to understand the contribution.

## Specific points

1. Section 4.1.2 add a citation to Mamba2
2. Section 4.2 it would be nice to have the problem written out mathematically as I find the current description a bit hard to follow
3. Section 5 is missing comment on hyperparameter tuning, especially for baselines
4. Fig 7 should be referenced in 5.1.1

questions:
1. Section 3.2 Where exactly does the StarCraft data come from? It would be nice to have a detailed description of the data collection and formatting process in the Appendix
2. How different is this method of encoding observations from previous work on StarCraft (regardless of the exact task, the encoding stage should be similar).
3. What is the "learned initialization" in Section 4.1.2
4. Is there some reason why baselines couldn't be matched on some axis? (parameter count is probably easiest)
5. What I understand in Section 3.2 is that the StarCraft task is a 3-class (idle, move to position, move to enemy) classification problem given sequential context, but still this doesn't exactly specify the setup for me, could you clarify what exactly is the task?
6. Section 7 Given the related work (especially Sun et al. 2023) is there some reason why none of these were used as baselines?
7. The limitations seem a bit unfortunate, but also foreseeable, is there some reason that the authors specifically designed their own benchmarks instead of working with something longer-term?
</content>
</thread_object>

</thread>
<thread>
creation date:
1730564750094

thread_type:
official_review

<thread_object>
creator: reviewer_ttpp

<content>
summary:
This paper investigates how to combine the inference efficiency of recurrent models and the training parallelism of multi-head attention transformers, especially for dealing with the varying size observation space. The authors propose a novel algorithm to cycle between cross-attention and inclusive scan to efficiently accumulate historical information. They evaluate their method on two benchmarks, Chasing-Targets gymnasium and StarCraft2.

strengths:
The proposed method, especially integrating inclusive scan algorithm, is novel and is an effective way to accumulate history.

The experiment results are good. The proposed method improves the performance and reduce the computation overhead.

weaknesses:
Please refer to the Questions part.

questions:
I’m not very familiar with this topic, so my questions are mainly about the application settings.

1.	The observations of Chasing-target and StarCraft2 are all low-dimensional. How to extend your method to more realistic and complex scenarios, e.g., vision-based navigation tasks?

2.	How can this method be applied to longer-horizon tasks?

3.	I would like to see more ablation studies on the selection of hyperparameter $\gamma$.
</content>
</thread_object>

</thread>
<thread>
creation date:
1731405134117

thread_type:
official_review

<thread_object>
creator: reviewer_avti

<content>
summary:
The paper discusses a approach to spatio-temporal modelling for tasks with irregular observation spaces. Experiments are conducted on multi-agent scenario's like robotic target-chasing simulations and StarCraft II gameplay analysis.


Motivation: The paper aruges that current architectures like Transformers and LSTMs works well in scenarios with fixed observation spaces but struggle with tasks where observation spaces vary over time, such as multi-agent interactions. The paper aims to combine the efficiency of recurrent models with the parallelism of attention mechanisms, focusing on varying-dimension observation spaces.

Proposed Method:  The paper introduce an algorithm using a 2D latent state, alternating between cross-attention on observations and a discounted cumulative sum over time, capturing long context.  

Experiments: Two benchmarks are used for evaluation: (a) Chasing Targets: Agents pursuing dynamic targets, evaluating the model's ability to infer intended goals in real-time; (b) StarCraft II: A environment with complex, temporally varying observations, where models infer combat unit assignments and predict unit movement.

Results: The proposed model achieves competitive accuracy with fewer parameters and faster inference compared to standard methods (e.g., transformers). The "resampling cycle," improves sequence modeling accuracy by conditioning observations on accumulated sequence data.

===================

After rebuttal: I've read the review by other reviewers as well as the rebuttal by the authors.

This statement by the authors summarizes the contribution of the paper well.
"Our results primarily set out to explore the trade-offs between efficiency and performance in the irregular observation case. This means results should be interpreted given the compute budget available to a practitioner. If only performance matters, practitioners should select a temporal transformer with sequential cross attention (X-Attn Sequential) for observation encoding. For efficiency and model parsimony, an inclusive scan should be selected with X-Attn Sequential."

It will be very helpful to revise the introduction and abstract accordingly.

strengths:
The paper investigates the use of various encoders to encode irregular observations into a fixed dimensional latent, and corresponding sequence modelling approaches.

weaknesses:
- The model's efficiency enables its use in real-time applications with irregular observation space, but studying the effect on longer sequence durations as well as diversifying the tasks (i.e., conduct experiments on more tasks) would further validate its utility in long-term planning tasks.
- There's some work in the literature like Temporal Latent Bottleneck (https://arxiv.org/abs/2205.14794) and Block Recurrent Transformers (https://arxiv.org/abs/2203.07852); which combines strengths of attention and recurrence.

questions:
See Weaknesses.
</content>
</thread_object>

</thread>
<thread>
creation date:
1732620181387

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We thank the reviewers for their targeted and thoughtful comments to improve our paper. We have added additional experiments and sections to the paper to include the additional requests such as including RetNet, more gamma ablations, and consistency with spatial encoder choice. Furthermore, we have added a direct comparison between our weighted summation and two gated recurrent unit variants, facilitated by directly swapping scan with GRU or GILR [1]. With these additional experiments, we have strengthened our introduction and contributions. To fit this within the page limit, we have moved the task decoding to the appendix as while it is still important for the understanding of the model, it is not pertinent to the core results of the paper.

[1] Parallelizing Linear Recurrent Neural Nets Over Sequence Length https://openreview.net/forum?id=HyUNwulC-
</content>
</thread_object>

</thread>
<thread>
creation date:
1733269430330

thread_type:
official_comment

<thread_object>
creator: authors

<content>
comment:
We would once again like to thank the reviewers for their valuable feedback which has helped us improve our paper. A summary of the key changes are as follows:

- Added evaluation of spatial encoding combinations, standardized usage of X-Attn Sequential
- Added learning rate sweep for chasing-targets temporal encoders
- Added gamma sweep for chasing-targets Scan 4x
- Included a new experiment, swapping the Scan module with GILR or GRU.
- Added RetNet and Transformer (per latent, “TT”) baselines to experiments
- Added discussion on motion forecasting datasets
- Added discussion on long-term forecasting challenges
- Improved clarity in various areas pointed out by reviewers

We believe that our contributions, introducing a framework for efficiently handling spatio-temporal challenges with irregular observations, a simple, yet effective method for accumulating temporal context on the back of extensive ablations and evaluations of efficient architectures alongside two challenging tasks to demonstrate model capability, are valuable to the community.
</content>
</thread_object>

</thread>


Metareview:
metareview:
The paper presents an interesting approach to sequence modeling with irregular observation spaces, exploring novel scan operations and cross-attention techniques for handling variable-dimensional observations. While the work demonstrates potential in multi-agent tasks and introduces innovative modeling strategies, it struggles with comprehensive validation and generalizability. Reviewers noted significant challenges in methodology, including unclear motivation and insufficient differentiation from existing approaches. The benchmarks, primarily based on custom multi-agent tasks, do not definitively prove the method's broader applicability, and performance improvements remain marginal. Despite the authors' efforts to address concerns through additional experiments and clarifications, fundamental limitations persist. The incremental nature of the contribution, combined with presentation challenges and limited experimental rigor, suggests that while the work shows promise, it is not yet sufficiently developed for publication at a top-tier conference.