Open Dreamer: Triển khai mã nguồn mở Dreamer 4 Pipeline bằng JAX/Flax với công thức huấn luyện đầy đủ
Một nhóm các nhà nghiên cứu AI đã phát hành Open Dreamer, một triển khai mã nguồn mở của pipeline mô hình thế giới Dreamer 4 được xây dựng bằng JAX và Flax NNX. Dự án bao gồm pipeline huấn luyện, khung suy luận cục bộ và bản demo trên trình duyệt tạo thế giới Minecraft theo thời gian thực. Mô hình sử dụng backbone transformer block-causal, đạt 57–58% mức sử dụng FLOPs của mô hình, và nhóm nghiên cứu nhấn mạnh rằng độ ổn định là thách thức chính, chứ không phải thông lượng.
Các nhà nghiên cứu Francesco Sacco, Diego Martí và Edward Hu, được tài trợ bởi Reactor, đã phát hành Open Dreamer, một triển khai mã nguồn mở của Dreamer 4 trong JAX/Flax NNX. Hai kho lưu trữ đã được phát hành: một cho huấn luyện (tokenizer, mô hình động lực học, rollout, chấm điểm FVD) và một cho suy luận cục bộ. Một bản demo trên trình duyệt phát trực tiếp một thế giới Minecraft được tạo ra và cho phép chuyển đổi giữa trò chơi thực và mô hình thế giới. Kiến trúc sử dụng một backbone transformer nhân quả theo khối (block-causal) cho cả tokenizer và mô hình động lực học. Tokenizer là một Masked Autoencoder dựa trên transformer đạt được nén khoảng 100 lần mà không cần mất mát KL hay mất mát đối nghịch. Mô hình động lực học (1,6 tỷ tham số) thực hiện dự đoán khung hình tiếp theo bằng cách sử dụng diffusion forcing, flow matching và shortcut models, đồng thời cũng dự đoán hành động tiếp theo. Huấn luyện chạy 200.000 bước với bộ tối ưu Muon, lịch trình WSD (warmup-stable-decay), tỷ lệ học đỉnh 3e-4 và hệ số suy giảm EMA (exponential moving average) 0,999. Nhóm đã đạt được mức sử dụng FLOPs (floating-point operations per second) của mô hình từ 57-58% trên GPU B200, với trạng thái mô hình phù hợp trong khoảng 24 GiB. Các kích hoạt (activations) là chi phí bộ nhớ chính, được giải quyết bằng song song dữ liệu và điểm kiểm tra kích hoạt (activation checkpointing). Các vấn đề về ổn định là thách thức lớn nhất; sáu bản sửa lỗi bao gồm bộ tối ưu Muon, trọng số EMA bắt buộc, độ chính xác hỗn hợp cẩn thận, trọng số mất mát v-space (velocity space), vận chuyển tối ưu barycentric theo minibatch (minibatch barycentric optimal transport) và giảm nhấn mạnh μ-parametrization. Kho lưu trữ không bao gồm vòng lặp huấn luyện bắt chước hành vi (behavior cloning) hay học tăng cường (RL), và điểm FVD (Fréchet Video Distance) chưa được công bố.
Nguồn: MarkTechPost —
bản gốc