Un grand modèle d'action récurrente : xLSTM permet une inférence rapide pour les tâches de robotique.

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

October 29, 2024

Auteurs: Thomas Schmied, Thomas Adler, Vihang Patil, Maximilian Beck, Korbinian Pöppel, Johannes Brandstetter, Günter Klambauer, Razvan Pascanu, Sepp Hochreiter

cs.AI

Résumé

Ces dernières années, on observe une tendance dans le domaine de l'Apprentissage par Renforcement (RL) vers de grands modèles d'actions entraînés hors ligne sur des ensembles de données à grande échelle via la modélisation de séquences. Les modèles existants sont principalement basés sur l'architecture Transformer, ce qui donne des agents puissants. Cependant, en raison de temps d'inférence lents, les approches basées sur les Transformers sont impraticables pour des applications en temps réel, telles que la robotique. Récemment, des architectures récurrentes modernes, telles que xLSTM et Mamba, ont été proposées, présentant des avantages de parallélisation pendant l'entraînement similaires à l'architecture Transformer tout en offrant une inférence rapide. Dans ce travail, nous étudions l'aptitude de ces architectures récurrentes modernes pour de grands modèles d'actions. Par conséquent, nous proposons un Grand Modèle d'Action Récurrent (LRAM) avec un xLSTM à son cœur qui présente une complexité d'inférence en temps linéaire et des capacités d'extrapolation de longueur de séquence naturelle. Des expériences sur 432 tâches provenant de 6 domaines montrent que le LRAM se compare favorablement aux Transformers en termes de performances et de vitesse.

English

In recent years, there has been a trend in the field of Reinforcement Learning (RL) towards large action models trained offline on large-scale datasets via sequence modeling. Existing models are primarily based on the Transformer architecture, which result in powerful agents. However, due to slow inference times, Transformer-based approaches are impractical for real-time applications, such as robotics. Recently, modern recurrent architectures, such as xLSTM and Mamba, have been proposed that exhibit parallelization benefits during training similar to the Transformer architecture while offering fast inference. In this work, we study the aptitude of these modern recurrent architectures for large action models. Consequently, we propose a Large Recurrent Action Model (LRAM) with an xLSTM at its core that comes with linear-time inference complexity and natural sequence length extrapolation abilities. Experiments on 432 tasks from 6 domains show that LRAM compares favorably to Transformers in terms of performance and speed.