Ra mắt Open Dreamer: Bản tái hiện mã nguồn mở của mô hình thế giới Dreamer 4 bằng JAX/Flax

Một nhóm nhỏ các nhà nghiên cứu AI (Reactor) vừa phát hành Open Dreamer, một bản triển khai mã nguồn mở của pipeline mô hình thế giới Dreamer 4, được viết bằng JAX và Flax NNX.

Những gì đã thực sự được phát hành

Hai kho lưu trữ đã được công bố. next-state/open-dreamer chứa pipeline huấn luyện: một bộ mã hóa video nhân quả (causal video tokenizer), một mô hình động lực học tiềm ẩn có điều kiện hành động, khả năng tạo rollout và chấm điểm FVD. reactor-team/open-dreamer chứa một bộ công cụ rollout cục bộ tối giản giúp tạo các khung hình từ tệp MP4 và tệp hành động tương ứng.

Sản phẩm thứ ba là bản demo trên trình duyệt được lưu trữ trên Reactor runtime. Nó truyền phát một thế giới Minecraft được tạo ra theo thời gian thực và cung cấp nút chuyển đổi Game ⟷ Dream, cho phép chuyển luồng từ trò chơi thực sang mô hình thế giới theo từng khung hình.

Mục tiêu đã nêu là tái tạo nghiên cứu Dreamer 4. Nhóm nghiên cứu đã cố tình tránh các phương pháp nằm ngoài bài báo nghiên cứu đó để giữ cho không gian tìm kiếm được thu hẹp. Họ bắt đầu với CoinRun, một trò chơi platformer 2D được tạo theo thủ tục có thể huấn luyện trên một GPU duy nhất, sau đó mở rộng pipeline hoạt động sang video lối chơi theo phong cách Minecraft/VPT.

Kiến trúc: một backbone, hai mô hình

Cả bộ mã hóa (tokenizer) và mô hình động lực học đều sử dụng chung một backbone transformer khối-nhân quả (block-causal). Backbone đó luân phiên giữa hai loại chú ý (attention). Các lớp không gian (space layers) truyền thông tin giữa các phần tử của một khung hình đơn lẻ. Các lớp thời gian nhân quả (causal time layers) truyền thông tin giữa các khung hình.

Bộ mã hóa là một Masked Autoencoder dựa trên transformer thay vì VAE. Nhóm nghiên cứu báo cáo mức nén khoảng 100 lần và lưu ý rằng thiết kế này không cần KL loss hay adversarial loss. Họ lập luận rằng việc che (masking) giúp không gian tiềm ẩn dễ khuếch tán hơn.

Mô hình động lực học thực hiện dự đoán khung hình tiếp theo và được huấn luyện bằng diffusion forcing, flow matching và các mô hình đường tắt (shortcut models). Nó cũng dự đoán hành động tiếp theo. Thay vì luân phiên giữa một mô-đun chuyển đổi và chính sách riêng biệt, quá trình rollout được gộp vào các khối theo từng bước thời gian gồm (hành động trước đó, trạng thái, chính sách). Chú ý không gian chạy bên trong mỗi khối; chú ý thời gian nhân quả kết nối các khối theo thời gian.

Quan trọng là, các token của mô hình thế giới không thể đọc được token của tác nhân (agent token). Do đó, thông tin về nhiệm vụ và chính sách chỉ có thể ảnh hưởng đến các trạng thái tương lai thông qua hành động tiếp theo.

Công thức huấn luyện, theo cấu hình

Các cấu hình Minecraft được phát hành giúp cụ thể hóa công thức này.

Mô hình động lực học có 1,6 tỷ tham số: 30 lớp block-causal, d_model 1920, 30 đầu chú ý và 3 đầu KV cho grouped-query attention. Cứ mỗi bốn lớp lại có một lớp chú ý thời gian. Mỗi bước thời gian mang 32 token đăng ký (register tokens) đã học, và packing_factor: 2 đóng gói các latents của bộ mã hóa lân cận vào mỗi token không gian của mô hình động lực học. Chú ý thời gian sử dụng cửa sổ trượt 192 bước.

Quá trình huấn luyện chạy trong 200.000 bước với Muon, lịch trình WSD và tốc độ học cực đại là 3e-4. Các mẫu shortcut/bootstrap được bật ở bước 100.000 với tỷ lệ batch là 0,25. Độ suy giảm EMA là 0,999.

Cấu hình bộ mã hóa phát ra 512 token tiềm ẩn mỗi khung hình với độ rộng nút thắt (bottleneck) là 16. Các khung hình thô 360×640 được đệm thành 368×640 để cả hai chiều đều chia hết cho các patch 16×16. Độ sâu bộ mã hóa là 12 tại d_model 1536; độ sâu bộ giải mã là 8 tại d_model 1024. Xác suất che MAE đạt tối đa 0,9 và LPIPS được áp dụng với trọng số 0,2 trên một nửa số bước thời gian.

Các hành động VPT được phân tích thành 27 kênh hành động nhị phân cộng với 121 lớp chuột phân loại, không có kênh liên tục.

