Resumo Técnico: FullFlow

Formulação do Problema

Modelos modernos de difusão e fluxo retificado de texto para imagem (T2I) possuem priores visuais ricos e capacidades generativas robustas, mas operam estritamente de maneira unidirecional: texto entra, imagem sai. Embora existam modelos unificados de visão e linguagem (VLMs), eles são tipicamente construídos condicionando um modelo de linguagem pré-treinado com um codificador visual ou realizando um pré-treinamento conjunto caro do zero. Essas abordagens frequentemente descartam os fortes priores de imagem já codificados nos backbones de T2I ou exigem retreinamento substancial que é computacionalmente proibitivo. A questão central abordada por este trabalho é se um modelo de fluxo T2I pré-treinado pode ser atualizado para um gerador bidirecional de visão e linguagem sem pré-treinamento multimodal completo ou sem sacrificar o prior de imagem original.

Metodologia

Os autores propõem FullFlow, uma receita eficiente em parâmetros que atualiza um modelo T2I de fluxo retificado pré-treinado (especificamente Stable Diffusion 3 e FLUX.1-dev) para um gerador bidirecional. A inovação central reside em tratar as duas modalidades de forma assimétrica, mas compatível, dentro de um único backbone compartilhado.

1. Extensão Multimodal no Espaço Nativo

• Modalidade de Imagem: As imagens permanecem em seu espaço latente contínuo de fluxo retificado nativo. O modelo preserva o processo de imagem contínuo pré-treinado.
• Modalidade de Texto: O texto é tratado por meio de um processo de inserção discreto construído sobre Fluxos de Edição. Diferentemente da difusão discreta baseada em máscaras (que introduz tokens fora do vocabulário incompatíveis com codificadores congelados), o FullFlow utiliza um processo forward baseado em deleção. Isso garante que, em qualquer nível de ruído, o estado de texto corrompido seja uma subsequência válida de tokens que pode ser decodificada e re-tokenizada para codificadores auxiliares (por exemplo, CLIP).
• Arquitetura: Um único backbone compartilhado (MM-DiT) processa ambas as modalidades. O modelo é condicionado em dois tempos desacoplados: $t$ para o tempo da imagem e $\tau$ para o tempo do texto.
  • Componentes Congelados: Os pesos do backbone e todos os codificadores de texto auxiliares (T5, CLIP) permanecem congelados.
  • Componentes Treináveis: Apenas adaptadores LoRA (aplicados às camadas lineares do transformador), cabeças de texto leves (para contagens de inserção e identidades de tokens) e novos caminhos de condicionamento de tempo são treinados. Isso equivale a treinar apenas ~5% dos parâmetros para SD3 e ~2,5% para FLUX.1.

2. Objetivo de Treinamento com Duplo Tempo

O objetivo de treinamento envolve prever conjuntamente a velocidade da imagem e as inserções de texto ausentes com base no par de tempos $(t, \tau)$.

• Programa de Corrupção: Os autores comparam um programa "alternado-limpo" (treinamento apenas em tarefas condicionadas ao ponto final) com um programa "corrupção-mista" onde $(t, \tau) \sim \text{Unif}([0, 1]^2)$. Eles constatam que uma abordagem híbrida — pré-treinamento com corrupção mista para aprender correspondência cross-modal, seguida de refinamento com alternado-limpo — produz os melhores resultados.
• Estabilização:
  • Balanceamento Adaptativo de Gradiente: Para abordar a incompatibilidade de escala entre a perda de imagem baseada em MSE e a perda de texto baseada em entropia cruzada/contagem, os autores ajustam dinamicamente o peso da perda de texto $\lambda_{txt}$ usando uma média móvel exponencial (EMA) da razão de gradiente.
  • Correspondência com Professor: Para evitar a degradação do prior de imagem pré-treinado durante o ajuste fino em novas distribuições de dados, a regressão de velocidade da imagem é substituída por um objetivo de "correspondência com professor no mesmo ruído". O professor é o mesmo backbone com camadas LoRA desativadas, não exigindo memória adicional.

