Open Dreamer: Open-Source-Implementierung der Dreamer-4-Pipeline in JAX/Flax mit vollständigem Trainingsrezept
Ein Team von KI-Forschern hat Open Dreamer veröffentlicht, eine Open-Source-Implementierung der Dreamer-4-Weltmodell-Pipeline, die mit JAX und Flax NNX erstellt wurde. Das Projekt umfasst eine Trainingspipeline, eine lokale Inferenzumgebung und eine Browser-Demo, die eine Minecraft-Welt in Echtzeit generiert. Das Modell verwendet einen Block-Causal-Transformer-Backbone, erreicht eine Modell-FLOPs-Auslastung von 57–58 %, und das Team betont, dass Stabilität die größte Herausforderung war, nicht der Durchsatz.
Die Forscher Francesco Sacco, Diego Martí und Edward Hu, gefördert von Reactor, haben Open Dreamer veröffentlicht, eine offene Implementierung von Dreamer 4 in JAX/Flax NNX. Es wurden zwei Repositorys veröffentlicht: eines für das Training (Tokenisierer, Dynamikmodell, Rollout, FVD-Bewertung) und eines für die lokale Inferenz. Eine Browser-Demo streamt eine generierte Minecraft-Welt und ermöglicht das Umschalten zwischen echtem Spiel und Weltmodell. Die Architektur verwendet einen block-kausalen Transformer-Backbone sowohl für den Tokenisierer als auch für das Dynamikmodell. Der Tokenisierer ist ein transformerbasierter Maskierter Autoencoder, der eine etwa 100-fache Kompression ohne KL- oder adversarialen Verlust erreicht. Das Dynamikmodell (1,6 Milliarden Parameter) führt eine nächste-Frame-Vorhersage mittels Diffusion Forcing, Flow Matching und Shortcut-Modellen durch und sagt auch die nächste Aktion voraus. Das Training läuft über 200.000 Schritte mit dem Muon-Optimierer, einem WSD-Schedule, einer Spitzen-Lernrate von 3e-4 und einem EMA-Zerfall von 0,999. Das Team erreichte eine Modell-FLOPs-Auslastung von 57-58 % auf B200-GPUs, wobei der Modellzustand in etwa 24 GiB passt. Aktivierungen waren der Hauptspeicherkostenfaktor, gelöst durch Datenparallelität und Aktivierungs-Checkpointing. Stabilitätsprobleme waren die größte Herausforderung; sechs Korrekturen umfassen den Muon-Optimierer, obligatorische EMA-Gewichte, sorgfältige gemischte Präzision, v-Raum-Verlustgewichtung, Minibatch-baryzentrischen optimalen Transport und die Entschärfung der μ-Parametrisierung. Das Repository enthält weder den Behavior Cloning- noch den RL-Trainingsloop, und FVD-Scores wurden nicht veröffentlicht.
Quelle: MarkTechPost —
Original