Workshop on World Models - Day 1 (WM@Booth 2026)

Sometimes they don't.

Check one, two, test. Test one, two.

Okay. Should get input.

I've got it muted on my end here. I just want to see if he'll be able to hear it.

Nick, are you able to sign in?

That's right.

And you can see what's going on up there because I can't see anything.

Yeah.

Check one, two, test. Test. This is a little loud. Check one, two. Mic test. Check. Check. Check. One, two. Test. Test. One, two. Check. Check. One, two. Mic test. Mic test. Check. Check. A little better.

Check. One, two. Mic test. Check. Check. One, two. Still going to be test one, two. That's pretty good.

What's that?

Where are you? Oh, you're hearing it out of Zoom. Okay.

You're in the eater room, right?

John.

What's that?

Not making decisions.

That was already.

We don't know.

They have no idea.

I already sub items.

Check one, two.

And then.

Chicago 1960.

One of the learning.

Careful.

Yeah. Are we streaming? Check. One, two. Check. Check. Check one, two. Check. Check one, two. Test one, two. Oh.

All right. There we go.

Check one, two.

Check. Check.

Check. One, two.

Oh, that's good. Okay, cool.

My goal is to place that is test.

Well, I don't know.

Foreign.

24 locations.

So you have Facebook here.

Table.

Yeah.

I think I know.

I.

Good morning everyone. Thank you for being here today. So, as you noticed, we're already two minutes late which is off to a great start. We are figuring out some last technical details. The slides will come up very soon. But meanwhile, yeah, I'm very happy to welcome everyone to the third world modeling workshop. In case you are not familiar with the previous, to put the links in the slide, there was one at the Flatiron and then one in Montreal last February. And every time we try to focus on a different application, different domain. This time of course because of the name Booth. If you're not familiar with it, it's a business school. So we focus a little bit more towards you know non-stationary signals, time series, economy, finance, business. And this will be very interesting because as world modeling researchers, we need to understand what type of properties those guys need and how to make world models useful. And at the same time they will learn about world models and I will let them know that they still have a few months if they want to get on board and you know not miss the train. So yeah very happy to have you here today and then we can get started with the slides whenever ready.

And meanwhile so just to give you a quick overview. So I have just couple of slides and announcement to make will be like 5 minuteish and then Mike will take it away to speak about Booth and why did they decide to fund and help us organize this event and then we'll directly go to the first keynote of Dave who is already there and then we have the first coffee break. So let's see once we can get started otherwise I can go without slides as well so we save time. So the first thing I want to say is that none of this would have been possible without everyone here. Okay. So even if you just came from ORD you know just grab a coffee and like okay I did not do anything yet still you are part of the event and without you there will be no such events right so we need to thank everyone here which includes the speakers the attendees the organizers of course I will put the name and their pictures later once the slides work and of course Booth who very generously provided funding so that we can have breakfast food. Very, very amazing building. So, I don't know if you had time to look a bit at the view, but it's pretty amazing. You have the river here, the lake there. Definitely take a look around for the architecture. Very beautiful scenery around. And all of this comes from Booth willing to help us organize the event. So, also huge thank to them and also Lambda who helped us not only on the financial side, but also we'll have a contributed talk on day three. And also remember that Lambda partners with you know financial firms to help them on compute. So I know there is at least one person maybe two from Lambda around I don't know if she's here yet but if you're interested on this definitely speak to them as well. So since we don't have the slides yet maybe let me just mention do we have some of the organizers here? If yes so I see Kawen here Bradford XY. Yeah, XY here and Brat phone maybe is not here.

Okay, so and then we have Katy and Lauren around. So again, if you need anything just ping us. Yes, Katy is there and Lauren there. So we are all the organizers if you have any question please come to us. And so with that first let's thank everyone for this.

Okay. So the second thing I wanted to mention which will be a bit easier with the slides is that there is as you know much much more interest now around world models and this started maybe a yearish ago. Ah great. So perfect timing. So I have a couple of lines that I got from Google trends yesterday night. So let me jump directly to that.

Uh okay the clicker does not work. Ah yes great. So okay so this we went over it already. Yes. Let me just mention that without any of those layers everything will crumble apart. So that's why everyone is to be thankful. Yeah so the first thing I want to mention quickly hopefully from this event everyone will learn something new by the end of the three days. Okay that's really something very important to me. If you did not learn anything new send me an email. It means we miss something on the program or attendees. So please keep that in mind. Make new connections very important. Also try to catch up with old ones. I know I saw a lot of faces of people who came here came to the second workshop even the first workshop. So you know reconnect with people see how the research is going and also maybe get a new paper ID. I think it's going to be very easy for most of the people here. But this is also very important and enjoy Chicago which is really really nice city. So if you don't if you don't fulfill all of those you know send me an email at the end of the three days and let's see how we can improve the program for the next version.

Okay. So now the thing I wanted to mention and I'm not pressing anything so I don't know what is changing on the side but anyway the next thing I wanted to mention is this those are two curves that hopefully will appear. Yes. So quick look at it because it may disappear. What this is is a trend of world models and LLM research on Google trends. Okay, so just a search term popularity on Google. So what you see is that world models started being quite popular roughly a year ago and of course the growth is quite rapid not at the scale of LLM yet. But what this means for us is that it's a perfect time to do research or to start doing research in world modeling. And so okay this is on its own and so if you are not doing research yet on world models it's not too late so this is very important thing to send in fact it's a perfect timing we have foundation papers that are coming out strong results that are coming out so there is things to work on but it's still not overcrowded and the ceiling for world model is really unbounded okay so you can really think of any application and try to see if you can apply world model there so this means that as researchers it's perfect time but for us we also need to keep doing those events for the community to talk to exchange and to brainstorm and so the two announcements I want to make is the first one which is in the next slide which is we already were able to secure the fourth version of the workshop which will happen okay I will spoil it already which will happen end of February in Aspen Colorado and with a focus on world models for physics. Okay. So, Aspen Center for Physics focusing on physics everything makes sense. If you like skiing also perfect sense to go there. So, this will be the fourth workshop. There is already a website out that I will send maybe later and call for paper and so on. So, this if you are interested in physics application for world model this will be the next thing to see. But this is not enough right? you see we start we start having more and more people so we cannot just keep having those small scale workshops. So that's why the second thing that we have to announce which of course is not going to be as grandio as with the slide but is that we are starting a conference on world model. Okay so it's going to be called the international conference on world modeling. There is already the website up you can check it up it's icwm.cc. And this will be much broader event more than 500 people at least. The first version will happen end of May in the Bay Area around San Francisco. And right now this is organized with Yann and Amir. So we have the full support of Amir and Lambda and of course we will have more sponsors to be announced.

Okay. So here we go. Thanks. So we are back back in track. So just to show you back the previous one. So this is the first workshop happening in Aspen. You can see the website and then the conference and the conference website. So still a lot of things to be figured out but this is something that is happening it's already on the move and so this is something you should look out if you are working on world models and I want to say because it's a new conference we can try to do things right from the beginning. So the thing we target nothing is guaranteed yet will depend on sponsors and such but students will get there for free. Okay no registration fees for students ideally even if you don't have a paper but let's see if you have a paper and you're a student you get there for free. Also we will rework the reviewing process to be slightly different than the other conferences. On average we expect it will reduce by 50% the reviewing load that is on people and also we will try to revamp a bit the whole world gen AI policy and those type of things. Last thing is that we want really financial transparency. So I know what the experience of everyone has been but if you went to whatever other I won't name names and you paid 1k to get just a five minute window to grab coffee you wonder where the money is going. So we want to have full transparency on where is the cost and where is the money going. Okay. So with that thank you very much. So I will now hand it to Mike if we get to the last slide.

Okay. I don't think the clicker is working but anyway sorry. So yeah please help me welcome Mike from Booth.

Thank you Randall. Fortunately I have no slides. And I really don't have much substance. So, my job is to welcome you here to Booth. So, I'm Mike Menace. I'm the deputy dean for faculty and on the accounting faculty here at Booth as well. And I want to welcome you here to Chicago Booth. For those of you who don't know, Booth is a business school within the University of Chicago. This is our obviously our downtown campus. Our main campus is obviously in Hyde Park. This is where we'll host evening weekend students and do events like this in our conference center here. So if you do get a chance to go down to Hyde Park to see our main business school campus, I welcome you to do so. I'm really glad we had a chance to host this. So our applied AI faculty area was the first new faculty area that Booth has introduced in decades. The last one was our behavioral science unit. So we have economics, we have finance, we have accounting, we have marketing, we have the areas that you'd expect to see in a business school, but several years ago there was a group of people who had the vision to say that we needed to be on the forefront of what's going on in AI and made the decision to bring in a new set of faculty in our applied AI group. We now have two people two faculty members both untenured assistant professors in our applied AI group that have started this or you are the first people to be able to do that. We have a number of other faculty who are what we call in applied AI so accounting and applied AI finance and applied AI that are doing a lot of very interesting research that are creating that intersection here. So we've made heavy investments in GPUs, heavy investments in faculty to try to make sure that we are keeping our foothold in this space as a business school.

The primary function of what we're trying to do and why we're here in this room with you is because a big chunk of what we try to do is how do we think about managing capital, human capital, physical capital, financial capital. And probably in case you haven't noticed, one of the biggest giant sucking sounds of capital right now is what we're doing with AI and all the investment that human capital, financial capital, and physical capital that's going into this space. So I think this is a great time and just fascinating research questions that are going on and a big chunk of what we have to think about is how we're investing and making optimal decisions in this space. So it's a great opportunity to connect both kind of the CS world with kind of the business economics world here too. So, just very glad that we could do this. And I also want to reiterate what Randall said and really thank the organizers. This has been a massive undertaking to get this pulled all together. I'm amazed at the set of speakers that we're able to have here and hear from and I'm really looking forward myself to being able to learn a lot from this. So, again, welcome. Thank you. And I'm gonna turn it over to people who actually know something about what we want to listen to. So Randall, I'll hand it back over to you. Thank you all.

Great. So now I'm very happy to introduce the first keynote of the day. So Professor Dave Donoho from Stanford statistics department. So Dave will be speaking about interesting parallel between Blackwell GPU and Blackwell the researcher and what he thinks about world models and what generative next token prediction model could bring to world modeling. So with that please let me welcome Professor Dave.

Actually sat in classes in this building while she was.

Okay. So, how do I get the.

Well, as I was just telling the dean, I've been in this building before I've even sat through classes in this building because years ago my wife was on the faculty here and then I made her a better offer. So, but it was really a nice place and they only have made it better. I think everyone knows that the current generation of dominant GPU hardware is called Blackwell, but probably very few people know anything about the Blackwell nomenclature or anything about the man's life. So, I thought I'd just fill in some of those details.

So this is a joint production with my brother Andrew who has used Claude very extensively to do all sorts of experiments. The experiments are all done on a dual DGX Spark system that allows you to have on your desktop instant response and inside there's GB10 Blackwell chips. So we're actually using the Blackwell chips for what we show. Also XY Han whose junior faculty here and Vardan Papon have given us extensive training and knowledge that was needed to complete the project and I'm sort of involving them in it as it goes along.

Okay, so let's talk about David Blackwell, the mathematician, Blackwell, the GPU, and Blackwell's world models theorem, and then get into some empirical work.

David Blackwell was a professor of mathematics and statistics at UC Berkeley. As his ultimate step, he started out at Howard University, which is an HBCU in Washington DC and really distinguished himself by his work. This was in the 40s. He came out to California as a sabbatical leave in the early 50s at Stanford University and then Berkeley snapped him up, made him the first black tenured professor at UC Berkeley. And throughout his career, he established many, many firsts. These included the first black mathematician elected to the US National Academy of Sciences and let's say in many ways he was the person who became known as the founder of a number of fields. So for example, essentially modern dynamic programming associated with understanding controlling of Markov processes is now called reinforcement learning and he did the canonical work on that in the late 40s and early 50s and wrote a book with Abe Gershik who was founder of the statistics department at Stanford.

He passed away in 2010 and was awarded the National Medal of Science posthumously in 2014 by Obama. The effort to get him this distinction which is the highest award that the US gives for science had been going on for some time. So I had been writing letters to Ken Arrow in support of this say for 15 years or something before this. So these things can take a while. There's a long queue of great scientists that the US has had in the past and we hope it continues in the future. It would be nice also to see national recognition for them from our federal government although that's more in doubt. Why is he in this talk? So in 1956 1957 Blackwell who had done all this work on controlling Markov processes proved a theorem showing exactly what a hidden Markov process observer must track in order to predict optimally. So if you think about next token prediction as the fundamental idea behind everything that's going on in the modern training of LLMs and many other systems then you're going to immediately go to the question well if you learn how to predict the next token accurately what are you really learning? So in essence, you've got to be doing whatever Blackwell's theorem says is going to emerge from this process. So he asked and answered this question many years ago. Also if you read Twitter and X very much then you'll know that there's claims that accurate next token prediction does not imply that you understand the world. You just know how to predict. So part of the reason that we're here is this fundamental question of if you understand how to predict, do you understand the world or do you have to do something extra in order to understand the world? Okay.

Well, so let's see what that means. Now, there's a second reason that he's in the talk. I'm not getting an advance of the slides.

So I was a faculty at UC Berkeley in 1986 to 1996 and my office was 401 and David Blackwell's office was 306 and he was a colleague of mine and also a mentor, a very sweet person at that age. I worked on projects with him. So for example just over coffee we started talking about if you have a two-person repeated game and but you don't know what the game is. It's just done like standard two-person games. You're going to play this game repeatedly. The only thing that you know about it is you have a state space of action. So, can you eventually achieve in long run play the value of the game? And we figured this out. I realized through writing this up that he was way smarter than I am. And in particular, he fixed everything that I considered a proof to be so much more elegant that I was just ashamed of how crude I was looking at everything. So he had in a repertoire of special laws of large numbers that he had developed over his research career and he would pull these things out and just finesse problems that I had, you know, done in a very messy way. In any event, we found that that Rodrigo Banuelos had beaten us to the punch and never published that. But I really enjoyed that project. Later I taught courses from his lecture notes and one was a course on game theory from which of course I learned a great deal.

So, you know, there's a lot since I mentioned Twitter X, there's a lot on Twitter X about DEI and the idea that you know, there are beneficiaries of DEI and they should be looked down upon and so on and so forth was kind of ridiculous. In this particular case, the obstacles that this man had to surmount were just unbelievable. And I learned about these from Jerzy Neyman's notes. He was the previous occupant of my office. The reason I'm mentioning all of this is that in fact he let's say understood at a very clear level all sorts of things despite all the distractions that one can pile up and the kind of let's say atypical thinking that people may have when they encounter someone who doesn't look like them. He's known of course for his work on the what's known as the Rao-Blackwell theorem which is you know standard first year graduate course. He was also a close friend of Abe Gershik on the right with whom he wrote a book on statistical decision theory and all the things about making decisions in time sequentially.

Ken Arrow is there on the left. Ken was my neighbor on the Stanford campus's house just downhill from me. So these are all people that are real to me. Abe Gershik had been a professor at Stanford before I was etc. I could fill you in if any of this were interesting. And this is a standard where they're essentially figuring out what is modern reinforcement learning. He has papers that have been foundational and you can see for example this one that I've used several times. That's his textbook and here's where he's working out a lot of the details in the early 60s in Annals of Math Stat connecting modern Markov decision processes and reinforcement learning.

So I'm a trustee of the Simons Laufer Math Institute, which used to be called MSRI, and this was the first conference for African-American researchers in the mathematical sciences, which was held at MSRI in Berkeley in 1996. And you can see there Blackwell in the center and on the right you can see Wilkins who was a member of the Manhattan project and on the left you can see Shabazz who was an awardee also at this particular event. There's the posthumous award to Blackwell's children and you can learn more about this event online. Of course, there's also a film that's been made by the SLMath MSRI called Journeys of Black Mathematicians, which is available to watch. And my wife and I financially supported it in part, but most of the money came from the Simons Foundation. Describes where African-Americans are doing math today. And there's a sequel that gives a lot more background information. At UC Berkeley, there's now a David Blackwell Hall, which is an undergraduate dorm.

Okay. The other Blackwell Blackwell GPU all the details are given here. This is from some press information that came out. It can in FP4 mode do 40 petaflops per second which is 5x the Grace Hopper GPU which was the basis of the A100 series that came before and it can hold models with 740 billion parameters and so on. So this is the most popular thing right now for building foundation models although new chips such as the Vera Rubin are on the way.

It's not a coincidence Nvidia on their own documents say they named the chip after him. The data center revenue that's being recorded by this new chip is really staggering. If you just look over a quarterly basis, you're seeing several tens of billions of dollars of revenue that's being booked because of the use of this chip. And so it's hundreds of billions of dollars being booked at an annual rate. So quite an amazing capitalistic monument to David Blackwell.

Some of the details for the generational jump and of course there will be future jumps and so on are indicated here. Supposedly up to 50x agentic AI performance and 35x cheaper cost per token.

Okay. So the relevance to this talk is given in this paper. It was presented at the Prague conference on decision and information 1957. So I mean we all go to a lot of conferences these days and let me just give you a little information about what it used to be like. So first of all this is a conference paper and as you can see here we have an example of it. What was required to get this example because it was not online was to write to the librarian of Charles University in Prague and ask them to go look for a physical copy of the volume that had been prepared in 1957. Scan and email to us. So this theorem has been out there and unknown to Claude until now, let's say. So presumably something from this conversation that we're having will leak out and appear in there. And he asks a question in this paper is if I have let's say a partially observed Markov chain. So the true dynamics are Markov but I don't get the ideal observables for that. So partially observed Markov chain and this is anticipating the full POMDP formalization that will come out that is used today behind linear control and nonlinear control and reinforcement learning. You just ask the following question. If I have such a thing, how do I calculate the entropy rate? So what is the entropy rate? It's basically the cross entropy loss of an ideal model in next token prediction. That's the entropy rate is several things but there's a chain of theorems and one of those theorems gives you that identity. So how do we calculate that? Well, then we have to understand the structure of what's going on.

And at the key is the theorem that's in this paper. I believe it was proved in 1956, but you know going to this conference and so on added in some extra delay about which I'll have some comments in a moment. So what Blackwell realized is that the understanding of the situation you're in when you face next token prediction is to know the full conditional distribution of all the tokens that correspond to the possible next token at this moment given all the tokens you've seen. There is such a conditional distribution and if you knew that that in a sense is the full object construct that's needed in order to optimally do next token prediction. But he observed more which is if that the sequence of those objects is actually a Markov chain itself or more generally Markov process depending on the details of the assumptions we're making. And so what he's basically saying is this is what you get from next token prediction. Maybe it's a world model. Maybe you know the world from that. Maybe you don't know the world from that but this is what you will get from next token prediction is a full understanding of this sequence of conditionals.

Okay. Now what could the relation of world model be? Suppose we think of ourselves as physicists. Lagrangian dynamics. What is that about? It's about I know all the positions of all the particles everywhere in the universe at one moment. I also know all the momenta of all the particles in the universe at that moment. Then what happens next right? That's the Lagrangian viewpoint. That's exactly what this is. So this is Lagrangian except it's moved from a deterministic calculation of the 19th or even 18th century into the framework of probabilistic calculation where you have conditional model. So whereas Lagrange, Leibniz, these kinds of people would say that a world model is to know all the particles know all their momenta and then just calculate what happens going forwards. Blackwell said no the right thing is to know the conditional distribution of what's happening next. It's probabilistic. So this fits in with of course the direction physics went over that same period where you went into quantum mechanics and everything became probabilistic. So roughly speaking the Blackwell theorem and the Schrodinger viewpoint are the same.

And Blackwell gives an example which is important for what comes next showing that it's quite interesting what can happen with this sequence of conditional distributions. Those condition if you have a large number of tokens like in a modern word embedding then that's a very high-dimensional object and yet what could be happening is that you're going to wander around in the simplex of probabilities according to a fractal and he did that in the same paper gave an example so that there was a variety of very interesting mathematical structure that could be exhibited.

Then Jim Crutchfield who was junior faculty at Berkeley when I was there I think now he's at UC Davis then developed things built on the Blackwell viewpoint and considered something he called mixed state presentation and sometimes called the Blackwell measure and so on in which we're looking at exactly where do all these conditional probabilities wander in high-dimensional space and you can look at a paper of Cosma Shalizi and Jim Crutchfield and learn more about this viewpoint.

