Treinamos um Modelo de Texto para Imagem em 24 Horas (Receita Completa)

Introdução: O Desafio de 24 Horas

Antigamente, treinar um modelo de difusão de texto para imagem competitivo custava milhões de dólares e semanas de GPU. Os tempos mudaram. Neste post, compartilhamos a receita completa de um speedrun de 24 horas usando 32 GPUs H200 (cerca de $1500 no total). O objetivo foi empilhar todos os truques que realmente funcionam e ver até onde podemos ir com um orçamento restrito.

Disponibilizamos todo o código no repositório PRX para você reproduzir, modificar e estender os experimentos. Este post resume as principais decisões arquiteturais e configurações de treinamento. Para uma análise detalhada de cada componente, confira nossos posts anteriores (Parte 1 e Parte 2).

Se você se interessa por ferramentas de IA que simplificam debugging, dê uma olhada no nosso guia sobre como construir um assistente de troubleshooting para Kubernetes.

A Receita de Treinamento: Um Stack de Técnicas Comprovadas

1. Treinamento em Pixel Space com X-Prediction

Adotamos a formulação x-prediction do artigo Back to Basics: Let Denoising Generative Models Denoise (Li & He, 2025). Diferente dos modelos de difusão latente que dependem de um VAE, essa abordagem treina diretamente no espaço de pixels. Isso simplifica o pipeline e remove o gargalo do VAE.

• Tamanho do patch: 32
• Dimensão do bottleneck: 256 na camada de projeção inicial de tokens
• Escala de resolução: Começar em 512px → fine-tune em 1024px (pulando 256px)

Em 512px, a sequência tem cerca de 256 tokens. Em 1024px, vira 1024 tokens. Com hardware moderno, isso é totalmente gerenciável.

2. Perdas Perceptuais (LPIPS + DINOv2)

Uma grande vantagem do treinamento em pixel space é poder aplicar diretamente perdas perceptuais clássicas. Adicionamos dois objetivos auxiliares à perda padrão de flow-matching:

• Perda LPIPS (peso: 0.1) — captura similaridade perceptual de baixo nível.
• Perda perceptual baseada em DINOv2 (peso: 0.01) — fornece alinhamento semântico mais forte.

Aplicamos essas perdas em imagens completas pooled (não patch-wise) em todos os níveis de ruído. Isso deu resultados consistentemente melhores que a abordagem do artigo original.

3. Roteamento de Tokens com TREAD

Para reduzir o custo computacional por passo, usamos TREAD (Krause et al., 2025). Ele seleciona aleatoriamente 50% dos tokens e os roteia do 2º bloco transformer para o penúltimo bloco, pulando as camadas do meio.

Escolhemos TREAD em vez de SPRINT pela simplicidade. Para compensar a queda de qualidade com CFG vanilla, implementamos um esquema de self-guidance que guia usando predições condicionais densas vs. roteadas.

4. Alinhamento de Representação com REPA + DINOv3

REPA (Yu et al., 2024) alinha as representações intermediárias do modelo com um teacher pré-treinado. Usamos DINOv3 como teacher e aplicamos a perda de alinhamento no 8º bloco transformer (peso: 0.5).

Como combinamos REPA com roteamento TREAD, calculamos a perda de alinhamento apenas nos tokens não roteados para manter o sinal consistente.

5. Otimizador: Muon

Trocamos Adam por Muon para todos os parâmetros 2D (matrizes). Biases, norms e embeddings permaneceram com Adam. A configuração:

Escala de Treinamento e Dados

Usamos três datasets sintéticos:

• Flux generated (1.7M amostras)
• FLUX-Reason-6M (6M amostras)
• Midjourney v6 (1M amostras, re-captionadas com Gemini 1.5 para consistência)

A escala:

• 512px: 100k passos, batch size 1024
• 1024px: 20k passos, batch size 512 (REPA desligado)

Mantivemos uma EMA dos pesos (smoothing=0.999, update a cada 10 passos, início no passo 0).