3. Modos de Inferência

Desacoplar $t$ e $\tau$ transforma a inferência em seleção de trajetória em um espaço generativo 2D:

• Texto $\to$ Imagem: Fixar $\tau=1$ (texto limpo), evoluir $t$ de 0 a 1.
• Imagem $\to$ Texto: Fixar $t=1$ (imagem limpa), evoluir $\tau$ de 0 a 1.
• Amostragem Conjunta: Evoluir simultaneamente $t$ e $\tau$.
• Geração Parcial de Texto (VQA): Fixar a imagem e os tokens da pergunta, restringindo o processo de inserção ao segmento da resposta. Isso permite Resposta a Perguntas Visuais sem um decodificador separado.

Contribuições Principais

1. C1: Converter um modelo pré-treinado de fluxo retificado T2I em um gerador bidirecional treinando apenas adaptadores LoRA e cabeças de texto leves, preservando com sucesso o prior de imagem pré-treinado.
2. C2: Introduzir um framework de duplo tempo onde as imagens permanecem em fluxo contínuo e o texto usa um processo de inserção discreto, permitindo geração condicional, conjunta e parcial de texto como trajetórias em um espaço unificado $(t, \tau)$.
3. C3: Demonstrar desempenho superior a uma linha de base Dual Diffusion-LoRA correspondente em SD3 e FLUX.1-dev tanto na retenção de texto-para-imagem quanto na qualidade de imagem-para-texto, enquanto suporta VQA downstream com recursos computacionais mínimos.

Resultados Experimentais

Os autores avaliam o FullFlow em SD3 e FLUX.1-dev usando contagem de parâmetros correspondente e rank LoRA contra uma linha de base Dual Diffusion reimplantada.

• Eficiência: O FullFlow treina em duas GPUs RTX A5000 em menos de 24 horas, reduzindo o pico de VRAM de ~84 GB (Dual Diffusion) para ~38 GB e aumentando o rendimento em ~8$\times$.
• Texto para Imagem (SD3): No tempo de parede correspondido, o FullFlow melhora o FID de 62,7 para 31,6 (e 29,3 em passos correspondidos) e o CMMD de 0,870 para 0,121, demonstrando forte retenção do prior de imagem.
• Imagem para Texto (SD3): O FullFlow alcança uma pontuação CIDEr de 99,4 (vs. 2,0 para a linha de base) no tempo de parede correspondido, e 13,4 em passos correspondidos.
• Tarefas Downstream: Em benchmarks de VQA (VQAv2, VizWiz, OKVQA), o FullFlow (com apenas 130M de parâmetros treináveis) alcança precisão competitiva, superando notavelmente a linha de base Dual Diffusion no VizWiz (50,7 vs 29,9).
• Transferibilidade: A receita transfere-se com sucesso para FLUX.1-dev, alcançando um CIDEr de 73,6 após 36 horas de treinamento.

Significado e Alegações

O artigo afirma que o FullFlow demonstra que capacidades bidirecionais robustas de visão e linguagem podem ser desbloqueadas a partir de modelos de fluxo T2I pré-treinados sem a necessidade de pré-treinamento multimodal completo. Os resultados sugerem que os priores T2I pré-treinados codificam uma estrutura visual-semântica mais rica do que comumente assumido, capazes de suportar tarefas inversas (imagem-para-texto) e geração conjunta através de adaptação leve.

Os autores posicionam este trabalho como uma "prova de conceito viável em termos de computação" em vez de uma alegação de desempenho máximo ponta a ponta. Eles enfatizam que seu método torna os modelos de fluxo retificado bidirecionais acessíveis à comunidade de pesquisa mais ampla, exigindo apenas GPUs comuns e uma pequena fração dos parâmetros tipicamente necessários para modelos unificados. Eles reconhecem limitações, incluindo o foco atual em resolução de 256$\times$256 e a falta de escalabilidade para conjuntos de dados maiores ou ajuste fino de instruções para diálogos complexos, notando esses pontos como direções futuras.