Open Dreamer: Dreamer 4 -putken avoimen lähdekoodin toteutus JAX/Flax-kehyksellä täydellä harjoitusreseptillä
Tutkijaryhmä julkaisi Open Dreamerin, Dreamer 4 -maailmamalliputken avoimen lähdekoodin toteutuksen, joka on rakennettu JAX- ja Flax NNX -kehyksillä. Projekti sisältää harjoitusputken, paikallisen päättelyvaljaat ja selaindemon, joka tuottaa Minecraft-maailman reaaliajassa. Malli käyttää lohkokausaalista transformattorin runkoa, saavuttaa 57-58 prosentin mallin FLOPs-käyttöasteen, ja ryhmä korostaa, että suurin haaste oli vakaus, ei suorituskyky.
Reactoriin sponsoroimat tutkijat Francesco Sacco, Diego Martí ja Edward Hu julkaisivat Open Dreamerin, avoimen toteutuksen Dreamer 4:stä JAX/Flax NNX:llä. Kaksi repositoriota julkaistiin: yksi harjoitukseen (tokenisaattori, dynamiikkamalli, rollaustiheys, FVD-pisteytys) ja yksi paikalliseen päättelyyn. Selaindemo suoratoistaa generoitua Minecraft-maailmaa ja mahdollistaa vaihtelun oikean pelin ja maailmamallin välillä. Arkkitehtuuri käyttää lohkokausaalista transformeri-runkoa sekä tokenisaattorissa että dynamiikkamallissa. Tokenisaattori on transformeri-pohjainen masked autoenkooderi, joka saavuttaa noin 100× pakkauksen ilman KL- tai adversariaalihäviötä. Dynamiikkamalli (1,6 miljardia parametria) suorittaa seuraavan kehyksen ennustamista käyttäen diffuusiopakottamista, virtausmatchingia ja pikakuvakemalleja, ja ennustaa myös seuraavan toiminnon. Harjoitus kestää 200 000 askelta Muon-optimointimenetelmällä, WSD-aikataululla, huippuoppimisnopeudella 3e-4 ja EMA-vaimennuksella 0,999. Tiimi saavutti 57–58 % mallin FLOP-käyttöasteen B200-GPU:illa, ja mallin tila mahtuu noin 24 Gigatavuun. Aktivaatiot olivat suurin muistikustannus, ratkaistuna datarinnakkaisuudella ja aktivointitarkistuspisteillä. Stabiilisuusongelmat olivat suurin haaste; kuusi korjausta ovat Muon-optimointimenetelmä, pakolliset EMA-painot, huolellinen sekalukutarkkuus, v-avaruushäviön painotus, minierän barysentrinen optimaalinen kuljetus ja μ-parametroinnin vähentäminen. Repositorio ei sisällä käyttäytymisklonointi- tai vahvistusoppimisen harjoitussilmukkaa, eikä FVD-pisteitä ole julkaistu.
Lähde: MarkTechPost —
Alkuperäinen