Designing Instance-Level Sampling Schedules via REINFORCE with James-Stein Shrinkage

Resumo Técnico: Projetando Cronogramas de Amostragem de Nível de Instância via REINFORCE com Encolhimento de James-Stein

1. O Problema

Atualmente, a maioria dos modelos de difusão e flow-matching (como Stable Diffusion e Flux) utiliza um cronograma de amostragem (sampling schedule) global e fixo. Isso significa que a mesma distribuição de passos de tempo (timesteps) é aplicada a todos os prompts e sementes de ruído.

O problema central é que esse cronograma universal não é otimável para todos os casos: diferentes prompts exigem diferentes níveis de detalhes semânticos ou espaciais, e diferentes sementes de ruído criam condições iniciais distintas. A falta de adaptação resulta em uma alocação ineficiente do "orçamento" de passos de inferência, prejudicando a fidelidade da imagem, a renderização de texto e o controle composicional, especialmente em regimes de poucos passos (few-step sampling).

2. Metodologia

Os autores propõem uma abordagem de pós-treinamento agnóstica ao modelo, que não altera os pesos do modelo de difusão (que permanece congelado), mas aprende a reescalonar o cronograma de amostragem para cada instância específica (prompt + ruído).

Principais componentes metodológicos:

• Política de Passo Único (Single-Pass Policy): Diferente de métodos autorregressivos que preveem o próximo passo um por um (o que aumenta o custo de inferência), os autores utilizam uma política baseada em uma Distribuição de Dirichlet. A rede de política (um encoder-decoder leve) recebe o ruído e o prompt e, em uma única passagem, gera todos os intervalos de tempo do cronograma. Isso garante um custo de inferência constante e mínimo.
• Otimização via REINFORCE: O aprendizado é formulado como um problema de aprendizado por reforço (RL), onde o objetivo é maximizar a recompensa esperada (ex: pontuações de alinhamento CLIP ou modelos de preferência humana).
• Baseline de Recompensa James-Stein (JS): Este é o núcleo matemático do trabalho. O uso de gradientes de política (REINFORCE) sofre com alta variância, especialmente em espaços de alta dimensão. Os autores introduzem um novo baseline de controle baseado no Estimador de James-Stein.
  • Eles tratam as recompensas como um modelo de "efeitos aleatórios", onde as recompensas de um prompt específico são "encolhidas" (shrunk) em direção à média global do lote (batch).
  • O baseline JS combina a média dentro do contexto (RLOO) com a média entre contextos (cross-context), adaptando o peso de cada uma de forma baseada nos dados. Isso provadamente reduz o erro de estimativa de gradiente e a variância em comparação com métodos tradicionais como o RLOO.

3. Principais Contribuições

1. Framework de Scheduler de Nível de Instância: Uma nova forma de adaptar o cronograma de amostragem por prompt e semente sem modificar o modelo base.
2. Baseline de James-Stein para RL: Uma técnica de redução de variância robusta e agnóstica que estabiliza o treinamento de políticas de longo horizonte.
3. Eficiência de Amostragem: Demonstração de que é possível alcançar qualidade de modelos destilados (como o Flux-Schnell) usando modelos não destilados (como o Flux-Dev) através de apenas 5 passos de amostragem otimizados.

4. Resultados

Os experimentos foram realizados em modelos de ponta como SD-XL, SD-3.5 e Flux-Dev, utilizando o benchmark Human Preference Dataset (HPD) v2.

• Alinhamento Geral: O método JS superou consistentemente todos os outros baselines (Default, RLOO, XCTX) em pontuações de preferência humana (HPSv2) em todos os orçamentos de passos.
• Renderização de Texto e Detalhes Finos: O cronograma aprendido melhorou significativamente a capacidade de renderizar caracteres legíveis e tarefas de contagem e composição (GenEval), onde o cronograma padrão falhava.
• Desempenho em Poucos Passos (Few-Step): Com apenas 5 passos, o modelo Flux-Dev com o cronograma otimizado atingiu uma qualidade comparável a modelos que passaram por um processo de destilação pesado, provando que o potencial de geração rápida já reside nos modelos pré-treinados, mas está "escondido" por cronogramas subótimos.

5. Significância

Este trabalho é significativo porque oferece uma nova alavanca de pós-treinamento para modelos generativos. Em vez de gastar recursos massivos treinando ou destilando novos modelos (o que é caro e altera os pesos), os pesquisadores podem simplesmente treinar um "scheduler" leve que desbloqueia o potencial latente de modelos já existentes. Isso abre caminho para sistemas de geração de imagem mais rápidos, precisos e adaptáveis com um custo computacional adicional quase desprezível.