ART for Diffusion Sampling: A Reinforcement Learning Approach to Timestep Schedule

Resumo Técnico: ART para Amostragem por Difusão

Declaração do Problema
Modelos de difusão baseados em pontuação geram amostras integrando uma dinâmica reversa no tempo aprendida. Uma escolha numérica crítica nesse processo é a discretização do horizonte temporal em uma grade de passos de tempo. Abordagens padrão tipicamente empregam grades uniformes ou cronogramas artesanais (por exemplo, EDM, DPM-Solver). Os autores argumentam que esses cronogramas fixos são frequentemente suboptimalizados dado um orçamento fixo de avaliações de função (NFE). Grades uniformes alocam esforço computacional uniformemente, mesmo quando a dinâmica reversa varia significativamente ao longo da trajetória, enquanto grades artesanais podem não se adaptar a distribuições de dados específicas ou complexidades da função de pontuação. O objetivo é redistribuir sistematicamente o esforço computacional ao longo da trajetória de amostragem para minimizar o erro agregado de discretização de Euler.

Metodologia
O artigo introduz o Tempo Reparametrizado Adaptativo (ART), um framework de teoria de controle para agendamento ótimo de passos de tempo.

1. Formulação de Controle (ART):

   • Em vez de selecionar pontos de grade diretamente, o ART trata a "velocidade do relógio" do processo de difusão como uma variável de controle.
   • Introduz uma variável de tempo reparametrizada $t$ mapeada para o tempo original de difusão $\tau$ via um mapeamento contínuo $\psi(t)$.
   • O controle $\theta(t) = \dot{\psi}(t)$ representa a taxa instantânea de variação do tempo de difusão.
   • O objetivo é escolher $\theta(t)$ para minimizar o erro local agregado de discretização de Euler, sujeito à restrição de que o orçamento total de tempo seja preservado ($\int_0^T \theta(t) dt = T$).
   • O proxy de erro local é derivado como proporcional a $|Q(x, \psi)| \theta(t)^2$, onde $Q$ depende da função de pontuação e suas derivadas.

2. Solução por Aprendizado por Reforço (ART-RL):

   • O problema determinístico do ART é intratável para solução direta devido ao espaço de estados de alta dimensão e à falta de uma solução de forma fechada para a equação de Hamilton-Jacobi-Bellman (HJB).
   • Os autores propõem o ART-RL, uma abordagem de aprendizado por reforço em tempo contínuo (CTRL).
   • Introduzem um problema auxiliar randomizado onde o controle $\theta$ é amostrado de uma política Gaussiana $\pi(\cdot | t, x, \psi) = \mathcal{N}(\mu, \lambda/|Q|)$.
   • Esta formulação permite o uso de algoritmos ator-crítico (adaptando especificamente o framework de Jia e Zhou [14]) para aprender a média da política $\mu$ e a função de valor via aproximação estocástica.

3. Ponte Teórica:

   • Uma contribuição teórica central é o estabelecimento de uma ponte bidirecional entre o ótimo determinístico do ART e a política Gaussiana ótima no problema auxiliar do ART-RL.
   • Suficiência: O controle ótimo determinístico $\mu^*$ eleva-se a uma política Gaussiana ótima cuja média é exatamente $\mu^*$.
   • Necessidade: Qualquer política Gaussiana ótima para o problema auxiliar deve ter uma média que recupere o controle determinístico do ART $\mu^*$.
   • Este resultado justifica parametrizar o ator como uma média Gaussiana e implantá-lo deterministicamente após o treinamento, fornecendo garantias rigorosas de que o cronograma aprendido é o ótimo verdadeiro para o problema de controle original.

4. Destilação:

   • Após o treinamento offline, a política estocástica é destilada em uma grade determinística, apenas temporal. Isso remove a necessidade de avaliações de rede de ator por passo ou cálculos de variância durante a inferência, garantindo custo computacional zero adicional comparado a cronogramas padrão.

Principais Contribuições

• Formulação: Propõe o ART, o primeiro framework de teoria de controle para agendamento de passos de tempo em difusão via reparametrização adaptativa de tempo.
• Teoria: Estabelece a equivalência bidirecional entre um problema de controle ótimo determinístico e um procedimento de aprendizado ator-crítico em tempo contínuo com políticas Gaussianas. Isso transforma aplicações heurísticas de RL em uma rota principial para encontrar ótimos determinísticos.
• Algoritmo: Desenvolve o ART-RL, um algoritmo ator-crítico que aprende cronogramas de tempo adaptativos de forma orientada por dados.
• Generalização: Demonstra que o cronograma aprendido pode ser destilado e transferido entre diferentes conjuntos de dados e contagens de passos de tempo sem retreinamento.

Resultados Experimentais
Os autores avaliam o ART-RL dentro do pipeline oficial do EDM em CIFAR-10, AFHQv2, FFHQ e ImageNet.

• CIFAR-10: O ART-RL supera consistentemente as grades Uniforme, log-SNR do DPM-Solver e cronogramas padrão do EDM em uma ampla gama de orçamentos de NFE.
  • Em orçamentos baixos (por exemplo, NFE=3), o ART-RL produz objetos reconhecíveis onde as linhas de base permanecem próximas ao ruído.
  • Em NFE=35, o ART-RL alcança um FID de 1,82, melhorando sobre o EDM (1,85) e o DPM-logSNR (1,89).
  • A significância estatística é confirmada via múltiplas sementes, com a pior execução do ART-RL superando a melhor do EDM.
• Generalização:
  • Interpolação/Extrapolação: Um cronograma aprendido em $K=18$ passos no CIFAR-10 mantém sua vantagem de desempenho quando interpolado ou extrapolado para outras contagens de passos ($K \in \{4, 9, 12, 15, 20\}$).
  • Transferência entre Conjuntos de Dados: O cronograma destilado do CIFAR-10 é aplicado ao AFHQv2, FFHQ e ImageNet sem retreinamento. Melhora consistentemente o FID sobre o EDM e o DPM-logSNR, particularmente em NFEs baixos e moderados.
• Análise Unidimensional: Em um cenário com uma função de pontuação analítica, o ART-RL reduz significativamente a distância de Wasserstein ($W_2$) comparado a cronogramas artesanais, confirmando que a melhoria decorre de uma melhor alocação de tempo e não de erros de estimação de pontuação.

Significância e Afirmações
O artigo afirma ser o primeiro a formular o agendamento de passos de tempo em difusão como um problema de controle via reparametrização adaptativa de tempo e o primeiro a fornecer uma ponte rigorosa bidirecional entre um ótimo determinístico e um procedimento de aprendizado ator-crítico.

Os autores enfatizam que seu método é orientado por dados e amortizado:

• Orientado por dados: O cronograma é aprendido a partir da função objetivo e da dinâmica dos dados, não sendo artesanal.
• Transferência Amortizada: O custo de treinamento offline (aproximadamente 1–2 horas em um Colab T4 para CIFAR-10) é uma despesa única. A grade determinística destilada resultante é reutilizável entre diferentes conjuntos de dados e contagens de passos com custo de inferência extra zero, igualando o tempo de execução de cronogramas padrão como Uniforme ou EDM.

O trabalho posiciona a reparametrização adaptativa de tempo como uma direção nascente, mas promissora, para o design sistemático de cronogramas em modelos generativos de difusão, oferecendo um ponto de partida de teoria de controle para pesquisas futuras.