Explorative Modeling: Unlocking a Third Pretraining Axis and End-to-End Generation — part 5/8

Cleaned-text code points [62048, 77719). Bibliography omitted here; source.md includes it.

Huge thanks to Flapping Airplanes for supporting Alexi as a fellow while completing this work. Massive thanks to Laude Institute for supporting this work, with a special shoutout to Braden Hancock and K. Tighe. Thanks to Soran Ghaderi for productive early discussions related to XMs. Thanks to Omead Pooladzandi and Samip Dahal for great feedback on XMs. This material is based upon work supported by the U.S. National Science Foundation Graduate Research Fellowship Program under Grant No. DGE 21-46756, U.S. DARPA ECOLE Program No. #HR00112390060, DARPA ITM Program No. FA8650-23-C-7316, NSF Molecule Maker Lab Institute, an AI Institute for Molecular Discovery, Synthesis Strategy, and Manufacturing funded by the U.S. National Science Foundation under Awards No. 2019897 and 2505932, the AI Research Institutes program by National Science Foundation and the Institute of Education Sciences, U.S. Department of Education through Award No. 2229873 - AI Institute for Transforming Education for Children with Speech and Language Processing Challenges, and NSF NAIRR award. Any opinions, findings, and conclusions or recommendations expressed in this material are those of the author(s) and do not necessarily reflect the views of the National Science Foundation, the Defense Advanced Research Projects Agency (DARPA), the Institute of Education Sciences, or the U.S. Department of Education. This research used the Delta and DeltaAI advanced computing and data resources, which are supported by the National Science Foundation (award OAC 2320345 and award OAC 2005572) and the State of Illinois. Delta and DeltaAI are joint efforts of the University of Illinois Urbana-Champaign and its National Center for Supercomputing Applications. Some of the computations in this paper were run on the FASRC cluster supported by the FAS Division of Science Research Computing Group at Harvard University.

## Appendix A Additional Experimentation

Figure A.1 reports the FID versions of the FDr 6 \mathrm{FDr}^{6} convergence plots in Figure 4 , where XRAE similarly converges much faster than the baseline.

Though we do not report them in detail, we have also observed these benefits in the discrete domain, where adding exploration to MDLMs improves both sample efficiency and generalization, mirroring the trends for image and video generation (Figures 4 and 9 ). Similarly, we have observed early evidence that exploration improves the data efficiency of autoregressive language models, which we leave to future work to report in detail.

### A.1 Exact Results for Exploration Scaling

For reproducibility, and to make future comparisons easier, Table A.1 reports the exact FID and FVD values behind Figure 7 .

### A.2 Comparing Exploration to Minibatch Optimal Transport Couplings

A natural alternative to exploration is to reduce mode blurring by computing a better coupling directly, most commonly with minibatch Optimal Transport (OT) [ 39 , 82 ] . Comparing the two, we find minibatch OT couplings actually hurt image generation performance, worsening FID from 46.3 46.3 to 54.5 54.5 for Small models and from 74.4 74.4 to 82.7 82.7 for Base models (Table A.2 ).

We attribute this to two problems that exploration avoids. First, minibatch OT is a biased approximation of the global coupling [ 40 ] , and this bias worsens as datasets grow, since each batch covers a vanishing fraction of the data. Second, the two select couplings on different grounds. Minibatch OT assigns pairs by geometry alone, computed within a batch regardless of what the model has learned; exploration selects by the model’s own current loss, so its coupling co-adapts with the model throughout training. Empirically, the model-aligned choice helps while the geometric one hurts, and exploration’s gains grow with scale rather than degrade (Figure 10 ).

### A.3 Reverse XM Language Models

We have successfully trained Reverse XM language models, where each generation searches the training data through a vector database rather than comparing against K K samples in a batch. Getting this to work relied on two ingredients. First, once a data point is matched and trained on, it is removed from the search pool for that epoch, which we track as a train coverage percent . This prevents nearest-neighbor search from collapsing onto the same few points, playing the role of the entropy term Reverse XM needs to avoid collapse (Section F.2 ), though as a coverage heuristic rather than a literal entropy bonus. Additionally, we found it important to set the train coverage percent to be less than the entire train set (e.g., 50 % 50\% ), or else models have high loss near the end of an epoch as they try to cover the extremes of a dataset. Second, the search happens in a hybrid of representation space and cross-entropy space, so that retrieved neighbors reflect the actual training loss. We strongly believe this can be further improved upon, and we have not yet open sourced this code, but plan to.

## Appendix B Additional Intuition

#### Generative Modeling with a For Loop.

