Resumo Técnico: DiffRGD

Problema

Os modelos de difusão tornaram-se fundamentais para a geração de imagens, mas controlá-los para tarefas específicas (ex: restauração de imagem, geração condicional) sem o custo de re-treinamento ou ajuste fino (fine-tuning) continua sendo um desafio. Os métodos de orientação existentes durante a inferência (ex: DPS, FreeDoM) direcionam o processo de amostragem injetando gradientes de um objetivo diferenciável diretamente no espaço latente. No entanto, esses métodos operam no espaço ambiente sem respeitar a estrutura Gaussiana subjacente do processo de difusão.

Essa falta de consciência de distribuição leva ao desvio de distribuição (distributional drift), onde as variáveis latentes se desviam do manifold Gaussiano original induzido pelo modelo. À medida que os passos de amostragem se acumulam, esse fenômeno "fora do manifold" degrada a qualidade da amostra limpa final. Tentativas anteriores de restringir as atualizações, como o MPGD (assumindo linearidade local) ou o DSG (projetando em uma hiperesfera de raio fixo), ou dependem de encoders perfeitos e inviáveis ou falham em aderir totalmente às propriedades Gaussianas específicas da distribuição latente, por vezes degradando para uma distribuição uniforme.

Metodologia: DiffRGD

Os autores propõem o DiffRGD, um framework de orientação durante a inferência do tipo plug-and-play que preserva explicitamente a estrutura Gaussiana latente. A ideia central é que, em cada passo de amostragem DDIM $t$, a distribuição latente é uma Gaussiana isotrópica centrada em uma média prevista $\mu_t$.

1. Decomposição Polar e Construção de Manifold

O método utiliza uma propriedade de decomposição polar das distribuições Gaussianas isotrópicas. Em qualquer passo $t$, uma amostra latente $x_t$ pode ser decomposta em um componente radial estatisticamente independente $r$ e um componente direcional $u$:
$$x_t = \mu_t + \sigma_t r u$$
onde $r$ segue uma distribuição chi escalonada e $u$ é uniformemente distribuído em uma esfera unitária.

Em vez de otimizar no espaço ambiente ou em uma esfera de raio fixo, o DiffRGD constrói um manifold esférico $\mathcal{S}_{t,r}$ induzido pela distribuição latente específica:
$$\mathcal{S}_{t,r} := \{ x \in \mathbb{R}^n \mid \|x - \mu_t\|_2 = \sigma_t r \}$$
Aqui, o raio $\sigma_t r$ é amostrado da escala natural da distribuição, garantindo que o conjunto viável respeite as propriedades Gaussianas do processo de difusão.

2. Gradiente Descendente de Riemann (RGD)

O problema de orientação em cada passo de tempo é formulado como um problema de otimização restrita:
$$x_{t-1} = \arg\min_{x \in \mathcal{S}_{t,r}} \mathcal{L}(\hat{x}_0(x, t-1), y)$$
onde $\mathcal{L}$ é a perda de orientação (ex: perda de reconstrução para restauração ou alinhamento de características para geração condicional).

Para resolver isso, o DiffRGD emprega o Gradiente Descendente de Riemann (RGD) no manifold esférico $\mathcal{S}_{t,r}$:

1. Projeção no Espaço Tangente: O gradiente euclidiano da perda é projetado no espaço tangente da esfera no ponto atual, removendo o componente radial para garantir que a direção da atualização seja tangente ao manifold.
2. Retração: Após realizar um passo ao longo do espaço tangente, uma operação de retração mapeia o ponto atualizado de volta para o manifold esférico $\mathcal{S}_{t,r}$, preservando o raio $\sigma_t r$.

Este processo é repetido por $K$ iterações internas (tipicamente $K=3$) por passo de difusão. Os autores provam que, sob suposições de suavidade $L$, esta otimização converge para um ponto estacionário.

Principais Contribuições

1. Geometria Consciente da Distribuição: Os autores identificam uma propriedade de decomposição polar que revela uma geometria induzida pela distribuição para a orientação de difusão. Eles propõem um framework de Riemann que respeita a geometria do manifold induzido pela distribuição, mitigando a degradação da amostra causada pelo desvio fora do manifold.
2. Análise Teórica: O artigo fornece uma análise de preservação de distribuição e prova a convergência do algoritmo de otimização de Riemann para um ponto estacionário. Na prática, a otimização geometricamente consciente proposta converge mais rápido do que métodos baseados em ADMM.
3. Desempenho Plug-and-Play: Demonstra-se que o DiffRGD é um método plug-and-play que não requer re-treinamento ou ajuste fino, alcançando resultados de estado da arte em diversas tarefas.

Resultados Experimentais

Os autores avaliaram o DiffRGD contra baselines (DPS, FreeDoM, MPGD, DSG, ADMMDiff) em uma única GPU NVIDIA GeForce RTX 4090 usando o framework HuggingFace Diffusers.

Restauração de Imagem

Avaliado em FFHQ 256×256 e ImageNet 256×256 para quatro tarefas: Inpainting, Super-Resolução (4×), Desfoque Gaussiano (Gaussian Deblurring) e Desfoque de Movimento (Motion Deblurring).

• Quantitativo: O DiffRGD alcançou melhorias estatisticamente significativas em PSNR, SSIM, LPIPS e FID sobre todos os baselines na maioria das tarefas. Por exemplo, em Inpainting, o DiffRGD alcançou um PSNR de 34.04 comparado a 32.38 do segundo melhor método (ADMMDiff).
• Robustez: O DiffRGD manteve um desempenho forte mesmo com menos passos de amostragem (100 passos), enquanto métodos como DPS e ADMMDiff mostraram degradação perceptível.
• Qualitativo: Os resultados visuais mostraram menos artefatos e melhor consistência de fundo em comparação ao DPS e MPGD, com qualidade comparável ou superior ao ADMMDiff, porém com menor tempo computacional.

Geração Condicional

Avaliado em CelebA-HQ 256×256 para orientação de Mapa de Segmentação, Esboço (Sketch) e FaceID.

• O DiffRGD superou consistentemente os métodos de inferência existentes (FreeDoM, DSG, ADMMDiff) em termos de métricas de alinhamento (mIoU, Sketch-$\ell_2$, FaceID-$\ell_2$) e qualidade de geração (FID, KID).
• Resultados qualitativos indicaram que o DiffRGD produziu resultados mais naturais e visualmente atraentes com melhor alinhamento às condições em comparação ao DSG, que frequentemente sofria de controle instável.

Força de Orientação (Guidance Strength)

Experimentos sobre a força de orientação ($\eta_t$) mostraram que o DiffRGD alcança um melhor equilíbrio entre alinhamento (mIoU) e qualidade de geração (FID) em uma gama mais ampla de intensidades em comparação ao DSG, que teve um desempenho ruim em todas as configurações.

Significância e Alegações

O artigo afirma que o DiffRGD aborda uma lacuna crítica na orientação durante a inferência: a falha dos métodos existentes em preservar a estrutura Gaussiana latente passo a passo. Ao formular a orientação como uma otimização restrita em um manifold esférico induzido pela distribuição, o DiffRGD evita os fenômenos "fora do manifold" que degradam a qualidade da amostra.

Os autores sustentam que a orientação geometricamente consistente é uma alternativa viável e superior à aplicação direta de gradientes no espaço ambiente ou ao uso de manifolds restritivos que não são conscientes da distribuição. O método oferece uma solução plug-and-play simples que melhora as métricas da tarefa enquanto preserva a qualidade de alta fidelidade inerente aos modelos de difusão pré-treinados, sem a necessidade de re-treinamento dispendioso.