Resumo Técnico: LiDAR (Lookahead Sample Reward Guidance) para Escalonamento em Tempo de Teste de Modelos de Difusão

Declaração do Problema

Modelos de difusão alcançaram um desempenho generativo notável em imagens, vídeos e linguagem; no entanto, amostras geradas frequentemente falham em se alinhar totalmente com a intenção humana. Embora existam funções de recompensa que capturam preferências humanas, integrar essas funções no processo de amostragem apresenta desafios significativos. O objetivo central do escalonamento em tempo de teste é amostrar de uma distribuição inclinada pela recompensa $p_\theta^r(x_0|c) \propto p_\theta(x_0|c) \exp(\lambda \cdot r(x_0, c))$.

Alcançar isso requer o cálculo da Recompensa Futura Esperada (EFR - Expected Future Reward), $r_t^\lambda(x_t, c)$, para partículas intermediárias $x_t$ para guiar o processo de difusão reversa. Métodos existentes enfrentam limitações críticas:

1. Rollout de Retrocesso (Backward Rollout): Amostrar diretamente de $x_t$ para $x_0$ para estimar a EFR é computacionalmente proibitivo devido à necessidade de denoising iterativo em cada passo de tempo.
2. Aproximações baseadas em Tweedie: Métodos como Gradient Guidance e Monte Carlo Sequencial (SMC) dependem de aproximações de Taylor de primeira ordem (fórmula de Tweedie) para estimar a EFR. Estes introduzem viés, particularmente quando a força do sinal de recompensa ($\lambda$) é alta, e frequentemente exigem backpropagation através de redes neurais (Gradient Guidance) ou sofrem de colapso de amostra em espaços de alta dimensão (SMC).

Metodologia: LiDAR Sampling

Os autores propõem o LiDAR (Lookahead Sample Reward Guidance), um framework que desacopla a dependência neural entre a partícula intermediária $x_t$ e a EFR, permitindo uma orientação de forma fechada e livre de derivadas.

1. Reformulação da Recompensa Futura Esperada

O artigo deriva uma nova formulação para a EFR (Teorema 3.1) que expressa a expectativa de recompensa usando apenas amostras marginais de um modelo de difusão pré-treinado e o kernel de perturbação direto, em vez de condicionar na partícula atual $x_t$ via redes neurais:
$$r_t^\lambda(x_t, c) = \log \mathbb{E}_{p_\theta(x_0|c)} \left[ \frac{p(x_t|x_0)}{\mathbb{E}_{p_\theta(x_0|c)}[p(x_t|x_0)]} \exp(\lambda \cdot r(x_0, c)) \right]$$
Esta abordagem de "rollout direto" (forward rollout) permite que a EFR seja computada para qualquer $x_t$ usando um conjunto pré-gerado de amostras marginais $\{\hat{x}_0^i\}$, eliminando a necessidade de backpropagation através da rede de score $s_\theta$ ou do modelo de recompensa durante o passo de orientação.

2. Amostragem de Lookahead Eficiente

Para mitigar o custo de gerar amostras marginais para cada prompt, os autores introduzem uma estratégia eficiente de amostragem de lookahead (amostragem de antecipação). Em vez de usar o processo de difusão completo e lento, eles empregam um gerador rápido $q(x_0|c)$, como um solver de ODE de poucos passos (ex: DPM-Solver) ou um modelo destilado, para gerar um conjunto de $n$ amostras de lookahead $\{\hat{x}_0^i\}_{i=1}^n$. Essas amostras são anotadas com valores de recompensa $r(\hat{x}_0^i, c)$.

3. Orientação de Forma Fechada Livre de Derivadas

Usando o conjunto finito de amostras de lookahead, os autores derivam uma expressão de forma fechada para o gradiente da recompensa empírica de lookahead (Teorema 3.3). O termo de orientação direciona a partícula atual $x_t$ para amostras de lookahead de alta recompensa e repele as de baixa recompensa:
$$\nabla_{x_t} \hat{r}_t^\lambda(x_t, c) = \sum_{i=1}^n (w_i^r - w_i) \frac{\hat{x}_0^i}{\sigma_t^2}$$
onde $w_i^r$ e $w_i$ são pesos softmax baseados nos valores de recompensa e na distância entre $x_t$ e as amostras de lookahead $\hat{x}_0^i$. Esta formulação não requer diferenciação neural.

4. Leis de Escalonamento Teóricas

O artigo estabelece leis de escalonamento para o método:

• Precisão do Lookahead ($\delta$): A distância de variação total entre a distribuição alvo e a distribuição gerada pelo LiDAR diminui como $O(1/\sqrt{\delta})$, onde $\delta$ é o número de passos de discretização no solver de lookahead.
• Tamanho da Amostra ($n$): A variância da estimativa de recompensa empírica diminui como $O(1/\sqrt{n})$, implicando que aumentar o número de amostras de lookahead melhora a precisão da aproximação.

Principais Resultados

Experimentos foram conduzidos no SD v1.5, SDXL, FLUX (flow matching) e UDLM (difusão discreta) usando ImageReward, CLIP, HPS e GenEval.

• Desempenho vs. Eficiência: No SDXL, o LiDAR atinge um desempenho no GenEval comparável aos métodos mais recentes de gradient guidance (DATE), mas com um aceleração de 9,5× no tempo de inferência.
• Eficiência de Memória: Ao contrário dos métodos de gradient guidance que exigem ~3× mais memória devido ao backpropagation, o uso de memória do LiDAR é quase idêntico ao da amostragem vanilla, pois a etapa de amostragem alvo (Algoritmo 2) não requer computação de gradiente.
• Robustez ao Reward Hacking: Enquanto métodos de gradient guidance frequentemente degradam em métricas como CLIP e HPS ao otimizar para ImageReward (reward hacking), o LiDAR melhora todas as métricas simultaneamente. Isso é atribuído à correlação entre o ImageReward e outras métricas nas amostras de lookahead, que naturalmente guia o amostrador em uma direção equilibrada.
• Melhorias Ortogonais: O LiDAR proporciona ganhos de desempenho adicionais quando aplicado a modelos já ajustados com Direct Preference Optimization (DPO), demonstrando sua ortogonalidade ao alinhamento baseado em treinamento.
• Escalabilidade: O método escala efetivamente com o número de amostras de lookahead ($n$) e passos do solver ($\delta$), e pode ser combinado com outras técnicas como SMC e Best-of-N.

Significância e Alegações

O artigo alega que o LiDAR aborda sistematicamente as limitações de eficiência e de reward-hacking dos métodos de escalonamento em tempo de teste existentes. Ao reformular o cálculo da EFR para depender de rollout direto e amostras marginais, os autores permitem uma orientação livre de derivadas que é computacionalmente eficiente e teoricamente fundamentada.

Os autores enfatizam que o LiDAR alcança um compromisso eficiência-desempenho superior em comparação com os baselines existentes. Especificamente, ele evita o overhead computacional do backpropagation neural e os problemas de colapso de amostra do SMC, ao mesmo tempo em que fornece estimativas de EFR mais precisas do que métodos baseados em aproximação de Taylor. O trabalho sugere que o desacoplamento do sinal de recompensa da dependência neural da partícula intermediária é um caminho viável para o modelamento generativo escalável e alinhado com humanos.