Resumo Técnico: NanoFLUX

Declaração do Problema

Modelos de difusão de Texto-para-Imagem (T2I) de grande escala, como a família FLUX (escalando até 17B de parâmetros) e o Qwen-Image (20B parâmetros), alcançaram uma qualidade visual de estado da arte. No entanto, sua escala massiva cria uma barreira significativa para a implantação em dispositivos (on-device), exigindo infraestrutura de GPU dispendiosa ou inferência baseada em nuvem. As soluções existentes para dispositivos costumam depender de modelos menores que carecem de qualidade comparável ou empregam técnicas de compressão (poda, quantização, redução de passos) que não abordam totalmente a redundância arquitetural dos Diffusion Transformers (DiTs) de grande escala e seus codificadores de texto associados. O desafio reside em comprimir esses modelos massivos para dispositivos móveis enquanto se preserva a qualidade de geração e se minimiza a latência de inferência.

Metodologia: NanoFLX

Os autores introduzem o NanoFLUX, um modelo T2I de 2,4 bilhões de parâmetros destilado do professor FLUX.1-Schnell de 17B. A compressão é alcançada através de um pipeline progressivo, impulsionado por destilação, visando os dois componentes mais pesados: o Diffusion Transformer (DiT) de 12B e o codificador de texto T5-XXL de 5B.

1. Compressão do Transformer (12B $\to$ 2B)

Os autores empregam uma estratégia de compressão progressiva de quatro estágios para reduzir o tamanho do DiT:

• Estágio C1: Poda de Cabeças de Atenção (12B $\to$ 5B):

  • Análise: A Decomposição de Valores Singulares (SVD) nas saídas das cabeças de atenção revelou que nem todas as cabeças são independentes; muitas contêm informações redundantes.
  • Ação: O número de cabeças de atenção ($H$) foi uniformemente reduzido de 24 para 16.
  • Destilação: Uma perda de nível de característica foi introduzida para garantir que as cabeças restantes preservem coletivamente a informação das cabeças podadas. O estudante foi inicializado com as primeiras 16 cabeças do professor.

• Estágio C2: Redução da Dimensão da Cabeça (5B $\to$ 3B):

  • Análise: A SVD foi aplicada novamente às características do modelo de 5B.
  • Ação: A dimensão da cabeça ($d_H$) foi reduzida de 128 para 96, retendo os principais componentes singulares. Isso resultou em um modelo de 3B de parâmetros.

• Estágio C3: Poda de Profundidade via Fusão de Blocos (3B $\to$ 2,5B):

  • Análise: A similaridade de cosseno entrada-saída foi calculada para os blocos do transformer. Os blocos de Fluxo Único (SS) mostraram alta similaridade (>0,85), indicando redundância funcional.
  • Ação: Em vez de deletar esses blocos, os autores fundiram uma cadeia de 15 blocos SS redundantes em um único bloco através da média de seus parâmetros. Isso reduziu a profundidade do transformer de 38 para 24 blocos.

• Estágio C4: Normalização de Camada Estática (2,5B $\to$ 1,8B):

  • Análise: Os coeficientes de Normalização de Camada Adaptativa (AdaLN) previstos pelo modelo mostraram baixa variância entre as amostras, sugerindo que o condicionamento dinâmico das projeções CLIP era redundante.
  • Ação: Os coeficientes dinâmicos de AdaLN foram substituídos por médias estáticas pré-computadas (Static-LN). Isso removeu aproximadamente 0,7B de parâmetros. O modelo foi ajustado brevemente com coeficientes congelados para recuperar a qualidade.

2. Downsampling Progressivo de Tokens

Para abordar a complexidade quadrática da autoatenção sem sacrificar o detalhe:

• Mecanismo: Um módulo de downsampling baseado em ResNet ($D$) e um módulo de upsampling ($U$) foram inseridos no transformer.
• Estratégia Híbrida: O modelo opera em um esquema de resolução híbrida. Os estágios iniciais de difusão (capturando a estrutura global) processam tokens com downsampling (reduzindo a contagem de tokens em 4x), enquanto os estágios posteriores (refinando detalhes) processam tokens de resolução total.
• Treinamento: Uma estratégia de "Downsampling Progressivo de Tokens" (PTD) foi usada para estabilizar o treinamento. Os módulos foram inseridos e treinados incrementalmente, movendo o ponto de downsampling mais cedo na rede ao longo das iterações.
• Resultado: 23 de 43 blocos do transformer operam em tokens com downsampling, reduzindo significativamente a latência enquanto mantém o processamento de alta resolução onde necessário.

3. Destilação do Codificador de Texto (5B $\to$ 330M)

• Alvo: Destilar o codificador T5-XXL de 5B para um modelo T5-Large de 330M.
• Método: Um processo de destilação em dois estágios:
  1. Correspondência de Embedding: Um MLP é adicionado ao estudante para corresponder à dimensão de saída do professor, minimizando o MSE entre os embeddings do prompt.
  2. Destilação por Bloco: O estudante é treinado para corresponder aos estados ocultos do prompt do professor em camadas específicas do transformer. Crucialmente, o método aproveita sinais visuais das camadas iniciais do denoiser (via atenção conjunta) para guiar o codificador de texto, mitigando a acumulação de erro.
  3. Estabilidade: Os gradientes são desconectados para timesteps de alto ruído ($t < \hat{t}$) para evitar supervisão instável, focando a destilação nos estágios de tempo posteriores e nas três primeiras camadas do transformer.

Principais Resultados

• Tamanho do Modelo: O modelo final NanoFLUX possui 2,4 bilhões de parâmetros (2B DiT + 330M Codificador de Texto), uma redução de 7$\times$ em relação ao professor FLUX.1-Schnell de 17B.
• Desempenho:
  • Em dispositivos móveis, o NanoFLUX gera imagens de 512 $\times$ 512 em aproximadamente 2,5 segundos.
  • Benchmarks quantitativos (One-IG, DPG, GenEval, HPDv3) mostram que o NanoFLUX mantém o desempenho comparável ao modelo professor, com apenas uma degradação menor em métricas específicas (ex: HPSv3 cai de 11,45 para 10,41).
  • Resultados qualitativos indicam que o modelo preserva a qualidade da imagem e a aderência ao prompt comparáveis ao modelo original de 17B.
• Eficiência: A estratégia de downsampling progressivo de tokens equilibra com sucesso a latência de inferência e a qualidade de geração, permitindo que o modelo rode eficientemente em hardware com recursos limitados.

Significância e Alegações

O artigo afirma que o NanoFLUX é o primeiro framework a estudar de forma abrangente a compressão de modelos de grande escala para geração T2I, visando especificamente modelos como o FLUX. Sua significância reside em:

1. Viabilidade de Geração On-Device: Demonstrar que a geração T2I de alta qualidade é viável em dispositivos móveis sem depender de infraestrutura de nuvem.
2. Insight Arquitetural: Fornecer um caminho prático para reduzir a escala de grandes modelos generativos ao identificar e explorar redundâncias arquiteturais específicas (cabeças de atenção, profundidade de blocos, dinâmica de normalização) em vez de depender apenas de quantização genérica ou redução de passos.
3. Generalização: Oferecer insights e um pipeline de destilação progressiva que podem se generalizar para futuros modelos generativos de grande escala.

Os autores posicionam este trabalho como um passo para ampliar a acessibilidade da IA generativa de alta qualidade para usuários com recursos computacionais limitados.