Open Dreamer: мировая модель на 1,6 млрд параметров генерирует Minecraft в браузере

Open Dreamer появился как попытка небольшой исследовательской группы Reactor воспроизвести результаты статьи Dreamer 4. Команда сознательно не добавляла ничего за пределами оригинальной работы, чтобы сузить пространство поиска. Исходный код опубликован в двух репозиториях, а браузерная демка размещена на инфраструктуре Reactor.

Что на самом деле делает Open Dreamer?

Мировая модель (world model) работает так: она «смотрит» несколько кадров видеоигры, запоминает действия игрока и генерирует следующие кадры, как если бы игра продолжалась. Игрок нажимает «вперёд», модель рисует, что будет дальше, без запуска настоящего игрового движка.

Конвейер Dreamer, мировая модель такого типа, состоит из двух частей. Токенизатор сжимает видеокадры в компактное представление (латентное пространство). Модель динамики берёт это сжатое представление и предсказывает следующий кадр с учётом действия игрока.

Обе части построены на одном и том же блочно-каузальном трансформере. Внутри него чередуются два типа внимания: пространственное (связывает элементы одного кадра между собой) и временное (связывает кадры друг с другом по порядку). Такая конструкция позволяет модели одновременно «понимать», что происходит на кадре, и помнить, что было раньше.

Токенизатор реализован не как VAE (вариационный автокодировщик, привычный генератор со случайным слоем), а как маскированный автокодировщик на трансформере. Команда отмечает, что такой подход не требует специальных потерь для регуляризации и даёт примерно стократное сжатие.

В браузерной демке есть переключатель «Game / Dream»: он в реальном времени передаёт поток от настоящей игры к мировой модели и обратно, кадр за кадром.

Что обнаружили: шесть уроков о стабильности

Команда прямо говорит, что стабильность обучения заняла больше всего времени. Главное наблюдение: большинство проблем возникают, пока функция потерь MSE (среднеквадратичная ошибка, базовая метрика того, насколько модель ошибается) продолжает снижаться. Числа улучшаются, а качество генерации падает.

Вот что зафиксировала команда:

• Muon вместо LaProp. Предыдущий оптимизатор LaProp давал случайные всплески потерь, которые становились всё чаще. Два запуска по примерно 400 часов на B200 каждый были потрачены впустую, прежде чем команда перешла на Muon.
• EMA-веса обязательны для инференса (инференс, это этап, когда обученная модель генерирует результат). Без экспоненциального скользящего среднего весов качество генерации нестабильно. Коэффициент затухания EMA выставлен на 0,999.
• Смешанная точность чувствительна к границам. Параметры хранятся в float32, основные вычисления в BF16, но нормализация и выходной слой модели динамики остаются в float32, иначе генерация ломается.
• Взвешивание потерь через x-prediction с v-space loss. Это даёт весовой коэффициент, похожий на оригинальный Dreamer 4, но с квадратом в знаменателе. Команда отмечает небольшое, но заметное улучшение.
• Баридентрический оптимальный транспорт в мини-батчах (способ лучше сопоставлять шум и латентные последовательности) сделал генерацию роллаутов стабильнее.
• μ-параметризацию протестировали и отбросили как ненужную, отчасти потому, что Muon и так держит гиперпараметры устойчивыми при разных размерах модели.

Почему data parallelism победил?

При 1,6 млрд параметров состояние модели (параметры, градиенты, состояние оптимизатора и EMA) занимает около 24 ГиБ и умещается на одном B200. Команда попробовала FSDP (шардинг состояния модели по нескольким GPU), тензорный параллелизм и параллелизм по последовательностям. Но узким местом оказались не веса, а активации: промежуточные данные, которые модель хранит при обработке 256 кадров на GPU.

Простой data parallelism (данные разбиваются по GPU, каждый GPU хранит полную копию модели) плюс activation checkpointing (пересчёт активаций вместо хранения) оказался быстрее и проще. Утилизация вычислений при этом составила 57-58%, что близко к заявленному ориентиру в 60% для здорового обучения трансформеров.

Загрузку данных решили пре-токенизацией: весь датасет заранее превращается в файлы формата arrayrecord, а затем подаётся через Grain с буфером предзагрузки на стороне GPU. Декодирование через ffmpeg не успевало кормить карты данными.

Что делать с этим прямо сейчас?

Если вы инженер или исследователь: конфигурации обучения опубликованы с точностью до гиперпараметров. 1,6 млрд параметров, 30 слоёв, 200 000 шагов, Muon с WSD-расписанием, конкретные настройки маскирования и LPIPS (метрика перцептуального качества, насколько картинка похожа на оригинал для человеческого глаза) с весом 0,2. Это позволяет воспроизвести обучение, а не гадать по статье. Особенно полезен урок про data parallelism: при модели до 24 ГиБ сложные стратегии шардинга не дали выигрыша.

Если вы автор или маркетолог: Dreamer, мировая модель, пока применима к играм. Но сам принцип (модель предсказывает, что будет, по действиям пользователя) через год-два дойдёт до генерации интерактивного контента и симуляций. Следить за этим направлением стоит уже сейчас.

Если вы работаете в РФ: JAX-код открыт, запустить можно на своих GPU. Из российских аналогов мировых моделей публичных открытых реализаций сопоставимого уровня пока нет, но инфраструктура JAX доступна.

Команда Reactor выложила рабочий конвейер с конкретными конфигами и шестью задокументированными граблями. Агентный цикл пока не реализован, FVD-скоры не опубликованы, но инженерная база для воспроизведения Dreamer 4 теперь доступна любому, у кого есть GPU и желание разобраться.

По данным Reactor (Open Dreamer)