One-Sided Quantile Coupling for Flow Matching

Resumo Técnico: Acoplamento de Quantis Unilateral para Flow Matching

Enunciado do Problema
O Flow Matching (FM) treina modelos generativos de tempo contínuo através da regressão do campo de velocidade de um caminho de probabilidade que conecta uma distribuição de origem simples (tipicamente uma Gaussiana padrão) a uma distribuição de dados alvo. A qualidade do acoplamento — o pareamento de amostras de origem e alvo — influencia criticamente a geometria das trajetórias de treinamento. O FM padrão utiliza um pareamento aleatório independente, o que frequentemente resulta em caminhos longos e conflitantes, aumentando a variância do campo de velocidade alvo e dificultando a eficiência do treinamento.

Embora acoplamentos estruturados inspirados em Transporte Ótimo (OT) possam retificar esses caminhos e melhorar a qualidade da geração, os métodos existentes (ex: OT-CFM) dependem da resolução de um problema de atribuição por mini-lote, o que exige a construção de matrizes de custo par a par e a resolução de problemas de atribuição (ex: via algoritmo de Hungarian ou iterações de Sinkhorn), levando a custos computacionais que escalam de forma pelo menos quadrática (ou cúbica para soluções exatas) com o tamanho do lote. Isso limita a escalabilidade e pode introduzir instabilidade. O artigo aborda a questão de se os benefícios geométricos do acoplamento estruturado podem ser alcançados sem o pareamento de dois lotes pré-amostrados, evitando assim o alto custo do transporte por lote.

Metodologia: Quantile Coupling Flow Matching (QC-FM)
Os autores propõem o Quantile Coupling Flow Matching (QC-FM), um esquema de acoplamento unilateral leve. Diferente dos métodos tradicionais que amostram tanto o lote de ruído quanto o de dados e depois os reapareiam, o QC-FM amostra apenas o lote de dados e constrói diretamente as amostras de origem pareadas.

1. Construção Unilateral:

   • Dado um lote de dados $\mathcal{B}_x$, o método projeta os pontos de dados em $k$ direções ortogonais aleatórias (fatias) definidas por uma matriz $U$.
   • Para cada direção, os pontos de dados são classificados dentro do lote. Esses postos empíricos são mapeados para quantis Gaussianos usando a CDF inversa da normal padrão.
   • O código latente de origem $e^{(i)}$ é construído para satisfazer essas restrições de quantil projetadas. Especificamente, a projeção da origem sobre as direções das fatias corresponde aos quantis Gaussianos derivados dos postos dos dados.
   • As dimensões restantes (o complemento ortogonal do subespaço da fatia) são preenchidas amostrando de uma distribuição Gaussiana condicional. Isso garante que a origem permaneça estocástica em direções não restringidas.

2. Extensões Híbridas de Gaussiana-Remanescente:
   Como o acoplamento QC completo (onde todos os pontos de origem são construídos via postos) desvia da priori Gaussiana i.i.d. padrão para $k \ge 2$ devido às dependências entre os códigos das fatias, os autores introduzem duas estratégias híbridas para manter os benefícios do QC enquanto preservam a priori:

   • QC-FM-Mixture: Aplica a construção QC a um subconjunto "âncora" do lote (ex: 20%) e pareia o restante das amostras de origem "remanescentes" com as amostras de dados restantes usando um pareamento aleatório independente (a abordagem Baseline). Isso cria uma mistura de acoplamentos estruturados e não estruturados.
   • QC-FM-Adjacency: Aplica o QC ao subconjunto âncora, mas pareia as amostras de origem remanescentes com base na geometria local induzida pelas âncoras. Os pontos de dados remanescentes são atribuídos à sua âncora mais próxima (no subespaço estruturado), e as amostras de origem remanescentes são pareadas dentro desses grupos. Isso propaga o viés estruturado sem duplicar alvos.

Principais Contribuições

• Algoritmo QC-FM: Um novo acoplamento estruturado local ao lote que impõe transporte monotônico unidimensional ao longo de fatias aleatórias. Ele elimina a necessidade de matrizes de custo par a par e solucionadores de atribuição, reduzindo o custo computacional dominante para $O(kB \log B)$ (ordenação) em vez de $O(B^2)$ ou $O(B^3)$.
• Análise Teórica:
  • Redução de Variância: Os autores provam que, para um quadro fixo, o acoplamento QC elimina a variância de regressão irredutível ao longo de cada fatia selecionada, tornando o fluxo ideal exatamente reto nessas direções.
  • Preservação da Priori: Eles quantificam que para $k=1$, a marginal da origem é exatamente Gaussiana. Para $k \ge 2$, o desvio da priori Gaussiana é confinado ao cópula dos códigos das fatias, e o custo de transporte deste desvio é limitado.
• Esquemas de Treinamento Híbridos: A introdução das variantes Mixture e Adjacency que equilibram a injeção de viés geométrico com a preservação do sinal Gaussiano explícito da Baseline.

Resultos Experimentais
O método foi avaliado nos datasets CIFAR-10, CelebA-64, FFHQ-64 e ImageNet-64 sob orçamentos de treinamento e arquiteturas equivalentes.

• Desempenho: O QC-FM-Mixture superou consistentemente a Baseline de acoplamento independente e o OT-CFM de Hungarian exato em todos os quatro datasets. As melhorias sobre a Baseline variaram de 6,3% (CIFAR-10) a 12,9% (FFHQ) em termos de FID-50K.
• Eficiência: Em termos de custo computacional, o QC-FM é significativamente mais rápido que os métodos baseados em OT. Com um tamanho de lote de 2048, o QC-FM-Mixture foi mais de 800 vezes mais rápido que o solver Hungarian OT e 35 vezes mais rápido que o solver Sinkhorn.
• Diagnósticos de Brinquedo (Toy): Em datasets sintéticos (Checkerboard e 8-Gaussians), as variantes do QC-FM produziram trajetórias aprendidas quase retas (Razão de Comprimento de Caminho $\approx 1$) e variância da velocidade alvo significativamente menor em comparação com a Baseline, confirmando os benefícios geométricos do acoplamento.

Significância e Alegações
O artigo afirma que preservar a estrutura de postos projetados é uma maneira simples, escalável e eficaz de injetar um viés geométrico útil em acoplamentos de Flow Matching sem resolver um problema de transporte por lote.

• Escalabilidade: O QC-FM oferece um meio-termo atraente entre a Baseline não estruturada e os métodos de transporte por lote caros. É computacionalmente leve e controlável via o número de projeções ($k$).
• Praticidade: O método não requer componentes aprendidos para inferência; a geração ainda começa a partir da Gaussiana padrão. O acoplamento é construído "on-the-fly" durante o treinamento.
• Limitações: Os autores observam modestamente que o QC-FM é um substituto local ao lote para o acoplamento estruturado, não uma aproximação do Transporte Ótimo global. Os quantis são estimados dentro de mini-lotes, e a teoria controla a variância por fatia em vez de garantir a retidão da trajetória global ou melhorias específicas de FID. A seleção do número de fatias ($k$) e da razão de âncora ($p$) atualmente depende de heurísticas e ajustes específicos para cada dataset.

Em conclusão, o trabalho demonstra que o acoplamento de quantis unilateral pode melhorar a qualidade da amostra e a eficiência do treinamento em Flow Matching ao estruturar o pareamento ruído-dados através da preservação de postos, oferecendo uma alternativa prática aos métodos de transporte ótimo computacionalmente proibitivos.