Then in this paper of Alexandra Jurgens and Jim Crutchfield, they calculate the Shannon entropy rate of hidden Markov process. This is exactly the problem that Blackwell solved in 1956-57. And but in this particular instance, they just do some very concrete calculations. And so if you really want to understand through some nice examples, you can go here and see two or three examples worked out in detail.

So just so we I'm throwing up this slide just so we understand a little bit better what could happen in a situation like this. So up here is a very simple Markov chain that has the structure if I'm not saying it wrong it's something like this. Every once in a while you chuck out a zero and then you know you're in state zero. But if you chuck out a one and print that out on the tape, you may be in state zero or one. And so you're never quite sure where you are until you get a zero. Then you resync and you know something. Okay. So this sequence of conditional probabilities can go from a to one over here where I know that I'm in state zero for sure and then depending on how often I've seen this indefinite state I could be many many times away and I'm less and less sure what's going on etc etc but then finally I see a zero again and I know I'm back where I am. Okay, so these are simple examples and there's variety of calculations that they do so that you can get a feeling for this and I'm just giving you this as an example so we know we're on solid ground on something.

Okay. So, interesting work also in the East Bay. I'm Shai Marsen-Reers and others. Some of them are at the Astera Institute now, which is I think one of these new labs. They asked can you somehow see a trace of this Blackwell measure actually hiding in the residual stream of a decoder-only transformer like GPT-2. Like can we see that this is really what the thing does. And then there's this incredible set of pictures that they show. So if you looked at the paper of Jurgens and Crutchfield, they would show you this a process like the thing you see in sub figure A. And then they would show you that it creates a fractal very much like the traditional one of these Sierpinski gaskets as the sequence of conditional distribution in this high-dimensional space. In this case, it's just, you know, it's three-dimensional space projected on two. So that's the way it looks. But that's because there's only three possible states. In higher dimensions, we'd have more. And they use Anthropic mechanistic interpretability techniques to resolve this kind of noisy version of the fractal. So the central one I think is like you know CGI drawn analytically from formulas and the one on the right has been extracted using essentially logistic regression and statistical estimation techniques and they say see there it is.

Okay. Any questions about what we're saying? Okay. Now, maybe I'll go back. So for the people at the Astera Institute, this is telling us like we know that we found this thing. We know that literally transformers have found a way to do optimal next token prediction. Even though Blackwell's theorem was buried in this dusty library at Charles University all those years and no one has had access to it before 2026, we know that somehow the transformers found out how to do the thing of the theorem. Of course, it's a theorem. It has to be true. So therefore, behaviorally it must be in there. Now one other thing is this is showing a linear embedding because the techniques are just based on linear models as we'll say later and so it's not just that there's some crazed embedding of the world here this is the simplest possible embedding of the world namely a linear one. Make sense? Okay now.

The Astera people did not have an interpretation for what they're doing in terms of world model but just we see these belief states. Now a separate paper that's of interest is by Keon Vafa who's now a junior faculty at UC Berkeley in statistics and then a variety of authors some of whom would be possibly well known in Chicago. John Kleinberg and Sendhil Mullainathan are two of the senior authors. So they looked at the question of finding world models from certain kinds of data sets and they considered primarily examples. One is taxi cab data showing ride routes through Manhattan, maybe some of the outer boroughs and then another was you and then there were a couple others that were a little more synthetic.

So the message line that they chose is that even if a model predicts essentially perfectly the world model implied could be pretty messy and wrong. And so next token prediction doesn't really give you an accurate world model. So that's their narrative which is kind of in conflict with what I just showed you about the recovery of the Blackwell measure. But let's see. Okay, so this is a paper you can go look up and I've heard Keon give this talk a couple of times. So I assume it's been videoed and so on. In fact, once I got in my feed a podcast explaining this. You know, there's those AI generated podcast where they just take papers that just hit archive and there's some Welcome back team. It's so good to see you again. Well, today we're going to talk about this paper. Okay, so this they got this treatment. I don't know how you get that treatment, but anyway, here's one of the data sets and experiments that they have in this.

So they have real data from taxi cab routes in Manhattan, which is a kind of a deformed messy grid. And they have training sequences that give an ID for a starting intersection and an ending intersection and then points on an eight direction compass rows for what you do at various intersections and then an end token. And then for evaluation, they give an origin and destination pair and just ask, you know, okay, go to town. So, you're supposed to learn the routes in Manhattan from seeing lots and lots of these routes. And then they say, what is the implicit world model behind this? That's the research program for this part of their paper.

Okay. One extra element here is that while of course there's a one-to-one between the IDs and the what intersection we're talking about, it's shuffled. So the ID itself doesn't tell you like you're at, you know, Fifth Avenue and 42nd. It just says you're at intersection 810. And so you don't get any clue about the world model from the actual ID. Whereas usually when people describe an intersection, the ID system they use is exactly like us. You're at Fifth and 42nd. You kind of know where they are without any other information. The key thing is that that's obfuscated in the origin destination pair spec.

Okay. So what they come up with is a model and they use it for predicting actual routes. And we'll get to the accuracy, but they also infer or back out of their model a map of Manhattan. And they remark that it's very noisy. And in fact, it's far more noisy than it would seem to be if you just added noise to some data about an actual map. They give information showing that for next token prediction, it's essentially correct. And they even do state probing of some sort. And it's essentially correct when it's trained the right way. But they propose new metrics by which you can see that it's not essentially correct. So they say next token prediction is not a good enough metric to understand if you got the world model.

A key point in their point of view is some research that originated at University of Chicago. So there's a paper by Nathan Myhill fundamental concepts in the theory of systems as you can see here contains essentially the theorem that they use which became known as the Myhill-Nerode theorem and it's talking about if I have a black box has it really understood a language and the question is whether if the language is generated by a deterministic finite automaton and it accepts exactly the same sequences, then it has to be actually that deterministic finite automaton or an equivalent one. There's various ways that two automata can be the same automaton and they come up with some metrics I'm just going to that they think do a better job than next token prediction. So that I'm not going to explain all of this. I don't think there's time now.

I want to get to some of the work that Andrew, XY, and Vardan and I have been doing. It's based on a framework for doing mind reading on LLMs, asking the LLM what it's thinking non-cooperatively. It's invasive. [laughter] It's invasive and not necessarily destructive, but could be destructive. So it's a little bit like planting electrodes inside someone's brain while they're awake and showing them stuff. We don't actually do that to humans, but we do it to laboratory animals. And if anyone is here from PETA and I just triggered them, step outside. I'm sorry. Okay. Anyway, there's an analogy here. So you could read this paper of the Anthropic team, which includes a lot of billionaires in it these days. They didn't start out that way when they wrote the paper. And then there's a variety of things downstream of that. There's the papers on logit lens, there's papers on tuned lens, there's papers on patch scope, so on and so forth. You can read all those fundamental papers. I taught this in a class at Stanford in the spring and took 20 lectures to get through all of it. So it's intricate and I'm not going to be able to tell you everything about it, but you're using the following fact that transformers output vectors in steps that are called layers. And then we can look at those vectors as predictors and we can do linear regression looking at those vectors as predictors. Say I have a model like GPT-2 where that vector is 768 elements long and we can use and I happen to know because I know a world model something about where I am in the world and I ask does this vector predict where I am in the world from that I will learn by fitting that model to enough data how to read out from the so-called residual stream where I am in the world and I'll know how accurately I can do it based on what the transformer is doing. I'm sorry for the brevity of this discussion. How much time do I have left? Okay.

Okay. So in our project, we're building up to the full Manhattan real data case, but the results I want to talk about are conclusive at this moment only in discussing a synthetic data set. In this synthetic data set, we do a kind of a baby version of unlike Manhattan which would be like 10 by 100 more or less in order of magnitude. It's this is complete. There's no.

I'm trying to advance the thing. There's a special thing about our tokenizer which is we use essentially a tight frame. We're not using tokenizers of the kind that you would use in LLMs normally because we understand the structure of the problem and I'll just leave out the details there because of time. Okay. And then the probes that we fit involved sorry let me just jump ahead here. So the left two plots I'm trying to infer essentially which avenue and which street am I on in this alternating one-way grid. And what you can see is that for this particular data set so this is like I'm on 7th Avenue, 8th Avenue, 9th Avenue etc. and the data points that you see. So this is a presentation of that Blackwell structure that we've talked about. And what you're seeing is that points are rendered by their position in this thing according to well each one is a linear combination. Its conditional distribution is some convex combination of all these possible places you could be like you could be in certain avenues and so on and so forth. So in this model what you see is that the cloud of points sometimes gets far away from the vertices. If you're at the vertex that means you know exactly where you are. And if you're a little bit away from the vertex, it means you're pretty sure. And the farther you get away from the vertex, the more equivocal you are about where you are at that moment. What you can see in this simple model is that at this data set size and scope, you kind of know roughly where you are because particularly if you're at, you know, 8th or 9th or West End Avenue or First Avenue, you know where you are very accurately because all those points actually just hit right on top of that point. But if you're over here, there's some kind of noise going on. So we did scaling studies where you add more data and then you can find that points migrate more towards the vertices. Okay, which suggests that in the alternating story you can infer what row or column you're from. Over here. What happens is there's a kind of a defined circular structure here. There's no alternating one-way and all intersections are four-way connected unless they're at the corners and stuff. And so what you see.

Here is, even if you have a lot of data, it turns out you have much more data in this model than over here. You don't get that close to knowing where you are. So, you actually don't know what street, you don't know what avenue you're on.

Please, but sorry. Could you say a little louder?

I'm so sorry. Let me come closer.

So I was just asking as I look at your diagrams, have you thought at all about the relationship between this and graph completion using cortical methods?

We can talk about it offline.

Okay. It sounds interesting. Thank you.

I'm juggling other things in my mind at the moment. So I can see how those words fit together, but I definitely have not thought about it previously.

Okay. So what we're seeing here is that you can recover your street and your avenue identity in a model with one kind of topology but not in the other. So in one you know your position and in the other you don't really know your position even if you have a lot of data. So okay, and another way to look at it is that we can look at the linear embeddings that we get of these. So for each street ID that's on the bar center plot there, we could consider all the representor vectors of these things and for each one of those, it's going to be a point up in a high-dimensional space but they all fit into a plane according to a simple SVD analysis and you see that it's a very accurate grid when we're doing the data that arise from alternate one-way streets and it's much less accurate of a grid otherwise. And in fact there's severe geometric distortion at the edge when it's a four-connected. So in the two-connected thing we learn where we are and implicit in the structure of the representors that teach us where we are is the actual geometry of the grid we're talking about, otherwise we're going to have distortion if we have a different connected set. Okay, and I could give you details of how we do the linear probes but I'll show you one thing that you learn in the connected fully connected set that we didn't learn otherwise which is, suppose we do the regression for each Y, we have a vector one-hot indicator of where we are in bearing space, how much farther we have to go from the place we currently are, how far we have to go to the right. So we can do a regression like this where these are the items the transformer has in its mind in the residual stream and these are the one-hots of we know in the world model how much farther we have to go. Can we do that regression? When we do that, the coefficients that represent each of those particular one-hots line up like this where if we have to go to the left, they all point this way. If we have to go to the right, they all point this way. So the set of regressions tells you that hiding in the mind of this LLM is very simply, it keeps track of the positive part of how far it has to go in one representor and it keeps track of the negative part of where it has to go in the other representor. And out of that we realized that there's just four representors intrinsic to this problem being learned by the whole LLM. And they are the positive part that you still have to go in X, the negative part in X, the same positive and negative in Y, and one single regression not with one-hots but just with the amount that you have to do picks up exactly that structure. So you can actually look at very simple linear formulas for predicting how far you have to go and you'll see that you get accurate reconstructions.

Okay, I've used a lot of your time and attention just to mention what we did. So we talked a bit about David Blackwell, winner of National Medal of Science and his understanding of Markov processes, hidden Markov chains, dynamic programming, reinforcement learning, and his understanding that next token prediction implies a very particular kind of understanding that you must develop in order to do it well. We mentioned that there are physicists doing what they call computational mechanics which is an outgrowth both of the traditional views of Landau in particle dynamics and Blackwell. We mentioned the work by the Astera Institute people, so Shai and others, showing that in LLMs for very particular Markov systems that they set up you can see these Blackwell belief states showing up. We presented some work by Keanu Vafa and others who are critiquing the idea that you can recover a world model if you do next token prediction. And in our project, we tried to understand how all of these four things come together. What we think we came up with is the following picture. If you have a completely symmetric grid where from every node you're connected to the maximal number of nodes you could have, then you don't learn position. The only thing that you learn is if you have to go from here to there, you just need to know the bearing, just the dx and dy. And that's the only thing that's learned by the model that would correspond to in the UK. Suppose I told you you're in the Euclidean plane and you're at these coordinates. You need to go to those coordinates. What's the one thing you would have to know? You would just want to know the vector that separates those two. You don't need to know where you are. You just need to know that vector. You need to go in the direction of that vector. And that's the only thing that matters to you. On the other hand, if I so-called break the symmetry the way physicists talk about this and now alternating streets go one way, then I have to know something. I actually have to know like, are you on a street where you can turn right or not? Okay. So therefore you have to know something and in that environment you do learn the thing that you need to learn. So the question of whether a next token prediction system is actually learning the world model in the sense of the world is like this and this is my place in the world has something to do with the symmetry breaking that's gone on and the idiosyncrasies that are present in the data. The more idiosyncrasies, the more the problem is going to lead to me understanding the world. The fewer idiosyncrasies, the more homogeneous things are, the less I understand about the world and the more I understand about some dimensionally reduced model of what the problem involves. Okay. Thank you very much for your time and attention.

Thank you very much. We have time for a quick question before the break and we have a microphone ready if needed. Otherwise you can ask Dave questions during the break as well. So we have 15 minute breaks now and then we have Adita and Janet. If you guys want to come to try your laptop that would be great. And let's thank Dave one last time.

All right everyone, we will be starting in one minute so you can slowly get back to your seats when possible. Thanks.

Yeah. Nice. It's on time.

Nice. Did you join the Zoom and everything?

Yes.

Okay. Perfect. And streaming is still good on YouTube, right?

I updated the website so it should be good.

Double check.

Okay. Perfect. I will just make you a sign when there is 10 minutes left. And then you can decide if you want to use the full time or leave time for questions. Okay.

Yeah.

Okay. Perfect.

All right. Everyone we are going to start the second part of the morning session. So please feel free to join your seat back. There will be a lunch break and another coffee break later as well. So it is my pleasure to invite Adita to give the next lecture. So the topic is actually very very important. It is about understanding how to handle multiple modalities. So if you've been working on world models at all and you read papers, you see most of the papers today are single modality or maybe they handle multimodality post-training. There is a huge need to understand how to do better multimodal world modeling and foundation modeling and so I'm very excited about the next talk and please help me in welcoming Adita. Thanks.

Are we good? All right. Thank you everyone for having me here. It's my first time visiting Chicago, so it's great to be here with this wonderful scenery outside. What I'm going to be talking about today is a new paradigm of research that I've been thinking a lot about over the last two years. So I wear two hats currently. So I'm a CS faculty at UCLA and at the same time I started this new company called Inception. We are what a lot of people like to call neolabs these days. So we are betting on a new paradigm for how generative AI would look like in the future and specifically the paradigm that we are very excited about is one where we try to develop a new class of language models based on diffusion. And what I'm going to do in today's talk is actually show some of the recent results we've had both in my research group at UCLA as well as at Inception in trying to bring this to the frontier of AI.

So to motivate why we would want to bet on this paradigm, I think one of the key things if you look at the last few decades has been the key importance of parallelism in bringing about advances in AI. And it started in the 90s when the GPU was invented and for a long time its primary use case was graphics until in the 2000s and very much so in the 2010s we started seeing deep learning take off and deep learning involved a lot of matrix multiplications for which there was no better hardware than the GPU to bring about that change. Now along with deep learning there was a specific kind of neural network architecture which has had a very important role to play which is the transformer. And this neural network architecture allowed us to train on arbitrary kinds of data, language, vision, you name it, in a very highly efficient manner. Now transformers are very general purpose in the kind of data you can apply them to. But as we think about 10 years later we now actually have lots of good transformers for various kinds of things and the bottlenecks these days are shifting towards inference. When we think about inference, we want models which do allow for parallel inference and that's what I'm going to focus on in today's talk and these are models which are trained in a certain way known as diffusion.

So most of you might have seen these kinds of visualizations where you have a diffusion model that starts with complete noise and then gradually uncovers structure into it to give something coherent. This is today the state-of-the-art for images. You see it in a lot of world modeling paradigms some version of diffusion being used. What we are going to do today in today's talk is think about how this paradigm actually can also be applied for text. So just to visualize how this might look like for something which is a discrete object like text, here you don't have a notion of a Gaussian noise that you would start training with. You have what are discrete tokens. So on the top you see the current paradigm which is token prediction or autoregressive models that we heard about even in the previous talk where you're trying to predict a new token conditioned on the previous ones. It is sequential. So the number of time steps it takes to get the full sequence is exactly the length of the sequence. In contrast when you think about a diffusion paradigm you have models. So at Inception we released these models called Mercury models. So these are diffusion language models and they operate in parallel. So here you start with a noisy sequence of tokens and these tokens can be random tokens that do not make any sense but you have trained the model in a way that allows you to then denoise these tokens in parallel. Right? So similar to image but now starting with some discrete noise.

So here is a simple example of how diffusion language models work. So this is a specific kind of diffusion language model which is the uniform diffusion language model. So you start with uniformly randomly sampled tokens which are sampled uniformly from your vocabulary and they could be gibberish. So maybe you start with 'the use SAT' which makes no sense whatsoever. But then what you do is you train the model to denoise it. So you look at your training data and you look at sequences which do make grammatical sense and try to match the output of the neural network with the ground truth and that's how you make this model uncover real structure from noise.

Now this is as of today, in the last few years a lot of advancements have happened. It's a drop-in replacement for your existing LLMs. So anything you use LLM for, whether it's prompting, supervised fine-tuning, reinforcement learning, all of those wonderful things you can also do with diffusion language models. It is also agnostic to emerging architectures. So there's a big wave of new paradigms trying to compete with transformers like the Mambas and the DeltaNets of the world. You can use those architectures as well and train them to denoise outputs that will still give you a diffusion language model.

Now what about inference? So this is where magic happens. So if you have let's say a prompt like 'which open math problem was proven by Andrew Wiles in 1994', how you would do inference with these models is again you would start with some randomly sampled gibberish tokens but now the model has been trained to denoise it and what's interesting is when you do a forward pass through this neural network you can actually output tokens all in parallel which means you might get actually more than one token matching what the correct answer is. So in this case we actually got two out of three tokens right, 'Fermat's' and 'theorem'. And what we can then do to get the correct answer is take the output of this step and use that as an input to the next step. So in the next step you pass this again through your transformer and at this point you actually have the correct answer 'Fermat's Last Theorem'. So in two steps you got three tokens in this example which is still faster than how autoregressive models would work where you'd get one token with every forward pass.

Now I won't go too much into details but one variant I would quickly mention is known as masked diffusion LMs. So what we talked about in the previous slides were uniform DLMs where you start with gibberish tokens. In masked diffusion LMs you start with masked tokens. So you replace gibberish by masks and denoising is equivalent to unmasking the masked tokens to get your correct answer. Again similar to uniform DLMs you can unmask multiple tokens in parallel.

Okay. So where is the field today when it comes to diffusion LMs? And this slide kind of summarizes how things have changed since diffusion itself was invented. So in 2019 the denoising diffusion was invented at my PhD advisor and also my co-founder Stefano's lab at Stanford and it quickly became the default for image and video generation at scale. And then many years later there were many advances in the field including some at our own research labs which prompted us to think that diffusion is now ready for language. So that's when we started Inception and at that point before the company was founded the biggest result was diffusion was matching GPT-2 in perplexity on text generation. So very small but sparks of what we now believe is going to be mainstream. A few months later we released Mercury which was the first commercial-scale diffusion LLM. This got a very good response as well in the community from both academics and industry experts. A few months later many competing companies like NVIDIA, Google, and Alibaba started their own diffusion efforts and they continue to release diffusion LMs. Many of them are wonderfully open-sourced and are spurring a lot of research in the community. And then earlier this year we released Mercury 2 which now incorporates reasoning and now it's also going into production in a lot of different use cases for voice agents, for search agents, in a lot of places where latency is a must-have.

