Která architektura je nejlepším základem pro world modely? Vědci porovnali RNN, Transformery a S4

Studie „Facing Off World Model Backbones: RNNs, Transformers, and S4“ (Deng, Park a Ahn, 2023) zkoumá, který typ sekvenčního modelu je nejefektivnější pro latentní world modely v modelově založeném posilovaném učení (model-based reinforcement learning, MBRL). Autoři představují model S4WM, který poprvé aplikuje strukturované stavově-prostorové sekvenční vrstvy (Structured State Space Sequence, S4) na vizuální latentní představivost ve vysokém rozlišení.

Související práce

• Latentní world modely / MBRL. Nejbližším základem jsou agenti s latentní dynamikou, jako PlaNet, Dreamer, DreamerV2 a DreamerV3. Tyto modely se učí kompaktní stochastické reprezentace stavů a predikují budoucí trajektorie v latentním prostoru pro optimalizaci politiky agenta. Studie staví na rodině rekurentních RSSM modelů, především na Dreameru, který je považován za dominantní baseline.

• Rekurentní neuronové sítě. RNN, LSTM a GRU modely dynamiky poskytují online rekurentní aktualizace stavů a jsou klíčové pro vizuální modelování světa. Jejich hlavním omezením je sekvenční trénink v čase, který snižuje možnost paralelizace a ztěžuje modelování velmi dlouhých kontextů.

• Transformerové world modely. Modely jako TransDreamer nahrazují rekurenci mechanismem pozornosti (self-attention), což umožňuje paralelní trénink a přímý přístup k dlouhé historii. Standardní kvadratická náročnost pozornosti v délce sekvence však činí dlouhé generování a modelování vizuálních kontextů výpočetně nákladným.

• Strukturované stavově-prostorové modely. S4 je model navržený pro závislosti na dlouhé vzdálenosti, který kombinuje rekurentní stavově-prostorovou interpretaci s konvoluční formou umožňující paralelní trénink sekvencí. Studie argumentuje, že S4 představuje slibný kompromis: lineárně časově náročnou inferenci a paralelizovatelný trénink.

• Odlišnost S4WM. Na rozdíl od použití S4 na konvenční benchmarky sekvencí, S4WM integruje S4 do latentního vizuálního world modelu s kódováním/dekódováním obrazu a stochastickou latentní dynamikou. Model je následně vyhodnocen v porovnání s RNN a Transformerovými základy v rámci společného rámce modelování světa. Uvedeným přínosem je první world model založený na S4, který generuje sekvence obrazů ve vysokém rozlišení pomocí latentní představivosti.

Stručné konstatování pro literární rešerši:

Dřívější world modely používaly převážně rekurentní latentní dynamiku nebo Transformery. RNN modely jsou efektivní při inferenci, ale trénují se sekvenčně, zatímco Transformerové world modely trénink paralelizují, ale mají kvadratickou časovou náročnost. S4WM zkoumá strukturované stavově-prostorové vrstvy jako alternativní základ, který si zachovává paralelní trénink a zároveň podporuje efektivní modelování dynamiky na dlouhé horizonty.