It is worth appreciating how simple end-to-end Explorative Modeling can be. Unlike diffusion, flow, or autoregressive models, which rely on many-step sampling procedures at inference to avoid mode blurring and often complicated masking/noising schedules (Section 2.1 ), end-to-end XMs generate with a single pass from the model and training XMs in the simplest case is just a short for loop with 3-5 lines of code (Algorithms 1 and 2 ).

#### Exploration as Building an Associative Memory.

A generative model turns an input, usually noise drawn from a multivariate Gaussian, into a sample. Explorative Modeling searches over this noise to find which inputs the model should tie to which samples. Once trained, each region of noise acts as a key that retrieves one mode of the data as its value, so XMs behave much like associative memories. A larger K K splits the noise into finer regions, so each mode gets its own instead of blurring together (Figure 2 ). This holds for every generative model, not just XMs. However, it does not mean that generative models simply memorize the training data. Large datasets act as a form of regularization: with far more examples than the model can store individually, it is forced to reuse parameters across them and learn the structure they share instead of the datapoints themselves, which results in generalization.

## Appendix C Approach Details

#### Batching Forward XM.

Algorithm 1 is written as a for loop for clarity, but in practice the K K explored generations can be folded into the batch dimension and computed as a single larger forward pass. Because accelerators process larger batches efficiently, this parallelizes exploration and speeds up training considerably compared to looping over candidates one at a time. We provide batched code for Forward XM in the released code.

#### Other Ways to Explore.

Drawing K K fresh candidates per step is not the only way to do Forward XM. For example, generations could be cached by their class or conditioning and reused as candidates for later samples with the same condition—on ImageNet, caching recently generated samples for each class would provide candidates essentially for free—though this works worse for continuous conditioning, where datapoints rarely share the exact same condition. We use fresh draws throughout, as this is the simplest approach to implement, is fairest when comparing across domains, and is less prone to collapse.

#### Gradients and Memory.

Only the best candidate receives gradients (Equation 1 ), which can be implemented in two ways. In the memory-saving mode, all K K candidates are forwarded without gradients, and only the best is re-forwarded with gradients to train on, keeping activation memory the same as standard training at the cost of one extra forward pass. In the FLOP-efficient mode, all K K candidates are forwarded with gradients and only the lowest loss is backpropagated, avoiding the extra forward pass at the cost of storing activations for all K K candidates. Switching between the two lets exploration adapt to whichever of FLOPs or memory is the bottleneck.

#### Exploration FLOP Cost.

For transformers, a training step costs roughly 6 ​ N ​ D 6ND FLOPs for N N parameters and D D tokens, where the forward pass costs 2 ​ N ​ D 2ND and the backward pass 4 ​ N ​ D 4ND [ 84 ] . Since only the best candidate is trained on, each additional explored candidate in Forward XM adds only a forward pass, so each additional mode explored costs roughly 1 3 \tfrac{1}{3} of a standard training step (XM- K K costs K + 2 3 \tfrac{K+2}{3} standard steps in the FLOP-efficient mode, plus one more forward pass in the memory-saving mode). Exploring more modes in Reverse XM is far cheaper, as additional data targets add no forward passes, only extra loss computations—essentially one matrix multiplication scaling with the data dimension—which is negligible compared to the network’s FLOPs.

#### Generating the K K Candidates.

For all continuous models in this work, both end-to-end XMs and the hybrid XMs built on Diffusion/Flow and Jumpy models (Section 4.1 ), each explored candidate uses a different input noise draw; for the hybrids, this means the same data sample, timestep, and condition (including condition dropping for guidance when the underlying recipe uses it), with only the noise varying. For XMDLMs, we instead learn K K discrete latent embeddings: each candidate samples one of these embeddings at random for every masked position, and the best-of- K K selection is over the resulting latent-conditioned predictions.

#### Implementing Reverse XM.

Reverse XM requires each generation to have a set of valid data targets to search over. With discrete conditioning, this comes almost for free through the dataloader: loading K K datapoints per condition lets each generation pick its best match within the batch, at no extra generation cost. The same can be done with continuous conditioning, just involving additional tricks to load similar latents together. At larger scale, the whole dataset can instead be indexed in a vector database, letting each generation search all training data in roughly logarithmic time, so the number of modes explored can in principle reach the dataset size (we take this approach for the Reverse XM language modeling experiments described in Section A ; see also Section 6 ).

## Appendix D Experimental Details

A lot of the experimentation paragraph and takeaway style is inspired by [ 85 ] . Some of the figure designs in this paper follow those of Mode Forcing [ 19 ] .

#### Model Sizes.

All image and video generation models are transformers following the standard DiT size conventions [ 56 ] , summarized in Table D.1 .

#### Image Generation.