Okay. So if you haven't seen how this looks at scale, this is just a visualization when we launched Mercury. So the same prompt given to three models and you can see that at the time how fast diffusion can generate tokens. So what took 6 seconds for Mercury to generate would take about 5x more for the other models.

So and this is now also going up in quality. So one of the questions we should think about really is like the speed is great but what is the quality at which we can achieve. So here in these charts what you see on the x-axis is speed measured as throughput. So you want to be more towards the right and on the y-axis we're looking at quality. This is a measure of quality that's independently verified by Artificial Analysis which does a lot of benchmarking of all kinds of LLMs out there. And what you see is that once we released Mercury 2, it was a reasoning diffusion LLM that improved upon Mercury 1 that we saw in the previous slide both on quality and speed. So if you compare it against a lot of the speed-optimized LLMs that you see from Frontier Labs. So for example, Claude has a version called Haiku series. GPT has the Mini series. You find that Mercury 2 gets you to similar quality to these LLMs while being much faster. A common question I get is 'is that the limit?' Not at all because this is where we are today. We don't see any signs of this stopping. So we do expect that we can translate these benefits to even higher quality and that's what we are working on.

Now why does this change how you should think about agents in the real world especially as we talk about world models and we think about a place where these models will interact with the physical world, real-time efficiency is key. So these are just some of examples that I have of how this changes the landscape of what AI can achieve when time is of essence. So here in voice, so every time you talk to a voice agent and that's an AI, they really care about how fast the response is. In fact no matter how good the response is if it misses the critical window in which the AI has to respond back the conversation is essentially churned because we all have been through this. If it's too slow to respond, we know immediately it's an AI and we want to transfer to customer support through a human operator. So what we can do in these cases is with a diffusion LM actually use reasoning to go up on the quality ladder but exploit the fact that these models are so fast that even with reasoning you actually get latencies which are much lesser than what you would get with frontier models. So basically while reasoning has provided us with this new axis to improve on quality it has typically come at the cost of speed but with diffusion LMs you can try to get the best of both worlds.

Search is another one of these use cases. So search pipelines are increasingly becoming LLM-driven. So every time you enter a query onto new kinds of search engine often there are a lot of steps that go in between before you get an answer. Maybe you write something vague. So an LLM is used to rewrite the query. Based on those rewritten queries, you try to retrieve some results and then again you use an LLM to re-rank those different kinds of results and summarize the outputs of them before you actually give the final output. So all of these pipelines which are now mainstream within the search community do add up and compound latency. So again, what you really require is a very fast LM that can do these steps with high quality and that's where diffusion offers a very compelling alternative.

Okay, so that was the first part of my talk. What I want to spend in the rest of the talk is actually focus on how we're taking these advances in diffusion LMs to the multimodal world. So and here again I'm going to be talking about some of the work, actually the entire series of work has been led by my PhD student Shuan Lee who you see here and we look at how diffusion LMs are now powering workflows with multimodal understanding as well as generation and reasoning and really cutting edge at the frontier of speed and quality.

So the series of work is called LaVIDA. So each of the papers I talk about will be an improvement over the first one that I'll discuss right now called LaVIDA itself and the first one tries to basically do multimodal understanding. So the setup is as follows. So you have a diffusion language model, it has been pre-trained on language and we now want to make it understand visual inputs. So in this particular image we want the model to be able to answer that what it's seeing is a dog. Now how you can do this generally is again you might start with a sequence that's completely masked and then try to unmask tokens one at a time or in parallel with a diffusion language model to get to your final output.

There are a few challenges that we observed when we try to do this. So a lot of this work was trying to address these challenges that did come across when we tried to follow this protocol. The first was of course that when we think about modern-day visual inputs, they're extremely high dimensional. So if we wanted to extend essentially the context length of diffusion language models to process extremely high dimensional inputs we had to essentially create views of them which were able to be efficiently processed by the language model. So in this case one strategy we found extremely useful and simple was to basically create both local and global views. So you take this image, you cut it out into four pieces and you also then resize the image. So the cutouts give you local features whereas doing a resize of it gives you global features that you can then pass through a vision encoder to get embeddings that eventually go through the DLM.

Another important observation we had in this work was that when we look at correspondences between vision and text, often they end up being sparse. And this is best illustrated through this example. Now out of all the tokens I could be unmasking in this sequence, clearly the missing token here is the most important token when it comes to being aligned and being able to answer the question being posed here. So it's really important during training of these DLMs to be able to see these kinds of sequences more often. But how a typical DM would be trained would be to mask and unmask tokens at random. So what this means is in this particular sequence the chance that I would land up at a training example where this specific token is masked is not guaranteed. Right? So in one epoch if every token is masked at random you have a 50% chance. In two epochs you get another shot and the probability rises up to 75%. And generally in any epoch you have the probability being essentially 1 - 1 / 2^n. Now can we do something better? So can we not rely on probability to get us to interesting sequences to train on? So in this work we introduced this technique called complementary unmasking. And what this technique basically does is it starts with a copy of the sequence where tokens are masked at random and maybe here we end up being unlucky. So we end up having some tokens which are masked which are not the most interesting tokens. So if you recall the most interesting token in this sequence is this middle token of 'dog'. And this is not masked. So the model is not training on to unmask interesting tokens. But what we could do is create a complement of it. So what we do is we create a second sequence which basically masks all the tokens which were unmasked in copy one and the opposite for the masked tokens themselves. So you get these two copies and you train the model on both these copies and you do repeat this process across all your sequences. This guarantees you 100% coverage of the tokens in every epoch. And this technique even though we developed it in the context of vision-language models based on diffusion, actually now it has become a standard for all modern-day diffusion language models. It's used in LLaDA, FastM, and a lot of the other models that are coming out in the open-source world.

So how do we bring some of these things together? So our overall training paradigm essentially involves two steps of training. So you have a pre-trained diffusion language model. You first learn to process the inputs in the ways I described to have both local and global views that gives you a set of vision embeddings that you then append to your text prompt. And then you use complementary masking to train this diffusion language model. In practice, what we find useful is a two-stage training process. The first stage we only train the vision encoders and the MLPs that convert your visual inputs to the relevant embeddings and in the second stage try to train the full model end-to-end.

Okay. How well does this work? It works pretty well. So on visual benchmarks, and here many of them test for different capabilities in science, in math, in general data analysis. And here we're comparing against autoregressive models like LLaVA-Next. You can see that LaVIDA, an 8 billion parameter diffusion model, is able to get high performance across benchmarks.

Okay. So what's also interesting beyond traditional benchmarks is because diffusion LMs can unmask tokens in parallel there are new use cases that you can enable with these models which are not possible by design with autoregressive models. So here's an example. So let's say I have an image and I want to generate a poem with a specific structure. For example, I want every line of my poem to start with a certain alphabet. Now with an autoregressive language model typically you would resort to prompting to get this kind of output. With a diffusion language model you can actually create your prompt such that this is guaranteed by design. So you could write a prompt to be 'write a poem with four lines' and already feed this to be the draft answer. So you already unmask tokens which you want to be set in a certain way, in this case the first token of every line, and that's how you can get an initial prompt and then when you pass it through this model you can get an output which exactly follows this kind of structure.

You can use it in a lot of creative contexts as well. So here is another example for how you can do visually grounded text editing. So here again, if you have a story that's associated with this kind of image and you want to make certain edits, with an autoregressive model, you again would have to resort to prompting. In this case, you can already feed this as a draft and replace any tokens you want to replace with masked tokens. So just by masking tokens you do want edited, you can actually guarantee an edit which is semantically consistent with what the user wants.

So and there are many such examples in the paper if you're interested. A second thing I wanted to talk about with these models is how do we actually improve their speed, right? So that's been the key focus of my talk so far as well that diffusion language models are not just about repeating what can already be done by existing autoregressive language models. They actually enable new capabilities and one of these main capabilities is the fast inference that you get with these models.

So now in any diffusion model including those which are visual in nature you typically would have a setting where you have to specify the noise schedule at which you unmask tokens right so basically which tokens you should be unmasking at which points in time and what rate these are things that you have to specify in a noise schedule. So in this case we experimented with different kinds of noise schedules. And here are the different shapes and what they correspond to in terms of the time mappings of these once you rescale them. And here what you find is basically that the convex schedule for the noising works extremely well. So what you're seeing in this chart is basically how many function evaluations you have to do before you get the final sequence. So if you're doing 100% function evaluations that basically corresponds to denoising one token at a time. So in this case you get no inference benefits in the last column but you will get very good quality. So what we wanted to see is are there noise schedules which actually have a smaller number of function evaluations and still preserve the quality and what we find is that the convex schedule is extremely good in this regard. So with the convex schedule you can actually get similar performance to that you would get with one token at a time and masking but needing only 25% function evaluations.

So this vastly accelerates the models. Now a second acceleration technique is how to use caching. So number of function evaluations even though it's correlated with wall clock time it is not exactly the same as wall clock time with modern-day generative models because of a lot of enhancements that have been done with autoregressive models in particular the use of caching mechanisms.

So what does KV caching basically do? Well, it basically exploits the fact that when you have a causal mask, as you do with autoregressive models, you can reduce the complexity of generating the nth token from quadratic to linear by storing the keys and values of all the previous tokens. So, attention in this particular case actually becomes a linear operator. Now if you wanted to apply the same kind of technique with diffusion language models you cannot naively because diffusion language models do not impose left-to-right order which means their default masking is actually bidirectional. So caching is not allowed by default. What we did in this work was to actually introduce a novel form of caching which we call as prompt caching. It is reminiscent of certain works called prefix LMs. So in this case what we do with prefix DLMs is basically have the prompt and the image inside the prompt attend to each other. So it is bidirectional inside the prompt but when it comes to the sequence that's actually being generated and which would be the noise initially and gradually refined to the final answer that sequence attends to all tokens. So you get some kind of partial caching but not full caching. But this is still a better trade-off than having no caching at all.

So this is again showing you results of how this works in practice. So you're looking at the same benchmark we had earlier, one of the image captioning benchmarks and here you're looking at full DLM with full bidirectional mask and prefix DLM. Regardless of your acceleration rate you find that prefix DLMs get you very similar quality but the latency is dropping by almost 3 to 4x so you can get similar quality but also the benefits of caching by caching only the prefix here.

Now how do these things come together? So here again is looking at LaVIDA comparing against LLaVA and another concurrent work called Mammada which was also diffusion based. So comparing against both AR and diffusion-based LMs, these inference acceleration techniques work exceedingly well. So for this specific image and prompt you can get an answer in just 2 seconds while other techniques would take much longer.

Okay. So that was the first part about multimodal diffusion language models that I want to talk about. In the remainder of the time, I also wanted to briefly go and discuss about how these techniques are now affecting multimodal generation and reasoning and really leading to a foundation for world models that are entirely based on the diffusion paradigm.

Okay. So again this is referring to two papers from this year. Again, both of them led by Shuan. So the first paper is called LaVIDA-O. Now the O here stands for Omni. So what we're basically doing differently from the previous one is we're allowing the outputs to be multimodal as well. Right? So in this specific example if we're asking to replace the animal in the image with a fish what we expect the diffusion language model to do is two things. First is understand that the animal being talked about here is actually the cat. So there is a linguistic component to understanding what's happening here. And then beyond understanding this, it also needs to then generate the corresponding image. So there's a visual generation component as well. Now this can get really complex. So in a follow-up work, we actually then post-trained these kinds of models with reinforcement learning and allowed them to reason. So this is a kind of prompt that becomes more solvable once you post-train these models with reinforcement learning. So here for example you see this image of a small boat and a few people on the boat and the prompt being asked here is to basically figure out who in this picture would be the best person to take people on a tour as well as provide bounding boxes for that person. So as you can see here in this case the model needs to realize that out of the three people the person steering the boat is this gentleman at the back and then not only recognize this individual but then be able to localize it and provide bounding boxes and then eventually render the bounding boxes in the image. So there's a lot of steps involved in getting the right answer and also being able to audit the process that it used to get these answers through these reasoning traces. So that's what we ended up doing in LaVIDA-R1 work.

So now what I'm going to do is I'm going to talk about both these two works in the sense of talking about two challenges. There are a few others that are listed in the papers but I want to focus on two challenges which were pretty interesting from the perspective of trying to address them in the diffusion context.

The first challenge was in trying to enable these kinds of multimodal outputs with these models. We realized that different modalities require very different amounts of modeling capacity to work in the real world. And here are just some statistics to drive home that point. So if you look at the open-weight models out there today, the largest open-weight model in the language space as of at least last week was Kimi K3 and this is about 3 trillion parameter models. Now if you look at the same image generation model it's about 80 billion parameters. So there's basically an orders of magnitude difference in the size of these different models. And that's why different modalities would require different capacities.

So when it came to designing multimodal LLMs, we enforced a similar behavior where we basically said that when we're trying to understand language, we would want to assign many more parameters than if you were trying to generate images. So all of this basically led to a design which we call as the Elastic Mixture of Transformers. Now what it basically does is three different things. So the first is where we actually build on prior work of Mixture of Transformers itself which proposes to share attention across layers but nothing else. So if you're trying to model different modalities, images, text, maybe even audio, video, you do not want to share all the layers across all modalities because that may not be the right approach to looking at different modalities in unison. You might want to think about modeling only a subset and Mixture of Transformers recommends looking at only attention sharing across the different layers. What we did on top of this based on our experiments was to then realize that actually when you're thinking about multimodal generation the generation bit itself also requires different capacity and this is what I was describing here where you see biggest language models are much bigger than the image generation models. So here we said we're going to break the layers into two parts. The first M layers which would look exactly like Mixture of Transformers in terms of sharing attention layers but then we stop sharing even these layers as we go to the last scale layers because that's where we want the model to start specializing to specific domains even more. And the last thing here again is as we do this kind of specialization we try to ensure that there is an asymmetry in how much modeling capacity is being assigned to different modalities because as we discussed images for instance would require much lesser modeling capacity than language itself.

So what does this mean from a downstream perspective? Well, it means depending on what task you care about, you're actually going to be activating many different neural networks within the same architecture. So if you care about image understanding, your inputs are images, your outputs are text. So you do not need to activate branches which would require the text part of this diagram. So you can only activate branches which do require images. In this case if you care about both image understanding and generation then you can choose a different set of output parameters. So in this case the yellow ones correspond to let's say image generation. Then you can activate only parameters which would correspond to image outputs here. And as you can see this would correspond to different number of parameters that are being opened here. And finally if you care about both images and video that's when you would activate the full network and take the full capacity of the neural network into account.

Now this is not just a way to model data better. These are actually a lot of implications on the efficiency of the whole paradigm. So this is comparing how this elasticity in modeling helps in the efficiency. So if you look at just training you find like a 3x speedup over the conventional paradigm. And during inference this just translates over. So depending on whether you're doing let's say text-to-image where you'll only activate the image branch or whether you're doing object grounding in which case you'll only activate the language branch at the output or math reasoning where again the language bit is the only one that's going to get activated, you can actually do much better compared to both the original design of Mixture of Transformers as well as a lot of the state-of-the-art LLMs out there which process both multimodal inputs and have multimodal outputs. So consistently for both training and inference you're seeing about 2 to 4x speedups.

Now these models are really competing with the frontier out here. So these are just some of the results. There are many many more tasks that they can do once you have both multimodal inputs and outputs. These are some on text-image generation. They can also do more intricate forms of understanding. So if I want to do localization of these different objects in the image I can get bounding boxes. I can render them through simply prompting. I can do a lot of editing of images out here to get very high quality transformations of input images based on what the user wants. So a lot of things in the multimodal domain are now possible with these kinds of diffusion language models.

Okay. So the second challenge that I want to talk about in bringing about these results relates more to reasoning. So there are a few moments in AI in the last five years which have grabbed eyeballs all over the world. Of course there was the ChatGPT moment but then very soon there was also the DeepSeek moment which really brought reasoning back to the core and reasoning primarily today is enabled by reinforcement learning. And I won't go into all the different objectives for post-training, but one of the key things that is needed to enable reinforcement learning is the ability to evaluate likelihoods, right? So this is also something which is actually a challenge for diffusion language models. So unlike autoregressive language models where you do have actually a very nice estimate for your likelihood, with diffusion language models you do not have that very tractably.

So a lot of the literature both in my research group and in the community has focused on trying to get efficient estimators to the likelihood and broadly there are a few things that one can play with. So what I'll focus in this specific talk is how you play with the weighting of the different terms that do come across when you're doing these kinds of likelihood estimation. So here basically on the left you have a likelihood on the right you're doing some kind of Monte Carlo estimate of the likelihood where you're trying to predict some tokens which were earlier masked in the input sequence and then you're trying to unmask them. While you're computing this you also try to assign some kind of weighting based on how you think these different terms should hold importance.

So there is a huge literature on trying to do these kinds of likelihood estimation. So the first paper trying to apply reinforcement learning for diffusion language models, this was led by my student Cenne who's here. What this basically did was show that you can start with a sequence that's completely masked and use a very simple mean-field kind of approximation where eventually the weight that you choose in this case was a simple 1/t kind of weighting. So in this case because you started with an all-masked sequence you would end up with every token having a uniform weight in this case. Now the problem with this kind of approach was that there's a big train and inference mismatch because when you're doing inference you will not always be having sequences which are completely masked. So while it works really well in terms of certain benchmarks when we tried to apply it for multimodal domains we found this to be a big challenge to use this kind of estimator.

Soon after D1 there was a work called UniGRPO which basically tries to change the input from being fully masked to something which is actually having tokens which are randomly masked and in this case again they use 1/t. Now this was trying to address a train-inference mismatch. So you would see lots of different kinds of sequences but you would have high variance because you're essentially choosing a random technique for masking.

Now what we ended up using in this work which we found very useful was to actually go back to what I was mentioning earlier in this talk which was this technique of complementary masking. So this is biased but just as a recall of what it does was here you would sample two time steps, one would be sampled uniformly at random the other would try to look at the complement of it and so in this case t2 would just be 1 minus t1 and based on these time steps you would generate two sequences specifically the first sequence would be generated with random masks but the second sequence would do the inverse of it. So whichever tokens were masked in the first sequence...

become unmasked in the second one and vice versa for the unmasked tokens. And doing this kind of sequence complementary masking, we end up getting every token to be either masked in one of the two copies and you get some kind of weighting. Now here there was one defect in practice that we saw which was that while trying to do reinforcement learning where stability is extremely key to training these models at scale, you ended up getting very different weights for different sequences. So the one overt weighting scheme was not something which was giving us good results even though we were getting a very diverse set of sequences. So what we then proposed in this case was complementary masking but using a uniform set of weights. And here basically every sequence and every masked token gets an equal weight. And this ended up giving us an estimator which was biased but also low variance and also had less mismatch between training and inference. So empirically it was able to cover for a lot of the demerits that we saw in previous schemes.

Yeah. And then applying this kind of reinforcement learning techniques helped us drastically improve across a bunch of benchmarks in math reasoning, question answering, image editing, question object grounding. So this is a fun example of what you can achieve when you train these models with reinforcement learning. So here we are trying to generate an image with the prompt of a light source that tried to display scandals in the early 20th century. If you look at a model that's been trained without reinforcement learning like LLaVA, it only looks at keywords like candles and 20th century and generates you an image which has scandals and a 20th century backdrop. But with reinforcement learning, it's able to reason the exact semantics of what's being asked about, reflect it in its trace and then actually generate an example which has a light bulb.

Okay. So this is my last slide where I want to focus on why we think diffusion can actually be a defining moment for the future of AI. So if I think about how AI will progress from where we are today to where we want to be, I think we really need to start thinking about not just pure intelligence as in a language model would do, but really intelligence that's universal and takes into account sustainability of how efficient it is. And universality would require us to reason over multiple modalities in parallel. And what I've been showing in this talk is in my group we have been thinking about how diffusion language models both now becoming mainstream at Inception and at UCLA how they've been trying to adapt to multimodal domains. And efficiency really requires us to go beyond the statistics and also deeper into the systems of how we can actually deploy these models and how we can think of techniques from all over the stack to make them efficient for pre-training, post-training as well as deployment during inference. With that, I'd like to conclude and thank all my students and collaborators for all the work they did here. Thank you.

Thank you very much. So we have time for a couple of questions and we have a microphone. Yeah, there's one.

Thank you. Really love the talk. Especially some of these visualizations, they're very nice. I had a question though regarding the decoding. How do you predetermine the sequence length that you need for the response? Is it using like a token for blank characters for example?

