Loading video...
Video Failed to Load
I spent the past month reimplementing DeepMind’s Genie 3 world model from scratch Ended up making TinyWorlds, a 3M parameter world model capable of generating playable game environments demo below + everything I learned in thread (full repo at the end)👇🏼
220,935 views • 1 year ago •via X (Twitter)
86 Comments

1/ understanding world models World models are neural networks that simulate physical worlds by generating videos. DeepMind’s Genie 3 proved that, just like LLMs, scaled-up world models exhibit emergent behavior: > Controllability: Pressing the right arrow makes the camera pan right. > Consistency: Fresh paint stays on the walls if you leave the room and come back > Quality: You can see your reflection in a puddle. How did DeepMind scale Genie 3 to unlock these abilities? Genie 3’s architecture is closed-source, but I found the key in an older paper: Genie 1. Before Genie, scientists thought that scaling world models meant using data of actions or 3D object structures. DeepMind realized that scaling raw video data is all you need. Desired behaviors like consistency naturally emerge from a massive world model, just like behaviors like grammar and syntax naturally emerge from massive LLMs. But world models need labeled actions for each frame (“when the player presses the right arrow key, move the camera right”). This means we can’t train on all the unlabeled videos on the internet, and we won’t have enough data to scale. Genie 1 bypassed this by training a separate model that infers action labels for each frame. This method is likely how DeepMind scaled training Genie to millions of hours of unlabeled YouTube videos and unlocked all the emergent capabilities mentioned above. I built TinyWorlds from scratch to help people understand how it works.

2/ building the dataset of worlds Before training TinyWorlds, I decided what video game worlds my model should generate by building the dataset. The set of worlds the model sees in training determines what worlds it generates. I created TinyWorlds' dataset by processing YouTube gameplay videos for: > Pong: the 2-player Atari game > Sonic: a 2D platformer > Zelda: birds-eye-view adventure game > Pole Position: 3D pixel racing game > Doom: 3D first-person shooter

3/ building the space-time transformer Normal transformers in LLMs understand language, which is 1D. TinyWorlds requires a model that understands video, which is 3D (height, width, time). This model also has to train quickly and learn using both actions and video. Space-time transformer understands video by handling space and time with 3 layers: > Spatial attention: tokens attend to all other tokens in the same frame. > Temporal attention: tokens attend to tokens in previous timesteps > Feedforward: tokens go through a vanilla neural network ST Transformer also lets actions influence video generation (so the user can interact with the world), either by: 1. Concatenating action and video representations 2. Using actions to scale and shift video representations When designing architectures, it’s nearly impossible to know what will be better ahead of time. So I always follow the rule ”An experiment is worth 1000 discussions”: I try both and see which works better. Since 2 performed the best, I used it for the final TinyWorlds ST Transformer. To make ST Transformer train quickly, I added innovations from LLM architecture: > SwiGLU: a modification of the feedforward network that speeds up learning > RMSNorm: makes learning more stable over long training runs > Position Encodings: tell each token where it is in the image

4/ designing the architecture Next, I adapted Genie's high-level architecture to TinyWorlds. I considered using either: > Diffusion: where we start with noise and slowly remove it until we have a completed video sequence. > Autoregression: where we predict small chunks of video one at a time until we have a full video. I chose Autoregression for TinyWorlds because it has: > faster inference: critical for real-time playable world models > faster training: hyper-optimized attention kernels make training faster > simplicity: diffusion is another layer of complexity aside from understanding world models I then built TinyWorlds using three modules that work together to make video prediction easy: > Video Tokenizer: turns video into small chunks that are easier to understand > Action Tokenizer: predicts the action taken between two frames > Dynamics Model: Using past video and action, predicts the next chunk of video

5/ quantizing video into tokens For the video and action tokenizers, we need a quantization method to produce tokens. Tokenizers represent videos as compressed data (like zip files). They learn by finding a set of small building blocks that makes reconstructing the video easy. Tinyworld’s method, Finite Scalar Quantization: > divides space into cubes > uses them to represent different chunks of each image FS Quantization is forced to make the tokens contain useful information from the original image because another neural network uses these tokens to reconstruct the original image.

6/ training the video tokenizer The first module, the video tokenizer, compresses videos into tokens using: > Convolutions to transform images into vectors representing each section of the image > ST transformer to let each vector share information > FS Quantization to turn the vectors into small video tokens Originally, my tokens didn’t compress the video small enough. To generate worlds, dynamics predicts future tokens, and the tokenizer converts them back to video. Large tokens make it much harder for dynamics to generate high-quality worlds. So I realized that the tokenizer should do almost all the compression. The dynamics model in video models should learn in the simplest space possible, because it has the most difficult job: predicting the future.

7/ training the action tokenizer The action tokenizer is the model that creates action labels and allows us to train on unlabeled video. > From raw video, it predicts the action that happened between two frames. > This lets dynamics learn to listen to actions without actually having action labels. > Most importantly, it means we can train on unlabeled video data like YouTube Originally, the tokenizer decoder ignored actions and failed to learn proper actions. To solve this, I: > Masked frames: this forced the image reconstructor to pay attention to actions > Added variance loss: this encouraged the encoder to label the video with all possible actions Ideally, each action token should correspond to a useful key like “right”, “left”, and “up”. At the small scale of tinyworlds, though, action tokens didn’t correspond to meaningful things. The best solution is to scale the model up. Alternatively, I could've used a few labels to supervise the action tokenizer on what to prioritize.

