Efficiently Scanning and Resampling Spatio-Temporal Tasks with Irregular Observations

Cont.

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?

Our results primarily set out to explore the trade-offs between efficiency and performance in the irregular observation case. A parameter-free inclusive scan is the cheapest possible method to accumulate the spatial-encoding over the sequence dimension. We hypothesize that embedding the notion of recency in the latent with the discount factor has enabled better model adaptation of target assignment changes throughout the sequence, which is why a standard RNN is significantly worse than our Scan method or a GRU, GILR or LSTM. We show in A.1.2 that a larger latent state improves performance, the parameter-free Scan method avoids a corresponding proportional increase in parameter count and can be performed efficiently in parallel.

What is Lx in 4.1.1? What does the first subscript mean in L_xt?

We have updated the paper to explicitly note that x denotes the layer index of the model.

Does the model ever output observation predictions? You mention "sample the current observation" but it's unclear what you mean.

We make predictions on properties of the environment that aren’t observed by the model. Since we use multi-head attention, with exception of the first layer, since we have accumulated knowledge from the scan operation, our next query includes historical features, so the attention on the key generated by the observation is now conditioned on the history.

What do you mean by "learned initialization" in 4.1.2?

The initial state of the recurrent neural network, we have adjusted the language to improve clarity.

What is the consequence of those PvE encoders from Figure 3 on the results?

We conducted an ablation study (A.1.1 and Fig. 5) on all the spatial encoding methods with a fixed temporal encoder and found X-Attn Sequential to have the highest accuracy, while X-Attn Fused was the cheapest. For consistency, we now use the X-Attn Sequential algorithm for all the models in the paper unless otherwise specified.

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.

We have added citations to relevant works in the introduction as well as tweaks to the main body and list of contributions to reflect the additional studies we have done in this revision and improve the quality of writing.

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.

While our approach is applicable across a broad set of domains, these multi-agent problems are particularly representative of the hardest aspects of the irregular observation problem, going beyond say missing data imputation in tabular data. We note that in multi-player, multi-unit games like Starcraft, the observations can be considered a tree-like data structure, with each player representing a branch, and units and their attributes sub-trees in the scene. When observations are not available, or units appear or disappear, the tree structure changes quite dramatically. Existing motion datasets are usually brief or small due to manual annotations required. Multi-agent modelling problems like those chosen are good examples of spatio-temporal modelling tasks with irregular observations, and also of interest in the machine learning community.

Critically though when the methods come around the organisation 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?).

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.

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. We have made this clearer in the text, and by adding an additional experiment where we directly replace the weighted cumsum with GRU or GILR. This experiment shows a more direct correlation of accuracy with temporal model sophistication and cost. We also changed to a Scan 3\times model in SC2 which now has the best training throughput, median memory usage with slight disadvantage in accuracy.

The methods in Figure 3 seem interesting, but I had difficulty finding their consequences in the paper later on.

We have improved the motivation and clarity around the choice of encoders and dynamics models. We conducted an ablation study on all the spatial encoding methods with a fixed temporal encoder and found X-Attn Sequential to have the highest accuracy, while X-Attn Fused was the cheapest. For consistency, we now use the X-Attn Sequential algorithm for all the models in the paper.

Why is it important that the benchmarks include PvP / PvE games?

This is a challenging problem as targets are often switched during the sequence, requiring the model to be adaptable to instantaneous changes in the signal. This is a potential reason why an RNN fails at this task.

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).

Yes, it is possible that some results are decoder dependent. We aimed to reduce this by evaluating the encoding methods with relatively simple intention prediction tasks, rather than more difficult tasks such as long horizon forecasting or using RL to train an agent on the task.

The main decoding ablation was the position decoding scheme. The target assignment can be seen in a similar vein to CLIP, correlating words and images, which is why we choose cosine-similarity of encoded agent and target features, although we have a one-way assignment which allows for a row-wise softmax.

TBC.

Cont.

Section 5 is missing comment on hyperparameter tuning, especially for baselines