Yeah, I think that's a great question. So unlike images, language does not have a predefined output length. So what ends up working really well in practice is what is known as block diffusion. So you do diffusion in blocks. So every block itself will generate tokens in parallel. But across blocks you can have a causal dependency where they would attend to the previous blocks.

I see. Is there a did you do any ablations over different block sizes and see does it depend on what did you find to be the ideal block size essentially?

Yeah. And that's actually one of the key choices. Not just because of the statistics again. So during training often it's common to use random block sizes because different sequences might lend themselves to different degrees of parallelization. But then at the systems layer the bigger the block size the more time it would take for the forward pass in terms of the overhead adds. So it's something you would ideally tune based on the application. Typically you would see block sizes in the community being anywhere from like 8 to more recently going all the way up to 256.

Thank you.

Yeah. So my understanding is this technique technology is used for multi-modal AI. You have a picture, image and also text image you mix up in the like article. So then you trying to convert the image into some language and then you know just input the converted language with your text into the transformer and then is that how it works?

Yes. So what we're trying to basically do is have any kind of input and any kind of output. Now how these things get combined is what is dictated by the transformer. So just like humans we can be reading inputs and also watching visual imagery and process all of it and at the output of these neural networks you would try to generate language in the form of reasoning as well as some consequence of that which could be an image output. Right? So the output itself is left open to being an arbitrary modality as well as the inputs.

Okay. But the input into the transformer is a language, right? Is a language token.

Yeah, it definitely is language but it can also actually be images. So if you want to have both of these, you can feed both of them into.

But the image you have to convert into the language first before you input into the.

Yeah. So you have to convert, you have to tokenize the image. Yeah. So you have to bring it in the same latent space as the language token.

Okay. Thank you.

Thank you.

All right. Thank you so much for the introduction and also thank you so much for inviting me here and letting me be part of this really nicely organized event. It's also my first time in Chicago and I really love it here. It's a beautiful city. Okay. So today I want to talk a little bit about world models from the perspective of robotics and what we actually need there. And for that can I just have a quick show of hands who worked with robots before in this room? Oh Alexander I see you now. Hi there. Okay a handful of people. So this is going to be I'm going to make it very hand-wavy. So maybe to convey what I'm trying to say to everyone. And if you have a specific question, please feel free to ask directly during the talk. You don't have to wait until the end. That maybe makes it a little bit more interactive as well. Okay, so today what I wanted to specifically focus on is something that Aditya actually also focused on which is latency. That is in robotics an even bigger problem than in let's say in the example that you brought which was like customer service. And just to have you understand this a little bit is like the robots that I'm going to talk about today you can actually send them a control command every millisecond. So it's a kilohertz that you have to serve a robot a command to and that is a latency requirement that is very different than again for customer service or for these GenAI chatbot kind of applications responding to user inputs. So from my perspective there are a lot of things you want to do on a robot that do not afford test time compute, not this kind of iteration that a lot of people are thinking about. But there are also applications in robotics especially reasoning where this is actually maybe okay you have maybe a few seconds it's okay to wait. So today I thought I gonna take three works from my lab that kind of cut off but I hope it's okay. But basically what I'm trying to do today is I give you three examples from my lab on work in robotics that we have done where we use some kind of world model. I'm going to use the term world model a little bit let's say liberally here. It's for example I going to use a simulator like a traditional simulator that I consider actually also a world model. And yeah so I'll give you a few examples in a few settings and based on those three examples I hope I can visualize for you of where in robotics we probably need world models and what kind of input and output they should actually have for these different tasks and what the latency requirement is as well. Okay so the first work I want to talk about is called Sim2Real and that's basically one of these scenarios where you have a robot that literally wants a command every second and so these are really this here for example closed loop policy and it needs milliseconds outputs to control this robot and have it react really with its fingers to really balance these tools without letting them slip. And then the second work I going to talk about is one where a robot actually has to navigate a messy environment. So imagine having a robot just like roaming around in your apartment which I going to assume is not free of clutter. Mine certainly has a lot of clutter and so a traditional path planner that you would normally apply employ on a robot that avoids collision at all costs is not really possible in that case. You want a robot that actually finds a path by maybe interacting with the environment and pushing things away to make some space. And that is probably fine. You can probably give that robot a few seconds to reason about. And then the last thing I want to talk about is a setting that is a little bit different. That is a setting where you kind of have to do creative physical problem solving. Think of a MacGyver robot. Okay, if you know MacGyver still I don't know. I'm probably also too old. You're probably all younger and you don't know MacGyver anymore, but basically the setting is you have a physical problem and you can do a few trials to figure out like how to solve it. And every trial provides you with some information that you then want to condition the next trial on. You want to basically make a smart choice on what to try next to hopefully solve the problem. So that's a little bit of a different setting that we are exploring here and there you obviously since you do multiple trials it's okay to think a few seconds about it. You would probably as a human do the same. Okay. So these are the three works I going to lay out and then at the very end I'm going to dish out some opinions I have on world models as well and then I hope I can instigate some discussion about it as well. Okay so let's start with this work. So I'm going to talk about it technically what it does and then again please feel free to ask questions because maybe I'm assuming you have some robotics knowledge that actually you don't and you need some clarification.

Okay. So let's talk about this work. First of all dexterous manipulation is really hard in robotics. And the typical front page news you see right now about what robots can do, they actually use a gripper. Okay? And if you use a gripper, you have a much smaller dimensionality of the action space actually that a robot has to do compared to a robot like that. It has a 22-dimensional anthropomorphic robot hand. So basically again it's a seven DOF hand plus the arm as well which is seven degrees of freedom. So very high dimensional problem that our model or policy has to output in. It is really hard because it has to continuously do multiple context switches. It has to do these kind of like low-level skills like first grasping then rotating the tool in hand and then actually using it as well. Sorry for the little lagginess there. Yeah but feel free to visit the website and actually look at the work. And then again at the lowest level we are sending milliseconds command to the robot every millisecond but the policy of course it's pretty slow still in robotics sense so it's at 60 Hz. Okay. All right. So one thing that is going on in robotics right now is that a lot of people are collecting data for robot especially the ones with the grippers through teleoperation right like so they basically have some kind of joystick or other device to do that and a lot of people invest a lot of money in people now collecting this data and again it's typically with a gripper if you do it with a dexterous hand it's really painful if you ever have the opportunity to try to teleoperate a dexterous hand with your own hands in some way. It's so hard. Like literally I tried it. My student forced me to do it and I was amazed how weird it is and I couldn't do the simplest thing. And so, the way this expresses itself here in this video is that you can see these people are trying to do these tasks with this complex multi-fingered hand, but really what they're doing is just using it as a fancy gripper. Okay, it's just pick and place because it's so hard to really do this kind of tool use that you've seen before in the videos. So, we really don't want to do that. We could use sim to real so basically train a reinforcement learning policy in simulation and then you know try to somehow bring it to the real world there is a gap between simulation and reality but we've seen quite a bit of success with the seminal work for example by OpenAI in 2018 before they shut down their robotics or for the first time now they're back in robotics but you know in 2020 actually they shut it down but this was like one of the headline results that they had. This looks very cool. Here in the second paper they solved even a Rubik's cube that again there was a high level policy that actually did the solving and then this one did the low-level control of this. And then we have other works like this one that finally turns around the hand and now has to fight against gravity a bit more and does also in-hand reorientation. So we could do that. But these prior works they usually learn one policy per skill right like per kind of low-level skill either in-hand rotation or spinning an object or stable tool use. But what we really wanted to do is or we asked the question how can we actually learn one single policy that can do all of these skills together. And so the key insight that we had in this work that I'm going to talk about is that you can actually unify manipulation of goal reaching. So essentially what that means is we want to learn a policy that takes an object goal pose. So a position and an orientation of an object as input along with the current state of the robot but also of the object. And so this work we called Sim2Real which is this object-centric policy for zero-shot dexterous tool manipulation together with Karol and our two students Kushal and Tyler. Okay so how does it work? Basically we have this unified objective where at train time we generate a large set of primitive objects in simulation and train a single RL policy that can manipulate each object to random goal poses. So this is kind of how this looks like. We sample these random goal poses and then we have this massively parallel reinforcement learning problem in simulation. I'm not going to talk about what algorithm we use and all of this. I can talk about that later maybe when if you are interested. So yeah we basically train it with a specific algorithm. And then kind of this is how this looks like here are the inputs the proprioception. So the state of the robot the current pose of the object and kind of a grasp bounding box that we want actually the hand to explore around and the sequence of goal poses. And then we learn this policy to do these tasks. Okay. So now that was training. What happens at inference? We basically at test time collect a human video demonstration. And then we extract a sequence of object goal poses. Notice that we don't care about what the person actually does. We just care about what the object actually does because that's what we want the robot to imitate, not what the person actually does. And then we track these goals in closed loop with this RL policy to do exactly the same. And again it's easy for humans to do these kind of things that we're showing but this capability hasn't been shown on a dex hand like this before. especially not where we have one policy that can zero-shot do use the same policy on all of these different tools that it hasn't seen in simulation. And also doing all these different tasks that again it also hasn't seen in simulation. Yeah, there's a question.

Yes, it's so we are using as input a camera only but also the state of the robot. So we can super easily figure out from the robot. It basically provides us with what how my joint is configured essentially my hand and my arm as well. That's the input and the object pose comes from a camera. Yeah. Okay. All right. So that's pretty cool. So here are some baseline comparisons, right? like ours is on the very left. And then we have a one baseline that basically figures out a grasp first and then doesn't do this in-hand manipulation and kind of fails because really it does need to rotate this brush. And then we have also a baseline that does kinematic retargeting. So basically tries to imitate what the person does. which is even though the hand looks anthropomorphic is still very hard even especially the thumb in anthropomorphic hands works very differently from how a human thumb actually works. It's the most complex finger to actually model. Yeah so if you are interested to look at this policy a little bit more feel free to play around with it on our website. you have this fun kind of little play thing there where you can basically manipulate this goal pose and then make the robot fail as well. Figure out what are the boundaries if you're interested. And what is nice is really that when these failures occurs also the policy shows very strong recovery behaviors when the hammers for example drops right like so sometimes it does make mistakes but it kind of recovers from it as well which is nice. So another thing we found is that this policies that we call Sim2Real can play it can kind of move these objects around in the way we want but it cannot really assemble because assemble means you have to be super precise right like you have to really insert for here for example this peg and this hole right like these are the kinds of things that you have to do under really tight tolerance and the tighter we go basically the less and less success this policy actually has. Okay, so here's some follow-up work that we did which we call Play to Perfect where we wanted to basically figure out if we can actually take this what we call here the play policy the Sim2Real and make it specialize it basically for all of these different much more complex assembly tasks that require much more precision. So right now Sim2Real already has these basic capabilities pick grasp and in-hand reorientation and now we wanted to basically do these assembly tasks as well. And so here is how this policy looks at inference time. It basically again gets the pose through a camera of the part of the object and of the fixture as well. And the goal pose as well and then again the robot state and some information about the object as well. And the policy that is used is this fine-tuned policy that outputs joint position targets. And it's display pre-training with the Sim2Real yields really effective assembly policy. So here are like a few of these examples that we did. So peg and hole or assembling like a whole I think it's I forget like what this object actually is but it has like a pillar and then a beam as well. And then here it's assembling a toy table one leg of it basically and screws it in. Yeah, and so one thing we found is that this using this play policy to then fine-tune it for these more specialized tasks works actually great. But there's also things you have to pay attention to when you do this. So what kind of play objectives you use actually matters then how good you are or how fast you at learning this new task. So first of all it definitely makes sense to fine-tune this play policy right like you need that because you definitely get better performance especially for tight tolerance tasks it's very fast to train even if compared to using a very dense rewards and it also compared to this dense reward policy it learns less brittle policies as well. The policies that actually come out of the dense reward, they have like weird things like balancing it on fingers and things like that. It would never work in the real world. Yeah. And so here are just a few examples that do all these different tasks. And again, the input is this whole scene is like seen through a visual sensor. And yeah, we're pretty happy with this. Again, I haven't seen basically this kind of performance on a Dexter hand yet. So we're very excited about this and continuing to work on this. Okay. So this was the first work. And so now let's move again. This was a work where I don't see how test time compute will help us at all. But I come back to my maybe controversial thoughts about this later at the when I wrap up. Okay. So let's move on to this second work from my lab that is again about this mobile robot that kind of is coming home to your home, right? Like and you have maybe you have kids that are letting stuff fly around all the time, right? And you have to clean up after them. Would be nice to have a robot that kind of carefully maneuvers around in your place and is able to touch things as well. That maybe could I think benefit from test time compute. So let's talk about how we did this. So what we call this work it's we call it MessyNav for messy navigation where we want to basically enable zero-shot navigation via manipulation of messy environments. So zero-shot here means we don't want to collect data necessarily in every apartment or small business or whatever this robot is going to encounter. And so, here is how our lab sometimes looks like. Maybe that's a metaphor for your home. I don't know. But this is how our lab sometimes looks like. And there are even robots lying in the way. And so we really need a robot that doesn't just you know, give up in these moments. But they do because classical navigation is like a car navigating on the street by all means wants to avoid collision with anything, right? Like which makes sense, but I think in a real world smaller environment like our homes or small businesses, warehouses and shops and things like that, we really want a robot to be able to touch things, push things, out of the way. So we're asking this question, what if this robot could decide which objects to move and how? And so here I have an example where we see some kind of like little cart with like a plastic box and a bunch of balls in it. But on the right side you have like another cart that actually has a bunch of blocks on it that maybe some kid, you know, built and now would be devastated if all of that stuff kind of falls apart, right? Like because the robot touched it. So, we really want the robot to be able to reason about this and then instead of destroying potentially this beautiful piece that this kid built, swipes actually or moves the left one, decides that that one is okay to move and it wouldn't result in any catastrophe. Okay. So, once it has done that, then suddenly the path is clear and can actually reach the goal for whatever it wants to reach. Okay. So going to the open world is tricky though, right? So, the robot must be able to operate across diverse environments and novel objects. So, we tested this thing actually in you know our robots in our robot test suites but also outside on Stanford campus during the day and during the night as well. And we tested with novel objects that we haven't seen or the robot hasn't definitely not specifically trained on. So there are scooters of course on campus. There are fragile objects on stools, piles of books and stuff that maybe like this plush animal that can be totally touched. And so what is interesting is the same object's affordance, let's say, or whether it can be interacted with or not can really change based on context. So this wheeled cart again by itself is totally fine to be pushed but as soon as we place a stack of cups or a laptop on it maybe we don't want to touch it. It becomes fragile and non-interactive. And we want our system to be able to deal with these open world challenges in this interactive navigation setup. Okay. So just walking you through this basically this MessyNav framework is like a zero-shot framework for open world interactive navigation to help us traverse these challenging situations and given the RGB image that it sees through its camera MessyNav first reasons on object interactability. Can the robot physically move the object? Are the objects fragile and stable or sentimental? And doing so helps it to identify that you know books on a wheeled suitcase is non-interactive while the wheeled metal cart with a book are interactive. Okay. And then once that happens basically we can plug all these labels basically into a planner that then plans like okay I can actually interact with this and that object because it's non-fragile and I want it to go in this particular space. So this space is free so I can push it in this way. And it does that kind of reasoning. Okay. There are a bunch of experimental results. So these results hold over all of these different VLMs. We are better if we actually explicitly reason about the affordance first instead of squishing everything into one of these models. And it matters. Yeah, and we also looked at all of these different scenes and how well it works. It doesn't work so well in the night. Okay so basically the takeaways are explicitly reasoning about the affordances of objects and task planning provides performance advantages over just zero-shot doing interactive navigation and then there's still a lot of challenging tasks like nighttime and geometric understanding that remains really challenging so I think maybe one takeaway here is that one of the things we noticed in this project is that current off-the-shelf VLMs definitely have very poor spatial understanding and this is something especially for mobile robots that is extremely important to do and yet we cannot use any of these VLMs to actually tell us something about oh can I push this object to the right is there something in the way what will actually happen will I have it pushed enough will I actually have space to move through honestly things are very difficult to actually figure out just from a single image and current VLMs are definitely not able to do that so again something that I'll discuss a little bit more at the end of this talk when I wrap up okay so let's move to the third project which is we call CausalPick and remember this is the MacGyver setup right like where you are given like a physical problem you have a bunch of tools or objects that you can use to solve this problem and you have a bunch of trials to kind of figure it out and every trial while you're learning something about whether you're already close and you just need to nudge change things a little bit or whether you have to completely switch your strategy and maybe do something else. Right? So that's the setting here. Okay so again this work is called CausalPick and we presented this at ICML last year. And so again the problem is we as I said is that we want to solve these physical tasks that are time extended and then give the agent or our robot in this case the ability to learn at inference basically from these trials. Okay. So here's again this problem formulation that we work with is a single intervention task where an agent chooses what action to take at the beginning of an episode and then observes the action's effect on the environment. If the agent doesn't succeed, it can attempt to solve the task again until it runs out of attempts. And how many attempts does it have? We're talking here about not millions like in RL. we're talking about three to five maybe something like that. Okay. Okay. So yeah so this is just kind of the model that shows this. These are the two benchmarks that we actually used that include a variety of single intervention tasks and they're called VirtualTools and Fire. They look deceptively simple. They are all 2D and they are actually very hard to solve. And these tasks are involving complex physical interactions between objects with which without complete information about the environment dynamics make it impossible for agent to plan kind of offline and then just open loop basically find the real solution. they really have to go in and make a trial to figure out what's actually going on. How heavy are these objects? What's the restitution coefficients? All of these things. And so if you don't believe me, you can go to our website and my student Kolada, she made like a little game so you can actually play these two challenges as well. And by the way, we are doing better than humans. I show you later in the results. Okay. So anyway so here is at a high level what this does CausalPick method does. So Causal includes a causality based action selection module. It's basically based on Bayesian optimization if you've heard about that and it proposes the best action to execute next based on a prediction from a Gaussian process that it uses. The Gaussian process is used to model a scoring function for the different actions and these predictions are generated by considering the rewards from the previous action that the agent has taken and the causal similarities between the candidate actions computed by this physics informed kernel that we are using. And so these similarities are derived from the individual effects of actions as predicted by a learned dynamics model. So aka the world model if you want to think about it like that. Okay so this kind of how this looks like. Let's dive in a little bit deeper. So given a set of randomly sampled candidate actions we use the predictions from the learned dynamics models or world model to compute the state changes of individual objects caused by a single action. So here for example so again here you have the single actions and now here for example this is maybe a little bit small but I hope you see it basically the actions that the agent can take is like picking one of these blue objects that you see there in the frames and placing them in the environments and then pressing play and gravity turns on or the object falls. The goal is to have any of the red balls to be in the green area. Okay, so in this case, you want one of these red balls to be in the green bucket just by using an object and letting it fall. Okay. So now you can see for candidate action one that I show here the dynamics model predicts that objects two to four so the red balls two to four will remain static while object one will move in the southeast direction and that is expressed as s dot of 1 one there in this kind of horizontal line. So oh actually I can maybe go here. So again it's maybe a little bit hard to see from the back but there's this purple arrow here on this object one that basically that the learned dynamics model predicts where this ball is going to move to. Okay. So and then we can sample another candidate action number two and objects one to three are predicted to remain static while object four is predicted to actually move. Right? like so the ball on the very right there that is predicted to move when I do this action when I take this blue object and let it drop basically. Okay and so then we can have a third action right like and it shows like that kind of the same thing happens as with the first action there. So now that we have like these action candidates, we're sampling more like hundreds, but we have all these action candidates and now we can compute pairwise similarities between these actions and we want the similarities to be based on the effect of what the action have. So how do we evaluate the individual effects of these actions? So we introduce causal reasoning in the physics informed kernel to assess the impact of an action on the environment and we evaluate action outcomes against a counterfactual baseline like the one on the screen where no action object is placed in the environment. So this one is actually our counterfactual baseline where basically we just turn on gravity no object is being placed. And then we want to measure the difference that an action makes against this counterfactual baseline. Right? So here is just turning gravity on versus actually placing an action an object. And now you can ask the question what made the red ball move in both of these cases and is it some kind of confounding factors arising from the dynamics of the environment. So here it's gravity right or is it a direct causal effect from the action of placing the object in this particular way. Okay. So now in CausalPick what we do is we find that the first and the last action here action number one and action number three they moved the same object right like the one on the very left with similar direction and magnitude as well. So that means their similarity is actually sorry their similarity is actually very high expressed here by this link where the causal similarity is like more towards yellow and green right and then if you compare the first and the third action they actually move different objects causing very different effects in the environment. So their similarity is actually minimal. And so now we can compute this for all our actions. We again we have actually we sample actually hundreds and compute this. And then we compute this matrix to form to build a kernel function that is used by the Gaussian process to approximate the score that actually each action has in terms of getting us closer to the goal. And this process is scored based on previous candidates as well. Okay so let's see how this looks like right like so this is let's suppose like this was the following this was a rollout that we observed we already did that trial and so now the physics informed kernel is actually able to reason about the following first the fact that the red ball moved to the left was directly caused by the action object and second compared to the counterfactual baseline we moved further away from the goal than we actually wanted remembered we want the red ball to be in that green cup. So now the question now you all have to wake up. I'm going to ask you is this going to be a quiz? Okay. So now would CausalPick try the following action of like imagine again it basically this red ball there is going to be on that gray block, right? So, who thinks CausalPick will pick that action next to kind of get to the goal of putting the red ball into the green after having observed this outcome on the very left. You think it will pick it? Okay.

