Um Modelo de Ação Recorrente de Grande Escala: xLSTM permite Inferência Rápida para Tarefas de Robótica

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

October 29, 2024

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

cs.AI

Resumo

Nos últimos anos, tem havido uma tendência no campo do Aprendizado por Reforço (RL) em direção a modelos de ação extensos treinados offline em conjuntos de dados em larga escala por meio de modelagem de sequência. Os modelos existentes são principalmente baseados na arquitetura Transformer, resultando em agentes poderosos. No entanto, devido aos tempos de inferência lentos, abordagens baseadas em Transformer são impraticáveis para aplicações em tempo real, como robótica. Recentemente, arquiteturas recorrentes modernas, como xLSTM e Mamba, foram propostas, apresentando benefícios de paralelização durante o treinamento semelhantes à arquitetura Transformer, enquanto oferecem inferência rápida. Neste trabalho, estudamos a aptidão dessas arquiteturas recorrentes modernas para modelos de ação extensos. Consequentemente, propomos um Modelo de Ação Recorrente Extenso (LRAM) com um xLSTM em seu núcleo que possui complexidade de inferência em tempo linear e habilidades de extrapolação de comprimento de sequência natural. Experimentos em 432 tarefas de 6 domínios mostram que o LRAM se compara favoravelmente aos Transformers em termos de desempenho e velocidade.

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.