Timestep Rescheduling in Diffusion Inversion

Resumo Técnico: Reescalonamento de Timesteps em Inversão de Difusão

Definição do Problema
A inversão de difusão é o processo de mapear uma imagem real $z_0$ de volta ao espaço latente Gaussiano ($z_T$) de um modelo de difusão, permitindo tarefas como reconstrução e edição de imagens. Embora os Modelos de Difusão Implícitos Determinísticos (DDIM) forneçam uma formulação determinística e eficiente para isso, o processo de inversão é inerentemente aproximado. A abordagem padrão envolve reverter analitamente a etapa de denoising, o que introduz desvios que se acumulam ao longo do tempo, levando a erros de reconstrução perceptíveis e fidelidade de edição degradada.

Métodos existentes tentam mitigar esses erros principalmente resolvendo problemas de ponto fixo para refinar a equação de inversão. No entanto, essas abordagens negligenciam amplamente um fator crítico: a influência do escalonamento de timesteps (a seleção e o espaçamento dos timesteps de difusão) na fidelidade da inversão. A maioria das pipelines utiliza timesteps uniformemente espaçados, o que os autores argumentam ser suboptimal. O artigo postula que a distribuição dos timesteps impacta significativamente o acúmulo de erros globais de inversão, embora essa relação permaneça subexplorada em comparação com seu papel na amostragem ou treinamento de difusão.

Metodologia: TRDI (Timestep Rescheduling in Diffusion Inversion)
Os autores propõem o TRDI, um escalonador de timesteps não uniforme projetado para minimizar o erro de inversão global sem exigir parâmetros adicionais ou sobrecarga computacional durante a fase de inferência. A metodologia é fundamentada em uma análise teórica do erro de inversão e implementada através de uma estratégia de duas etapas "coarse-to-fine" (do grosseiro ao refinado).

1. Insight Teórico:
   Os autores derivam que o erro de inversão induzido por um salto de timestep $\Delta t$ pode ser reformulado como um problema de ponto fixo escalonado. A magnitude do erro $\delta$ é proporcional a um coeficiente de escala $c_\alpha(t, \Delta t)$ multiplicado pela discrepância entre previsões consecutivas do modelo ($\Delta \epsilon_\theta$).

   • Observação Chave: O coeficiente de escala exibe uma tendência parabólica em relação ao índice do timestep. Os erros concentram-se tanto em timesteps pequenos (etapas iniciais de difusão) quanto em timesteps grandes (etapas tardias), com um declínio no meio. Além disso, tamanhos de passo maiores ($\Delta t$) induzem inerentemente erros maiores.

2. Reescalonamento Global (Ajuste Grosseiro):
   Para abordar a distribuição parabólica do erro, os autores introduzem um mecanismo de reescalonamento global. Eles modificam a sequência original de timesteps uniformes $t_k$ usando uma transformação de lei de potência controlada por um hiperparametro $\gamma$:
   $$t_k = t_1 + (t_K - t_1) \left( \frac{k-1}{K-1} \right)^\gamma$$

   • Quando $\gamma > 1$, os timesteps são esticados em direção aos passos iniciais (alocando mais resolução para a região de alto erro).
   • Quando $\gamma < 1$, os timesteps tornam-se mais densos em direção ao final.
     Esta etapa prepara o escalonamento para a otimização de grão fino ao ajustar globalmente a alocação de passos.

3. Reescalonamento Local via Programação Dinâmica (Ajuste Fino):
   Para minimizar ainda mais o erro acumulado, os autores empregam uma abordagem de programação dinâmica (DP) de janela deslizante.

   • Mapa de Custo: Um mapa de custo $E[k, t]$ é inicializado para rastrear o erro acumulado mínimo para atingir o índice de passo $t$ na $k$-ésima etapa de inversão.
   • Recursão: Dentro de uma janela de largura $2d+1$ ao redor do passo atual, o algoritmo busca recursivamente pela transição localmente ótima entre timesteps adjacentes que minimize a soma do erro acumulado anterior e o coeficiente de erro do passo atual.
   • Backtracking: A sequência ótima de timesteps é recuperada através de backtracking do passo final para o primeiro.

Todo o processo de reescalonamento é realizado offline (ou como um passo de pré-processamento) com base no cronograma de ruído conhecido, o que significa que adiciona uma sobrecarga de tempo de execução negligenciável ao processo real de inversão ou edição de imagem.

Principais Contribuições

• Formulação Teórica: O artigo deriva uma conexão teórica entre saltos de timestep e erro de inversão, enquadrando o erro como um problema de ponto fixo escalonado dependente do cronograma de ruído.
• Escalonador Adaptativo: Propõe uma estrutura de reescalonamento coarse-to-fine (TRDI) que aloca adaptativamente passos maiores para regiões de baixo erro e passos mais finos para regiões de alto erro, equilibrando efetivamente eficiência e precisão.
• Melhoria Plug-and-Play: O método é projetado como um aprimoramento "pronto para uso" que não requer parâmetros extras, treinamento ou sobrecarga computacional significativa, sendo compatível com pipelines de inversão existentes.

Resultados Experimentais
Os autores validam o TRDI através de experimentos extensos em reconstrução de imagem (dataset MSCOCO) e edição de imagem (PIE-Bench) usando vários modelos de difusão (Stable Diffusion v1.5, SDXL e SDXL Turbo) e métodos de inversão de linha de base (DDIM, ReNoise, NPI, GNRI).

• Reconstrução: No SD v1.5, a integração do TRDI melhorou consistentemente as métricas em todas as linhas de base. Por exemplo, com a linha de base DDIM, o PSNR aumentou de 20.07 para 20.21, e o LPIPS diminuiu (melhorando a qualidade perceptual) de 193.97 para 187.85. Resultados qualitativos mostraram correções em detalhes estruturais como janelas de trens e formas de frutas que as linhas de base não conseguiram reconstruir com precisão.
• Edição: No PIE-Bench, o TRDI melhorou significativamente a preservação de estrutura e a preservação de fundo, mantendo alta similaridade CLIP. Notavelmente, para modelos de poucos passos como o SDXL Turbo (4 passos), onde a degradação do fundo é tipicamente severa, o TRDI alcançou melhorias substanciais na preservação do fundo (por exemplo, reduzindo a Distância Estrutural em 17.43% para o DDIM).
• Estudos de Ablação: O estudo confirmou que um fator de escala $\gamma \approx 1.05$ e uma largura de janela $d=8$ forneceram o melhor equilíbrio entre fidelidade estrutural e editabilidade semântica.

Significância e Alegações
O artigo afirma que o escalonamento de timesteps é um fator crítico, porém anteriormente negligenciado, na fidelidade da inversão de difusão. Ao revelar a relação parabólica entre o tamanho/índice do timestep e o acúmulo de erro, os autores demonstram que uma alocação estratégica e não uniforme de passos computacionais pode reduzir significativamente os erros globais de inversão.

Os autores posicionam o TRDI não como um substituto para formulações de inversão mais fortes (como iterações de ponto fixo ou inversão estocástica), mas como um complemento leve de nível de escalonador (retrofit). Ele permite que pipelines de inversão determinísticos existentes alcancem resultados superiores de reconstrução e edição sem aumentar os custos computacionais ou exigir o retreinamento do modelo. O trabalho enfatiza que otimizar quando dar os passos é tão importante quanto como dá-los.