Maybe it drops high enough like goes up the ramp and skips it.

Okay. I saw some other hands. Okay. Who thinks CausalPick will pick something different? Okay, a little bit more hands. Okay. And then a lot of people very confused. Okay, that's great. Okay. Okay, so here is Oh, sorry. Yeah, sorry. I messed this up like so bad. Okay, sorry for this. Actually, I was going to I was going to say like the first one on top, that's the action that CausalPick has just done and observed and the ball is moving to the left, right? And so now my question was actually is it going to pick that one? the one where you have the diamond actually falling down, right? And anyway, it's let's just assume you had answered that question. So would it basically try this action and actually CausalPick would not pick it because it actually has a very similar effect, right? Like so essentially the effect of this new action that was sampled and what it had observed before actually very similar and give you very low reward. The ball is like going it's a failure basically, right? Like so it's going in the wrong direction even further from the goal than you want it. And it's not going to pick the same action again and make the same mistake, right? Like it's definitely going to explore. Yeah, sorry I messed that up really bad. But, yeah, so it's not going to pick that one. And so, it's going to explore much more smartly. And so, here we have actually comparison of a person trying a particular puzzle. Again, you want the red ball touching the green ground. Here in this case you have our model on the very left and then you have a very commonly used kernel function in Bayesian optimization that is the radial basis function which is basically not physics informed. So your hand so and you see like basically you have an attempt of three for our model. The human kind of gets stuck like trying it found like a this person found like a close solution and is kind of like trying the same thing and just tries to get a little bit better. and then finally realizes has to switch strategy and then this RBF kernel just gets completely lost. But there was a question.

Yeah. Yes. Okay. So the physics or the question was like what do we assume about the weight distribution in the objects and the blue objects that are kind of the tools that we're using here, right? And the way this benchmark is built is it uses a very simple physics engine, right? And actually in this physics engine the distribution of mass is basically uniform in these objects. But our method actually doesn't care about that even if let's say in the diamond there was like the center of mass was actually in one of the very left corner of the diamond then it would spin right. And so because it would actually observe this effect, it would actually play and interact with it and see like, oh, this is like this is spinning, maybe it would actually kick the ball potentially to the right, right? Like then it would make a different decision because the causal effect would probably be very different, right? So I think there's the benchmark that you have to see which we're agnostic to, but the method itself, the idea is that you have to go you cannot observe everything. For example, the distribution of mass in this object, you have to actually try it. And then you have to make decisions based on that.

There's another question.

Or the other way of asking the question is that is that input of the models image based which means that if we change the environmental setting or the physics setting then we need to like kind of refine train model again. Yeah.

So there's no physical parameters of these objects are input to the model. What we do is we actually do detect the objects. So it's not pixel based. It's more like object based, right? Like it detects where all the balls, where's the tool and all of this. So it's object based. So you could think of like oh there's some kind of like detector or something running. And that's what is actually being fed into the model and used as a comparison. So yeah, so maybe we did retrain one dynamics model. So we the only thing we learn here is the dynamics model that does the prediction and we did learn I believe I'm not 100% sure but I do think we trained a different model for the two different benchmarks but the test environments that we tested on were actually held out from the training. So these are completely new physical scenarios and the physics dynamics model the world model is trained on scenarios that are then it is not tested on those training scenarios. Okay. So this virtual tools benchmark for example has like I'm forgetting the exact number but it has like 50 different scenarios and so we train on 40 of them and then we test on 10 of them something like that. Does that make sense?

So the world model kind of observe what's going on the first try and then observe problem after that.

So it tries to estimate scores of actions right like an actions here are which tool you pick that's a discrete action and then where to place it which is a continuous action. So it samples from all these possibilities and then what this model tries to estimate from these rollouts or trials that it does here is a scoring function for these actions. And the scoring is estimated based on the similarity of the effect the new sampled action has to the trials it had observed before. Does that make sense?

Yeah. Okay. More questions. There's a question over there.

Yeah, we tried that obviously. I mean, this again, this is like a year old and now you can throw Sora at it. It's not good at these kind of physical reasoning tasks. But you know, I do think it will be better and maybe we can think about more having like in context learning kind of with these videos that are coming out. It's not good yet. So, because they're not good at physical reasoning. Yeah, more questions. There's a question.

If a task is impossible, so is this model able to figure out if a task is impossible?

I don't think so actually. It would try indefinitely if the task was impossible. So, we've only confronted it with possible tasks. Yeah, it would maybe it would potentially have let's say if you compare the scores that it basically computes for actions of a possible model versus the scores of an impossible environment. Let's say the absolute score would maybe be higher because you never get close to the actual solution, right? Like in an impossible environment. But we haven't made that test actually. So this is just a hypothesis I have.

Yeah, there's another question there.

Let's say there's a task where in order to complete the task you need to first move the red ball away from the target location.

Yes. This is exactly the problem that we face. Euclidean distance is not enough for determining whether you are actually close to solving these. This is why these actually are harder than they look even though they're only in 2D and all of this. And so we had to come up with a notion of closeness or almost like you almost did it. And so the way we figure out how close you were is we actually look at the entire sequence of what these objects are doing over the entire rollout to get a notion of this. But also remember we're using Bayesian optimization based on this Gaussian process that is based on this similarity function that we constructed here. And so Bayesian optimization is a global optimizer over this again over the scoring function. And it does sometimes if it's stuck somewhere it does have a probability to pick an action it has never tried and because there's a very high uncertainty about this scoring function that it actually tries to estimate. Yeah, five minutes. Okay, perfect. Because that's almost I'm almost done here. Okay, so let's maybe go to my opinions.

I don't know what's so funny about having opinions, but anyway, we're Yes, I know. So it's doing great. It's doing better than humans, you know. That's why you should actually try it. Because I feel like people don't believe me. Okay, so basically CosPick is this novel approach that integrates a probabilistic approach with a physics-informed kernel using Bayesian optimization that enables multi-action belief updates from a single rollout and smart environment exploration derived from this causal reasoning about the effect. And yeah, so we show that it outperforms all the state-of-the-art baselines in both benchmarks while holding high alignment with humans on the physical dynamics. We actually did a user study as well, which is how we got these human numbers for both of these benchmarks too. Yeah. Okay. So I gave you a little bit of a pot of the work in our lab that are very different from each other. Okay, like super different. And my main point was like, and there's a lot of people who are very excited about using test-time compute with world models and I'm confused as a roboticist. I'm very confused about this because definitely, especially at the left end of this error, you cannot do test-time compute even if you have a very fast model. It's probably going to be a big model that is running on the cloud somewhere, and that network latency is also going to bite you again, you need like milliseconds. And so at the very least you would need to somehow run it on board, right, like to reduce that. So I'm also, I'm also love to hear from all of you if you have something like that already, right, like maybe I'm wrong. So but yeah, anyway, let's move to my opinions. I'm wondering really like when do we actually need world models and if we do, how should they actually look like, what's the input, what's the output. And so right now if you are following a little bit the robotics field maybe on Twitter or whatever and you see all the cool results, then you probably know that right now the main dominant paradigm is actually imitation learning using something called vision language action models. And what this basically is is you have a VLM and you slap an action head on top of it, typically something with flow matching. And it's kind of amazing what models or what robots can do if they're equipped with these kinds of models. So here are just like two examples. One is from Generalist, one of the kind of flashy startups out there. That shows very interesting results. And one is I took a recording from Maria Ba's talk. She is with Google DeepMind like Alexander is as well. And she shows this amazing comparison of a person trying to put a belt on around two gears and really failing hard at it and then the robot kind of figured it out. But again, yeah, it's I would say it's very impressive what comes out and more recently we have seen also something called I see it as a VLA but with a video backbone and these are called world action models right so instead of a VLM you have a video model as a backbone but basically does kind of the same thing and one cannot like the paper that actually coined this world action model term is DreamZero from Nvidia and they're not yet as flashy and you know the results that are coming out of these are like very interesting and all these papers that are using these models are showing like oh we're better than the baseline blah blah blah but all the startups are of course using VLAs mostly most of them. Okay so one thing to notice is that this is both of them even though this is a video backbone it still needs teleoperated data very painful to collect please try it and you will understand how painful it is and yeah but companies invest a lot in it they buy it from China and so on right like that's what's currently happening and yeah so that's one way on how world models are used right now in robotics. One thing to note though is like again these teleoperated data points that are being collected they're all successful and they're very carefully curated as well especially for the post-training not the pre-training but the post-training especially and also for the video model part you have to train with successful examples if you want to be successful at the task and so it takes a little bit away the advantage that I see in world models which is that you should be able to train them also on failure data so that it gets a good understanding of what the physics actually are like but it's unclear to me how you can use these world action models for example more in a world model mode that can explore versus just as a VLA with a video backbone but that only is a policy that outputs something successful. Okay. So then okay this whole teleoperation and collecting data for robots is very painful and so especially if you remember for dexterous hands if you want to do things like this you cannot teleoperate what this robot is doing. Okay. And so one interesting thing could be like what if we do learning in imagination right and we actually learn a policy from a world model. And so let me show you the state of current video models, let's say, right? So here is first of all the ground truth that you've already seen, right? Like this robot hammering a nail. This actually even has sound, but I don't play this right now. And then here is the generated version of it. First of all, the hand morphs becomes different. The hammer becomes different. Look at the nails. They make no sense, right? And actually these physical differences they really matter if you want to learn something from it. Here's like the brushing action again. The hand becomes different. What is this brush? So it's kind of hilarious and you know I mean you understand how this is going. So this actually morphs from a left hand into a right hand for drawing with a marker and does something very strange. And then here is a video from a person putting this lid on top of it. And you saw that this lid actually at some point had openings on both sides as well. And again, like if you want to do something intricate with a dexterous hand and learn in imagination, this is not good enough. You do need to have a better approximation or like an approximation of the world or prediction of the world in order to make it useful. And I understand, you know, this is going to get better and all of this, but we're definitely not there yet. Okay. So basically this part I just want to say you could learn in imagination to distill everything into a policy so that you don't have to do test-time compute especially for these time critical parts. But we're not there yet. That's kind of the idea from this slide. But then what about slower planning with world models, right? Like so here what we observed here is that in this work is that they're really bad at spatial reasoning. So I'm actually wondering if world models in the way like World Labs, a startup, actually thinks about that is explicitly geometric actually makes more sense for robotics because you can actually reason about open space that you can push things into and things like that which these image based world models and the video based models don't have like a sense of and I'm wondering if fundamentally they can actually learn this kind of spatial maybe they can you know you tell me but I feel like in robotics, you actually could put a depth sensor, for example, on it. In robotics, you have the luxury of putting all sorts of sensors on the robot. And so, you could maybe do that. And then, I feel like we need, if you wanted to use world models, kind of in the sense that most people think about them, we need some way to have memory, right? like especially in these settings that I discussed before where you have like multiple trials and you try to learn something from them to make a smart next action decision. And it's really not clear to me how to exactly do it. Maybe in-context learning would somehow work. It doesn't yet work. And also you probably want much larger scale memory. For example, again, if you have a robot coming into your house, or in an office or in a warehouse or something like that, in a small business, you probably needed to remember like almost like a pretty large map of what is going on, in this world to make again the right decisions. And we have some work in this direction where we use like scene graphs and all sorts of things, but it's not really a world model. But that's the last slide. Thank you so much for your attention and these are the students in my lab that have done all this work.

So thank you very much and also I think the Phoenix Foundation will thank you which is inside joke for people who know MacGyver but anyway I guess just us. So we're a little bit over time so if you have questions for Randall use the lunch break we have a bit more than an hour. I will see you back at 1:15 for the afternoon session and let's give a last thank you to all the speakers of this morning.

All right everyone, we will start in five minutes. Five minutes. Thanks.

All right everyone, welcome back for the afternoon session. So we will be trying to start exactly on time to maintain the schedule and I'm very happy to introduce Ilun coming from Arvad University and we will speak about topics that will remind you of the talks we had this morning about multimodality and understanding what is stochasticity in the world and what should we model or not and so with that please help me welcome Illen.

So it looks like it's plateauing is that concerning

So, so model do you think it's a model? So, you it's a model size issue that caused the plateau rather than like a fundamental gap. That's the hypothesis.

Yes.

How accurate does the reconstruction need to be of the real world function from 2D?

The robots cannot?

Uh yes, I'll talk about it now. Yeah, the verifier in general is some type of metric you want to optimize whether it's some type of 3D consistency or downstream task. But you do have to separately train a verifier for all of these. and there's one other idea which is this idea that during your generation procedure you don't have to generate everything all at once in the same noise level. You can imagine that so let's say I want to generate a plan from one location to another one. I can first generate the start of the plan and then I can gradually sequentially refine the later parts of the plan. So we don't have to necessarily generate everything all at once. So this combination allows us to do some type of backtracking search for efficient video reasoning. and the baseline of this approach is that we just have an approach which simply just from the starting state it just generates multiple different videos and you just choose the best one. So that's just essentially the first simplest search algorithm I talked about. What we do in this paper instead is when we're trying to generate a video plan again we generate during the intermediate generation. So here our entire plan is length time step t but we just generate until time step k and then after time step k we have some type of verifier that just checks is the video up to this point plausible like is it physically accurate? Do I think it's reaching my goal? It's a separate verifier that we learn. And if we assume we have access to this verifier, what we do now is that instead of generating the entire video, we generate these subsequences. We choose the sequence subsequence that seems to be making the most progress. And then we just repeatedly generate different suffixes until we get one that finally solves the final task we want to solve. So let me give an example of this. So here this is a simple puzzle where I want to move this yellow block to this red block. if I just simply sample many different samples from our video model, it always typically samples this yellow block moving downwards which leads to a dead state. We have a verifier that predicts feasibility of different states to the final goal. So using that verifier, we can determine when we have an issue here and we just regenerate until we get a plan that allows us to reach this final red goal state. and we find that this type of backtracking search is pretty useful when it comes to longer path planning tasks. So once we go to a path length of 41, we find that if we do this backtracking search with 10 or 20 video iterations, we find that it works quite a bit better than simply just sampling from the video model 20 times. And we find a pretty large boost in performance especially on OOD scenes where the original system is only around 1%. but with this search process you get a much larger boost around 27%. Yeah and then this just shows that it's pretty useful across a couple different video architectures. Here's another so here's a different task. So the task here is we want to navigate to the goal region green star shown here. and if you just directly generate from the video model originally it generates this path that kind of veers off at some point. and now what you do is this approach is the verifier detects whether you're making intermediate progress and if you're making intermediate progress you just backtrack and then try to continue to do a generation until you have a generation that can reach the target star. Yeah and we find that this improves the performance on these navigation tasks. and we can also do this in this more manipulation setting. So here we're using this video generation approach. We're given a starting image. the system will generate a sequence of videos depicting how to solve this task. and now what we find is that if you just simply generate multiple video plans and choose the best one each of these video plans have some errors. but this video shears like if you just generate from the video model it keeps on generating these plans that do not successfully finish dropping the objects. But if you do this if you have a separate verifier that detects how much progress is being made and then you backtrack then you can eventually get a video plan that is able to drop those objects in the tray.

Yeah. Did McDonald's give you any funding to

And here's just a couple videos showing the generated videos of this approach for picking up objects. So you can see it is a bit blurrier because it is a very, it's also a very small model. Not only is it single step, it's also very small and we trained it from scratch. But we find that it's still able to generate some fairly realistic actions for picking up objects.

So what makes this a world model compared to just vanilla RL or something? So I think to me a world model means you're predicting the future states of the world. Right. So here we're predicting some type of pixel state. In contrast I guess with the RL policy you're just predicting actions.

That's my definition. I guess it predicts the future states. And because you can predict the future states now you can search and optimize to find the right sequence of actions to reach the goal that you wish to reach.

Yeah. Thanks.

Okay. Great. And now I'm just going to show two other examples of using some type of world model to try to get actions. So one idea is one really nice thing about VLMs and all of these language models that we've seen is that they just have a ton of high level knowledge. So anything you can ask a VLM with they can kind of give you some type of sensible answer. But if you actually try to take some of these VLMs and ask them about physical dynamics, this is a setting that none of the models are very good at. They all tend to hallucinate a lot. So let's say I just give this simple image here where I give the model this robot and I give it a rope and I ask it to bend the rope in a U-shape. It turns out that this is pretty challenging for a lot of these models. They tend to mess up a lot. Even when you see these models sometimes succeed. They actually just try a ton of times and then at the very very end does it finally make some type of U-shape. But roughly these models are very bad at precise low-level control. So, it can very easily tell you that there's a white robot arm over the table and beneath it is a curved piece of brown tape, but it has a really challenging time telling you, oh, you should move your gripper right 0.3 and you should move it up 0.4. So, this is a thing that the models just don't have a very good sense of.

So one thing we can do is we can again use this idea of world models as a way to really convert VLMs into low-level actions. An idea here is we're going to take an instruction and then we're gonna, so given this image here, we're going to construct a simulator. So this simulator is not going to be a learned world model actually, this simulator will just be, so we're just going to do a 4D reconstruction and we're going to actually put it in some type of physics simulation. So we're going to get a simulator and then we're going to have the VLM propose some abstract actions. So to the VLM it will say something like push this much, like it will say like push this distance, like push one, push 0.5, push 0.25. So we'll just predict different coarse actions and we're just going to simulate each of these actions in the constructed simulation. And then once we have all of these simulations, we're going to have the VLM take the outputs of all of these simulations and ask it, well, here's how the world would have changed if you predicted each of these actions. And based off this, how would you update your action? And if we repeat this process a bunch of times, now the VLM, even though it doesn't really understand the low-level continuous actions, it's able to use your simulations as a tool to really kind of ground its understanding and imagine and really get the precise low-level actions for solving a task. So this is what allows you to take the high level knowledge of the VLM and really try to ground it into low-level actions.

And here's just a couple different examples of this. So these are zero shot tasks that we did. And then here's actual constructed simulation for each of these tasks also.

Okay. And then we have a new variant of this which we call the world action planner where we're now not using a 4D simulator. We're actually using the action conditioned world model but it's roughly similar in style. So here we have this model, this model gets an instruction. For instance, you want to place the white mug on the left plate and you want to place the yellow and white mug on the right plate. So we just want to push one mug over here and push the other mug here. What we're going to do is we're going to take this instruction and we're just directly going to ask a VLM agent to predict what are coarse possible actions to take in this environment. And the way we get the VLM to predict coarse actions here is we ask it to point to different coordinates in each of these images of the scene that corresponds to where it wants to move the object. But essentially, we have the VLM predict some very coarse actions to execute in the environment. And then once we have these coarse actions, we're going to take each of these predicted actions and we're going to simulate it in our action conditioned world model. And again once we have all these simulations, similar to the previous one, we're going to have the VLM analyze what each of these generations looks like in accordance with these predicted actions. And then once we have this analysis, we're going to have the model predict a set of actions again, but now instead of just trusting the VLM to do your actions, we're now going to do some type of local search. So once we have the initial VLM predicted actions, we're going to do some type of, we're either going to just tile different possible surrounding actions and simulate each of them in the world model or we can do some type of gradient based optimization with respect to the world model and a reward model or some type of cross entropy maximization but we're just going to use this prediction of the VLM as some coarse initial guess that then we're going to do a lot of continuous action refinement to eventually get a final action that we execute in the environment. So one way to view this is you have this type of LM agent but now you also have this world model and then you have some search and this combination allows you to get these detailed actions for solving the task.

