Flow-Factory: A Unified Framework for Reinforcement Learning in Flow-Matching Models

1. Problema

O alinhamento de modelos generativos avançados (como modelos de difusão e flow-matching) com preferências humanas através de Aprendizado por Reforço (RL) tem se mostrado promissor, mas a implementação prática enfrenta barreiras significativas:

• Códigos Fragmentados e Acoplados: Os algoritmos de RL proliferaram com formulações específicas de Equações Diferenciais Estocásticas (SDE) e procedimentos de treinamento. Isso resultou em repositórios isolados onde a lógica do algoritmo está fortemente acoplada à implementação do modelo, dificultando a transferência de algoritmos entre arquiteturas ou a comparação controlada de múltiplos métodos.
• Ineficiência de Treinamento e Gargalos de Memória: O ajuste fino (fine-tuning) via RL exige passagens forward repetidas para amostragem de trajetórias e cálculo de recompensas. Em modelos de grande escala com componentes congelados (codificadores de texto, VAEs), o consumo de memória GPU é excessivo, limitando o tamanho do batch e a escalabilidade em hardware comercial.
• Falta de Flexibilidade em Recompensas: As estruturas atuais geralmente suportam apenas recompensas pontuais (pointwise). No entanto, algoritmos emergentes exigem recompensas grupais (groupwise, para avaliação de ranking) ou combinações de múltiplas recompensas com estratégias de agregação diferentes, o que muitas vezes resulta em carregamento redundante de modelos.

2. Metodologia

O Flow-Factory é apresentado como um framework unificado e extensível que resolve esses desafios através de três princípios de design principais:

A. Arquitetura Modular Baseada em Registro (Registry-Based)

O framework desacopla modelos, algoritmos, recompensas e agendadores (schedulers) em quatro componentes padronizados:

1. BaseAdapter: Para operações do modelo (codificação, forward pass, checkpointing).
2. BaseTrainer: Para lógica do algoritmo (amostragem, cálculo de vantagem, otimização).
3. BaseRewardModel: Para computação de recompensas.
4. SDESchedulerMixin: Para amostragem estocástica e cálculo de log-probabilidade.
   Essa abordagem reduz a complexidade de integração de $O(M \times N)$ para $O(M + N)$, permitindo que novos componentes sejam combinados apenas através de configurações YAML, sem reescrita de código.

B. Otimização de Memória Baseada em Pré-processamento

Para mitigar o uso excessivo de memória:

• O framework pré-computa e armazena em disco todos os embeddings de condição (prompts, pooled embeddings, latentes do VAE) antes do treinamento.
• Durante o treinamento, esses embeddings são carregados diretamente, permitindo que os componentes congelados (codificadores de texto, VAEs) sejam completamente removidos da memória da GPU.
• Isso resulta em um layout de memória focado apenas no transformer alvo, permitindo batches maiores e eliminando operações de codificação redundantes.

C. Sistema Flexível de Múltiplas Recompensas

O sistema oferece interfaces unificadas para:

• Recompensas Pontuais e Grupais: Suporte nativo para modelos que avaliam amostras individualmente ou rankings dentro de grupos.
• Deduplicação Automática: Um carregador inteligente (MultiRewardLoader) garante que modelos de recompensa únicos sejam carregados apenas uma vez, mesmo que referenciados por múltiplas configurações.
• Agregação de Vantagem Configurável: Suporte para estratégias de ponderação e normalização (ex: estilo GDPO), facilitando o treinamento multi-objetivo.

3. Algoritmos Suportados

O framework implementa unificadamente algoritmos de ponta, incluindo:

• Flow-GRPO e Variantes: Estende a amostragem determinística (ODE) com injeção de ruído estocástico (SDE) para permitir exploração em RL. Suporta variantes como MixGRPO (mistura de SDE e ODE) e GRPO-Guard (reponderação de perda).
• DiffusionNFT: Otimiza um objetivo contrastivo diretamente no processo de flow-matching, sem necessidade de estimativa de verossimilhança tratável.
• Advantage Weighted Matching (AWM): Alinha a otimização de RL com o objetivo de pré-treinamento de flow-matching ponderando a perda de correspondência de velocidade pelas vantagens calculadas.

4. Resultados Experimentais

Os experimentos foram conduzidos utilizando o modelo FLUX.1-dev com três algoritmos (Flow-GRPO, DiffusionNFT, AWM) e modelos de recompensa (Text-Rendering, PickScore).

• Reprodução de Resultados: O Flow-Factory conseguiu recuperar com fidelidade os ganhos de desempenho reportados na literatura original, demonstrando que o desacoplamento da lógica algorítmica não compromete a estabilidade ou a qualidade.
• Eficiência de Treinamento: Em um ambiente com 8 GPUs H200:
  • Redução de Memória: Diminuição de 13,0% no uso de pico de memória por dispositivo (de 61,08 GB para 53,14 GB) ao descarregar componentes congelados.
  • Aceleração: Aceleração de 1,74x no tempo por passo (redução de 144s para 82,68s) devido à eliminação de operações de codificação redundantes.
• Qualidade Visual: Comparativos qualitativos mostraram que todos os modelos ajustados via RL apresentaram melhoria na qualidade visual e melhor alinhamento com preferências estéticas humanas em comparação ao modelo base.

5. Significado e Contribuições

O Flow-Factory representa um avanço significativo na engenharia de sistemas de IA generativa:

• Padronização: Estabelece um padrão modular que permite a rápida prototipagem e comparação justa de novos algoritmos de RL em modelos de flow-matching.
• Acessibilidade: Ao reduzir drasticamente os requisitos de memória e o tempo de treinamento, torna o ajuste fino de RL em modelos grandes mais acessível para pesquisadores que utilizam hardware comercial.
• Futuro-Proof: A arquitetura baseada em registro garante que o framework possa se adaptar facilmente a novos modelos fundamentais e paradigmas de RL que surjam no futuro, sem a necessidade de grandes reescritas de código.

Em suma, o Flow-Factory remove barreiras de engenharia, acelerando a pesquisa comunitária no alinhamento de modelos generativos via RL. O código está disponível publicamente no repositório do projeto.