All image generation experiments train class-conditional models on ImageNet 256 × \times 256. Aside from the RAE experiments, we follow the SiT setup [ 29 ] , with the exception of not using a horizontal flip augmentation. Images are encoded by the VAE [ 55 ] into 4 × 32 × 32 4\times 32\times 32 latents, which with a patch size of 2 gives a 16 × 16 16\times 16 grid of patches, or 256 tokens per image (this is the token count behind the data scaling in Figure 5(a) ). All models train with a batch size of 256 and a learning rate of 1 ​ e − 4 1\mathrm{e}{-4} , using 10k steps of linear warmup followed by cosine decay to 1M steps (or to 3M steps for the longer runs in Figure 5(b) ), along with gradient clipping of 1.0 and weight decay of 0.01. For Diffusion models we use the Flow Matching formulation, and following the RAE and JiT papers [ 45 , 86 ] we sample with 50 Heun steps [ 87 ] by default, which we found was enough to converge. For the FLOP scaling curves in Figure 5(b) we use Jumpy models [ 30 ] , as they were more stable for longer runs and more performant. These experiments all report validation-set generative metrics, as they better measure overfitting and generalization.

#### RAE Image Generation.

For the RAE experiments (Table 1 and Figure 4 ) we follow the RAE setup [ 45 ] , building directly on their codebase and only adding exploration, including their batch size of 1024. We rely primarily on FDr 6 \mathrm{FDr}^{6} [ 43 ] for these experiments, as FID at this level of performance is overfit and no longer tracks true sample quality [ 43 , 44 ] . For guided results, both the RAE baseline and XRAE use AutoGuidance [ 26 ] : the RAE baseline uses a guidance scale of 1.42 with the released S model at epoch 14 as the guiding model; for XRAE-2, the best FDr 6 \mathrm{FDr}^{6} uses a scale of 1.5 with the same released S model at epoch 14, while the best gFID uses a scale of 1.35 with an XM-2 L model at epoch 9.

#### Video Generation.

Video generation experiments use the Something-Something V2 dataset [ 49 ] at 128 × 128 128\times 128 resolution, as higher resolutions required too much compute. We always model 10 frames, passing in frames 0, 1, and 9 as conditioning (simulating goal-conditioned world modeling [ 18 ] ), and use a 3D video transformer [ 88 ] with Base model size and a patch size of 4, also chosen due to limited compute, as these experiments aim for fair comparisons rather than state-of-the-art results. Models train with a batch size of 256, a learning rate of 1 ​ e − 4 1\mathrm{e}{-4} , and weight decay of 1 ​ e − 2 1\mathrm{e}{-2} .

#### Individual Results in More Detail.

Figure 7(a) uses Small models, while Figure 7(b) uses Base models. Figure 5 uses Base models for class-conditional image generation, where the baseline is the optimally tuned SiT recipe [ 29 ] . Figure 9 uses 4-step Jumpy models, plotting FVD over the course of training rather than the best performance reached. Figure 10(a) measures image generation improvements across Jumpy model sizes from Small through XLarge, and Figure 10(b) does the same for XLarge models across training steps. Figure 11 uses Base-sized image generation models.

#### Behavior Cloning.

All Explorative Policy runs use XM-10. All Behavior Cloning models are CNNs rather than transformers, and we follow most of the Diffusion Policy setup [ 17 ] . However, we use a newer version of robomimic [ 89 ] , so results are not perfectly comparable to those originally reported; we therefore reproduce Diffusion Policy’s results under our setup (Table 2 ).

#### Goal-Conditioned World Modeling.

All Explorative World Model runs use XM-10. We evaluate on the single-task Maze2D setting from Table 1 of the Diffuser paper [ 18 ] , training for 1M steps instead of 2M and using a goal radius of 0.45. As with Behavior Cloning, for a fair setup, we reproduce Diffuser’s results under our setup (Table 3 ). Our end-to-end explorative world models’ NFE exceeds one only with a recurrent block [ 65 ] : U-Maze recurs over the full depth 3 times (NFE 4), Large over the middle 3 times (NFE 1.9), Medium not at all.

#### Language Modeling.

Our MDLM [ 27 ] experiments use a context length of 256 and pretrain for 210k steps with a batch size of 64. Models are xxs-sized: 6 transformer blocks, 6 attention heads, and an embedding dimension of 384. We report generative perplexity using the Qwen2.5-1.5B base model (not post-trained) [ 90 ] as the evaluator, as it is a much better oracle than the commonly used GPT-2 [ 91 ] , though GPT-2 results confirmed the same trends. Results are primarily across an entropy range of 5.0 to 5.7 and the sampling step counts shown in Figure 8 . We also confirmed that exploration’s gains hold at larger scale as well as for infilling tasks.

## Appendix E Related Works

### E.1 End-to-End Generative Modeling