So what we find is that this approach, when the task you want to solve is exactly the same as a training task, it doesn't have too much of a benefit compared to just learning a policy. But we find really dramatic benefits when it comes to trying to solve compositional tasks. So if we take a task that contains two steps instead of the one step the model was trained on or even if it consists of a novel combination of language instructions we find that if you take an end-to-end policy like Type 0.5 or you take some type of world action model like Cosmos Policy we find that these approaches because they really rely on behavior cloning they really have a hard time solving these new instructions that they have not been explicitly trained on. In contrast if we do our approach here which is world action planner we get much better success. I guess it's around 70 to 80%. The closest one is actually using a vision language model directly to try to predict the actions because they again have some of this prior but we find that that loses a lot of the accuracy of the low-level actions.

And this is also pretty true when you try new different layouts of objects. So these are layouts of objects that these models were not explicitly trained on. And again, we find that this world action planner works quite a bit better than ever, and where I guess the biggest striking difference is it turns out that a lot of these imitation learning methods just really struggle whenever you just even change the scene slightly. It's just too out of distribution for the models. They tend to memorize a lot of the actions exactly, the mapping from images to actions. But it seems like any type of planning just generally improves their performance quite a bit.

Oh, and then here's a final example. So this one we actually learned an action conditioned world model but this action conditioned world model is not in image space. It's actually in the locomotion states of a robot. So you have a system that gets a sequence of actions and it just predicts the future proprioception of the unitree dog here. And what we find is again given this dynamics model, if I just specify different reward functions I can adapt the behavior of the robot. So I can make it so it walks in a particular way with joints in a particular angle. I can make it so it's a particular height or I can make it walk up an incline. So actually let me show you the behavior. So here you have the original model perform model behavior. So original model just kind of walks along the ground.

And now you can add some type of reward function that you want to optimize. And now you can change that calf angle of the model.

You can also make the angle larger.

You can change the height of the robot.

And you can also ask the model to balance itself.

And you can also try to step up.

Yeah, and then there's also additional objective functions you can specify. But I guess one really powerful thing about world models, even if let's say they're in the low dimensional control space, is that it just allows you to really flexibly adapt your behavior in a more robust way than just directly learning a policy does.

Okay, great. So those are some examples of trying to connect world models with lower level actions. So the final part I'll talk a bit about is how you can also combine world models with much higher level reasoning. So trying to solve longer horizon tasks or trying to do some type of spatial reasoning.

So if we want to solve a really long horizon task, for instance, let's say I want to just set up my kitchen. It's pretty hard to imagine just directly using a policy to solve this task. Because it's just so long horizon, the policy doesn't have some type of memory. It's just generally hard for it to reason to solve these tasks. And in contrast it might seem really natural to use a vision language model to predict each of the high level steps towards solving a task but the issue is the vision language model doesn't actually see the, isn't able to reason about the underlying physics of the world so it can't really make very detailed physical plans for solving a long horizon task either. So we find that we can use a combination of a vision language model as well as a video model to build some type of planning system for solving these long horizon tasks where we use a vision language model to propose these high level plans and then we use a video model to synthesize and simulate what would be the effects of certain actions proposed by the vision language model.

So let me give a concrete example. So imagine I have the start image here and I have the instruction I want to put the fruits in the top drawer and now the visual language model is like loosely grounded. It can see the starting image of this scene but it can't really reason so well about the physics. So it will predict two different actions. It will predict either place the banana in the top drawer or it will predict open the top drawer.

And this is where the video model comes in now. So now the video model will, taking a possible action predicted by the VLM. So given this image and the instruction place the banana in the top drawer it can predict this instruction where it just picks up the banana but now because there's no drawer that's open the video model is able to simulate dropping the banana down. And now in contrast if I have the instruction open the top drawer the video model is able to predict the instruction of opening the drawer.

And now once we have both this video generation as well as this final goal put the fruits in the top drawer. Now we're going to have the VLM taking each of these videos and just estimate how much progress is being made. And for the one where it just drops banana it estimates very little progress. And in contrast for the one where actually you open the drawer because you need to open the drawer to place the fruits in the top drawer it estimates more progress. So you can imagine that this combination of this vision language model as well as a video model allows you to do this closed loop imagination and planning to find the right language and video instruction for solving a particular task.

And here's what this kind of looks like. So the first step is what I just described earlier. You have the image. You predict these two actions. You choose the one that's the most promising. And once you have this image, you can just repeat the process. So you just again predict different actions, choose the most promising one. And then at that final frame, you again do this. And if you just repeatedly do this, you can eventually choose a sequence of actions that allows you to exactly place all the fruits in the drawers and finally close the drawer. Right? And this is the one that this step is the most interesting one where after you've put the banana in the top drawer, you could have feasibly imagined closing the top drawer. You can imagine putting the coke can in the drawer or placing the banana in the drawer or placing the apple in the drawer. But you can see that by simulating each of these using a vision language model and video model, you can choose the appropriate one for solving this task.

And here's the actual execution of this. So now here we want to put the fruits in the top drawer and this is a corresponding real robot execution.

And one thing to note is that in general when it comes to these planning tasks, it's hard for you to predict the exact precise actions for solving the task. So here for instance, it messed up once and then we just replan and then it finally closes the drawer and puts the fruits inside.

Here's another example of this planning procedure. So here I have this set of objects and my goal is to arrange all these objects in a line. And you can imagine that if you just directly have a VLM try to solve this task, it can kind of give you these instructions such as move the red circle to the left of the yellow hexagon or move the green circle closer to the red star. But it's hard for you. So let's say if I told this video generation model, move the red circle to the left of the yellow hexagon. Well, on one hand, this is left, this is also left, this is also left, right? It's not clear what moving left to the hexagon means, right? So if I just predict high level instructions, it's hard for me to construct a detailed video plan for actually arranging this in a line. But if we actually do this entire planning procedure I talked about earlier, you actually get this fully synthesized video. So you can see that it's able to move the objects pretty precisely so that they form a line. And you can see that this, we were doing this a couple years ago. So the video generation models weren't so reliable. So there's some small shifts of the objects also.

But you can also convert this into corresponding real robot execution.

And overall what we find is that for these longer horizon tasks we find again that this type of planning is much more effective than some type of end-to-end, oh I guess this was just an end-to-end video generation model, so kind of like a world action model, and we also find it's better than just some end-to-end VLA model which is shown here. I guess this is a couple years ago so the terms are a bit different but roughly we find that planning is much more effective than either end-to-end video generation or end-to-end policies.

Okay and the final thing I'll talk about that's a little different than what I've talked about earlier is this idea that we can also use these video models as a way to kind of improve the spatial understanding of VLMs.

And if you take an image and then you ask a VLM, for instance, if I sit on the couch on my right and face the chairs, will the kitchen be on the right or left? These types of spatial reasoning problems are things that are very very challenging for existing VLM models to accurately do.

And what we propose in this work is that we can again use the world model as a tool or a visual imagination to help the VLM understand how to solve the task. So when I want to solve this task, the VLM is going to ask the world model to simulate different actions and then for each of these simulated actions, I'm going to give the simulated actions to the VLM as additional evidence to help it solve the task. So hopefully this is a little more concrete. So here my instruction is I have this image and it says if I sit on the couch and I face the chairs, will the kitchen be on my right or left? So this is the instruction we wanted to solve. And what the system will do is it's going to take in the starting image.

And from the starting image it's going to simulate what the world would look like if I were to do the sequence of actions. So it's going to simulate what the world looks like if my robot moves like this or moves like this or moves like this. So it's going to simulate each of these video traversals in the environment. And once it simulates each of these video traversals, we're going to have the VLM choose which of these video traversals seem the most helpful for solving the task. And we're going to choose the most helpful one. And we're going to use that as an in-context example to teach the VLM to solve this particular task. And you can see that from all these imaginations, you can see that this particular one basically allows you to directly solve this task because it allows you to simulate what it would look like if you sat on the couch and faced the chairs. So essentially you can use the model to simulate the spatial traversals and use that as context to solve a task.

And yeah roughly what we do to actually do this is we have an initial image, a model predicts certain actions to execute in the environment and then we just sequentially refine trajectories until, or like we render each of these trajectories and we choose the one that's most helpful for the VLM.

Here's a final illustration of this. So here's just a different problem. So here it says if I'm at the green point and facing the battle, will the whiteboard be on my right or left? Again, what we do is we predict a certain set of actions. Either move, turn left, turn right.

Based off each of these imagined actions, the VLM is going to select the one that's the most promising. From this imagined action, we're just going to repeat this process of simulating different actions.

And you can see that this final traversal directly solves the question. So the VLM selects that one and uses that as evidence to solve the final task.

So overall what we find is that this type of mental simulation in combination with VLMs is a pretty effective tool to really improve the spatial reasoning of VLMs in the physical world. So overall I've talked a bit, I talked about essentially about three different streams of work. So I've talked a bit about how we can directly use video generation as a way to directly reason in the visual space and try to solve tasks. I've talked a bit about how we can combine video models with lower level actions to do some type of model predictive control to get more precise low-level actions and to try to generalize a bit to novel scenarios. And finally, I've talked a bit about how we can use video models at a higher level to do more long horizon planning and as a simulator to try to improve the spatial understanding of these current systems. Yeah, and I mean overall I think there's still a lot of room to go with these world models. I think that on one hand the physics is still somewhat inaccurate. Also the speed is a very very large issue which especially is very problematic for combining video and action representations but overall I think they're a pretty interesting tool for trying to build systems that are able to generalize a little better out of distribution or outside the task that you've trained on. Yeah, thanks everyone for listening and happy to take any questions.

I guess a lot of the instructions were, examples were done with images. I mean in principle you could imagine all of these also working with audio, right? Because instead of predicting the images you just predict audio. I think the question is how accurately can you, like for instance if I wanted to do this video action reasoning right, how accurate is audio gonna be in enabling you to find the right context which it might be for depending on the inclusion task you want.

Yeah, sure. So I mean the first one, right, is you're just using the video model to imagine the future, right? And you're doing some type of search on top of the future images that are generated in order to let's say control a robot or to reason in the visual space. The second pillar is talking about how you can use this visual simulation in combination with some type of lower level action proposal to find the right actions to solve a task. And the final one's talking a bit about how you can use this video generation in combination with a higher level reasoner and using that final simulation to find the right level of high level steps for solving a task. So this is just talking about different ways in which you can use simulation as a way to get more accurate task solving.

You could imagine actually having all three of these at the same time. So you can imagine that this video and action reasoning is for lower level control, right? And then you have a separate model that's combined with this language for higher level long horizon tasks and then maybe you imagine that this first system is the one that actually proposes your initial actions. So you could imagine all of these three at the same time, or you could imagine these are just different levels of temporal resolution also. So this is very fast, this is in the middle and this is very long horizon.

So, I think this might be a question you get a lot, but I was just kind of curious. So, like in one of the things you were showing like having basically doing like several videos and like sort of like different scenarios and then having it like picked towards that. How are you like handling essentially like the memory bottlenecks around that like and like how much hardware are you ending up having to use just to be able to keep up with that?

Oh, so bottleneck's about storing all the past videos or?

Well storing all the videos but eventually as you like are doing more tasks this would like kind of balloon into a lot of usage right memory usage.

Yeah. Yeah. I guess it's really so yeah depending on, so I mean you could imagine that when you're doing this tree search procedure for instance you could just store the successful tree up until your current stage in which case maybe the memory usage wouldn't be so huge but in general yeah I mean like if you want to generate multiple videos at once this is going to be very expensive in terms of memory. So this is why we've been kind of interested in trying to build these more efficient models like that are a bit smaller so that maybe the memory usage is a bit less.

Okay that makes sense thank you.

Thanks.

Thanks.

I just said that's right.

Hi everyone, thanks for inviting me. Excited to talk about world models and some concrete algorithms but more so I also want to give a broader overview of where I see the field of AI right now, how world models fit in with that and what I think will matter the most for the next couple years. So, you know, where are we in AI really from the perspective of matching human capabilities? Obviously, there's been tremendous progress for digital agents, text understanding, textual reasoning, solving long horizon digital tasks, right? And we have certain basic visual understanding. There are all kinds of classification models working really well, but you know, I would argue we don't have a really deep understanding of the physical world yet. And more so the textual models and visual models are very much like detached for the most part, right? And so obviously as the field progresses, you know, we want to go towards one model that can do everything, has all the capabilities at the very least of a human brain and then beyond, right? And so I think it's easy to forget with all this recent progress that there are some pretty fundamental capabilities still missing from AI. And so that's a broad physical understanding. I think even more so motor control is very obvious if you look at how robots fail and how limited they are to narrow scenarios at the moment. And you know beyond that obviously there are a lot of interesting directions of ongoing research. Which I think to me something I would really like to see become like mainstream AI is as we move towards these more multimodal models, you know, how can we get the capabilities of let's say a textual reasoning model without special casing text right because it's not special enough of a modality to deserve special treatment and so I think you know at that point these algorithms won't be anchored to like human reasoning performance and a little bit improvement on top, but they can actually go and detach from sort of the human baseline and go much further. And then, you know, there are all kinds of other things like training efficiency. We're still like about a thousand times, require a thousand times more energy for training a big model than the human brain. You know, even in terms of flops and energy. So, I'm going to focus on the first two. I think the timing for these is actually great because we have a lot of the ideas of how to make progress there. Whereas let's say for latent reasoning of course there are some algorithms there. It's active area of research. But I would say we still haven't really found the right algorithms to solve that problem. And then training efficiency right it's even worse. We need new algorithms and we need new hardware as well. So that is like even longer term on the road map of AI for me. So let's talk about the physical understanding and motor control right if we think about AI as a field and the progression of algorithms is very much determined by how much compute is available and so I think it's just like important to understand of course we all have heard of the lesson and so on but it's important to understand that the correct algorithm to use the best algorithm for a given problem and situation very much depends on how much compute you have available right and so the entire field of AI, I would say if you squint a little bit, is well explained by how much compute was available, right? Before neural nets started to work, the best you could do is just manually program heuristics. Very cheap to run, you know, not very general at all, but at least you can do something in narrow cases that a human programmer understands really well. And then you know it's still too expensive to train neural networks but at least we could, that was at a time when neuroscience was still ahead of AI and people understood the sort of types of filters visual cortex learns and they were just like handprogrammed into the neural network. And so then you have these hard-coded filters and you learn a little read out on top and suddenly you know of course you need more compute to actually train that but or you need compute at all to train that but you get a more general model out of it right much more general than anything you could just like program with if and else statements and then of course you know with AlexNet we managed to do end-to-end supervised learning and beat like more classical algorithms there. But it's important to understand, you know, first of all, it was unlocked by writing the entire neural net training in CUDA from scratch. There's no deep learning frameworks at the time. So massively more compute unlocked. But it's also important to understand that obviously it's still a very narrow objective function, right? You're taking millions of pixels as input and you're predicting just like a thousand numbers of what class is in the image. And so, you know, it took another three years for things like CPC to come around and work and actually get state-of-the-art generalization performance on ImageNet through an unsupervised pre-training phase, right? And so from then on, of course, LLMs became a thing and train on the entire internet and then now we're and then started adding more modalities to them sort of in a post-hoc way combining different models, training adapters and so on, right? And I think from this trajectory it's quite clear that you know the field will converge towards just predicting everything and to me that's what a world model really means or you know omni model there are a lot of sort of early steps in this direction now and so you know if so you get the most general model if you predict all your tokens right but it's also important to understand that you just truly need more compute like you could have tried CPC in 2014 and you just wouldn't have gotten very good results even if you got all the algorithmic details right right because it just takes 10 times more training and so so you know I think it'll be a gradual progress but that is where most of the progress will come from and so then the question is you know if we want to predict everything what objective functions do we want to use for that right and so I put some examples here obviously there's a lot of work in this space but broadly speaking you could either predict the future like next token prediction, you know, your future inputs or you could like predict your current inputs after they've been corrupted, right? To make it a non-trivial task, you then apply corruption. So that's how diffusion models work and masked autoencoders and so on. And you could do the same thing, you know, either you predict in input space or you could predict in latent space, right? So there are these latent prediction algorithms like DINO and so on that rely on corruption and losses in latent space but ultimately it is also just predicting the same information and then you know you could predict forward in time you could sort of try to get the best of both worlds I think that's also very reasonable you know there are different sort of regularizers to make sure these are like interesting objective functions that actually extract bits from the data. And so yeah, like either you corrupt or you mask your input or you have a causal mask and you shift your labels or you have certain regularizers like slow prediction targets or contrastive negatives to make sure your learning algorithm doesn't collapse, right? So which one should we use? I would really argue it doesn't matter that much at the end of the day. And I was actually a big fan of sort of these more fancy algorithms for quite a while but you know as we've experimented with these things in the last couple of months it seems like the scaling trajectories are actually quite similar and it doesn't, you know, ultimately ours, you know, the real world is so high dimensional and we have to like really compress to make sense of it and turns out compression is just like bites you in the tail. You know, if you have, you want to train a model that's good for multiple tasks, you really don't want a compressed representation that was sort of trained to discard all the information except what's useful for one particular task. And so, you know, no bottlenecks. And then should we predict pixels? Isn't that just like way too high dimensional? I actually say, you know, pixel is actually quite low dimensional as we are scaling up these models more and more and so you know I've come around to thinking that you know we just predict your data and do it at scale and it's actually you know we don't need to get so fancy about the algorithms and so you know we've seen this work pretty well and so this is on the side of physical understanding right and if we look the scale obviously as I said you need scale for generality and I would argue I mean this is hypothesis but I think the scaling ceiling for physical understanding is actually much higher than for just textual understanding as well right because we've already started overfitting the internet and then come up with synthetic tasks in language world but for physical understanding you know there's more videos uploaded to YouTube every day then you can even than I can even push into my data loader. So I think we'll not run out of that data anytime soon. And there's a lot of information there that the model wants to extract. And so you know it'll be exciting to see where this goes. But basically for context you know motor control physical understanding not working that well yet. But these models are also just like tiny, right? And often they're not trained on the right tasks that actually help. You know, often it's like you take a VLM that was trained with question answering tasks and so it abstracts semantic features, but it doesn't really have physical features. And so, you know, then you fine-tune it a little bit and it overfits to specific task. But you know I think the biggest improvements in motor control will actually come from like pretty aggressively pushing the physical understanding right and so for context you know I think models like broadly start to generalize for the real world like complexity scale of data around 10 to the 25 flops and that's like roughly where I mean these are all estimates of course these numbers aren't public but yeah where like GPT-3 and then the top video models are trained at whereas you know obviously frontier LLMs are trained like two orders of magnitude above that and I think you know we just have to push pretty aggressively on that front and that's when we'll start seeing the benefits of like multimodal pre-training for like omni models and world models and so so then you know on top of that physical understanding will help a lot right but how do we get motor control you know I think it's quite clear now that imitation learning is only an initial initialization we just need RL so for one if you have humans collect demonstrations they will fail in all kinds of ways and there may be some recovery behaviors in those demos as well and it's great because the agent can learn that but you know once you deploy the policy it'll do its own thing. It'll fall into a different distribution. It'll observe a different distribution of states. And so, you know, the policy will just fail in ways that the human data collectors never failed. And so, you know, you need RL to do trial and error and for the policy and the agent to discover, you know, its own ways of failing and learn how to correct from that. I think another reason RL is so necessary is that demos at scale will just not be clean. You know, we know that from pretty much all scalable training runs. There's always, you know, the bigger your data set, the harder it is to ensure it's high quality. And so you know RL allows you to once you have an initialization that comes from imitation learning to amplify the right modes and suppress the worst ones and then obviously longer term you want to have let's say in robotics you know superhuman performance then RL is the only way to get that self-improvement right so we need that we need the physical understanding to get generality and then we need RL to get robustness. And practically, you know, doing that purely in the physical world is going to be too expensive. It's not how humans learn as well. And so, you know, with a world model, of course, you can start to unroll the model with itself. It can predict actions and future observations and so on and the RL algorithm starts to look a lot more like a LLM post training recipe than sort of the old school like real world robotics like RL experiments and so world models unlock this imagination training right and so with that context I'm now going to talk about some of the details that went into our work on DreamerV4.

So probably some of you have already seen that but I want to go into some of the algorithmic details here. Obviously this is a world model agent that predicts actions and observations together and it's training to do all of these tasks inside of its own mind. And you know it's doing that in some token space. We decode the tokens here just for visualization.

