Resumo Técnico: Mage-Flow

Declaração do Problema

Modelos generativos visuais de larga escala alcançaram capacidades notáveis na geração de texto para imagem e na edição de imagens baseada em instruções. No entanto, esses sistemas dependem cada vez mais de backbones massivos (variando de 6B a 80B de parâmetros), criando barreiras significativas para o implante prático, ajuste fino (fine-tuning) e pesquisa sob orçamentos de computação restritos. Além disso, os modelos de código aberto existentes frequentemente herdam arquiteturas de tokenização (VAEs) otimizadas para reconstrução ao nível de pixel que se tornam gargalos computacionais em altas resoluções, onde os custos de codificação e decodificação podem rivalizar com o próprio backbone de difusão. Existe uma lacuna entre o lançamento de pesos abertos e o fornecimento de um sistema que seja verdadeiramente prático para implantação em desktop local, uso interativo e pesquisa sistemática sem exigir dezenas de bilhões de parâmetros.

Metodologia

Os autores apresentam o Mage-Flow, um stack generativo compacto de escala 4B projetado para eficiência através de co-design de sistema, em vez de simples escalonamento de parâmetros. O stack compreende três componentes principais:

1. Mage-VAE: Tokenizador Leve de Alta Fidelidade

Para abordar a latência e os gargalos de memória dos VAEs padrão em altas resoluções, os autores projetaram o Mage-VAE, um tokenizador latente leve.

• Arquitetura: Utiliza uma abordagem de estilo difusão de um único passo tanto para codificação quanto para decodificação. O decodificador é um modelo de difusão de pixel totalmente convolucional destilado para um único passo, evitando blocos de atenção global. O codificador é construído como o dual arquitetônico do decodificador, também operando como um modelo de difusão de um único passo condicionado aos pixels.
• Treinamento: O modelo é treinado em três estágios: (1) Pré-treinamento de flow-matching de múltiplos passos onde o codificador prevê latentes âncoras e o decodificador reconstrói pixels; (2) Destilação de decodificador de um único passo usando perdas de reconstrução, DMD e GAN; (3) Ajuste fino conjunto com uma regularização de latente âncora. Esta regularização alinha o posterior do Mage-VAE com a distribuição latente de um VAE público forte e computacionalmente caro (FLUX.2-VAE), garantindo que o espaço latente permaneça pronto para geração e compatível com geradores de difusão existentes.
• Eficiência: O Mage-VAE alcança fidelidade de reconstrução comparável ao FLUX.2-VAE enquanto reduz os custos computacionais de codificação e decodificação (MACs) em aproximadamente 12× e 22×, respectivamente.

2. Transformer de Difusão Multimodal de Resolução Nativa (NR-MMDiT)

O backbone generativo é um NR-MMDiT de 4 bilhões de parâmetros treinado com flow matching retificado.

• Empacotamento de Resolução Nativa: Ao contrário do treinamento convencional baseado em buckets que restringe os lotes a buckets de resolução/proporção de aspecto fixos, o Mage-Flow utiliza um esquema de empacotamento de resolução nativa. Sequências de imagens de comprimento variável (com resoluções e proporções de aspecto arbitrárias) e sequências de texto são empacotadas em um único lote usando kernels de comprimento variável do FlashAttention e embeddings rotativos 2D por amostra. Isso permite que o modelo aprenda de tamanhos nativos heterogêneos em uma única atualização, melhorando a flexibilidade e a generalização para proporções de aspecto extremas (ex: 512×2048).
• Backbone Compartilhado: O mesmo backbone NR-MMDiT é usado tanto para geração de texto para imagem quanto para edição baseada em instruções, com formatos de condicionamento ajustados para cada tarefa (prompts vs. imagem de origem + instruções).

3. Otimização de Nível de Sistema

A infraestrutura de treinamento é otimizada para remover o overhead limitado pela memória.

