Explorando o Espaço de Design da Retropropagação de Recompensa para Flow Matching

Resumo

Alinhar modelos de correspondência de fluxo texto-imagem com preferências humanas via retropropagação direta de recompensas é eficiente em termos de amostras, mas é prejudicado por duas patologias bem conhecidas: as ativações não podem ser armazenadas ao longo de toda a trajetória de amostragem na escala moderna dos modelos, e os produtos Jacobianos encadeados entre etapas inflam o gradiente da recompensa à medida que ele retrocede aos índices iniciais. Métodos baseados em conectores, como o LeapAlign, abordam esses problemas substituindo a trajetória completa de retropropagação por um caminho curto e fixo, destacando uma separação útil entre amostragem e otimização. No entanto, a qualidade do gradiente resultante depende da precisão com que esse caminho curto aproxima a execução completa, especialmente em intervalos longos. Propomos o FlowBP, uma estrutura unificada de trajetória substituta que trata a própria trajetória de retropropagação como objeto de design. O FlowBP mantém uma execução armazenada em cache sem gradiente para amostragem e, em seguida, constrói um substituto leve de retropropagação a partir de velocidades armazenadas em cache e reenviadas seletivamente. Essa visão separa quatro escolhas: a entrada do modelo de recompensa, o conjunto ativo, os pesos de integração e o acoplamento de ponte, e recupera métodos anteriores de gradiente direto como configurações particulares. Dentro dessa estrutura, instanciamos três variantes: FlowBP-Sparse usa reconstrução esparsa de Euler, FlowBP-Bridge adiciona acoplamento de ponte controlado e FlowBP-Lagrange eleva a ordem da quadratura de salto. Todas as três limitam a memória pelo tamanho do conjunto ativo e restringem o encadeamento de gradientes a no máximo um fator Jacobiano. Em SD3.5-M, FLUX.1-dev e FLUX.2-Klein-base, em métricas de preferência, qualidade e composicionalidade, as três variantes melhoram em relação às linhas de base de gradiente direto na maioria das métricas.

English

Aligning text-to-image flow matching models with human preferences via direct reward backpropagation is sample-efficient but hampered by two well-known pathologies: activations cannot be stored across the full sampling trajectory at modern model scale, and chained Jacobian products across steps inflate the reward gradient as it travels back to early indices. Connector-based methods, such as LeapAlign, address these issues by replacing the full backward trajectory with a short pinned path, highlighting a useful decoupling between sampling and optimization. However, the quality of the resulting gradient depends on how accurately this short path approximates the full rollout, especially over long intervals. We propose FlowBP, a unified surrogate-trajectory framework that treats the backward trajectory itself as the design object. FlowBP keeps a no-gradient cached rollout for sampling, then builds a lightweight backward surrogate from cached and selectively re-forwarded velocities. This view separates four choices: the reward-model input, active set, integration weights, and bridge coupling, and recovers prior direct-gradient methods as particular settings. Within this framework, we instantiate three variants: FlowBP-Sparse uses sparse Euler reconstruction, FlowBP-Bridge adds controlled bridge coupling, and FlowBP-Lagrange raises the order of leap quadrature. All three bound memory by the active-set size and limit gradient chaining to at most one Jacobian factor. Across SD3.5-M, FLUX.1-dev, and FLUX.2-Klein-base on preference, quality, and compositional metrics, the three variants improve over direct-gradient baselines on most metrics.