Flow-OPD: Destilação On-Policy para Modelos de Flow Matching

Resumo

Modelos existentes de Flow Matching (FM) para texto-para-imagem sofrem de dois gargalos críticos sob alinhamento multitarefa: a esparsidade de recompensa induzida por recompensas escalares e a interferência de gradientes decorrente da otimização conjunta de objetivos heterogêneos, que juntos geram um "efeito gangorra" de métricas concorrentes e um hacking de recompensa generalizado. Inspirados pelo sucesso da destilação on-policy (OPD) na comunidade de grandes modelos de linguagem, propomos o Flow-OPD, o primeiro framework de pós-treinamento unificado que integra destilação on-policy em modelos de Flow Matching. O Flow-OPD adota uma estratégia de alinhamento em duas etapas: primeiro, cultiva modelos professores especializados em domínios por meio de ajuste fino GRPO com recompensa única, permitindo que cada especialista atinja seu teto de desempenho de forma isolada; em seguida, estabelece uma política inicial robusta por meio de um esquema de partida a frio baseado em Flow e consolida de forma contínua as expertises heterogêneas em um único estudante por meio de uma orquestração de três etapas de amostragem on-policy, rotulagem de roteamento de tarefas e supervisão densa em nível de trajetória. Introduzimos ainda a Regularização de Âncora de Manifold (MAR), que utiliza um professor agnóstico a tarefas para fornecer supervisão de dados completos, ancorando a geração a um manifold de alta qualidade, mitigando efetivamente a degradação estética comumente observada em alinhamento puramente baseado em RL. Construído sobre o Stable Diffusion 3.5 Medium, o Flow-OPD eleva a pontuação GenEval de 63 para 92 e a precisão OCR de 59 para 94, resultando em uma melhoria geral de aproximadamente 10 pontos em relação ao GRPO puro, preservando a fidelidade da imagem e o alinhamento com preferências humanas, e exibindo um efeito emergente de "superação do professor". Esses resultados estabelecem o Flow-OPD como um paradigma de alinhamento escalável para a construção de modelos generalistas de texto-para-imagem.

English

Existing Flow Matching (FM) text-to-image models suffer from two critical bottlenecks under multi-task alignment: the reward sparsity induced by scalar-valued rewards, and the gradient interference arising from jointly optimizing heterogeneous objectives, which together give rise to a 'seesaw effect' of competing metrics and pervasive reward hacking. Inspired by the success of On-Policy Distillation (OPD) in the large language model community, we propose Flow-OPD, the first unified post-training framework that integrates on-policy distillation into Flow Matching models. Flow-OPD adopts a two-stage alignment strategy: it first cultivates domain-specialized teacher models via single-reward GRPO fine-tuning, allowing each expert to reach its performance ceiling in isolation; it then establishes a robust initial policy through a Flow-based Cold-Start scheme and seamlessly consolidates heterogeneous expertise into a single student via a three-step orchestration of on-policy sampling, task-routing labeling, and dense trajectory-level supervision. We further introduce Manifold Anchor Regularization (MAR), which leverages a task-agnostic teacher to provide full-data supervision that anchors generation to a high-quality manifold, effectively mitigating the aesthetic degradation commonly observed in purely RL-driven alignment. Built upon Stable Diffusion 3.5 Medium, Flow-OPD raises the GenEval score from 63 to 92 and the OCR accuracy from 59 to 94, yielding an overall improvement of roughly 10 points over vanilla GRPO, while preserving image fidelity and human-preference alignment and exhibiting an emergent 'teacher-surpassing' effect. These results establish Flow-OPD as a scalable alignment paradigm for building generalist text-to-image models.