A Large Recurrent Action Model: xLSTM enables Fast Inference for Robotics Tasks

Thank you for your feedback, which helped us to considerably improve our paper. We are happy that you find that our work demonstrates substantial research effort.

Lower performance of DT? We want to clarify that DT and xLSTM have the same number of parameters in all our experiments to ensure a fair comparison. Only Mamba has a slightly higher number of parameters (see x-axis in Figure 2). With respect to the scaling laws, xLSTM results in slightly better performance than DT, but also comes with linear inference time. A similar performance advantage has been observed in the original xLSTM publication [1].

Why does xLSTM outperform Mamba? It is true that both xLSTM and Mamba are designed for linear-time inference complexity. Line 300 tries to say that we observe performance differences at larger model scales such as 206M (see Figure 2b, last point), which indicates that xLSTM (and also other backbones) may continue to benefit from more parameters and more data. To provide preliminary empirical evidence for this, we present additional model performances at the 408M model scale in Figure 14 (only 1 seed for computational reasons), and find that performance continues to improve slightly. One other benefit of xLSTM over Mamba may be the state tracking abilities of sLSTM blocks [2], which can be useful for partially observable observation spaces such as Darkoom. Therefore, we conduct an additional ablation to study the effect of the ratio of mLSTM-to-sLSTM blocks in Appendix D.3. However, we believe that more research is necessary to fully investigate the difference between modern recurrent architectures, and hope our work is a step in this direction.

Discussion on DeMA [3]: Thank you for bringing up DeMa [3]. We added the paper to our related work section.

Thank you for the rebuttal and the inclusion of detailed experimental results. I appreciate the large-scale experimental settings and the analysis of the xLSTM architectures within the RL domain. I hope the authors will consider releasing the collected datasets and corresponding codes soon. I have raised my score, good luck!

Thank you for your constructive feedback. We appreciate your positive assessment and are glad that you find our work insightful and a valuable study.

Differences between xLSTM/Mamba/Transformer: We agree with the reviewer that a discussion on the differences between xLSTM, Mamba and Transformers can be useful and added a discussion to the paper. Important differences include:

• The complexity difference between Transformers and recurrent backbones such as xLSTM and Mamba. While the self-attention in Transformers is quadratic, the recurrent backbones are linear. This can have advantages for inference speeds and memory requirements (see Figures 6, 7), in particular for long sequences. Note that we conduct our comparison on a high-end data-center GPU with 40GB of RAM. In contrast, applications on the edge, like robotics, may have to deal with less powerful accelerators, making linear-time inference more attractive.
• Transformers have other advantages. For example, for tasks where exact recall is required, self-attention is typically superior (see Figure 5 in [1]), which can be important for decision-making tasks [2]. xLSTM, in particular, can be better for tasks that require state-tracking [3], which Mamba and Transformers cannot model.

Normalized vs. Raw scores: Throughout the paper, we presented the normalized performances (e.g., in Figure 2b). This is to make the scores from different domains more comparable. Did the reviewer mean that we should also add tables for the raw-scores to the paper? If so, we added the raw performance scores for all 432 tasks for the 206M model scale to Appendix F of our manuscript.

Thank you for spotting our mistakes on the bottom border of Table 5 and the scientific notation of the learning rate. We corrected them!

Regarding your questions:

1. Ratio of mLSTM to sLSTM blocks: We agree with the reviewer that other ratios would be interesting to study. The rationale for the ratios we chose is that we adopted the ratios of the original xLSTM publication [1] (see Appendix B.3). To get a better understanding of the effects of this design choice, we conduct additional ablation studies on the ratio in Appendix D.3. We compare different ratios (1:0, 0:1, 1:1, 3:1, 7:1) and different placements of the sLSTM blocks, both on the 432 tasks used in the main paper, and on the Darkroom environment. While we do not find benefits on the 432 tasks at the 16M model scale (see Figure 29), we do find that different ratios (such as 3:1) can be beneficial on Darkoom (see Figure 30). Therefore, the effectiveness of sLSTM layers depends on the task at hand. For more complex tasks with long horizons or partial observability (as is common in real-world applications), sLSTM may be a useful tool for practitioners.
2. Replacing Transformer blocks with xLSTM blocks: We agree with the reviewer that combinations of Transformer and xLSTM blocks may be interesting and effective. While we cannot provide empirical evidence for this combination, it is possible that they may lead to performance improvements, by combining the benefits of both architectures. Similarly, combinations of SSM and Transformer layers have been found to be effective in Griffin [4] and Jamba [5]. However, the combination comes at the cost of losing the linear inference time benefit.
3. Visual encoder: The visual encoder is trained jointly with the downstream policy. We use a CNN architecture similar to [6], and do not make use of ViT-style patchification of images (as used in [7,8]). Consequently, every image is a single embedded “token”. In preliminary experiments, we found that patchification led to similar results as using a separate image encoder, but comes at the cost of increased sequence length. The image encoder is randomly initialized, but initializing from a pre-trained model and fine-tuning may also be effective [9]. We describe the image encoder in Appendix B.3.

