Сравнение архитектур мировых моделей: RNN, Transformer и S4

Работа «Facing Off World Model Backbones: RNNs, Transformers, and S4» (Deng, Park, Ahn, 2023) посвящена выбору наиболее эффективной архитектуры последовательностной модели для латентных мировых моделей в обучении с подкреплением (model-based RL). Авторы предлагают S4WM — первую мировую модель, использующую структурированные слои пространства состояний (Structured State Space Sequence, S4) для высокоразмерного латентного воображения.

Связанные работы

Латентные мировые модели и model-based RL

Ближайшую основу составляют агенты с латентной динамикой, такие как PlaNet, Dreamer, DreamerV2 и DreamerV3. Они обучают компактные стохастические представления состояний и прогнозируют латентные будущие траектории для оптимизации политики. В статье рекуррентные модели типа RSSM (особенно Dreamer) рассматриваются как доминирующее baseline-семейство.

Рекуррентные нейронные сети

RNN, LSTM и GRU обеспечивают рекуррентное обновление состояний в реальном времени и долгое время были центральными в визуальном моделировании мира. Их главный недостаток — последовательное обучение во времени, что снижает возможность параллельной обработки и затрудняет работу с очень длинными контекстами.

Трансформеры в мировых моделях

Трансформерные модели динамики, включая TransDreamer, заменяют рекуррентность механизмом самовнимания, что позволяет распараллеливать обучение и напрямую обращаться к длинной истории. Однако квадратичная стоимость внимания от длины последовательности делает длинные воображаемые траектории и визуальное моделирование протяжённых контекстов вычислительно дорогими.