Большая рекуррентная модель действий: xLSTM обеспечивает быструю инференцию для задач робототехники

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

October 29, 2024

Авторы: Thomas Schmied, Thomas Adler, Vihang Patil, Maximilian Beck, Korbinian Pöppel, Johannes Brandstetter, Günter Klambauer, Razvan Pascanu, Sepp Hochreiter

cs.AI

Аннотация

В последние годы в области Обучения с Подкреплением (RL) наблюдается тенденция к использованию больших моделей действий, обученных в автономном режиме на масштабных наборах данных с помощью моделирования последовательностей. Существующие модели в основном основаны на архитектуре Трансформера, что приводит к созданию мощных агентов. Однако из-за медленного времени вывода подходы на основе Трансформера непрактичны для приложений в реальном времени, таких как робототехника. Недавно были предложены современные рекуррентные архитектуры, такие как xLSTM и Mamba, которые обладают преимуществами параллелизации во время обучения, аналогичными архитектуре Трансформера, при этом обеспечивая быстрый вывод. В данной работе мы изучаем пригодность этих современных рекуррентных архитектур для больших моделей действий. В результате мы предлагаем Большую Рекуррентную Модель Действий (LRAM) с xLSTM в ее основе, которая обладает сложностью вывода линейного времени и естественной способностью экстраполяции длины последовательности. Эксперименты на 432 задачах из 6 областей показывают, что LRAM сравнима с Трансформерами по показателям производительности и скорости.

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.