Tối đa hóa tính toán và bức tường bộ nhớ

Nhóm nghiên cứu báo cáo hiệu suất sử dụng FLOPs của mô hình đạt 57–58%, so với mức chuẩn 60% cho việc huấn luyện transformer ổn định. Lý do là lập luận về đường mái (roofline). Trên B200, điểm giao thoa giữa giới hạn băng thông và giới hạn tính toán nằm ở mức 292 FLOP/byte. Việc nạp 256 khung hình mỗi GPU đẩy khối lượng công việc vượt qua điểm giới hạn đó.

Việc phân mảnh (sharding) đi ngược lại với kỳ vọng. Với 1,6 tỷ tham số, trạng thái mô hình — tham số, gradient, trạng thái bộ tối ưu và EMA — chiếm khoảng 24 GiB, vừa vặn trên một GPU B200. Các kích hoạt (activations) mới là chi phí thực sự. Nhóm nghiên cứu đã thử song song hóa dữ liệu, FSDP, song song hóa tensor và song song hóa chuỗi, sau đó quyết định sử dụng song song hóa dữ liệu đơn thuần kết hợp với checkpointing kích hoạt.

Việc nạp dữ liệu được giải quyết bằng cách tiền mã hóa toàn bộ tập dữ liệu thành các tệp.arrayrecord, sau đó sử dụng Grain với bộ đệm nạp trước phía GPU. Việc giải mã bằng ffmpeg không đủ nhanh để cung cấp dữ liệu cho các GPU.

Phần về độ ổn định mới là giá trị cốt lõi

Nhóm nghiên cứu khẳng định rõ ràng rằng độ ổn định chiếm phần lớn thời gian của họ. Quan sát chính của họ: hầu hết các vấn đề về độ ổn định xảy ra mặc dù hàm mất mát (loss) vẫn giảm. MSE cải thiện một cách mượt mà trong khi chất lượng tạo hình lại suy giảm.

Sáu bản sửa lỗi đã được ghi lại. Muon thay thế LaProp, vốn thường xuyên gây ra các đột biến ngẫu nhiên, qua hai lần chạy, mỗi lần khoảng 400 giờ trên B200. Trọng số EMA được coi là bắt buộc đối với suy luận khuếch tán. Độ chính xác hỗn hợp (mixed precision) rất nhạy cảm với ranh giới: các tham số giữ ở float32, BF16 bao phủ hầu hết các kích hoạt matmul và đầu vào chú ý, còn float32 được giữ cho chuẩn hóa và đầu ra flow của mô hình động lực học.

Về trọng số loss, họ sử dụng x-prediction với v-space loss, giúp giảm xuống một số hạng trọng số tương tự như của Dreamer 4 nhưng với mẫu số bình phương. Họ báo cáo một sự cải thiện nhỏ nhưng đáng kể. Vận chuyển tối ưu barycentric theo minibatch giữa các chuỗi nhiễu và chuỗi tiềm ẩn giúp việc tạo rollout ổn định hơn. μ-parametrization đã được thử nghiệm và đánh giá là không cần thiết, một phần vì Muon giữ các siêu tham số ổn định hơn trên các quy mô mô hình khác nhau.

Một kết quả khác từ giai đoạn CoinRun: một đợt quét iso-FLOPs cho thấy quy mô tối ưu tính toán ở mức khoảng N∝C0.56 và D∝C0.44.

Những gì không có trong gói phát hành

Kho lưu trữ không bao gồm vòng lặp huấn luyện bắt chước hành vi (BC) hoặc RL; một vòng lặp tác nhân BC/RL đầy đủ của Dreamer 4 được liệt kê là một mục trong lộ trình phát triển mở. Công việc về chính sách CoinRun được mô tả trong bài đăng đã không được sử dụng cho Minecraft và không được phát hành.

Bài đăng cũng không công bố điểm FVD, mặc dù các tập lệnh scripts/eval_fvd.py được phát hành kèm theo một bộ công cụ dựa trên I3D được cấu hình cho 4 khung hình ngữ cảnh và chân trời 240 khung hình.

Những điểm chính cần lưu ý

Hãy xem bài đăng trên blog và bản demo, kho lưu trữ huấn luyện, kho lưu trữ suy luận và Reactor trên X. Mọi công lao cho nghiên cứu này thuộc về các nhà nghiên cứu của dự án.

Asif Razzaq là CEO của Marktechpost Media Inc.. Là một doanh nhân và kỹ sư có tầm nhìn, Asif cam kết khai thác tiềm năng của Trí tuệ nhân tạo vì lợi ích xã hội. Nỗ lực gần đây nhất của ông là ra mắt Nền tảng Truyền thông Trí tuệ nhân tạo, Marktechpost, nổi bật với việc đưa tin chuyên sâu về tin tức học máy và học sâu, vừa đảm bảo tính kỹ thuật vừa dễ hiểu đối với đông đảo khán giả. Nền tảng này tự hào với hơn 2 triệu lượt xem hàng tháng, minh chứng cho sự phổ biến của nó đối với độc giả.