COT-FM: Emparelhamento de Fluxo com Transporte Ótimo por Agrupamento
COT-FM: Cluster-wise Optimal Transport Flow Matching
March 11, 2026
Autores: Chiensheng Chiang, Kuan-Hsun Tu, Jia-Wei Liao, Cheng-Fu Chou, Tsung-Wei Ke
cs.AI
Resumo
Apresentamos o COT-FM, uma estrutura geral que remodela o caminho de probabilidade no Flow Matching (FM) para alcançar uma geração mais rápida e confiável. Os modelos FM frequentemente produzem trajetórias curvas devido a acoplamentos aleatórios ou por lotes, o que aumenta o erro de discretização e reduz a qualidade da amostra. O COT-FM corrige isso agrupando amostras-alvo e atribuindo a cada cluster uma distribuição de origem dedicada, obtida pela reversão de modelos FM pré-treinados. Esta estratégia de dividir e conquistar resulta em um transporte local mais preciso e campos vetoriais significativamente mais retos, tudo sem alterar a arquitetura do modelo. Como uma abordagem plug-and-play, o COT-FM acelera consistentemente a amostragem e melhora a qualidade da geração em conjuntos de dados 2D, benchmarks de geração de imagens e tarefas de manipulação robótica.
English
We introduce COT-FM, a general framework that reshapes the probability path in Flow Matching (FM) to achieve faster and more reliable generation. FM models often produce curved trajectories due to random or batchwise couplings, which increase discretization error and reduce sample quality. COT-FM fixes this by clustering target samples and assigning each cluster a dedicated source distribution obtained by reversing pretrained FM models. This divide-and-conquer strategy yields more accurate local transport and significantly straighter vector fields, all without changing the model architecture. As a plug-and-play approach, COT-FM consistently accelerates sampling and improves generation quality across 2D datasets, image generation benchmarks, and robotic manipulation tasks.