We have added a section in the supplementary on varying the latent state size for chasing-targets with the temporal transformer encoder. While we explored the learning rate on Scan $4\times$ (5e-3, 1e-3, 5e4, 1e-4), we didn’t feel it was particularly pertinent to include this. We did not explore different learning rate schedules.

Fig 7 should be referenced in 5.1.1 Done

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

We have added a footnote with a link to the source of the raw data. As mentioned in the reproducibility section, the StarCraft II data collection and conversion tool is made by us and publicly available so is not currently linked for anonymity.

How different is this method of encoding observations from previous work on StarCraft (regardless of the exact task, the encoding stage should be similar).

The most prominent SC2 model, AlphaStar, used an LSTM for temporal recurrency. AlphaStar-Unplugged uses Transformer XL with up to 2*10^10 frame memory, finding that in pure behaviour cloning LSTM is detrimental compared to no memory. In both models, units are initially encoded with a transformer, then “summarised” into a single feature vector with an averaging operation, and merged with other contextual observation features. This method isn’t particularly sophisticated and we have shown that for a simpler task, chasing-targets, that memory and a larger capacity latent state greatly improves performance.

What is the "learned initialization" in Section 4.1.2

The initial state of the recurrent neural network, we have adjusted the language to improve clarity.

Is there some reason why baselines couldn't be matched on some axis? (parameter count is probably easiest)

We matched by the number of layers - 2. We have added this note in the first paragraph of the results section. Models such as the transformer naturally have significantly more parameters to project Q,K,V + Linear FeedForward compared to our method which has zero. This is illustrated in our new 5.2 section that replaces the weighted sum with GILR and GRU, each of which doubles the parameter count (0.9->2->4M), which comes with an associated performance increase. We believe this is a good illustration of scaling between parameters in sequence modelling.

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?

The task is to predict what a player is doing at a given point in time: attacking an enemy unit (assignment to unit), moving to a position (position prediction + position prediction is valid) or idle (no assignment and position prediction is invalid), given the observations of unit positions and their attributes and the context of the minimap. We have revised the wording in Section 3 to improve clarity.

Section 7 Given the related work (especially Sun et al. 2023) is there some reason why none of these were used as baselines?

We have added RetNet to represent this family of linearized transformers.

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?

We note that our benchmarks are longer duration than many existing (eg. multi-agent tracking). However, as mentioned in the discussion and limitations, models such as these are not really suited to long horizon forecasting. This is a common concern with autoregressive state space models [2]. Our core focus is on architectures that efficiently encode complex irregular observations into a latent and how to pair these with efficient state space models, so it is likely that our approach will inherit these weaknesses.

[2] Appendix B.1 Long-term Forecasting with TiDE: Time-series Dense Encoder https://openreview.net/forum?id=pCbC3aQB5W

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. [...] It seems that some domains like sequential decision making in healthcare, or some kind of variable-sized text prediction task would be equally fine [...]

We chose these benchmarks because we were unable to find existing public datasets with irregular observation spaces. Multi-agent modelling problems like those chosen are good examples of spatio-temporal modelling tasks with irregular observations, and also of interest in the machine learning community.

We have been unable to find medical datasets with an irregular observation space like our multi-agent interaction tasks. Most settings here (eg. EEG, Imaging and Tabular data) all exhibit regular shapes (1D, 2D or 3D signal processing). Text based tasks are often treated as a 1D signal, GPT generates token by token until [EOL], BERT reduces a sequence to 1 token, so creating a benchmark here feels somewhat contrived. As mentioned by the reviewer, Shared Autonomy is a domain that is conducive to what we are working on, however common benchmarks from Waymo Open Motion Dataset, ArgoVerse2 and nuScenes are all very short sequences, a few seconds of history to produce a few seconds of prediction. We use our own environments to test modelling over longer sequences, although as mentioned in the limitations, we observe that the amount of context needed saturates around 15 timesteps in (~10sec real time), while motion prediction usually works with 1-5sec history. We have included a mention of motion forecasting datasets in the Discussion and Limitations section.

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.

