Free Lunch for Stabilizing Rectified Flow Inversion

Resumo Técnico: Free Lunch for Stabilizing Rectified Flow Inversion

1. Problema e Motivação

Os modelos generativos baseados em Rectified Flow (RF) emergiram como alternativas de alto desempenho aos modelos de difusão tradicionais, oferecendo geração de alta qualidade e suporte a tarefas de inversão (mapear dados observados de volta para representações latentes) sem necessidade de treinamento adicional. A inversão é crucial para tarefas downstream como reconstrução de imagens e edição.

No entanto, os métodos de inversão existentes (como a inversão baseada em RF "vanilla") sofrem de erros de aproximação que se acumulam ao longo dos passos de tempo. Devido à instabilidade inerente das equações diferenciais ordinárias (ODEs) em altas dimensões, pequenas perturbações no espaço latente podem levar a grandes erros de reconstrução. Isso resulta em:

• Campos de velocidade instáveis e desviados.
• Degradação da qualidade de reconstrução e fidelidade na edição.
• Acúmulo de erros que afastam a trajetória de inversão da distribuição de prior (ruído gaussiano), levando a regiões de baixa densidade de dados.

O objetivo do trabalho é estabilizar o processo de inversão e edição corrigindo esses erros de velocidade sem exigir treinamento adicional ou aumentar significativamente o custo computacional (avaliações de função neural - NFEs).

2. Metodologia

Os autores propõem duas técnicas principais que atuam como uma "refeição grátis" (free lunch), pois são plug-and-play e não requerem treinamento:

A. Proximal-Mean Inversion (PMI) - Para Inversão e Reconstrução
O PMI visa corrigir o campo de velocidade perturbado guiando-o em direção a uma média mais estável.

• Mecanismo: Em cada passo de tempo, o método calcula uma média móvel das velocidades previstas desde o início até o passo atual.
• Correção Proximal: Aplica uma atualização proximal que guia a previsão de velocidade atual ($v_{t_k}$) em direção a essa média global ($\bar{v}_{t_k}$).
• Restrição Teórica: Para evitar que a trajetória corrigida se desvie para regiões de baixa densidade da distribuição de dados, cada passo de correção é restrito dentro de uma esfera gaussiana. O raio dessa esfera é derivado teoricamente com base em uma análise de estabilidade, garantindo que os pontos de amostragem permaneçam em regiões de alta densidade.
• Eficiência: A correção é realizada via uma expansão de Taylor de primeira ordem, resultando em uma solução de forma fechada que não adiciona NFEs extras.

B. mimic-CFG - Para Tarefas de Edição
Para tarefas de edição, onde é necessário equilibrar a preservação da estrutura original com a capacidade de edição, os autores introduzem o mimic-CFG.

• Inspiração: Baseia-se na estrutura do Classifier-Free Guidance (CFG), mas aplicado à correção de velocidade.
• Mecanismo:
  1. Calcula a média móvel das velocidades passadas.
  2. Projeta a velocidade prevista atual sobre a direção dessa média histórica.
  3. Interpola entre a velocidade projetada (que representa estabilidade/estrutura) e a velocidade original (que representa a direção de edição).
• Equilíbrio: O parâmetro de interpolação ($w$) permite ajustar o trade-off entre a preservação do fundo (estrutura) e a fidelidade da edição.

3. Contribuições Principais

1. PMI (Proximal-Mean Inversion): Um método livre de treinamento que corrige erros de inversão em modelos RF através de correção de gradiente em direção a uma média móvel, estabilizando o campo de velocidade e reduzindo erros acumulados.
2. mimic-CFG: Uma estratégia leve de correção de velocidade para edição que imita a lógica do CFG, interpolando entre a velocidade atual e sua projeção na média histórica, equilibrando consistência estrutural e eficácia de edição.
3. Análise Teórica e Prática: Derivação teórica das condições de estabilidade (raio da esfera gaussiana) e demonstração experimental de que o método reduz a necessidade de passos de amostragem sem sacrificar a qualidade.
4. Desempenho SOTA: Estabelecimento de um novo estado da arte (SOTA) no benchmark PIE-Bench, superando métodos existentes com menos NFEs.

4. Resultados Experimentais

Os métodos foram avaliados no PIE-Bench (700 imagens, 10 categorias de edição) utilizando o modelo Flux.1-dev.

• Reconstrução e Inversão:

  • O PMI melhorou consistentemente métricas como PSNR (Relação Sinal-Ruído de Pico), SSIM (Similaridade Estrutural) e reduziu MSE (Erro Quadrático Médio) e LPIPS (Similaridade de Patch Perceptual) em comparação com solvers base (Euler, Heun, RF-Solver, FireFlow).
  • Melhorias notáveis foram observadas no solver Heun (aumento de 1.22 no PSNR) e RF-Solver (aumento de 1.55 no PSNR).
  • A avaliação "sem prompt" (unconditional) mostrou que o método estabiliza a trajetória mesmo na ausência de condicionamento textual.

• Edição:

  • A integração de PMI e mimic-CFG melhorou a similaridade CLIP (fidelidade semântica da edição) e a preservação do fundo (medida por PSNR e SSIM nas áreas não editadas).
  • Os métodos alcançaram desempenho comparável ou superior a baselines com menos passos de amostragem (NFEs reduzidos). Por exemplo, o FireFlow com a correção alcançou resultados superiores com apenas 18 NFEs, enquanto baselines exigiam mais passos para resultados inferiores.

• Análise de Hiperparâmetros:

  • O peso de interpolação $w$ no mimic-CFG foi otimizado em $0.94$, oferecendo o melhor equilíbrio entre preservação de fundo e qualidade de edição.
  • O uso de norma $L_1$ na regularização proximal mostrou-se ligeiramente superior à norma $L_2$.

5. Significado e Impacto

Este trabalho resolve um problema fundamental na aplicação de modelos Rectified Flow: a instabilidade na inversão que limita a qualidade de reconstrução e edição.

• Eficiência: Ao estabilizar a trajetória de inversão, o método permite obter resultados de alta qualidade com menos passos de inferência, reduzindo o custo computacional.
• Versatilidade: Sendo livre de treinamento e plug-and-play, pode ser integrado a qualquer modelo baseado em RF ou solver ODE existente.
• Fundamentação Teórica: A derivação teórica do raio de correção baseado na análise de instabilidade de ODEs de alta dimensão fornece uma base sólida para evitar regiões de baixa densidade, algo que métodos puramente heurísticos não garantem.

Em resumo, os autores demonstram que é possível estabilizar a inversão de modelos de fluxo retificado e melhorar tarefas de edição com um custo computacional mínimo, estabelecendo novos padrões de desempenho e eficiência na área de geração de imagens.