And so yeah, this was work back at Google with Wilson and Timothy. And you know, we use the world model here as a neural simulator. And for one, we leverage the representation learning and then we see pretty strong benefits from that, but also the ability to imagine the future to do RL inside of the world model. And so, you know, the biggest challenge here was really to get the models to be fast enough, right? Because most video models just predict in blocks. And so if you predict an entire block of video, you can diffuse everything jointly. And that's great, but it doesn't allow you to interact with it in a fine grained way. So to do this imagination training well you actually have to generate one frame at a time so that you can then predict the next action based off that and then you can predict the frame afterwards. So you know previously in DreamerV3 we're using like RNN based architectures and they were about a thousand times faster than just like a video diffusion transformer and so switching to these video diffusion transformers out of the box was just completely infeasible to do RL in, right? And so it took a while but we figured out a bunch of things that together made this feasible and we ended up somewhere in the middle. So roughly like a 30 to 100x improvement in sampling speed compared to like diffusion forcing transformers. Yeah which like roughly in between the RNNs we had before and the sort of architectures used for frontier video models. And so you know we did that through architectural improvements but also objective function improvements. Right? And then we train this agent that learns complex skills purely inside of this world model. So you know now we have a really scalable architecture and algorithm for world model agents. So we can pre-train on web data. We can train on data without actions and fine-tune actions into it and so on. And we get a simulator that's accurate enough to actually learn successful long horizon behaviors inside of it.

Right? And so, you know, it's the first algorithm that got diamonds in Minecraft purely offline without any environment interaction and outperformed OpenAI's VPT policy while using 100 times less data.

So, you know, at a high level we have like a little tokenizer here and then a big sequence model transformer on top and that is just one unified model that interleaves a bunch of modalities, right? And recently I've been working more on robotics and there are all kinds of modalities. There's like tactile sensing and joint states and all of that. And so you really can't afford to train adapters and separate models and give special treatment to different modalities, right? So that was why we try to really just have one general architecture that can sort of flexibly, you can just throw modalities at it, right? And so we pre-trained the world model. First you train the tokenizer of course, right? And then you train the world model and you can do that just without actions. You just leave out those tokens. You have to be a bit careful with your rope embeddings, but besides that, it just works really well. And then you know you fine-tune that on the actual experience data set you have, which in this case was a data set of contractors, 2,500 hours of playing Minecraft. And you know you have rewards either in your data set or you annotate them and you train your reward model. Training reward model, fine-tuning that into the world model backbone is very straightforward because it already has good understanding of the world. And then you do RL post training in imagination.

And so you know to get this to work you need large model capacity you need this fine grain action conditioning and you still need fast sampling for RL right and the first step towards enabling that is shortcut models we played around a lot with meanflow approaches as well ultimately shortcut is more stable and then if you tune everything right we found it to consistently outperform meanflow based approaches but both work pretty well. And the idea here is as a diffusion model starts from noise and then gradually fills in more detail to come up with the next frame, right? That's a process that goes over many frames, but there's actually a deterministic mapping once the model is learned, right? Because the sampling process is deterministic. And so that means you can distill it into the neural network and you can just go from the noise and train it to just predict the outcome directly, right? So then one way to do that naively would be to just do very fine grain samples and collect this paired data set effectively of noise and data points and train on that. But that's not really...

Very efficient because the sampling process is so slow. And we can actually make this much more efficient if we do some sort of hierarchical distillation. And the way shortcut models do that is actually at training time. That has several benefits. You know, it's not like you pre-train this model, you just look at the loss, but sampling is really slow and then at the very end you distill it to actually get some generations.

You can actually get, you know, long horizon video generations throughout the training run and that's very valuable for logging. It's also valuable because you can just choose when to stop the training run based on like the results you're getting and so, you know, we've like completely switched towards just training time few step sampling.

The way it works is just that, you know, sometimes you train it with a small step size, right? The model gets conditioned on the noise level but also on the step size of how much it should denoise. And so for the smallest step size you just train it like a diffusion model and then for larger step sizes you ask the model itself to compute its own targets, make two smaller steps and that becomes the prediction target for the longer step. And then this bootstraps all the way towards just going all the way from noise in a single step to the data.

Now for video generation there are a couple details to extend that, which is what we call shortcut forcing. It's basically shortcut models and diffusion forcing where you have, you know, a whole sequence and you have a different noise level and step size at each time step within the video. You know, I'm not going to walk through these equations here but it's basically what I described intuitively.

Oh, we're back. Okay, just when the equations come up. How convenient. Let's just skip over them. So, you know, so that alone gave us about 50-fold or 30-fold improvement. Sorry. You know, you go down from like 64 sampling steps to just one or two, but that does not give you like a 60x speed up because now other things start to become the bottleneck like accessing the KV cache for example, which sort of closer to the problems people run into with language models and fast inference there.

And so, you know, we put together a bunch of architectural details. I mean, it's all based on a pretty standard, like, modern transformer recipe. You know, I'm not going to talk too much about that, but overall, you just have a lot more tokens in video to deal with and to attend into. And so it's even worse when you have a diffusion model. You might need like two or three, two or four forward passes. You have to access the entire KV cache each time. And then, like I said, doing this fine grain prediction frame by frame, there's just a lot of tokens to attend into in your KV cache because like one video frame will probably be like 256 tokens, something like that. And of course, you want like decent temporal context as well.

And so, you know, one thing we found in this work, but also consistently since then, and I talked to a couple people, I think this has popped up over and over. Actually, if any of you guys have intuitions there, I'd love to hear them, too. But, we found that actually the temporal connections can be very lightweight. So, you don't need a lot of KV heads. You don't even need full attention. You can get away with just axial attention quite well. Where let's say the top left patch of the image attends to like the top left patch of the previous image and of the image before and so on. And so, that drastically reduces the computation of the attention operation, right? And then GQA of course helps on top of that.

And I think, you know, my hypothesis is that maybe, you know, generating an image and having a good like image generation model itself is already a bulk of the work. You know, a lot of the layers we use are actually just spatial attention layers and only every four blocks there's one with the temporal attention in there. And so that seems to be like good both in terms of generation speed but also in terms of generation performance.

And so I think we don't really understand why yet but maybe there's just not that much information to share over time. Once you have a good like latent space or representation of all the different objects and so on like their motion vectors and so on. Maybe that's just like pretty low dimensional. Anyways, it could also be a thing where if we go up two orders of magnitude in compute, maybe, you know, the attention can learn its like sparse patterns better. Even if you started with like a full quadratic attention.

And yeah, so we use the same architecture for a tokenizer and world model. We can see that here with three spatial attention layers and one causal one and there's one transformer on the interleaf like actions and then noise level and denoising step size and the representations.

Yeah, we're moving towards smaller and smaller tokenizers so with less and less compression so that, you know, you really just have one big model to scale on top. In DreamerV4 we're still using like MAE to help the representations but at scale it doesn't seem to really matter that much.

And then to train the policy we want that to be in the same model but we actually have to be careful because, you know, there are a lot of video generation models and people say, you know, they don't learn physics really well and if you look at videos generated by like Veo and Sora you can see that. But that's because there's inherent bias in there where you train a text to video model. You ask it for a certain like interaction outcome and, you know, pick up this bottle and then it will always predict a trajectory where at the end of it I'm holding the bottle or the person is holding the bottle. Even if that violates like fine grained details in the physics along the way, right? Because the model is just trained with a hindsight labeled data set and it's trained to generate a trajectory where this is the outcome.

So you don't want that if you care about having a causally correct model of the world that you can train policies inside of. What you want to do instead is, you know, an unbiased more objective one where the task goes into the policy part of the model. So that's a different set of the tokens is one of the modalities and you make sure that, you know, the world model like generation tokens prediction tokens don't attend to the task tokens. And so what happens now is the world model is just conditioned on the fine grained actions, the motor actions that come from the policy. And so they will just faithfully predict the outcome of those actions. And so that is, you know, that and even just having data that covers enough of the scenarios those two things together basically solve the physics problem. Like I think that plus scale, I mean it keeps getting better of course the more data you train on and the bigger the models are and so on but there's no inherent bias in it anymore at that point.

And so, you know, I think, yeah, it's just if you care about self-improvement inside of video models, you have to switch towards something where the task doesn't go into the generation.

Yeah, to train the policy and the reward model, it's just very simple multi-token prediction. Basically, next token prediction, but you predict like a whole chunk of the next 16 actions. That helps with representation learning.

And then for RL we're using PMPO which is an interesting RL algorithm because effectively it only uses the sign of the advantage. So normally in RL if you just do policy gradient you do samples, you do rollouts from your policy and then you upweigh or downweigh the actions proportionally to how well, how good the outcome was. Right? But that becomes like pretty finicky depending on what your reward function is and the scale of that and if you have multiple tasks. Some of them are short horizon tasks, some of them are long horizon tasks. So the, you know, some of future rewards will actually be very different even if your rewards are in a nice range.

And so, you know, I think over like the last 10 years, people have invented all kinds of fancy RL algorithms with like streaming normalization and so on built in, advantage normalization and whatnot. But we found this PMPO algorithm, prior work from people at DeepMind, we found that to work very well. And so all you do is just look at whether the rollout was better than the average or worse than the average and you split all your states into these two sets. And then you normalize each set individually and on one you just maximize log probability of your actions and on the other one you minimize. And so there's equal pull up or down or at least you can choose it with this hyperparameter alpha very well.

And then you can regularize the policy towards a prior which all of these are basically KL terms or log prob terms. They're all measured in nats or bits and so they're very easy to trade off. And these hyperparameters are like 0.5 and this is like 0.1 or 0.5 depending how strong you want the prior to be. But you can sort of have an intuition for what value to set up front and you don't really have to tune it much.

And so yeah, how much time do I have left? Okay, cool. Yeah. And so then the prior is important because we've trained the world model on offline data. And so if we train this algorithm, this agent to improve within the world model, it'll start to find exploits. And so you have to keep it sort of somewhat close at least to the actions in the data set. So it doesn't like go off to parts of the world model where it just doesn't really know what the outcome will be, right?

And so you regularize it towards the prior and then we did some experiments as well where after you do that and you get a good policy out of it, you run it in your actual eval and you feed that data back, you fine tune the world model on it and then the next round of improvement you can dial down the scale regularizer from like one to 0.5 and then the next time to 0.2 and then the fourth round you just set it to zero and you just let the RL run free. And so the world model has learned all the relevant counterfactuals by then. And we're surprised to see that that only takes like three, four rounds.

And so, you know, even though I'm a big believer in imagination training, I also believe in real world feedback to make sure the world model like corrects itself over time, right? And training an agent that exploits the world model and then deploying that is a pretty effective way of collecting the right data for that. Yeah, which is all pretty intuitive, but it's nice to see that actually works.

And so, yeah, for DreamerV4, we did a lot of experiments in Minecraft, and the reason is that it's a benchmark with a lot of diversity. I actually think, you know, video games have a bad reputation maybe these days because of all this like benchmark maximizing on Atari, which I sort of agree with, but there are very complex benchmarks like Minecraft out there and I think they will actually be great challenges for like the next 5 years. And it just makes sense to iterate in like a well-controlled environment rather than let's say in the physical world.

And so yeah, Minecraft is complex as a 3D game. I mean, most of you, I'm sure, have seen it. Every single world is randomly generated. So the model cannot overfit to the locations of resources and objects or anything like that. And there are over 500 items. A lot of these are long horizon tasks. If you want to obtain a certain item, you have to get all the other resources that lead up to it first. There's fine grain mouse and keyboard control because you need to craft in the inventory and drag little icons around and so on, right? And so we have a simulator so we can run like reproducible evals and there's like open data sets for it as well. And of course a lot of gameplay on YouTube as well.

Yeah, so just, you know, to put this into context, in DreamerV3 we did some Minecraft. It was very low resolution and we're training from scratch just online directly against this game. Tried to be data efficient and hill climb on it. There was no inventory crafting. So was using this MineRL competition action space and then OpenAI they did from pixels although at 128 by 128 which is the input resolution you can't really see anything. Which is why they needed so much RL with PPO to sort of memorize like one crafting trajectory that ends up working.

And then in DreamerV4 we use a resolution where, you know, as a human you can actually see everything and so yeah it learns to do these like long horizon tasks from high resolution video and low-level mouse keyboard inputs.

Another comparison point here, obviously this model is trained at much higher scale. So DreamerV4 was trained at about 10 to the 23 flops and, you know, much better generations and higher resolution than what these RNN based models were able to do before.

In terms of the actual task we're tackling. So, Minecraft diamond challenge. I think it's a great test bed, but also I want to sort of make clear that there is a lot more to do in Minecraft than just to get diamonds. It's like a relatively early game item if you actually play the game. And so, you know, this is not to say that like Minecraft is done as a benchmark now, right? And like we didn't even get that high success rate on diamonds even though we got some.

So yeah, there's a sequence of all these items the agent has to obtain along the way. So you need to find trees and break them to get logs and then open your inventory to craft the planks and so on. And each of these subtasks there is a reward for it, a plus one. So it's a very sparse reward. And it takes like about 20,000 actions, a sequence of 20,000 actions to go from randomly dropped into like a world you've never seen before, get all the resources to actually get the diamond, right? And so DreamerV4 does this without any training in the actual game.

So yeah, we condition the agent on these subtasks and then at inference time we just run through them.

Right? So first of all for fast sampling the shortcut forcing works really well. You know, you get massively better generations even, you know, we ended up using four sampling steps for the DreamerV4 results compared to like classical diffusion forcing where you really need to go at least to like 32 steps or maybe, yeah, 16 to 32 steps to get sort of any reasonable video out of it. So that's a massive improvement.

And then here are all the ablations for the efficient transformer. You know, you can see like all these like changes and details. We try to be like rigorous about that and it's like basically a trade-off of capacity and modeling quality. So at the end the model is so fast that we can actually increase the spatial tokens further and still get like reasonable throughput during training and during RL as well. And so then that improves performance further.

So once you have this world model trained with all these improvements it runs at 20 FPS on a single GPU. So you can actually just, you know, play the world model and it's compared here to previous models, Lucid and Oasis. And, you know, it's not perfect, but it's learned a lot of the game mechanics, and you can consistently build like a structure without it like just dreaming up stuff that you've never built. Or, you know, just derailing this entire city somehow now in this one here.

And then yeah, so we wanted to test the ability or the potential for pre-training on web data. We weren't allowed to put the actual YouTube pre-training results in the paper. But at least what you can see is, you know, you need very few actions to actually like fine-tune the grounding into this video model. Now I think people in robotics are surprised. People in general for modeling are like, you know, obviously if you pre-train on video and then you give it extra conditioning signal it only makes it easier, the model can collapse the distribution towards the smallest hit of futures. So it's very easy to fine-tune in.

And we found that as well and were able to confirm that experimentally that even just with 100 hours of action labeled video and everything else, the remaining 2,400 hours had all the actions removed, you get basically the same action following performance. And then we did some out of distribution tests as well, the sort of two separate worlds in Minecraft like the overworld and the underworld. And so we trained with no actions for the underworld and still got like quite good action following in there. So there's pretty strong generalization and I think this makes like very strong case for the pre-training to transfer well.

Yeah, we have some experiments here on real world physics as well and like smaller robotics data sets where, you know, obviously it's not at the visual quality of like a frontier video model but there's no physics bias in there. And so you can see as the arm moves around the ball like will rotate correctly and flip correctly and so on, can pick up the object and stuff is like following the laws of physics.

Yeah, and in terms of, you know, the world understanding or physical understanding, we're comparing to, you know, pre-trained models that were using a lot more compute like Gemma 3, and just imitation learning or behavioral cloning from scratch. And you can see that with the world model representations, you get, you know, significantly better success rate and even more so significantly faster task success. And so, you know, that's just because the policy works better. If it's like a bit flimsy or like a bit, I guess like sloppy in its actions. If it keeps going for long enough, maybe it'll still do it eventually. But with good representations, you latch on to the right features, you can actually, you know, be much more efficient here. Despite this world model of course being trained with like less compute than Gemma 3.

And then yeah, once you add RL, you can speed things up further and this is with purely offline. So if you do like real world interactions, of course, you can keep going. But I think yeah it's sort of interesting to see that of course if your tasks are easy it doesn't really matter but then the harder the task the further down in the sequence of all the tasks the more you get this benefit.

So here are some videos of the agent training in its own imagination and decoding it back into image space for us to look at along with the reward predictions and the value predictions. And you can see the rewards are pretty accurate. Sometimes it doesn't know the exact frame. So it's like spreading out its reward prediction here over two or three frames. So these two frames sum up to probability of one. It just doesn't know the exact timing of it. So overall the reward model is very accurate and so then yeah we're able to go all the way to diamond with this and yeah even without any sort of YouTube pre-training data and just training purely offline in the world model.

So now to close up, I'll actually show you the rollout of the agent playing the game. And so this is at the frame rate and resolution that the agent actually sees. So these are the inputs to the agent and is the agent just, you know, collecting its logs so that I can start crafting and then, you know, doing the like inventory interaction, building crafting table, placing that.

Not sure what it's doing here. Oh, it's making sticks maybe. Okay. There we go.

Yeah, that's quite cool to see it actually like interact through this like human interface and be able to do these like long horizon tasks.

Right. So, now it has this pickaxe. It can start digging down and collecting stone. Once it has enough stone, I'll make a new crafting table and make a better tool.

See what else we have. Need some more sticks first. Okay.

And then it's found some iron there, you know. And I think in this instantiation of the world, it's like the iron was like relatively easy to find, for example. And I think that's why it works, you know, with a certain success rate. Obviously, yeah, it'd be interesting to see how far this can be pushed.

Now I have to wait for his iron to smelt.

Iron pickaxe.

And then now it's trying to find diamonds. So it needs to dig around quite a bunch and go a bit lower as well because they're closer to the sort of core of the earth.

There we go.

And then Oh no. Crisis averted.

So it tries to stay alive but, you know, at some point, you know, so yeah, just to summarize, you know, I think like obviously this is like a good step forward on Minecraft and just like computer use and then training agents to learn long horizon behaviors and world models, figuring out a really scalable recipe for that that's applicable to the real world and to robotics and so on.

But at the same time, Minecraft is far from solved. And, you know, even after, you know, if you're like a well-trained human, you can go basically 100% success rate on diamonds. And then there are so many more items. And I think there are people like online that literally play this game for like thousands of hours to try and unlock everything the game has to offer. And so, you know, there are lots of sort of challenges for hierarchical control and so on. Or even, you know, unsupervised exploration in Minecraft would be super cool. Just an agent that discovers all these items by itself without any external rewards. So, I think it's a great test bed. There's a lot to do there for many years to come.

And then, you know, sort of tying things back to the intro, we have an algorithm now for strong physical understanding. We can pre-train it on web data and we can do more multimodal data mixes and so on to get the physical understanding for motor control and robotics. And so, you know, we get the generality that way and we have a recipe for RL post training to get reliability and self-improvement. So I think it's an exciting time for robotics and because all the pieces are in place now. Cool, thank you.

So, we have time for a couple of questions. If you can please go to the microphone to make sure the stream captures everything, that would be great. There is one on the right and one on the left. Yeah.

I can just repeat the question also. Well, I forgot what I want to say. I mean the difference between the world model like VLA versus the model representation.

Yeah. You mean like between taking like a VLM and using it for training policy versus a world model?

I saw you compared a VLM representation and also the world model representation. You mentioned world model representation beat the other one, right? And I'm wondering, I mean, I don't know, I couldn't tell the difference to me is the same. Sorry to say that.

You mean the empirical performance or the intuitions of...?

I mean the VLM and the world model representation what's the difference? I don't see the difference.

Yeah. But do you want to understand the difference on the like empirical performance or the intuition of what's the algorithm...?

The mechanism.

The mechanism, okay. Yeah, so VLMs or VLAs is basically you take a VLM that's been pre-trained by predicting text, right? For the most part VLMs are pre-trained by predicting text. Even though they have some image inputs the way they learn to perceive the images is by answering questions about them. Right? So they learn, if it's trained on scalable web data, they learn semantic understanding of images which is better than no understanding of images, right? But they lack the precise physical understanding, right? So for example, they're not very good at even understanding the precise location of an object, right? Or the precise shape of an object because they're only trained to answer textual questions about what they see, right?

