Який «хребет» найкращий для світових моделей? Порівняння RNN, Transformer та S4

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

Пов’язані роботи

• Латентні світові моделі / model-based RL. Найближчою основою є агенти з латентною динамікою, такі як PlaNet, Dreamer, DreamerV2 і DreamerV3. Вони навчають компактні стохастичні представлення станів і прогнозують майбутні латентні траєкторії для оптимізації політики. Стаття позиціонує рекурентні моделі типу RSSM, особливо Dreamer, як домінантну базову родину.

• Рекурентні нейронні мережі. Моделі динаміки на основі RNN, LSTM та GRU забезпечують онлайн-оновлення рекурентного стану і були ключовими для візуального моделювання світу. Їхній головний недолік — послідовне навчання в часі, що знижує паралелізм і ускладнює роботу з дуже довгими контекстами.

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

• Структуровані моделі простору станів. S4 — це структурована модель простору станів, призначена для довгострокових залежностей. Вона поєднує рекурентну інтерпретацію стану простору зі згортковою формою, що дозволяє паралельне навчання послідовностей. Автори стверджують, що це робить S4 перспективним «золотим серединою»: лінійна за часом рекурентна інференція та паралелізоване навчання.

• Унікальність S4WM. На відміну від використання S4 для стандартних бенчмарків послідовностей, S4WM вбудовує її в латентну візуальну світову модель із кодуванням/декодуванням зображень та стохастичною латентною динамікою, а потім порівнює з RNN та Transformer у спільному каркасі моделювання світу. Заявлений внесок — перша S4-світова модель для генерації високорозмірних послідовностей зображень через латентну уяву.

Стисле позиціонування для огляду літератури:

Попередні світові моделі переважно використовували рекурентну латентну динаміку або Transformers. Моделі на основі RNN ефективні під час інференції, але тренуються послідовно, тоді як Transformer-моделі паралелізують навчання, але мають квадратичну часову вартість. S4WM досліджує структуровані шари простору станів як альтернативний хребет, що має зберегти паралельне навчання, підтримуючи ефективне моделювання довгострокової динаміки.