We have improved the motivation and clarity around the choice of encoders and dynamics models. Our core focus is on determining appropriate encoding strategies and complementary dynamics models, with efficiency in mind. Baselines and encoding/ dynamics models were chosen primarily for computational complexity reduction and our core goal was to explore the compute performance trade-offs for a range of architectures tackling the problem of spatio-temporal modelling with irregular observations.As such, each of the baselines represent common modelling meta-categories, Recurrent, Transformer and SSM. Given the lack of prior work in the irregular observation handling case, we selected baselines that we felt were appropriate points of comparison. We have added RetNet to represent these linearized transformers (as mentioned in Q6). We have also added citations and additional context for spatio-temporal transformer. Furthermore, we conducted an ablation study on all the spatial encoding methods and found X-Attn Sequential to have the highest accuracy. For consistency, we have now used that algorithm for all the models in the paper.

The last major issues I see is that the results in Figure 5 and Table 1 feel somewhat inconclusive. [...] the fact that it does not appear to have particular advantages in computation that was claimed as a contribution ("efficiently address sequence modelling" at the end of Section 1), makes it hard to understand the contribution.

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. We have made this clearer in the text, and by adding an additional experiment where we directly replace the weighted cumsum with GRU or GILR. This experiment shows a more direct correlation of accuracy with temporal model sophistication and cost. We also changed to a Scan $3\times$ model in SC2 which now has the best training throughput, median memory usage with slight disadvantage in accuracy.

Section 4.1.2 add a citation to Mamba2

Done

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

While we moved section 4.2 to the appendix due to space constraints, we have added formulas that represent the decoding algorithms.

TBC.

observations of Chasing-target and StarCraft2 are all low-dimensional

There is already substantial work on tasks with fixed-dimensional observations like vision-based navigation tasks. Our work focuses on a less commonly studied setting of variable dimension observations. These are often low-dimensional, but with complex data structures that are often seen in real-world settings. For example, in our starcraft data, the observations not only include variable numbers of units and positions, but also their health and the unit type (as a learned embedding). This problem setting is very realistic, for example in inventory management, digital twin simulations or tabular data structures with missing data. That being said, high dimensional observations could be trivially included in the proposed architecture, by tokenising the scene (eg. with image chips or object tokens from a DETR-like vision encoder).

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

Although our benchmarks do not explicitly test longer horizon modelling, the method can still be applied to these tasks directly. There is no particular limitation on the input sequence length, like any other recurrent model. Training requires O(L) memory whereas inference is O(1). However, as mentioned in our response to Reviewer aVTi, autoregressive models are known to perform worse in longer horizon prediction tasks, so our approach is likely to inherit these weaknesses.

I would like to see more ablation studies on the selection of hyperparameter (Gamma)

We have added gamma=1.5 and 3 and created a separate figure in the supplementary material plotting all the curves (1,1.5,2,3). gamma=2 retains the highest performance, followed by 1.5, 3 and 1. We have also added baselines to GILR, a state-space model that can be formulated as an inclusive scan with an input-conditioned gamma parameter.

...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.

We note that our benchmarks are longer duration than many existing (eg. Waymo Open Motion Dataset). However, as mentioned in the discussion and limitations, models such as these are not really suited to long horizon forecasting. This is a common concern with autoregressive state space models [2]. Our core focus is on architectures that efficiently encode complex irregular observations into a latent and how to pair these with efficient state space models, so it is likely that our approach will inherit these weaknesses.

[2] Appendix B.1 Long-term Forecasting with TiDE: Time-series Dense Encoder https://openreview.net/forum?id=pCbC3aQB5W

There's some work in the literature like Temporal Latent Bottleneck [...] and Block Recurrent Transformer…

We thank the reviewer for pointing us to another strategy for sub-quadratic transformer modelling, we have included discussion of this in our Related Work section. While we have not included the mentioned block+recurrency methods in our comparison studies, we have added RetNet to represent linear attention models of this form.

We recommend reviewing the plots within the revised manuscript which have been updated, but here are tables with new ablations mentioned in the comment above. Other new ablations such as latent state size are plots in the appendix of the revised paper.

Spatial-Encoder with Temporal-Only Transformer

Replacing Scan with Gating (4-Block)

Updated SC2 Evalauation

Recurrent Chasing-Targets Assignment