Modelagem Exploratória: Como um Loop de Treinamento de Três Linhas Libera o Treinamento End-to-End para Modelos Generativos

O deep learning tem uma lei de ferro: o treinamento end-to-end supera pipelines montados manualmente. A AlexNet provou isso, e classificação, detecção e segmentação de imagens vivem segundo essa regra desde então. Os modelos generativos são a exceção teimosa. Os modelos autorregressivos e de difusão mais fortes aprendem, durante o treinamento, a prever apenas um único pequeno passo e, na inferência, desdobram esse passo centenas ou milhares de vezes. Treinamento e inferência nunca amostram da mesma forma, e o erro de cada passo alimenta o seguinte — a notória exposure bias (viés de exposição).

Um artigo recente de pesquisadores da UIUC e de Harvard, Explorative Modeling (XM), argumenta que a peça que faltava é menor do que qualquer um esperava: um laço for de três a cinco linhas dentro do objetivo de treinamento. A implementação completa em PyTorch é open source, com scripts de treinamento para ImageNet condicionado a classes em 256×256, world models de vídeo e modelos de linguagem. Este guia apresenta a ideia e, em seguida, coloca o código oficial para rodar.

Por Que Modelos Generativos Quebram no Treinamento End-to-End: é a Média

Classificação é tolerante porque cada entrada geralmente tem uma única resposta certa, então aprender um mapeamento determinístico resolve o problema. Geração não é assim. Peça a um modelo que desenhe um cachorro e todo cachorro plausível é uma resposta válida — cada um é um modo separado, seu próprio pico na distribuição dos dados.

O problema começa com perdas de reconstrução como o MSE. Quando uma entrada é pareada aleatoriamente com muitos alvos válidos, a previsão que minimiza a perda é a média desses alvos. Para praticamente qualquer dataset real, essa média não está na variedade (manifold) dos dados. Ela cai entre os modos e não se assemelha a nenhum deles. O sintoma é o borramento de modos: três agrupamentos de pontos previstos como uma única mancha no meio, uma foto de cachorro que renderiza como papa, uma frase que degenera em “o” repetido para sempre.

Todo modelo generativo mainstream contorna isso fragmentando a geração em pedaços. Autorregressão prevê um elemento por vez. Difusão remove uma lasca de ruído por vez. Cada alvo de treinamento é fatiado até conter essencialmente um único modo, de modo que a perda de reconstrução não tem nada para tirar média. Essa decomposição protege a qualidade — e mata o treinamento end-to-end.

A pergunta dos autores é direta. Em um modelo generativo, só há duas coisas que você pode decompor: como ele gera e como ele treina. Se decompor a geração destrói o caráter end-to-end, por que não decompor o treinamento em vez disso?

O Mecanismo Central: Exploração Best-of-K

O XM decompõe o próprio loop de treinamento. A cada passo, o modelo não se compromete com uma única amostra e a força contra o alvo. Ele gera K candidatos, mantém o mais próximo dos dados ground-truth e propaga o gradiente apenas por esse. O pseudocódigo do README oficial:

# 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

Por que isso dissolve o borramento de modos? Imagine umaanka dardos. Se você tem um único palpite avaliado pela distância, a estratégia racional é palpar o centroide — o ponto onde quase nenhum dardo realmente cai. Se você tem K palpites e apenas o mais próximo conta, a estratégia ótima inverte imediatamente: espalhe seus palpites para que cada um reivindique um agrupamento diferente. Os modelos respondem da mesma forma. Ruídos de entrada distintos reivindicam cada um um modo diferente, em vez de se aglomerarem em direção à média. Aquilo que você pode explorar, você pode capturar.

Os autores nomeiam essa capacidade há muito ignorada de expressividade generativa, e a rastreiam até o objetivo de treinamento — aumente parâmetros e dados quanto quiser e ela nunca cresce por conta própria. Esse enquadramento também explica a dependência do campo em relação à guidance: o classifier-free guidance essencialmente empurra as previsões para longe da média borrada. Um modelo que não borra em primeiro lugar não precisa do empurrão.

Inserindo o XM em Modelos de Difusão e Flow

A integração com modelos generativos existentes é deliberadamente enxuta — explore sobre o latent (o ruído, no caso de difusão) e propague o gradiente apenas pelo melhor candidato. O exemplo oficial:

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

No repositório, a exploração é uma única flag de linha de comando, --xm_best_of_k K; K=1 reproduz a baseline sem exploração. As duas direções de exploração também se combinam:

Forward XM mantém o alvo real fixo e busca entre suas próprias gerações a mais próxima — orientado a recall, construído para cobrir todos os modos.
Reverse XM mantém uma geração fixa e busca nos dados reais a correspondência mais próxima — orientado a precisão, com quase nenhum custo computacional extra; a contrapartida é o possível colapso em poucos modos.

Rodando o Código Oficial

O repositório é organizado. model/ contém as implementações de DiT, flow matching, Jumpy e modelos de linguagem; job_scripts/<modality>/ separa os scripts de treinamento por modalidade; e slurm_executor.sh submete qualquer um deles a um cluster Slurm. Para começar:

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

Datasets de vídeo (SSv2, Kinetics-400) precisam do ffprobe no PATH — consulte primeiro data/vid/README.md. FID e FVD rodam online entre os passos de treinamento via --run_online_evaluation, então os números de qualidade chegam sem precisar esperar o fim de uma execução.

Lições Práticas

Trate K como um terceiro eixo de escalonamento. No artigo, os ganhos da exploração crescem com a escala: à medida que os dados escalam, o ganho sobe de 7% para 36%; à medida que o tamanho do modelo escala, de 13% para 23%; e triplicar o cálculo aproximadamente dobra o ganho de eficiência. Em pequena escala, parâmetros e dados são o gargalo e a exploração mal aparece. Quanto maior a execução, mais o K compensa.
Comece com K=2. O Reverse XM adiciona quase nenhum custo computacional, então valide por aí primeiro; adicione o Forward XM quando precisar de cobertura completa dos modos. A exploração multiplica o custo do forward pass, então não comece com um K enorme.
Fique de olho nos casos de uso end-to-end. Usado como modelo end-to-end independente, o XM iguala o Diffusion Policy — que precisa de 100 forward passes na inferência — com um único forward pass em behavior cloning, e supera o Diffuser em world modeling com 16–256× menos computação de inferência. Controle sensível à latência e cenários embodiment são onde isso vale mais a pena.
Reproduza você mesmo os números de destaque. O artigo reporta eficiência de 4,1× em FLOPs, 6,2× em eficiência de amostras, 47% de eficiência de parâmetros e um FID de 1,43 sem guidance no ImageNet. Rode um sweep com --xm_best_of_k 1 e outro com --xm_best_of_k 5, e verifique a curva de ganho contra a sua própria carga de trabalho.

Artigo e Código

Artigo: arXiv:2607.27372
Código: github.com/alexiglad/XM
Página do projeto: explorative-modeling.github.io
Explicação do autor: Explorative Modeling no blog do autor

Por uma década ajustamos modelos generativos com dois botões: fazer o modelo maior e alimentá-lo com mais dados. Este artigo adiciona um terceiro botão tão simples que parece suspeito — palpite várias vezes, fique com o melhor palpite. Mas todos os experimentos apontam na mesma direção: os dois primeiros botões chegam ao fim do curso eventualmente, e o terceiro mal começou a girar.