Ein großes wiederkehrendes Aktionsmodell: xLSTM ermöglicht schnelle Inferenz für Robotikaufgaben.

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

October 29, 2024

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

cs.AI

Zusammenfassung

In den letzten Jahren gab es im Bereich des Reinforcement Learning (RL) einen Trend zu großen Aktionsmodellen, die offline auf umfangreichen Datensätzen über Sequenzmodellierung trainiert werden. Die bestehenden Modelle basieren hauptsächlich auf der Transformer-Architektur, die leistungsstarke Agenten hervorbringt. Aufgrund langsamer Inferenzzeiten sind Transformer-basierte Ansätze jedoch für Echtzeitanwendungen wie Robotik ungeeignet. In letzter Zeit wurden moderne rekurrente Architekturen wie xLSTM und Mamba vorgeschlagen, die ähnliche Parallelisierungsvorteile während des Trainings wie die Transformer-Architektur aufweisen und schnelle Inferenz ermöglichen. In dieser Arbeit untersuchen wir die Eignung dieser modernen rekurrenten Architekturen für große Aktionsmodelle. Daraus resultierend schlagen wir ein Large Recurrent Action Model (LRAM) mit einem xLSTM im Kern vor, das eine Inferenzkomplexität in linearer Zeit und natürliche Sequenzlängenextrapolationsfähigkeiten bietet. Experimente an 432 Aufgaben aus 6 Domänen zeigen, dass LRAM in Bezug auf Leistung und Geschwindigkeit vorteilhaft im Vergleich zu Transformers abschneidet.

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.