8/ training the world generator Lastly, I trained the dynamics model that predicts the next frame. > In training, it predicts masked tokens > In inference, we add masked frames and it autoregressively decodes them. When I first trained the dynamics model, loss plateaud early and the output videos were low-quality. I realized the dynamics model had to be much larger than the tokenizers because dynamics contains most of the intelligence of the world model. Once I scaled it up, the dynamics improved dramatically For inference, the user chooses an action token to condition dynamics on. This lets them interact with and control the generation of the world. So after I finished training, I played the tiny game worlds: > I raced on the roads of Pole Position > I walked in the landscapes of Zelda > I explored the 3D world of Doom I was impressed that such a small model generated decent videos and paid attention to actions, but generations were often blurry and inconsistent. To improve TinyWorlds, I'd scale it up to >100B parameters, train on more data, and try more architecture changes (see repo). Incorporating diffusion might also improve generations (and is likely what all frontier world models use). The Bitter Lesson strikes again :)

9/ Finally, here’s tinyworlds: It’s a minimal codebase to help people understand world modeling. Try it out yourself and make a PR, there are many easy + impactful additions to make (such as Mixture of Experts, Muon, and scaling). Thank you to @runpod @LukePiette for top-tier gpus and @MajmudarAdam, @suryasure05, @evanliin, and @karpathy for inspiring this project!

Hey! I lead one of the biggest open source world model groups, wanna find some time to chat? This is what we're building (and open sourcing)

Very cool!

Very cool! Any chance you could share the datasets or models on hf for the community?

It’s already shared :) thank you HF

Ts hella tuff, looking forward to diving in this weekend

Thank you goat

next step: train tinyworlds on tinytpu

Lmfaooo unironically

Incredible work. Fantastic writeup. It's phenomenal that you open sourced everything, including the data. What did you train on and how long did it take?

Thanks! Trained on the data I mentioned in the third post, u can download it via the repo The development took a long time, the final training runs took ~a day each in total depending on the complexity of the world

Which GPUs did you train it on? How many?

Final runs were 2-4 H200s, ~1 day each for 5 runs

awesome work!!

Thanks!

Wtf is in the majmudar gene @MajmudarAdam

incredible banger

inspiration

LETS GOOOO

fucking cracked

👉🏼👈🏼🫵🏼

🐐🐐🐐

🫵

Im sure this is not how Genie 3 work. You prob need a pre-trained few-step diffusion model and post train with self-forcing and language-based actions.

Yea definitely agree genie 3 does way more than genie 1 I mentioned in the post they very likely use diffusion transformer instead of autoregression I’d guess they initialize weights from VEO3, and also have a post post training dataset that is actually labeled with actions + only videogames

It’s not fully autoregressive. I believe the MLLM part is autoregressive to intake condition from language space, this includes camera and movement (yes, through language), and a few step diffusion process roll out the frames (that’s why you see latency).

Hmm might be true Autoregression also has latency, it just depends how many tokens u sample per step (similar to maskGIT inference if ur familiar)

Autoregressive video model is not scalable enough, considering the fact of pretraining being used, MLLM + diffusion adaptor is more likely

Wym, autoregression is way easier to train and everyone already has large scale AR infra because of LLMs I agree they probs use diffusion but it’s probs for generation quality not scalability

Video is a different modality.

Can still use the exact same infra (except dataloading) in theory

Great work!! Do check out Jasmine :)

Very cool Out of curiosity what’s the range of compute this took

iirc the runs were 2-4k gpuh each (on an A100 cluster)

Nice!

bro this is so sick

thank you king ur an inspiration

wait this is so cool wtf

thanks king

Cracked

🥹🙏

@Scobleizer It's amazing. All the different rabbit holes different people go down

@Scobleizer Indeed

Impressive. Genie3 people have publicly said they use an autoregressive model (not diffusion, or not purely diffusion).

Interesting, haven’t seen this before, where’d you hear this?

Pretty sure it was this MLStreetTalk podcast, but there was a similar one from Google DeepMind podcast also

Really surprised me, because AR has so many issues with long rollout often. But they make the point that you have to be AR, because you need to be reactive to the user history, you don't have a set fixed target.

Interesting, I guess both ways you’re autoregressive over time, even if you predict in chunks (which both can do)

So you have the same “longer rollouts go out of distribution” issue either way

holy shit

🙏

3 million parameters is crazy

maximal compression indeed

Some absolutely insane work, fire!!!

@AliesTaha Tysm urs too🙏🏼

knew you were goated from when i met you at the kscale event 🫡

@aliuahma deep cut lmao thank you

this is incredibly goated

ty🙌🏼

incredible

🥹

are you @MajmudarAdam's brother?

@MajmudarAdam I wonder

😂

Great work! 3D->2D and time is kind of true when we sit in front of a screen. Real world we use 2 x 2D (Human vision two eyes) add tactical movement data to feel real 3D space and time. So we are at the beginning of these world models. Still comput is key to explore. Lets play…

Agreed, world modeling is one of the end applications of deep learning, debatably the end application

incredible work & lovely name lol

haha thank you credit goes to tinygpu and tinytpu lmao

HOLY SHIT insane work

thank you king

Impressive work on TinyWorlds! Reimplementing DeepMind's Genie 3 from scratch with only 3M parameters while maintaining playable game environments is remarkable. The efficiency gains in world model architecture are fascinating.

Thanks!

@artmarryscience

this is so sick

Thank u

hot

thank u u🫵

so cool