Whereas if you take a video model and, you know, predicting forward, let's say the camera rotates around, the camera perspective changes, now the model is forced to predict how objects look like from the other side. It's forced to learn understanding of 3D structure and precise location in 3D space. If you see, you know, a hand picking up an object, it's actually forced to internally learn or at least estimate how strong is my hand, how heavy is this bottle, how much friction is there on the surface here, right? So there's a lot more detailed physical understanding and I think that's why the generalization is better and why people also in robotics are starting to see quite consistently there's like about two orders of magnitude compute efficiency that you can get out or data efficiency if you start from the pre-trained model and you just fine-tune with robot data from fine-tuning models that have like a precise sort of dense representation learning loss on images, right?

Even you can take like a text to image model and fine-tune that into a policy and it works significantly better than if you take a VLM. Now I don't really like the like one versus the other kind of thinking that is quite common in the community. I actually think, you know, we should train one big model that can do all of these tasks and should be able to answer questions but should also be able to predict the video ahead. And like these models once you have good video model is very easy to fine-tune like semantic segmentation into the model or depth prediction and all of these things seem to fall out of the video pre-training loss.

Yes, I have a very specific question. Thank you for the talk. For DreamerV4, how many TPUs you used and how many hours of training? So total TPU hours for a single training run.

Yeah. So single training run is 10 to the 23 flops. So we're using TPUs. So, you know, depending on what exact like chip you use, but it's on the order of like a couple hundred GPUs for like two weeks, something like that.

So, is it 256 TPUs or 1,024 TPUs?

That's a very precise question. What's the background here? I just want to know like how many TPUs you used for training and how many hours to calculate...?

Like very very close to 10 to the 23 flops.

Okay. Yeah, with like 50% MFU. So you can from those two numbers you can like find the GPU in your own and like work out the exact days and number of chips.

I'll do the computation. Thank you.

I think it's very exciting that you're able to prevent your RL from exploiting the world model by a few iterations. You said like three or four iterations of training a policy and then feeding and then running it in the simulator real world and then feeding that back into the world model training. I was wondering if you're seeing that in the real world or is that with simulated data only or...?

So that experiment was in Minecraft and we haven't tried it for a physical world yet. So, you know, I can't tell you yet. Hopefully at the next talk, I'll be able to answer that question better. I think, you know, given how easy it is to fine-tune actions into these models, with the version of pre-training, I don't think it'll need to be a lot of rounds, but that's just my intuition.

Yeah. You also maybe adding one thing. You sort of want the policy to exploit your world model but just a little bit. So it works best if you set the KL strength such that it's actually exploiting a bit more. So you're not optimizing for like the best downstream performance. But just like a little bit more aggressive than that. And so then you can get the most effective corrective data. Of course, if it's totally exploiting, you'll deploy the policy, it'll just like do some stupid actions and not actually get like interesting data. Yeah.

Yeah. Hey, I wanted to ask a few questions on the parallels to robotics. You said you started working on this much more now. I was also surprised by the 100 hours that you need just for fine-tuning a world model. I'm wondering though if this is a function of the simplicity of Minecraft. I know you said like it's actually challenging and I totally believe you especially in terms of the long horizon but visually it's like very different. The physics are very simple, right? Like the picking up doesn't actually require any kind of physical interaction and things like that, right? So I'm wondering what is your prediction on how many hours in robotics terms of teleoperation data maybe you would need or action generated data. And then my second question is where do you get the reward in the real world because for Minecraft that's also easy it's a game right like you get it. And so I'm wondering how do you imagine this in robotics like you have a person there sitting next, well it's an imagination but like then you need a learned value function. I had the same question for Yun actually I think there were lots of learned value functions. How do you get that?

Great questions. So yeah, first question is how much robot data do you need after like video pre-training? I think the answer, well there's like two parts to the answer. One is it's not that hard to collect 10,000 or even 100,000 hours of robot data.

Have you done it?

Close to.

Sure. So you paid someone or you bought data.

Done it. Yeah, I get it.

No, you've done it.

We've collected a bunch of data. But it's a complex data mix. So there's like a lot of like, we buy human data for example, right? Which has pose annotations which helps with like mid-training phase. But I think everything you describe, the intuitions are all correct but they're all about pre-training because the object interactions, learning those more detailed complex physics you can actually do that from observation, right? And it's possible that you will need some sort of counterfactual correction to get all the counterfactuals right but yeah you need massively more scale and bigger model and more diverse data to even be able to accurately predict the interactions of real world physical objects. But the action conditioning I think still remains quite simple because...

Do you collect failure data also?

Yeah. Yeah. You do want broad coverage.

Mhm.

Yeah. Exactly. But yeah, basically it's just the mapping of, you know, this number means how much is this elbow of the robot going to move versus like this pinky. And so that's like pretty easy one-to-one mapping. The model already knows how the hands can move.

I disagree on that one.

Okay, we should talk more about it later then.

But okay, I get where you're going, but maybe can you comment on the reward function as well?

Yeah. Yeah. I mean, there's different ways to do it either during data collection or...

You need labeled data as well on...

Yeah. I think, you know, my take on rewards is they are just expressing, I mean, maybe in the long term we'll have agents that, you know, are sort of unsupervised in that regard, right? I didn't talk about this today but I sort of have this like road map of how stuff will play out over like the next maybe 10 years where, maybe 15 years even where robots, you know, there are intrinsic objective functions that robots can follow to like autonomously explore to like find interesting scenarios in the world and also learn how to control the world in sort of a general sense. So they set themselves goals and try to achieve them let's say. So I think these are very fundamental principles and that'll come one day but, you know, the ordering is like determined by how much compute is available and what the economic incentives are around all this. So I think this stuff comes like much later. And so then if we're just talking about like task rewards that's I think just inherently like human preferences. So there's just not much you can do like you just annotate it, right? And it's okay and it's like not that expensive compared to like compute or other things.

Okay.

Yeah.

Thanks. We are at time unfortunately but if you have questions please use the break to ask questions to Yann and everyone please again for the... We now have 15 minutes coffee break and after that we have the lightning talks and yeah speaking so let's come back in 15 minutes.

If you are giving a lightning talk, please come with your laptop to get ready to start again. So we have four people for the lightning talk. Please come to the podium. Thanks.

All right, everyone. If you can please come back to your seat. We're going to start the lightning talk.

Everyone. So we are starting the last part of the afternoon program. So at first we'll have four lightning talks. Those were selected based on the submission and the best submissions to the workshop. So we have first Elun, then Ryan, then Teodros and Ryan again. So each one is 5 minutes and then we have a very short one, two minutes for questions if you have and they will go one after the other and after that we will have Yann giving his keynote. Okay so with that please, thanks Elun for being the first speaker.

Thank you Randall. Can... Okay. Hello. My name is Elun. I'm a PhD student at NYU working with Professor Yann LeCun. So today I'm going to present our work LP World Model where we are trying to incorporate sparse representation in a JEPA style world model. So basically when you're doing, you know, this like JEPA world model you always encode your representations in a latent space and you learn a latent dynamics on your latent space. So in this case you can always choose to have a dense representation or a sparse representation. So for example, one example of a dense representation is this like VAE model method where you match your features to an isotropic Gaussian distribution. So here's like one dimensional examples of this like Gaussian distributions. You will have a dense code.

What we're doing in our project is LP World Model is instead of matching your features to a dense Gaussian distribution we're matching it to a rectified Gaussian or rectified Laplace distribution. So you have one example of this rectified Gaussian distribution in one dimension where you have a delta mass at zero and a Gaussian tail. And you will get exact zeros in your features. So you might wonder like why do we want like a sparse representation? Why do we want to model dynamics in a sparse coordinates? There are many reasons for this. One of the concrete reasons which we will formalize more rigorously in our paper is this idea of like if you have a Lipschitz nonlinear dynamics then you can find a one-hot sparse encoding such that your latent dynamics is just linear. But you know, basically there's this like thing we have in more details in our paper.

So the idea is that in practice, you know, you cannot really have like a one-hot encoding of your input because it's like impractical and infeasible. But somehow we still think, we still hope that we can learn a representation such that it simplifies the dynamics in the embedding space. So how do we actually do it, like to learn a sparse representation? So this is like a diagram of our LP World Model approach where we have like an encoder F that, you know, encode your observation OT, passes through ReLU as well. And also a predictor that takes your sparse vector ZT and your actions and you predict, you output a prediction with a ReLU attached to it. So that you also predict a sparse vector.

And the training objective is just a standard MSE loss plus like RDM reg which we cover in the next slides. So we called it the rectified distribution matching regularization. Basically the idea is we sample, you know, random projection vector C. And then under these like random projected directions we match our rectified feature Z to like a rectified Laplace samples. So the loss we're using here like the L is defined using the two-sample distance. So the intuition is you're getting the empirical CDFs for your projected features and your projected target distribution and you just match those like empirical CDF differences and you can achieve the distribution matching.

So we tested this in this like PushT environment where we train on a trajectory of like this dot agent pushing this T-shaped object to a target location. And to test the hypothesis of whether, you know, sparsity can somehow simplify the dynamics what we do is on the x-axis we've listed a bunch of different types of predictors ranging from on the left is the expressive transformer with like adapted layer norm action conditioning all the way down to all the way on the right is like a linear, it's a linear time invariant state space model. So it's like a completely linear dynamics.

So what you can see is for the very expressive types of predictors like the gray box, the sparse LP World Model is like on par with the dense model but for like a linear predictor with just one extra layer of nonlinearity. So that's the blue and orange bars there. We can see that the LP World Model actually substantially outperforms the dense model. And, you know, we think this is like evidence that sparsity can just like reduce the predictor complexity that's required for planning or basically, you know, given a dynamics you can actually get away with a much simpler predictor.

And we also hope that a sparse representation somehow gives you more interpretability. So we have this environment we call Peacewise where the red dot agent is moving in this with background force field, like four of this background force fields. And we saw that the sparsity pattern like the support pattern of your feature vector is very similar within the region for the dots that are red dot that are moving within the region versus outside of the region. So we take this as like, you know, one of the evidence that, you know, sparsity is kind of tracking the underlying dynamics as well. So this is joint work with my co-authors and these are the two relevant papers for this talk. Yeah. Thank you.

So we have time for a quick question while Ryan will set up. Okay. So Ryan, if you want to come next. So thanks again for the talk.

Okay. So, hello everyone. I'm Suzin Wu from AXE Research. Well, not Duke. They obviously make a mistake although I hope so. I have some association with Duke. So, I'm presenting my paper called Selection Induced Optimism in Social World Models. So, basically in the current world models there are always some planning stage where you optimize over plausible rollouts you predicted based on your world models. But our paper shows that in some social settings, in some certain settings you may actually, the optimization process may actually end up introducing bias into the final decisions instead of making the overall decision better.

So we kind of, we test this in a synthetic classroom where it's a social world model setting. So basically there's a case presented to the LLM showing that there's AI authorship flag, the student is using AI for their homework. So LLM is then asked to choose one of the four actions here. The LLM can either choose to take no actions, to clarify with the students about their homework, to ask another human to review this homework, and just then impose automatic extensions. That's the four actions they can choose from. And they're also asked to predict the immediate downstream outcomes of each actions.

So based on these downstream outcomes and immediate outcomes, we kind of calculate the utility to see the impacts of the language model choosing one of the four actions. We want to see like whether this planning over the utilities calculated from this predicted outcomes actually improve the decision or just simply amplifying the errors. So to compare the language model predictions with a baseline we establish an oracle using programmatic rules and probabilistic rules already set in this synthetic environment. So they start from the same partial observations and from then we constructed an oracle as well as the language model predictions. We then compare these two to show the language model predictions against a baseline.

Basically, we choose to use GPT-4o and Gemini 1.5 Flash. It used to be two of the more advanced models like back in those days, you know. So, yeah, we, so among the 18 cases we test here using the language model, you can see there is positive X-val values. This means that among the 18 cases we test in this synthetic classroom setting, there is more error on average. A positive X-val basically means the chosen action's error is greater than the mean error of all valid actions under this synthetic environmental rules.

So we then run an ablation test to see where this error is being introduced. We start from the continuous oracle where you basically have the numerical utilities constructed from the rules established in the synthetic classrooms and we then convert these numerical values into categorical options that we ask the language model to choose from. So you can see in this step there's already some errors being introduced. The match is going down, the regret is going up and there's a very interesting bias that we'll talk about later. So even before the language model is added to this decision pipeline, there's already some errors. And when you add the LM planner, it's just basically amplify the error further.

There's an interesting bias towards A0, which is no action, which means the language model tend not to choose no action among all the cases where the Oracle believes A0 is actually the optimal outcome. So we try to test the robustness of this against A0 action bias by varying the aggregation rules, varying the utility functions, varying the mapping functions and you can see none of the A0 actions is recovered which means there's a very strong bias of language model against taking no actions even though it might be the optimal actions according to the programmatic oracle we established here.

So we then look further into this bias towards a zero. A very interesting thing is that when you add one prompt change, one sentence that explains well this partially observed case can be resolved by choosing one of the four actions, like added into the experimental prompt, you can see that the bias towards a Z is reversed. Now the clarify prompts make the language model choose six out of the six cases where a Z is actually the better option according to the oracle. The mean regret and the utility mean absolute error also go down as well, which show quite a huge improvement, although the X valid is still positive, where like we call it selection induced optimism. So there's still selection induced optimism in this case even after the prompt change is being introduced to this experiment.

So what can we learn? Well, so what can we learn from this experiment? First of all, you need to audit the whole experimental decision stack starting from the continuous dynamics all the way to the optimization and planning stage of the experiment. You need to evaluate decisions, not just simulations, because like valid trajectories can, planning over valid trajectories can sometimes it cannot actually improve the decisions in the end. And treat prompts and schemas as part of the model because slight changes to the experimental prompts in this experiment reverse the whole results. And compliance is not authority because valid actions can still be misranked. And finally, models can share framing errors. In this case, the two model families share move in the same direction that share the same positive selection induced optimism here. So that's all about my research. I'm also presenting a poster. So if you have any questions, feel free to reach out to me in the poster session. Thank you very much.

All right. Thank you. And now please help me welcome Teodoros to speak about latent action world models.

Hey, my name is Theodoros. I'm presenting CLAW: Learning Continuous Latent Action World Models via Adversarial Latent Regularization. So the question we are trying to answer in this project was: can a robot learn how an environment changes under action without ever directly observing the actions themselves? So typically, a world model would take in a current observation or a history of observations and an action and then tries to predict the next observation. But then the problem with this is action labels are robot specific and then they are expensive to scale. However, raw videos are abundant and then we can observe how the world translates from one observation to the other within the videos. Then the key idea behind what we try to do in this project is: can we jointly learn latent actions from visual transitions and world models that predict their effects? So typically the way latent action world models are designed is they take in sequential frames and then we would use an encoder and then encode to a latent action. Then the latent action and the initial frame are fed to a decoder and then this decoder tries to predict the future frame. The problem with this is there are easy solutions that would solve the optimization over the pixel space. So typically a future information leakage through ZI, through the latent representation of the action happens and then an encoding of the future frame is created in the latent space and then that gets decoded. That means the transition is not encoded in a state, it's a representation of the future frame that gets encoded. So trying to address this, we propose an adversarial training system where we take the standard way of encoding and then decoding of a latent action world model and then we added a parallel network which we call an adversary and then the goal of the adversary is taking the latent action and then trying to recover the future frame which it shouldn't be able to. So basically because we don't want the latent actions to represent the future frame, what we do is during training, during optimization we invert the gradient in the opposite direction. So incentivizing it not to encode any information about the future frame. So that leads to CLAW. So we created this end-to-end pipeline where sequential frames are encoded to latent actions and then a conditioning frame is given to the latent action diffusion model and then that gets used to predict the future observation and then we train this end-to-end to be able to recover latent actions as well as a dynamics model that could do world modeling. So we went on after training this, we went on and experimented on different environments. I try to focus on the robotics environments here. But we tried planning as well as latent action learning experiments on the top. As you can see, the model that is trained solely from videos was able to plan on the right side and then tries to close the slide. And then at the bottom you can see a model that is pre-trained on human data and the latent actions that were used to post-train a robotics task and we show that we perform better than the current state-of-the-art. Thank you. I'm going to present at the poster station. If you have more questions and want to see more experiments, I'll be at the poster station. Thank you.

Thank you. And now let me welcome the last speaker. Ryan will be speaking about Latent Flow Waypoints for Latent Planning.

Okay. So today I'm going to talk about a very simple method for generating intermediate waypoints for guidance when planning with latent world models. And as some initial background, there have recently been a bunch of different methods that have developed world models trained in latent space to plan. These models use a terminal goal matching cost which is effective but often provides little signal for useful intermediate states during long horizon planning. So we wonder is there a way to leverage the latent space to produce more useful intermediate waypoints for planning. Let's first consider a simple setting where we want to control a robot arm to pick and place a cube. So we're given a start and a goal frame and we must choose our sequence of actions to move it from start to goal. So we might start by embedding our start and end frames using a frozen world model encoder. And now that we have our start and end frames as latent states, what might a simple method be for producing a plausible latent trajectory between these points that could provide effective intermediate guidance? And so a simple one might just be linearly interpolating between these two states. But unfortunately, this isn't actually very good. If you notice when we decode the intermediate points found through linear interpolation back to image frames, they lose coherence in the middle of the trajectory, which is exactly where you would want additional guidance. And in practice, it doesn't really improve success rate very much to use them at all. But what if we just model the error term on top of simple linear interpolation, which is essentially the method we propose. So during training, we use the trajectories from training the original world model. And for each trajectory, we encode the start and end frames as well as intermediate frames using a frozen world model encoder. And then to predict each intermediate frame, we first compute the prediction from the linear path between start and end. And then we train a model just to predict the residual to the true intermediate latent. And so at inference time we can generate waypoints recursively by using the last produced waypoint as a new starting point. And for planning we incorporate this with an additional weighted cost function that encourages the model to pass near the waypoints at some point during the plan. So to evaluate we consider a few different tasks for navigation and manipulation. For navigation we use two rooms where an agent must navigate from one room to a target in another room through a doorway. For manipulation, we use push T where we have to push a T-shaped block into a target pose as well as a cube task earlier where we have to pick and place a cube using a robot arm. We can revisit our previous experiment and ask whether our waypoints are visually coherent now. And we find that when we decode our waypoints back into images, not only are they more coherent than those from linear interpolation, but they also seem to track the oracle trajectory, which gives us some initial signal they might provide a more effective intermediate guidance for our model during planning than those found through linear interpolation. But we can confirm this by just using them directly during planning. And we can compare to some additional baselines. For example, in addition to the linear interpolation, we might also snap each linear waypoint to the nearest training trajectory frame to test whether or not linear waypoints suffer from being off manifold for the model. So, when we use our waypoints for guidance, we find that the latent flow waypoints improve performance in all tasks, but slightly less so for push T. And they're particularly beneficial in the cube task where they lead to a 13.8 percentage point increase in success rate. We can also just compare directly with Oracle guidance by embedding intermediate frames from the true trajectory. And we find that we can approach their performance with purely generated waypoints from our trained model. This also extends to more difficult planning scenarios such as when we want to plan under a limited budget and we can actually plan more efficiently just by using these waypoints for intermediate guidance. So we show this in particular by reducing the number of CEM optimizer steps and the number of samples available to the model and we find consistent benefits with the weakest results in push T. Finally, we can extend this to solving more difficult tasks by filtering our cube dataset to create buckets which involve higher cube displacement. We find that we can provide consistent benefits over final goal planning in all scenarios even though in both cases performance degrades as tasks become more difficult. Future work is looking a little bit more into better ways to place our intermediate waypoints as currently they're just done through kind of evenly spaced points along the latent path. Additionally, we might want more adaptive methods as we get towards more long horizon and more difficult tasks. Finally, there might also be better ways to handle cases where we have multiple valid solution trajectories and we want to model all of them as opposed to just collapsing down to one particular way of correcting the linear interpolation. So, thank you. Any questions?

All right. So, let's thanks again all the speakers of the lightning talks. And we now have the last keynote of the day and after that there will be a panel discussion which hopefully will be a little bit heated. So definitely stay around and I am very happy to introduce Professor Yann LeCun from NYU and Meta AI Labs. So first of all of course as everyone I could speak two hours to introduce Yann but the main thing you probably need to know is that he receives the Turing Award and he has been one of the few people to really fight for open source and that's something I really wanted to mention. I don't see a lot of people mention that but this is very important for most of us. I won't say all I don't know but most of us and so at least for this you know if nothing else we are really thankful to him so with that please help me welcome Yann.