Explorative modeling: Train on the best of K guesses

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 lef