Open Dreamer: Dreamer 4 Boru Hattının JAX/Flax ile Açık Kaynak Uygulaması ve Tam Eğitim Tarifi

Reactor sponsorluğunda araştırmacılar Francesco Sacco, Diego Martí ve Edward Hu, Dreamer 4'ün JAX/Flax NNX'te açık bir uygulaması olan Open Dreamer'ı yayınladılar. Biri eğitim (tokenizer, dinamik model, rollout, FVD skorlama) ve biri yerel çıkarım için olmak üzere iki depo yayınlandı. Bir tarayıcı demosu, oluşturulmuş bir Minecraft dünyasını akışla aktarıyor ve gerçek oyun ile dünya modeli arasında geçiş yapmaya izin veriyor. Mimari, hem tokenizer hem de dinamik model için blok-nedensel bir transformer omurgası kullanıyor. Tokenizer, transformatör tabanlı bir Maskeli Otomatik Kodlayıcı (Masked Autoencoder) olup, KL veya çekişmeli kayıp (adversarial loss) olmadan yaklaşık 100 kat sıkıştırma sağlıyor. Dinamik model (1,6 milyar parametre), difüzyon zorlaması (diffusion forcing), akış eşleme (flow matching) ve kestirme modelleri (shortcut models) kullanarak bir sonraki kareyi tahmin ediyor ve ayrıca bir sonraki eylemi de tahmin ediyor. Eğitim, Muon optimizer, WSD çizelgesi, 3e-4 tepe öğrenme oranı ve 0,999 EMA sönümü (EMA decay) ile 200.000 adım boyunca çalışıyor. Ekip, B200 GPU'larında model durumunun yaklaşık 24 GiB'ye sığmasıyla %57-58 model FLOPS kullanımı elde etti. Ana bellek maliyeti olan aktivasyonlar, veri paralelliği ve aktivasyon denetim noktalaması (activation checkpointing) ile çözüldü. En büyük zorluk kararlılık sorunlarıydı; altı düzeltme arasında Muon optimizer, zorunlu EMA ağırlıkları, dikkatli karma duyarlılık (mixed precision), v-uzayı kayıp ağırlıklandırması (v-space loss weighting), minigrup barycentric optimal taşıma (minibatch barycentric optimal transport) ve μ-parametrelendirmenin (μ-parametrization) önemini azaltma yer alıyor. Depo, davranış kopyalama (behavior cloning) veya pekiştirmeli öğrenme (RL) eğitim döngüsünü içermiyor ve FVD skorları yayınlanmamıştır.