Loading video...

Video Failed to Load

Go Home

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

anandmaj's profile picture
anandmaj1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

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 :)

anandmaj's profile picture
anandmaj1 year ago

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!

Louis Castricato @ Tokyo Game Show's profile picture
Louis Castricato @ Tokyo Game Show1 year ago

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)

anandmaj's profile picture
anandmaj1 year ago

Very cool!

clem 🤗's profile picture
clem 🤗1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

It’s already shared :) thank you HF

ludwig's profile picture
ludwig1 year ago

Ts hella tuff, looking forward to diving in this weekend

anandmaj's profile picture
anandmaj1 year ago

Thank you goat

surya's profile picture
surya1 year ago

next step: train tinyworlds on tinytpu

anandmaj's profile picture
anandmaj1 year ago

Lmfaooo unironically

Lon Lundgren's profile picture
Lon Lundgren1 year ago

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?

anandmaj's profile picture
anandmaj1 year ago

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

Lon Lundgren's profile picture
Lon Lundgren1 year ago

Which GPUs did you train it on? How many?

anandmaj's profile picture
anandmaj1 year ago

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

Vincent Weisser's profile picture
Vincent Weisser1 year ago

awesome work!!

anandmaj's profile picture
anandmaj1 year ago

Thanks!

Sigil Wen's profile picture
Sigil Wen1 year ago

Wtf is in the majmudar gene @MajmudarAdam

adammaj's profile picture
adammaj1 year ago

incredible banger

anandmaj's profile picture
anandmaj1 year ago

inspiration

surya's profile picture
surya1 year ago

LETS GOOOO

justin's profile picture
justin1 year ago

fucking cracked

anandmaj's profile picture
anandmaj1 year ago

👉🏼👈🏼🫵🏼

evan's profile picture
evan1 year ago

🐐🐐🐐

anandmaj's profile picture
anandmaj1 year ago

🫵

Bell Ch's profile picture
Bell Ch1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

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

Bell Ch's profile picture
Bell Ch1 year ago

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).

anandmaj's profile picture
anandmaj1 year ago

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)

Bell Ch's profile picture
Bell Ch1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

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

Bell Ch's profile picture
Bell Ch1 year ago

Video is a different modality.

anandmaj's profile picture
anandmaj1 year ago

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

Mihir Mahajan's profile picture
Mihir Mahajan1 year ago

Great work!! Do check out Jasmine :)

anandmaj's profile picture
anandmaj1 year ago

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

Mihir Mahajan's profile picture
Mihir Mahajan1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

Nice!

surya's profile picture
surya1 year ago

bro this is so sick

anandmaj's profile picture
anandmaj1 year ago

thank you king ur an inspiration

Krish Shah's profile picture
Krish Shah1 year ago

wait this is so cool wtf

anandmaj's profile picture
anandmaj1 year ago

thanks king

Jason Ma's profile picture
Jason Ma1 year ago

Cracked

anandmaj's profile picture
anandmaj1 year ago

🥹🙏

Mike Boysen (JTBD/acc)'s profile picture
Mike Boysen (JTBD/acc)1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

@Scobleizer Indeed

Michael Churchill's profile picture
Michael Churchill1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

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

Michael Churchill's profile picture
Michael Churchill1 year ago

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

Michael Churchill's profile picture
Michael Churchill1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

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

Jibraan's profile picture
Jibraan1 year ago

holy shit

anandmaj's profile picture
anandmaj1 year ago

🙏

sway's profile picture
sway1 year ago

3 million parameters is crazy

anandmaj's profile picture
anandmaj1 year ago

maximal compression indeed

ali's profile picture
ali1 year ago

Some absolutely insane work, fire!!!

anandmaj's profile picture
anandmaj1 year ago

@AliesTaha Tysm urs too🙏🏼

ali's profile picture
ali1 year ago

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

anandmaj's profile picture
anandmaj1 year ago

@aliuahma deep cut lmao thank you

chi's profile picture
chi1 year ago

this is incredibly goated

anandmaj's profile picture
anandmaj1 year ago

ty🙌🏼

mikael haji's profile picture
mikael haji1 year ago

incredible

anandmaj's profile picture
anandmaj1 year ago

🥹

Schrodinger's Cockatoo's profile picture
Schrodinger's Cockatoo1 year ago

are you @MajmudarAdam's brother?

anandmaj's profile picture
anandmaj1 year ago

@MajmudarAdam I wonder

anandmaj's profile picture
anandmaj1 year ago

😂

sc's profile picture
sc1 year ago

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…

anandmaj's profile picture
anandmaj1 year ago

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

aiden's profile picture
aiden1 year ago

incredible work & lovely name lol

anandmaj's profile picture
anandmaj1 year ago

haha thank you credit goes to tinygpu and tinytpu lmao

rajan agarwal's profile picture
rajan agarwal1 year ago

HOLY SHIT insane work

anandmaj's profile picture
anandmaj1 year ago

thank you king

Ak afridi's profile picture
Ak afridi1 year ago

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.

anandmaj's profile picture
anandmaj1 year ago

Thanks!

Ayush ᯅ's profile picture
Ayush ᯅ1 year ago

@artmarryscience

omkaar's profile picture
omkaar1 year ago

this is so sick

anandmaj's profile picture
anandmaj1 year ago

Thank u

Jason Gao's profile picture
Jason Gao1 year ago

hot

anandmaj's profile picture
anandmaj1 year ago

thank u u🫵

Toms's profile picture
Toms1 year ago

so cool

Related Videos