Ponderação de Timestep Reduz Custos de Consulta ao Modelo de Recompensa para RLHF de Difusão

Um recente artigo da CMU e da Universidade de Maryland, intitulado *Selective Timestep Weighting and Advantage-Based Replay for Sample-Efficient Diffusion RLHF*, relata uma melhoria de seis vezes na eficiência de amostragem para RLHF de difusão ao reponderar os timesteps de desnoising e reproduzir trajetórias de alta-vantagem. Isso poderia reduzir significativamente os custos de consulta ao modelo de recompensa para equipes que trabalham com Stable Diffusion, SDXL ou FLUX. O artigo trata os pipelines de estilo DDPO existentes como um aumento plug-and-play em vez de uma arquitetura de substituição.

A intervenção aborda uma falha de atribuição de crédito no DDPO padrão, que trata a geração de imagem como um Processo de Decisão de Markov e propaga a recompensa da imagem final uniformemente em cada passo de desnoising. Como o feedback humano ou do modelo de recompensa está disponível apenas no output completamente desnoidoso, os timesteps iniciais e finais recebem pesos de perda idênticos, diluíndo os gradientes. Os autores corrigem isso ao reponderar cada timestep durante a otimização da política pela magnitude ao quadrado da mudança latente por timestep, usando-o como um proxy para a variância da vantagem do erro TD. Além disso, eles mantêm um buffer de replay baseado em vantagem que hard-mines trajetórias amostradas historicamente com as maiores vantagens, permitindo que o modelo aprenda com rollouts passados em vez de gerar continuamente novos e consultar o modelo de recompensa. Ambas as modificações são integradas no otimizador de política; o agendamento de inferência permanece inalterado.

Os autores relatam que seu método alcança a paridade com - e depois supera - as linhas de base DDPO e variantes de GRPO usando um sexto das amostras sob hiperparâmetros idênticos, resultando em uma redução aproximada de 83% nas avaliações do modelo de recompensa. A generalização para prompts não vistos é preservada, indicando que o ganho de eficiência não é à custa do overfitting na distribuição de treinamento. Para equipes de produção, o apelo está do lado da economia de treinamento: nenhuma cirurgia no modelo, nenhuma latência adicional de inferência e sem ajuste de hiperparâmetros é necessário.

O artigo é um pré-impressão sem código vinculado publicamente, e a avaliação não relata o tempo de execução real, horas de GPU ou custos em dólares - apenas contagens de amostras em relação a linhas de base acadêmicas. A heurística do latente quadrado é uma aproximação da verdadeira quantidade de importância do timestep, e o artigo deixa a lacuna de aproximação não quantificada. O buffer de replay introduz overhead de memória e efeitos potenciais de gradientes obsoletos que não são caracterizados. Essas lacunas são importantes porque a literatura mais ampla de RLHF de difusão documenta patologias persistentes - incluindo trajetórias instáveis, hacking de recompensa e tempos de inferência longos - que a ponderação do timestep por si só não resolve.