• Kernels Fundidos: Os autores fundem cadeias de operadores dominantes (normalização, ativação, conexões residuais, embeddings rotativos) dentro dos blocos repetidos do Mage-VAE, do codificador de texto Qwen3-VL e do NR-MMDiT em kernels CUDA customizados.
• Impacto: Isso reduz o tráfego de memória de ativação e o overhead de lançamento de kernel, aumentando a Utilização de FLOPs do Modelo (MFU) de ~33% para ~77% e alcançando um aceleração de treinamento end-to-end de 2.5×.

Pós-Treinamento e Destilação

• Diffusion-NFT: Os modelos base são alinhados usando Diffusion-NFT, uma técnica de pós-treinamento que realiza o ajuste fino consciente de negativos no processo direto para melhorar o seguimento de prompts, a renderização de texto e a qualidade estética sem exigir estimativa de verossimilhança.
• Variantes Turbo: As variantes Mage-Flow-Turbo de 4 passos e Mage-Flow-Edit-Turbo são produzidas via destilação de poucos passos. Isso combina o Decoupled-DMD (separando a augmentação de guia livre de classificador da correspondência de distribuição) com orientação perceptual adversária baseada em Modelos de Fundação de Visão congelados (DINOv2, CLIP).

Principais Contribuições

1. Stack Generativo Mage-Flow: Um fundamento compacto de escala 4B para geração e edição eficientes, co-projetado com um tokenizador leve, um backbone de resolução nativa e uma infraestrutura de kernel fundido.
2. Mage-VAE: Um tokenizador leve de alta fidelidade usando codificação/decodificação de difusão de um único passo e regularização de latente âncora, reduzindo os custos de tokenização em uma ordem de magnitude enquanto mantém a compatibilidade com VAEs públicos fortes.
3. NR-MMDiT: Um backbone treinado com flow matching retificado e empacotamento de resolução nativa, permitindo treinamento de resolução flexível e inferência de guia livre (CFG) empacotada eficiente.
4. Família de Modelos Completa: O lançamento das variantes Base, alinhada por RL e Turbo de 4 passos, tanto para geração quanto para edição, demonstrando que um modelo de 4B pode alcançar desempenho competitivo contra sistemas de código aberto significativamente maiores.

Resultados

• Desempenho: O Mage-Flow alcança resultados competitivos ou superiores em benchmarks de geração padrão (GenEval, DPG-Bench, TIIF-Bench, CVTG-2K, OneIG, LongText) e benchmarks de edição (ImgEdit-Bench, GEdit-Bench, TextEdit-Bench) comparado a modelos significativamente maiores (ex: FLUX.2, Qwen-Image, HunyuanImage-3.0).
• Eficiência:
  • Velocidade de Inferência: Em resolução 1024² em uma única GPU NVIDIA A100, o Mage-Flow-Turbo gera uma imagem em 0.59s, e o Mage-Flow-Edit-Turbo edita uma imagem em 1.02s.
  • Memória: O uso de memória de GPU pico permanece em torno de 17–20 GB, o mais baixo entre os modelos comparados.
  • Treinamento: A infraestrutura de kernel fundido proporciona um aceleração de treinamento de 2.5× sobre os baselines usando FLUX.2-VAE.
• Qualitativo: Os modelos demonstram fortes capacidades de fotorealismo, arte estilizada, renderização de texto bilíngue (Inglês e Chinês) e diversas tarefas de edição (substituição de fundo, manipulação de objetos, edição de texto) através de resoluções e proporções de aspecto flexíveis.

Significância

O artigo argumenta que uma geração e edição visual forte não necessita de backbones de dezenas de bilhões de parâmetros. Ao co-projetar cuidadosamente o tokenizador, o backbone e o sistema de treinamento, uma família de modelos compactos de 4B pode entregar capacidades de alta resolução, interativas e amigáveis à pesquisa. O Mage-Flow fornece um fundamento prático para implantação em desktop local e pesquisa downstream, preenchendo a lacuna entre modelos de código aberto de alto desempenho e as restrições de orçamentos de computação realistas. O trabalho destaca que a eficiência do sistema (tokenização e fusão de kernel) é tão crítica quanto o escalonamento arquitetural para a próxima geração de geradores visuais eficientes.