Alexi Gladstone | Explorative Modeling -- Unlocking a Third Pretraining Axis and End-to-End Generation
Pangram verdict · v3.3
We believe that this entire text is human-written.
AI likelihood · overall
HumanArticle text · 1,380 words · 1 segments analyzed
Website: https://explorative-modeling.github.io/ GitHub: https://github.com/alexiglad/XM TLDR: We introduce Explorative Modeling, a new paradigm for generative modeling that acts as a third pretraining axis when added to existing generative models, and also enables end-to-end generation. Increasing exploration monotonically improves existing models across images, video, and language, and the gains grow with scale (7%→36% with data, 13%→23% with parameters). Concretely, Explorative Models (XMs) reach 6.2× sample efficiency, 4.1× FLOP efficiency, and 47% better parameter efficiency. Exploration also enables scaling generalization, and scaling how end-to-end existing models are. As end-to-end generative models, XMs match diffusion on control tasks with up to 256× less inference compute. Let me start with a question that sounds simple. If I ask a model to “generate a dog”, how many correct answers are there? It turns out there are a lot… likely billions or more images that we could count as dog images. So what happens if we train a neural network to directly predict dog images? The model sees thousands of different valid dogs during training, and the single prediction closest to all of them is their average. That’s what the model learns to output, and the average of thousands of dogs looks nothing like a dog, it’s a brown blur. A real dog from the data → What the model predicts Figure 1: Training a model to directly predict images gives you the average of them all, a brown blur. This is why direct regression doesn't work for generative modeling. To make this concrete, let’s play a game. I’m going to throw darts at the board below, and each dart will land somewhere random on the rings. Your job is to guess where my next dart will land, and the further off you are, the worse your score. Figure 2: The dartboard for our game. So where should you guess? It turns out the guess that minimizes your error is the exact middle of the board.1 We trained a model to play this game, and sure enough, it guesses the middle every time (this is the optimal prediction here)! Figure 3: A model playing our game (blue) guesses the middle of the board. This is terrible though… the middle is almost never where a dart actually lands. The “optimal” guess is a spot that no darts ever land. This is the core problem of generative modeling. When a prediction has many valid answers, the best single prediction is their average, and the average of data is generally a bad answer that looks nothing like the real data.2 And this problem isn’t special to dartboards or dogs, it shows up with any kind of data. When we trained a model to directly generate three piles of 2D points, it predicted a single dot in the middle of them, and when we trained one on text, all it could say was “the”. Real data What the model predicts Real data What the model predicts Figure 4: Direct prediction collapses to the average on any kind of data. But wait. ChatGPT writes coherent text, and image models generate really amazing images. Clearly this problem has been solved somehow, right? It has, and every scalable generative model today solves it the same way, by breaking generation into many small steps during training, so each step has roughly one right answer. When a step has one right answer, there’s nothing to average, and the blur disappears. Let’s look at how this works. Autoregressive models (like LLMs) predict one piece at a time, which in our game means never guessing the dart’s exact position all at once. Instead, you first guess only how far left or right the dart lands, and then given that, you guess how far up or down. Once you know the dart landed on the far right, there are only a couple places it could be. Step 1: pick a left-right spot Step 2: pick up-down, given left-right Figure 5: Autoregression predicts one sequence element at a time. In our game, once the left-right position is chosen, the up-down prediction only has two small spots left to choose from. Diffusion models do this differently. They start from pure random noise and take hundreds of tiny steps toward the data. Early on, their guess could still become any dart, but every step narrows the possibilities, so no single step ever faces many valid answers at once. Start: could become any dart Partway: fewer areas remain Near the end: pinned down Figure 6: Diffusion takes small steps from noise to data. Blue shows the darts the model's guess (purple) could still become, which narrows with every step. It turns out this is basically how every modern generative model works, by breaking the generation process into smaller pieces that can be predicted well. This includes LLMs, image and video models, and even newer few-step models like MeanFlow and consistency models. We refer to this idea of breaking generation into pieces as factoring generation. This approach of factoring generation works, but it’s also evil for a couple of reasons. The first is that models get trained on a single step, yet run for hundreds or thousands of steps at inference, so their own imperfect outputs get fed back in as inputs, errors compound, and generations slowly drift away from anything the model saw during training. This problem is called exposure bias (I wrote a whole blog on why it’s evil), and it’s why video models melt into mush after ten seconds and why LLMs get less coherent over really long generations, directly hurting performance and generalization. The second evil builds on the first, because that mismatch between training and inference means these models are never end-to-end, where an end-to-end model runs at inference exactly the way it was trained. End-to-end learning is what kicked off the deep learning revolution with AlexNet, and the lesson has held ever since… letting models learn everything directly from data beats hand-designing parts of the pipeline, and a model that runs the way it was trained is never forced into out-of-distribution territory. Nearly all of deep learning has gone end-to-end by now except generative modeling, and factoring generation is exactly what’s blocking it. So ideally we’d stop factoring generation, but factoring is also the only trick we know that handles the many-answers problem. The natural question then is whether we could factor something else instead, and it turns out a generative model only has two processes, how it generates and how it trains. If generation is off the table, that leaves the training loop. So what does factoring training look like? To answer this, let’s go back to our game, except this time I’ll give you twenty guesses instead of one, and only your closest guess counts. It turns out that with twenty guesses, guessing the middle becomes a terrible strategy. This is because you can now spread your guesses over the spots where darts actually land, lowering your error far more than the middle ever could. In other words, the winning strategy is to use your guesses to explore different answers. And this is exactly what happens. When we train a model this way with twenty guesses (middle panel below), its guesses spread across the board! 2 guesses 20 guesses 200 guesses Figure 7: When only the closest guess counts, guesses spread across the board instead of averaging. Take a second to appreciate what just happened here. The darts land in the exact same places as before, but because we changed how guesses are scored, the best possible prediction moved from the middle of the board onto the spots darts actually land. This reveals something important, which is that the training objective alone controls what the best prediction is (the loss minimizer), and by changing it, we moved the loss minimizer from the average of the data onto the data itself. This is Explorative Modeling. At each training step, the model explores K possible matches between what it generates and the real data, and only the best match gets trained. We call models trained this way Explorative Models (XMs). In the simplest case, this is literally, beautifully, a for loop: losses = [] for i in range(K): generation = model.generate() # e.g., from a different random noise losses.append(loss_fn(generation, data)) min(losses).backward() # only the best generation gets gradients