Resumo Técnico: Melhorando o Fluxo Retificado com Condições de Contorno

Declaração do Problema

O Fluxo Retificado (RF), também conhecido como Correspondência de Fluxo (Flow Matching), emergiu como um poderoso framework de modelagem generativa que aprende um campo de velocidade para transformar uma distribuição de ruído em uma distribuição de dados alvo via uma Equação Diferencial Ordinária (ODE). Apesar de seu sucesso em tarefas como geração de texto-para-imagem e vídeo, os autores identificam uma limitação crítica nas implementações padrão de RF: o campo de velocidade aprendido frequentemente falha em satisfazer as condições de contorno teoricamente derivadas.

Especificamente, sob interpolação linear entre o ruído ($t=0$) e os dados ($t=1$), o campo de velocidade ideal deve satisfazer:

1. Contorno à Esquerda ($t=0$): $\mathbf{v}(\mathbf{x}, 0) = \mathbb{E}[\mathbf{X}_1] - \mathbf{x}$
2. Contorno à Direita ($t=1$): $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$

A parametrização de redes neurais padrão para o campo de velocidade não impõe inerentemente essas restrições. Como visualizado no artigo, a velocidade prevista em $t=1$ frequentemente desvia-se do mapa identidade esperado. Essa violação é particularmente prejudicial durante a amostragem estocástica (por exemplo, usando dinâmica de Langevin). Como a estimativa da função de score (via fórmula de Tweedie) envolve uma divisão por $(1-t)$, uma velocidade imprecisa próximo a $t=1$ faz com que a estimativa do score divirja. Isso leva a uma amostragem instável, resultando em imagens excessivamente suavizadas ou artefatos "cartunescos", especialmente ao usar tamanhos de passo ou escalas de ruído maiores.

Metodologia

Para abordar isso, os autores propõem o Modelo de Fluxo Retificado com Contorno Forçado (Boundary RF Model). Esta abordagem incorpora explicitamente as restrições de contorno no campo de velocidade através de duas variantes distintas:

1. Modelo de RF com Contorno Baseado em Máscara (Mask-based)

Esta variante impõe ambas as condições de contorno ($t=0$ e $t=1$) estruturando o campo de velocidade como uma combinação linear de termos de contorno e a saída de uma rede neural treinável:
$$\mathbf{v}(\mathbf{x}, t) = g(t) \cdot (C - \mathbf{x}) + f(t) \cdot \mathbf{x} + h(t) \cdot m_\theta(\mathbf{x}, t)$$
Onde:

• $m_\theta(\mathbf{x}, t)$ é uma rede neural padrão (ex: U-Net ou DiT).
• $C = \mathbb{E}[\mathbf{X}_1]$.
• $f(t), g(t), h(t)$ são funções escalares projetadas para desaparecer ou ativar em contornos específicos (ex: $g(1)=f(0)=h(0)=h(1)=0$ e $g(0)=f(1)=1$).
• Os autores sugerem uma escolha "Standard Cosine": $g(t) = \cos(\frac{\pi}{2}t)$, $f(t) = \sin(\frac{\pi}{2}t)$, e $h(t) = \sin(\pi t)$.

Esta construção garante que as condições de contorno sejam satisfeitas por design, independentemente da saída da rede neural.

2. Modelo de RF com Contorno Baseado em Subtração (Subtraction-based)

Motivado pela observação de que impor apenas a condição de $t=1$ gera benefícios empíricos significativos, os autores propõem uma variante mais simples de contorno único:
$$\mathbf{v}(\mathbf{x}, t) = \mathbf{x} + m_\theta(\mathbf{x}, t) - m_\theta(\mathbf{x}, 1)$$
Esta formulação satisfaz inerentemente $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$ porque o termo $m_\theta(\mathbf{x}, 1)$ se cancela em $t=1$. Embora esta abordagem ofereça simplicidade de design e flexibilidade, ela incorre em um custo computacional, exigindo uma passagem direta adicional para calcular $m_\theta(\mathbf{x}, 1)$ para cada avaliação.

Impacto na Amostragem Estocástica

