Recentemente, pesquisadores da Universidade da Califórnia em Irvine e outras instituições utilizaram estratégias como mascaramento tardio, MoE e expansão hierárquica para reduzir o custo de treinamento de modelos de difusão para $1.890. Isso é uma ordem de magnitude menor que o método mais barato anterior (Wuerstchen) de $28.400, enquanto modelos como o Stable Diffusion têm custos ainda mais altos.
Os pesquisadores experimentaram várias abordagens para reduzir essas despesas enormes. Por exemplo, o modelo de difusão original requeria cerca de 1000 etapas do ruído à imagem, que agora foi reduzido para cerca de 20 etapas ou menos. À medida que os módulos básicos nos modelos de difusão foram gradualmente substituídos de Unet (CNN) para DiT (Transformer), algumas otimizações baseadas nas características do Transformer também foram implementadas, como quantização, pular alguns cálculos redundantes na Attention, pipeline, etc.
Este estudo treinou um modelo de difusão com 1,16 bilhão de parâmetros do zero por apenas $1.890. Isso representa uma melhoria de uma ordem de magnitude em comparação com o estado da arte, dando esperança às pessoas comuns de poderem experimentar o pré-treinamento. Mais importante ainda, as técnicas de redução de custos não afetaram o desempenho do modelo, com 1,16 bilhão de parâmetros produzindo resultados muito bons. Além da percepção visual, as métricas do modelo também são excelentes, com pontuações FID muito próximas ao Stable Diffusion 1.5 e DALL·E 2.
Os segredos para economizar dinheiro incluem principalmente:
-
Estratégia de mascaramento tardio: Usar um misturador de patches para pré-processar antes do mascaramento, incorporando informações de patches descartados em patches sobreviventes, reduzindo significativamente a queda de desempenho causada por alto mascaramento.
-
Ajuste fino: Realizar um pequeno ajuste fino (desmascaramento) após o pré-treinamento (mascaramento) para mitigar artefatos de geração indesejados causados pelo uso de máscaras.
-
MoE e expansão em camadas: Usar camadas MoE simplificadas baseadas em roteamento de seleção de especialistas para aumentar os parâmetros e a capacidade expressiva do modelo sem aumentar significativamente os custos de treinamento. Também considerou métodos de escalonamento em camadas, aumentando linearmente a largura dos blocos do Transformer.
Em termos de configuração experimental, os autores usaram duas variantes de DiT: DiT-Tiny/2 e DiT-Xl/2, com tamanho de patch de 2. Todos os modelos foram treinados usando o otimizador AdamW com decaimento de taxa de aprendizado cosseno e alto decaimento de peso. O front-end do modelo usa o codificador automático variacional (VAE) de quatro canais do modelo Stable-Diffusion-XL para extrair características de imagem, e também testou o desempenho do VAE de 16 canais mais recente em treinamento em larga escala (versão econômica).
Os autores usaram o framework EDM como configuração de treinamento unificada para todos os modelos de difusão, usando pontuações FID e CLIP para medir o desempenho dos modelos de geração de imagens. O codificador de texto escolhido foi o modelo CLIP mais comumente usado.
O conjunto de dados de treinamento usou três conjuntos de dados de imagens reais (Conceptual Captions, Segment Anything, TextCaps), contendo 22 milhões de pares de texto-imagem. Como o SA1B não fornece legendas reais, legendas sintéticas geradas pelo modelo LLaVA foram usadas aqui.