Exemplo de Código: Configurando o Loop de Treinamento

Abaixo um snippet Python simplificado mostrando como configurar os componentes chave. O código completo está no repositório PRX.

# config.py (simplificado)
import torch
from prx import Trainer, TREADRouter, REPAAlignment, PerceptualLoss

# Configuração do modelo
model_config = {
    "patch_size": 32,
    "bottleneck_dim": 256,
    "num_layers": 24,
    "hidden_dim": 1024,
}

# Componentes de treinamento
router = TREADRouter(
    route_ratio=0.5,
    start_block=2,
    end_block=-2  # bloco penúltimo
)

repa = REPAAlignment(
    teacher_name="dinov3",
    loss_weight=0.5,
    apply_at_block=8
)

perceptual_loss = PerceptualLoss(
    lpips_weight=0.1,
    dino_weight=0.01,
    apply_at_all_noise_levels=True
)

# Otimizador
param_groups = [
    {"params": model.muon_params, "optimizer": "muon", "lr": 1e-4, "momentum": 0.95},
    {"params": model.adam_params, "optimizer": "adam", "lr": 1e-4, "betas": (0.9, 0.95)}
]

trainer = Trainer(
    model=model,
    router=router,
    repa=repa,
    perceptual_loss=perceptual_loss,
    param_groups=param_groups,
    batch_size=1024,
    steps_512=100000,
    steps_1024=20000
)

trainer.run()

Resultados e Limitações

Após 24 horas, o modelo produz imagens coerentes e esteticamente agradáveis com forte aderência ao prompt. O fine-tune em 1024px afia detalhes sem quebrar a composição. No entanto, alguns artefatos permanecem:

• Falhas de textura ocasionais
• Anatomia estranha em prompts complexos
• Performance degrada em prompts muito difíceis ou fora da distribuição

Esses problemas são consistentes com subtreinamento e diversidade limitada de dados, não uma falha estrutural na receita. Com mais computação e cobertura de dados mais ampla, a mesma configuração deve continuar melhorando de forma previsível.

O Que Isso Significa para a Área

Este speedrun demonstra que o treinamento de difusão se tornou notavelmente acessível. Combinando treinamento em pixel space, roteamento eficiente, alinhamento de representação e orientação perceptual leve, um modelo significativo é alcançável em um único dia com um orçamento que seria inimaginável há alguns anos.

Limitações e Cuidados

• Viés de dados sintéticos: Todos os dados de treinamento foram gerados por IA. Imagens do mundo real podem se comportar de forma diferente.
• Dependência de hardware: A receita assume GPUs H200 com grande memória. Resultados em hardware de consumo variam significativamente.
• Sem filtro de segurança: O modelo não foi ajustado para segurança ou mitigação de viés. Use com cuidado em produção.

Próximos Passos

Esta corrida de 24 horas é apenas o começo. O trabalho futuro focará em:

1. Escalar com mais computação e conjuntos de dados mais amplos
2. Melhorar a qualidade e diversidade das legendas
3. Explorar mais esquemas de self-guidance

Convidamos a comunidade a construir sobre este trabalho. Todo o código é open-source no repositório PRX. Se você está construindo seu próprio pipeline de difusão, também pode achar útil nosso guia sobre API nativa de popover para tooltips para criar melhores interações de UI em seus apps de demonstração.

Conclusão

Treinar um modelo de difusão de texto para imagem competitivo do zero em 24 horas não é mais ficção científica. A combinação de treinamento em pixel space, roteamento TREAD, alinhamento REPA e otimizador Muon funciona sinergicamente para entregar resultados fortes com um orçamento modesto.

A principal lição: a pesquisa em difusão se democratizou. Você não precisa de milhões de dólares para contribuir com trabalho significativo. Comece com o código do PRX, ajuste os botões e veja o que pode descobrir.

Obrigado por ler. Junte-se à discussão da comunidade no Discord e fique ligado na próxima rodada de experimentos. 🚀