Once again, thank you for your helpful comments and positive assessment of our work.

[1] “xLSTM: Extended Long Short-Term Memory”, NeurIPS 2024
[2] “When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment”, NeurIPS 2023
[3] “The Illusion of State in State-Space Models”, ICML 2024
[4] “Griffin: Mixing Gated Linear Recurrences with Local Attention for Efficient Language Models”, ArXiv 2024
[5] “Jamba: A Hybrid Transformer-Mamba Language Model”, ArXiv 2024
[6] “IMPALA: Scalable Distributed Deep-RL with Importance Weighted Actor-Learner Architectures”, ICML 2018
[7] “Multi-Game Decision Transformers”, NeurIPS 2022
[8] “A Generalist Agent”, TMLR 2022
[9] “RT-1: Robotics Transformer for Real-World Control at Scale”, ArXiv 2023

Thank you for your helpful feedback. We are glad that you appreciate the experimental details we provide and that you find our data pipeline useful for the community.

Decision Transformer vs. Behavior Cloning: We agree with the reviewer that it is important to validate whether our findings translate to the behavior cloning setting. It is true that many recent works focus on behavior cloning without return-conditioning, as used in Decision Transformers. Therefore, we conduct an additional ablation study on the behavior cloning setting at the 206M parameter scale. Indeed, we find that the same trends hold when training models without return-conditioning (see Appendix D.2 of our updated manuscript). Thank you for the suggestion!

Influence of context length: It is true that most sequence lengths in Figure 23b tend to result in similar performance. We believe that this is an artifact of the Meta-World environment. Episodes in Meta-World last for 200 steps and, therefore, the effect of longer context may not be visible. To further investigate the effect of the context length, we conduct an additional ablation study for models trained on all 432 tasks and with varying context lengths in Appendix D.1. Indeed, we find that performance improves both for DT and xLSTM as the sequence length increases. For DT, the average normalized score increases from 0.3 with C=1 to 0.7 with C=50 (see Figure 24). Similarly, for xLSTM the normalized score increases from 0.4 to 0.8 (see Figure 26). Furthermore, we find that domains with longer episode lengths, like DMControl or Atari, benefit more from increasing context length than Meta-World (see Figures 25, 27). This result highlights the fact that large action models can strongly benefit from increased context length, even on the simulated environments we consider in this work. We believe that this effect may be even bigger for complex real-world applications that require longer-term interactions.

mLSTM to sLSTM ratio: The reviewer is correct that we consider the ratios [7:1] and [1:0] for xLSTM. For example, [7:1] indicates that there is 1 sLSTM block for every 7 mLSTM blocks. We adopted the ratios from the original xLSTM publication [1] (see Appendix B.3 for a discussion on this choice). While mLSTM is fully parallelizable, sLSTM enables state-tracking [2], which can be useful for logic tasks or partially observable environments. To get a better understanding of the effects of this design choice, we conduct additional ablation studies on the ratio in Appendix D.3. We compare different ratios (1:0, 0:1, 1:1, 3:1, 7:1) and different placements of the sLSTM blocks, both on the 432 tasks used in the main paper, and on the partially observable Darkroom environment. While we do not find benefits on the 432 tasks at the 16M model scale (see Figure 29), we do find that different ratios (such as 3:1) can be beneficial on Darkoom, while others hurt performance (see Figure 30). Therefore, the effectiveness of sLSTM layers depends on the task at hand. For more complex tasks with long horizons or partial observability (as is common in real-world applications), sLSTM may be a useful tool for practitioners.

Removing the action condition: Thank you for bringing up the “copycat problem” paper [3], which we now refer to in our work. We believe this paper studies a similar problem as we encountered, even though we came up with a different solution. In preliminary experiments, we found that robotics domains were more strongly affected by this problem than discrete control domains (Procgen, Atari). To provide further empirical support for this, we conducted an ablation study on this in Appendix D.1.2. We find that Meta-World and DMControl are strongly affected, but Atari and Procgen are not (see Figure 25). Furthermore, we also found that the robotics domain Composuite is also not affected. Therefore, we compute the average MSEs between subsequent actions in the datasets (similar to [3], see Table 7) and find that the MSE for Composuite is considerably higher. Therefore, we conclude that removing actions from the agent’s context is particularly effective for domains where actions change smoothly (such as Meta-World, DMC). However, note that the final performances on domains “unaffected” by the copycat problem are still considerably worse when training with actions than when training without actions in the context (0.35 vs. 0.7).

