After Diffusion & FLOW Models: Equilibrium Matching (MIT, Oxford, Harvard)

Key Takeaways

Hello community. So great that you are back. Nano banana. What a beautiful image generation model. But you know what? Today we go the next step because we do have a research paper on equilibrium matching. Equilibrium matching is a glimpse into the future of image generation with AI. And you say of course because we here at the channel Discover AI for the latest in AI research. Those are our two papers. University of Oxford October 3rd 2025 and MIT and Harvard University October 2nd 2025. And at first we have to understand the first paper and then we will jump immediately to paper number two. So here we go diffusion models and the manifold hypothesis that is connected now with adaptive geometry. And you say yes of course no because if you look at the isotropic smooning of the score function that identifies the manifold structure you say we have seen this before. Of course in my video just days ago where we talked about the unified theory of agentic reasoning I told you that this work here by Jingua and picking university they have now a unified reasoning manifold framework for understanding the reasoning in LLMs. And I also showed you this is based on a representation entropy that connected here the information tier with the LLM geometry. This was done by Chingua University. And of course you remember that we talked about the real core in the understanding was this paper. This was the paper here from 2013 testing the manifold hypothesis and this was the most important paper here if you want in AI in image generation. And now what we understood in 2013 in mathematics in pure 50 pages of pure mathematics about separable hilbert spaces now we implement in a code for the next image generator. So let's have a look how does it play out. Now, if you're not familiar with manifold hypothesis, easy. It just tells us that we have highdimensional data, 10,000 dimensional vector spaces and the images in the audio in the real world data set are concentrated on a lower dimensional submanifold embedded within the ambient space. And this explains why our machine learning models can generalize despite the course of dimensionality. Now, if we looked at diffusion models, it was rather easy. You have about a dozen image uh videos on the diffusion model explaining them. So very short generate samples from data by reversing a forward noising process that transforms data into noise via an SDA and then approximated the score function with a ner log PT that ends it at a time through a score matching and then we have a loss and you are familiar with this. Now the conjecture is now that the diffusion model excel because this smoothening adapts to the lowdimensional manifold structure remember assumed by the manifold hypothesis allowing us therefore the interpolation along the manifold rather than the full ambient space. So computational much simpler but mathematical to prove much more difficult and if you want here there are some beautiful mathematical lameter for you to follow having understand paper number one let's move now to the main paper of this video MIT and Harvard say now okay so what about okay diffusion and flow matching models learn non-equilibrium dynamics but what is the next step and finally a paper that is challenging absolutely love it So generative modeling with implicit energy based models and they call it equilibrium matching. Why? What's happening here? Now let's take a step back and you know what we have today are either flow matching or energy. So let's start with flow matching. Flow matching models learn to match a conditional velocity along a linear path connecting the noise and the image samples. And during the sampling the flow matching starts from a pure gshian noise and iteratively den noises the current sample using a velocity predicted by a particular function f. This process is governed by differential equation framework in which the predicted velocity is treated here as the time derivative of the desired sampling path and integrated over the total length. And you know this because in this video we applied here the flow to the reinforcement learning methodology based on Gflow net architectures. So if you're still using GPO or DPO have a look at this video because now we have a new reinforcement learning methodology with flow reinforcement learning. But now now there's the next step equilibrium matching. It is simple. We are standing on the shoulders of giants. Equilibrium matching combines now the advantages we found and the energy based model and the flow-based model and we build now something much more beautiful. If you take only one sentence, let it be this sentence. Equilibrium matching learns a time invariant gradient field that is compatible with an underlying energy function eliminating the time the noise conditioning and fixed horizon integrators. Have a look at this. If we have a look at the different time slices here, the flow matching you are familiar. So we have two if you want valleys down here in now a three-dimensional uh manifold and everything is now pointing exactly the way we wanted here with the flow matching. But equilibrium matching look at this. So what is happening here? I was thinking how can I explain this? I think the simplest way is here the diffusion and the flow models are like a GPS. It gives you a precise turn by turn direction. If you are somewhere positioned, you have to be have to have a defined start. Of course, then the system tells you, hey, take five steps north, three steps west, and you got it. No. So direction change at every single point in your journey. What you do at the beginning, let's say if you're high up in the plateau is different what you do near the end. So if you're approaching here, the energy minimum. Now EQM is like a static map. So we have here dynamic static map much more easier much computational faster a static map of slopes always guiding to stable spot energy minimum offering the freedom in the path choice. This is amazing offering freedom in the path choice. If you want it's like a magical compass and you have given here a topographical map and a compass and the model learns here of course. AI model, a magical compass that always points downhill and you just have here the indicator. So this is now a static map. How beautiful. But how do we create this map? How do we create this compass in this new methodology? Well, it's simple. But there's a profound mathematical trick. But at first, let me try to explain it in a simplified version. To train this model, we show that start with a corrupted image. No, it's a mix between the real image here, the real picture and some pure noise. You know, it's a mixture and the model's job is to look at this and predict now the compass direction where the model has to go. You're familiar with this if you are here on a curved gradient. Now, most models are trained to predict the direction towards the real image. This is like learning the velocity velocity to get there. And now the trick quotation mark is we do the exact opposite. And we'll explain later why. But let's do this simple explanation first. So EQM trains the model to predict the direction away from the real and towards the noise. And it seems backward I know but believe me this is the key. If we look at the energy surfaces in 3D you will understand it immediately. So by learning now the uphill direction the model implicitly defines here a landscape maybe an energy landscape. So to generate an image now we simply follow the compass in the reverse direction. We go downhill. Of course we have to make the valleys flat. No we need the compass to stop pointing anywhere if we are in the local deepest valley because if it would keep pointing in any other direction we would just continue. We would overshoot and climb out of the other side. No. If we have here local or absolute minimum, we want to stay there. So this is where EQM introduces here another trick idea, a dimmer switch function calls C of Y. So if Y equals Z, the image is pure noise. Y equal 1, the image is a pure real image, a real photo. Another dimmer switch is designed in a way to be strong when the image is mostly noise. So if we are far away from the valley somewhere up in the Austrian mountains and smoothly dials down to zero when the image is now pure data when we're down in the valley. So the most important rule is if we are really down at the valley our C function equals to zero because this new methodology learns now a static landscape and a compass. the inference our generating image task is now so much easier becomes incredible flexible also because instead of just taking a step downhill you can make it even faster because you know from theoretical physics that we have not only the normal uh gradient functions but we have accelerated gradient with a momentum and maybe you know an Nesterov accelerated gradient here with having a momentum and you can imagine here the acceleration and also It avoids getting stuck here in these small local minimums, these small bumpy parts here in the energy landscape that we are have to pause to go here to the local or absolute minimum. There's another benefit. You know what we can do in this new thing? We can just add the two compass direction together and follow a combined path. We don't have to retrain them all. We just can't do here. Just add the directions. And why is it possible? Because the main idea is we're just learning. We the AI is learning an energy landscape. And adding the gradients is equivalent to adding the energy landscape. It is mathematically such an elegant solution. Please have a look at the paper, read the paper yourself. You're going to take some time or you're a genius and you understand it immediately. Congratulation. And it is real simple to implement. Whereas doing this with a diffusion model is absolutely complex and often requires a complete retraining effort. But you know what? With EQM, it's a native feature. Isn't this beautiful? Okay. Again, the main idea, what is it? Equilibrium matching's core idea is to train a generative model to learn an uphill direction field that always points away from the valley downhill data. And by scaling this direction with a dimmer switch that turns off at pure data points here down in the valley, it creates a stable energy landscape. It builds us here this beautiful energy landscape. And you might say, hey, but why is this method so easy? It sounds so simple, so triple, so so beautiful. Well, now part two, there is a complete mathematical framework behind everything I just told you. But you know what? I give you the result. The result is the equilibrium matching after all the mathematics. A is theoretically guaranteed to learn this data manifold. You remember the manifold hypothesis we started with the video and produce samples from this manifold using gradient descent. An equilibrium matching achieves here one point. Never mind inception distance on the image net. Forget it. It outperforms existing diffusional flowbased counterparts in the generation quality. Flow matching at its best. Have a look at this. But we want to be better. We go here to equilibrium matching. Therefore, here you have a GitHub. Everything is available update four or two or four days ago. Beautiful. We have an MIT license. Isn't this beautiful? You have all the files that you need. If you want to have a look at the paper, look at chapter number two where they explain the flow matching because they build upon this. Then switch over to the training of the equilibrium matching model. You will understand here this function C that I told you that controls here the gradient magnitude. And finally to learn here the explicit energy function. The complexity here there is a simple function defined but of course you have to go to the appendix because there you have the complete analysis and there the orers tell us hey we provide the mathematical justifications for equilibrium matching and if you have a weekend that you say I want to learn for two days straight just pure mathematics this is your paper okay so you see summary there's a easy way to understand it or if you really want to have a deep dive if you understand why it is necessary to have a PhD in theoretical physics and mathematics or maybe just in mathematics those ideas are not easy to prove. You can have this idea like this here flipping here the direction target direction is not the difference the exact opposite here of the FM velocity why now we understand it no because if the data x is to be a low energy minimum you are at the minimum you can't go deeper no you have to go in the opposite direction the gradient must point away from it towards the mountain towards the higher energy regions like the noise epsilon and by learning now this field we create a landscape where they're descending. The gradient here takes you from the high energy, the noise to the low energy to your data points. And then the magnitude controller here of the target gradient if you want. This is the secret sauce. It's so easy. It's a knob that you can have different magnitudes. But it's really a simple Scala function that controls here the magnitude of the target gradient and it only must satisfy one exceptional uh if you want boundary condition C of 1 equals zero. So this ensures if you're down in the valley the needle just stops. This is it. You have reached your goal. So this ensures that when the interpolated sample is pure data the target gradient is zero. You achieved your goal. By training the model to match this, we explicit teach it that the gradient of all data points should vanish. Beautiful. And now a sentence that broke my brain because there is a hidden information theory behind this sentence. Can you feel it? This forces the data manifold to become a set of local minima in the learned energy landscape. We started the video here with the manifold approximation and the manifold hypothesis on everything and a reduced subspace in the dimensionality. And now we build it in a particular way to become a set of local minima. This is mathematically non-trivial. This is if you come from theoretical physics and you work in quantum field theory non-trivial. But it is beautiful. So with this objective EQM learns a single time invariant underscore time invariant vector field f of x that acts as the gradient of an implicit energy function efx where the ground truth samples are the stable stationary points. Great. And here you have now the equilibrium matching. It's so easy. Look it couldn't be easier. fast, gorgeous, the next generation because inference is now a simple optimization problem that we know exactly how to solve. We have now the ultimate flexibility because since EQM learns here a gradient field, generating a sample is no long integrating here a velocity over a fixed path like we do today. It is simply about finding a minimum of the energy landscape. And this is a standard optimization problem. So beautiful. But if you made it to the end of the video, you know, let's come back and let's put this here where it belongs. I told you at the beginning, now we have the flow miles and we have the energy miles. You know what? If we look now at this idea, we suddenly understand that flow and energy are just two sides of the same coin. distinguished by whether the learned vector field is a conservative a gradient field or not. But now we have a third perspective and this third perspective will bridge those other two side of the same coin. So you see we can build now a unified perspective and this unified perspective EQM could really inspire an entire new class of models that simply moves you between the two regimes and you can decide when to use what. You don't have to build separate models. You don't have to train separate models. It's just three sides of the same coin in a higher dimensional mathematical space. Let me explain. Now all the diffusion and flow-based models that we have, you know, this our current champions here for image quality. But you know, we both understand that it's a non-equilibrium process. Meaning that the rules of the movement, the velocity fields in high dimensional space change at every single step at every single time. And this rigidity leads now to constraint, fixed sampling step, predefined schedule, and more. It's the newer model, but it's complex. Our good old friend from hundred of years ago, maybe a little bit incorrect, the energy based models, no EBMS, the master flexibility. Good old friends. They learn here a timeless static energy landscape where the data points are the valleys, lower energy, what a surprise, and everything else is in the hill and the mountains and the high energy. To generate now a sample, you place a ball anywhere in this threedimensional landscape and let it down downhill via an optimization. Our typical gradient descent. This is an equilibrium process. This is a non-equilibrium process. But this is the big but here. All of this is highly unstable, difficult to train and only full of problems. But if we understand now that there is a bridge function where we just take here the advantage here the best parts of model of the diffusion flow based models and everything that is good about energy based model and we build now the next generation equilibrium matching. It is really kind of a revolutionary unification of these two let's call it kingdoms. It proposes here a way and a mathematical pure way but I have not yet time I need a weekend to really understand the mathematics behind this to build a generative model that learns a static energy landscape like an EBM but it is trained with the stability and the scalability of modern flow matching models. The best of both worlds. Absolutely. But you know what? We can unlock EBM superpowers at scale because naturally EQM enables your tasks that are too complex for the diffusion or real complex and real time intensive real expensive for the diffusion models to do for the training such as the dnoising partially corrupted images or inherent out of distribution detection and simple powerful compositional generation. Isn't this beautiful? combine the best of both worlds. And maybe this is really a conceptual shift that might redefine how we approach generative sampling. So there you have it, two brand new studies. One from University of Oxford Department of Statistic October 3rd and MIT and Harvard University October 2nd, equilibrium matching. What a beautiful idea. I hope you enjoyed it. I hope I was able to communicate the beauty and the simplicity of these ideas but also the complexity if you really want to publish this and you really want to code this because you have to prove that all your assumption are correct and you can build a mathematical construct that allows you to argue I am able to do this in the limit or I can have a convergence of this time series. I can build here a time invarant gradient. I have my energy fields. Therefore, you need a lot of mathematics, a lot of theoretical physics. But if you are interested in this, you see, you can have everything from pure mathematics to theoretical physics, you can if you understand why we are doing something. It is rather easy then to build the code and implement it. So this was it for today an outlook for the next image generation AI systems and maybe they will use equilibrium matching. Subscribe and I see you in my next video.