Explorative Modeling: How a Three-Line Training Loop Unlocks End-to-End Training for Generative Models

AI Guide
Explorative Modeling: How a Three-Line Training Loop Unlocks End-to-End Training for Generative Models

Deep learning has an iron law: end-to-end training beats hand-assembled pipelines. AlexNet proved it, and image classification, detection, and segmentation have lived by it ever since. Generative models are the stubborn exception. The strongest autoregressive and diffusion models learn to predict only a single small step during training, then unroll that step hundreds or thousands of times at inference. Training and inference never sample the same way, and every step's error feeds the next — the notorious exposure bias.

A recent paper from researchers at UIUC and Harvard, Explorative Modeling (XM), argues that the missing piece is smaller than anyone expected: a three-to-five line for loop inside the training objective. The full PyTorch implementation is open source, with training scripts for class-conditional ImageNet at 256×256, video world models, and language models. This guide walks through the idea, then gets the official code running.

Why Generative Models Break Under End-to-End Training: It's the Averaging

Classification is forgiving because each input usually has one right answer, so learning a deterministic mapping does the job. Generation is not like that. Ask a model to draw a dog and every plausible dog is a valid answer — each one a separate mode, its own peak in the data distribution.

The trouble starts with reconstruction losses such as MSE. When one input is randomly paired with many valid targets, the loss-minimizing prediction is the average of those targets. For almost any real dataset, that average does not lie on the data manifold. It falls between the modes and resembles none of them. The symptom is mode blurring: three clusters of scatter predicted as a single blob in the middle, a dog photo that renders as mush, a sentence that degenerates into "the" repeated forever.

Every mainstream generative model dodges this by shattering generation into pieces. Autoregression predicts one element at a time. Diffusion removes a sliver of noise at a time. Each training target gets sliced until it contains essentially a single mode, so the reconstruction loss has nothing to average over. That decomposition protects quality — and kills end-to-end training.

The authors' question is blunt. A generative model has only two things you can decompose: how it generates and how it trains. If decomposing generation destroys end-to-end-ness, why not decompose the training instead?

The Core Mechanism: Best-of-K Exploration

XM decomposes the training loop itself. At each step the model does not commit to one sample and force it against the target. It generates K candidates, keeps the one closest to the ground-truth data, and backpropagates only through that one. The pseudocode from the official README:

# Before exploration: generate one, backprop straight away
y = model(sample_latent())          # one output from noise/mask
loss = recon_loss(y, x)             # scored against real target x
loss.backward()

# After exploration (Forward XM): sample K, keep the best
losses = []
for _ in range(K):                  # explore K candidate outputs
    y = model(sample_latent())      # generate one candidate
    losses.append(recon_loss(y, x)) # score each against x
min(losses).backward()              # gradient flows through the closest one only

Why does this dissolve mode blurring? Picture a dartboard. If you get one guess scored by distance, the rational strategy is to guess the centroid — the spot where almost no darts actually land. If you get K guesses and only the closest one counts, the optimal strategy flips immediately: spread your guesses so each one claims a different cluster. Models respond the same way. Distinct input noises each claim a different mode instead of crowding toward the average. Whatever you can explore, you can capture.

The authors name this long-overlooked capacity generative expressivity, and trace it to the training objective — scale up parameters and data all you want and it never grows on its own. That framing also explains the field's dependence on guidance: classifier-free guidance essentially pushes predictions away from the blurry mean. A model that does not blur in the first place does not need the push.

Dropping XM Into Diffusion and Flow Models

Integration with existing generative models is deliberately thin — explore over the latent (the noise, in the diffusion case), and backprop only the best candidate. The official example:

t = sample_timestep()               # sample a noise level
losses = []
for _ in range(K):                  # explore K candidate noises
    z = randn_like(x)               # one candidate noise
    x_t = add_noise(x, z, t)        # noise the data to level t
    losses.append(diffusion_loss(model(x_t, t), x, z))
min(losses).backward()              # gradient through the closest candidate only

In the repository, exploration is a single command-line flag, --xm_best_of_k K; K=1 reproduces the no-exploration baseline. The two exploration directions also combine:

  • Forward XM holds the real target fixed and searches your own generations for the closest one — recall-oriented, built to cover every mode.
  • Reverse XM holds one generation fixed and searches the real data for the closest match — precision-oriented, with almost no extra compute; the trade-off is possible collapse onto a few modes.

Running the Official Code

The repository is tidy. model/ holds the DiT, flow matching, Jumpy, and language-model implementations; job_scripts/<modality>/ splits training scripts by modality; and slurm_executor.sh submits any of them to a Slurm cluster. To get started:

git clone https://github.com/alexiglad/XM.git
cd XM
conda create -n xm python=3.12
conda activate xm
pip install -r requirements.txt

# Key configuration
export HF_HOME=/path/to/cache     # dataset/model cache
export HF_TOKEN=...               # ImageNet requires accepting the license
wandb login                        # logging via W&B

# Class-conditional ImageNet 256x256 training (XDiffusion = DiT + exploration)
bash job_scripts/img/pretrain_class_conditional/xdit.sh

# Or submit to Slurm (fill the TODOs in example_h100.slurm for your cluster)
bash slurm_executor.sh example_h100 job_scripts/img/pretrain_class_conditional/xdit.sh

Video datasets (SSv2, Kinetics-400) need ffprobe on the path — check data/vid/README.md first. FID and FVD run online between training steps via --run_online_evaluation, so quality numbers arrive without waiting for a finished run.

Practical Takeaways

  • Treat K as a third scaling axis. In the paper, exploration gains grow with scale: as data scales up, the gain climbs from 7% to 36%; as model size scales up, from 13% to 23%; and tripling compute roughly doubles the efficiency gain. At small scale, parameters and data are the bottleneck and exploration barely shows. The bigger the run, the more K pays.
  • Start at K=2. Reverse XM adds almost no compute, so validate there first; add Forward XM when you need full mode coverage. Exploration multiplies forward-pass cost, so don't open with a huge K.
  • Watch the end-to-end use cases. Used as a standalone end-to-end model, XM matches Diffusion Policy — which needs 100 forward passes at inference — with a single forward pass in behavior cloning, and beats Diffuser on world modeling with 16–256× less inference compute. Latency-sensitive control and embodied settings are where this is worth the most.
  • Reproduce the headline numbers yourself. The paper reports 4.1× FLOP efficiency, 6.2× sample efficiency, 47% parameter efficiency, and a guidance-free 1.43 FID on ImageNet. Run one sweep at --xm_best_of_k 1 and one at --xm_best_of_k 5, and check the gain curve against your own workload.

Paper and Code

For a decade we have tuned generative models with two knobs: make it bigger, feed it more data. This paper adds a third knob so plain it looks suspicious — guess several times, keep the best guess. But every experiment points the same way: the first two knobs run out of travel eventually, and the third one has barely started turning.

Scroll to Top