Ao impor $\mathbf{v}(\mathbf{x}, 1) = \mathbf{x}$, o Modelo de RF com Contorno garante que o numerador na fórmula de Tweedie ($t\mathbf{v}(\mathbf{x}, t) - \mathbf{x}$) se aproxime de zero conforme $t \to 1$. Isso evita que a função de score $\nabla \log \rho_t$ divirja, estabilizando os amostradores estocásticos (amostradores SDE) e permitindo uma geração robusta mesmo com escalas de ruído ou passos maiores.

Principais Contribuições

1. Identificação de uma Limitação: O artigo destaca que campos de velocidade não restritos no Fluxo Retificado violam as condições de contorno teóricas, levando a trajetórias de ODE imprecisas e amostragem estocástica instável.
2. Parametrização Inovadora: A introdução do Modelo de RF com Contorno, que oferece duas variantes práticas (Baseada em Máscara e Baseada em Subtração) para impor essas condições com modificação mínima de código.
3. Estabilização de Funções de Score: Demonstração de que o reforço do contorno estabiliza diretamente a função de score próximo a $t=1$, mitigando os problemas de divergência que assolam a amostragem estocástica no RF vanilla.
4. Validação Empírica: Experimentos abrangentes mostrando melhorias consistentes na qualidade de geração e estabilidade de amostragem em diferentes conjuntos de dados e escalas de modelos.

Resultados Experimentais

Os autores avaliaram seus modelos no CIFAR-10 e ImageNet (256$\times$256 e 512$\times$512) utilizando arquiteturas de Diffusion Transformer (DiT).

• Melhorias Quantitativas:

  • No ImageNet 256$\times$256, o Modelo de RF com Contorno Baseado em Subtração alcançou um FID de 6.32 (amostragem ODE) e 6.38 (amostragem SDE), comparado ao baseline de RF vanilla de 6.87 e 6.79, respectivamente.
  • O modelo Baseado em Máscara alcançou um FID de 6.63 (ODE) e 6.18 (SDE).
  • Melhorias também foram observadas no Inception Score (IS), Precisão e Recall em ambos os datasets.
  • A variante Baseada em Subtração alcançou consistentemente os maiores scores de IS.

• Robustez de Amostragem:

  • Em experimentos usando o Amostrador de Overshooting com variando escalas de ruído ($c$), os modelos de RF com Contorno demonstraram desempenho significativamente mais estável e menos degradação na qualidade da amostra em comparação ao RF vanilla conforme a estocasticidade aumentava.
  • Resultados qualitativos mostraram que os modelos de RF com Contorno produzem imagens mais nítidas com menos artefatos e melhor fidelidade de objetos, particularmente em configurações de amostragem estocástica onde o RF vanilla tende a suavizar excessivamente.

• Escalabilidade:

  • Os benefícios do Modelo de RF com Contorno foram consistentes ao escalar para modelos maiores (DiT-L, DiT-XL) e resoluções mais altas (512$\times$512).

Significância e Alegações

O artigo afirma que o Modelo de RF com Contorno oferece uma solução simples e prontamente implementável para uma limitação fundamental do Fluxo Retificado. Ao impor explicitamente as condições de contorno, o método:

1. Melhora a Qualidade de Geração: Supera consistentemente o RF vanilla em FID e outras métricas.
2. Permite Amostragem Estocástica Estável: Resolve a divergência da função de score próximo ao tempo terminal, permitendo estratégias de amostragem mais robustas.
3. É Fácil de Adaptar: Pode ser integrado em modelos de RF existentes (como os que usam U-Net ou DiT) com mudanças arquiteturais mínimas.

Os autores posicionam este trabalho como um refinamento necessário para o uso prático do Fluxo Retificado, particularmente em cenários que exigem amostragem estocástica de alta qualidade. Eles sugerem que, embora o foco seja o RF, o conceito central de impor condições de contorno poderia potencialmente ser estendido para modelos de difusão com adaptações menores. Trabalhos futuros são propostos para explorar o reforço de contorno adaptativo e a aplicação para geração de texto-para-vídeo.