We hope to have answered all your questions and believe our manuscript has become substantially better due to your suggestions.

[1] “xLSTM: Extended Long Short-Term Memory”, NeurIPS 2024
[2] “The Illusion of State in State-Space Models”, ICML 2024
[3] Fighting Copycat Agents in Behavioral Cloning from Observation Histories, NeurIPS 2020

I appreciate the detailed response from the authors, and agree on the value of this work as large-scale empirical research. However, as noted in my initial review, demonstrating xLSTM's relevance for robotics requires evaluation on learning frameworks that are more widely used in the field.

The new Appendix D.2 is a step in this direction, but it does not provide direct comparisons with any of the models cited by the author in the section (L1951): RT-1, RT-2, and Octo. The added experiments are still restricted to offline RL, removing return tokens but keeping state and reward tokens, which fundamentally differs from behavior cloning. Most of the SoTA models in robotics have no access to reward and take only recent observations (e.g., 1–4 image frames) as input.

Overall, while I value the empirical contributions of this work for xLSTM, its claims on robotics are rather debatable. I therefore consider this work borderline and cannot champion it in its current form.

Thank you for your helpful feedback on our work. We are glad that you find that our empirical studies bring some useful practices to the community. We address your open points in the following.

When and why to use recurrent models? We agree with the reviewer that an additional discussion on LSTM-based versus Transformer-based models can be useful, and we added the discussion to our updated manuscript. Regarding the inference speed, we want to highlight that we do compare against FlashAttention in our experiments. While FlashAttention speeds up the computation of the self-attention, the computational complexity remains quadratic with the sequence length in terms of speed and memory. In contrast, for xLSTM and Mamba the complexity is linear. We observe benefits in inference speed (Figure 6) and importantly also in memory consumption (Figure 7). Consequently, we believe that xLSTM can be a more efficient alternative to Transformer, for example for applications on edge devices. Furthermore, xLSTM enables state tracking [1] via the sLSTM block, which Transformers and Mamba cannot. This property can be particularly useful for partially observable environments. To provide more insights on this, we conduct additional ablations on the ratio of mLSTM to sLSTM blocks in Appendix D3 of our updated manuscript. On the other hand, Transformers may be particularly effective for applications that require exact recall of memories [2]. Consequently, the “best” choice among architectures is a question the community is still exploring, and may depend on the task at hand.

Inference time comparisons: We agree with the reviewer that batch size 1 is commonly used for inference of real-world robotics models. For B=1, xLSTM is almost twice as fast as the Transformer, which makes use of highly optimized FlashAttention (Figure 6). Importantly, the inference speed and memory consumption of xLSTM does not change with the sequence length (Figure 7). Note that we conduct this comparison on a high-end data-center GPU with 40GB of RAM, and the Transformer goes OOM for larger context sizes. In contrast, applications on edge devices, like robotics, may have to deal with less powerful accelerators. Therefore, the recurrent backbone enables processing significantly longer sequence lengths without increased requirements in terms of speed and memory. This is important because this property enables handling high sample rates as prevalent in real-world applications, or multi-episodic context as used for in-context RL (see Lines 72-78). To highlight the benefits of longer context, we conduct additional ablation studies in Appendix D.1 of our updated manuscript.

Differences of recurrent vs. Transformer. We agree with the reviewer that sparse reward tasks are interesting to study the difference between recurrent backbones and Transformers. We do conduct preliminary experiments on Dark-Room, which exhibits sparse rewards and a partially-observable observation space. We find that xLSTM compares favorably to Transformers on Dark-Room (see Figure 4 and 17 of our updated manuscript). We also agree that additional experiments in this direction could be of value. While we consider them out-of-scope for this work, we aim to investigate this in future work. However, for a similar comparison of vanilla LSTM to the Transformer, we refer to [2]. Furthermore, we refer to the xLSTM publication [3] Figure 5, for a comparison of xLSTM and Transformer on associative recall tasks. We added a discussion on both references to our updated manuscript.

Our manuscript improved considerably by addressing and incorporating your comments: thank you.

[1] “The Illusion of State in State-Space Models”, ICML 2024
[2] “When Do Transformers Shine in RL? Decoupling Memory from Credit Assignment”, NeurIPS 2023