پروژه Open Dreamer: بازسازی معماری Dreamer 4 با JAX و Flax

اگر قصد دارید مدل‌هایی بسازید که بتوانند دنیای اطراف خود را شبیه‌سازی کنند، احتمالاً با دیوار بلند «ناپایداری آموزش» برخورد کرده‌اید. در حالی که مدل‌های جهانی با دقت بالا اغلب در تئوری مورد بحث قرار می‌گیرند، اما دستور پخت‌های واقعی آموزش که برای پایدار نگه داشتن آن‌ها در مقیاس بزرگ لازم است، همچنان دست‌نیافتنی بوده‌اند. تیم پژوهشی Reactor با انتشار Open Dreamer، دقیقاً همان دستور پخت مهندسی را منتشر کرده است که برای عبور از این دیوار لازم است؛ این پروژه در واقع یک بازسازی کامل و متن‌باز از خط لوله مدل جهانی Dreamer 4 است.

مدل‌های جهانی (World Models) — شبیه به تخیلات یک انسان که می‌تواند پیش از انجام کاری، نتیجه آن را در ذهن خود پیش‌بینی کند — تلاش می‌کنند فیزیک و منطق یک محیط را شبیه‌سازی کنند تا عامل‌های هوش مصنوعی (AI Agents) بتوانند بدون تعامل مستقیم با دنیای واقعی، آینده را «رویاپردازی» یا پیش‌بینی کنند. طبق مستندات این پروژه، در حالی که پژوهش‌های اولیه Dreamer 4 تئوری‌های لازم را تثبیت کرده بود، اما مهندسی عملی — به‌ویژه نحوه جلوگیری از «فروپاشی مدل» (Model Collapse) در طول آموزش — تا پیش از این انتشار، یک جعبه سیاه باقی مانده بود.

همان‌طور که در تحلیل‌های پیشین ما درباره‌ی امنیت و پایداری مدل‌های بازمتن اشاره کردیم، فاصله میان مقاله پژوهشی و استقرار واقعی همیشه در جزئیات مهندسی است. تیم Reactor برای اطمینان از بازسازی دقیق، بر روی یک فضای جستجوی محدود تمرکز کرد. آن‌ها ابتدا این خط لوله را روی بازی دوبعدی CoinRun آزمایش کردند تا صحت آن را تایید کنند و سپس آن را برای پردازش ویدیوهای بازی Minecraft مقیاس‌بندی نمودند.

معماری فنی

این سامانه بر پایه یک ترنسفورمر (Transformer) با ساختار block-causal مشترک برای هر دو بخش توکن‌ساز و مدل دینامیک بنا شده است. مشخصات کلیدی عبارتند از:

• مدل دینامیک: ۱.۶ میلیارد پارامتر در ۳۰ لایه block-causal با d_model برابر ۱۹۲۰.
• توکن‌ساز: یک خودرمزگذار ماسک‌شده (Masked Autoencoder) مبتنی بر ترنسفورمر که بدون نیاز به زیان‌های KL یا زیان‌های متخاصم (Adversarial Loss)، به فشرده‌سازی تقریباً ۱۰۰ برابری دست می‌یابد.
• مکانیزم توجه: ترکیبی از لایه‌های مکانی (Space Layers) برای داده‌های درون‌فریمی و لایه‌های زمانی علی (Causal Time Layers) برای انتشار بین‌فریمی، که از یک پنجره لغزان ۱۹۲ مرحله‌ای استفاده می‌کند.

بهینه‌سازی محاسبات و حافظه

به گزارش تیم Reactor، آموزش این مدل ۱.۶ میلیارد پارامتری روی واحد پردازش گرافیکی (GPU) مدل NVIDIA B200 نشان داد که گلوگاه اصلی حافظه، نه وضعیت مدل (Model State)، بلکه «فعال‌سازها» (Activations) هستند. وضعیت مدل حدود ۲۴ گیگابایت فضا اشغال می‌کرد، اما تیم مجبور شد از موازی‌سازی داده‌های ساده (Plain Data Parallelism) در کنار «نقطه-گذاری فعال‌سازها» (Activation Checkpointing) استفاده کند تا بار حافظه مدیریت شود.

برای اینکه پردازشگرهای گرافیکی همواره با داده تغذیه شوند، تیم Reactor رمزگشایی کند ffmpeg را دور زده و کل مجموعه داده را پیش‌از-اینکه توکن‌بندی شود، به فایل‌های .arrayrecord تبدیل کرد. این استراتژی باعث شد بهره‌وری عملیات اعشاری مدل (MFU) به ۵۷ تا ۵۸ درصد برسد که بسیار نزدیک به استاندارد ۶۰ درصدی برای آموزش‌های سالم ترنسفورمر است.

نبرد برای پایداری

تیم Reactor صراحتاً اعلام کرده است که پایداری بزرگ‌ترین چالش آن‌ها بوده است؛ چرا که منحنی‌های تابع زیان (Loss Function) اغلب گمراه‌کننده بودند و در حالی که MSE بهبود می‌یافت، کیفیت تولید مدل در واقع کاهش می‌کرد. برای حل این مشکل و جلوگیری از واگرایی، ۶ مداخله فنی خاص اجرا شد:

۱. جایگزینی LaProp با Muon برای توقف جهش‌های تصادفی در طول اجراهای ۴۰۰ ساعته روی B200.
۲. اجباری کردن استفاده از وزن‌های EMA برای استنتاج (Inference) مدل انتشار.
۳. استفاده از استراتژی دقت ترکیبی (Mixed Precision) حساس به مرز (استفاده از float32 برای نرمال‌سازی و خروجی جریان یا Flow Output).
۴. اعمال x-prediction با استفاده از یک زیان v-space.
۵. پیاده‌سازی انتقال بهینه باری‌سنتریک (Barycentric Optimal Transport) در سطح مینی‌بچ بین نویز و توالی‌های نهان.
۶. استفاده از زمان‌بندی WSD با نرخ یادگیری حداکثری 3e-4 در طول ۲۰۰,۰۰۰ گام.

این تلاش مهندسی، بحث را از معماری تئوریک به «ماده تاریک» آموزش هوش مصنوعی منتقل می‌کند: انتخاب‌های دقیق بهینه‌ساز و مرزهای دقت که مانع از واگرایی مدل‌های چند میلیارد پارامتری می‌شوند. برای توسعه‌دهندگان، این بدان معناست که سد ورود به پژوهش‌های مدل جهانی فرو ریخته است. دیگر نیازی نیست برای یک خط لوله به سبک Dreamer، ابرپارامترها را حدس بزنید؛ زیرا پیکربندی‌های دقیق Minecraft اکنون عمومی شده‌اند.

گام بعدی شما

• پژوهشگران علاقه‌مند می‌توانند خط لوله آموزشی را در مخزن next-state/open-dreamer بررسی کنند.
• تست استریم بلادرنگ مدل جهانی Minecraft را از طریق دموی runtime Reactor تجربه کنید.
• ابرپارامترهای منتشر شده را برای پیاده‌سازی مدل‌های پیش‌بینی‌کننده در محیط‌های شبیه‌سازی شده به کار ببرید.

اما داستان سخت‌افزاری این تحول حتی شگفت‌انگیزتر است — به تحلیل ما درباره‌ی تراشه‌های Blackwell و بهینه‌سازی‌های حافظه در مقیاس بزرگ مراجعه کنید.