Técnica de Treinamento Aprimorada para Modelos de Consistência Latente

Improved Training Technique for Latent Consistency Models

February 3, 2025

Autores: Quan Dao, Khanh Doan, Di Liu, Trung Le, Dimitris Metaxas

cs.AI

Resumo

Os modelos de consistência são uma nova família de modelos generativos capazes de produzir amostras de alta qualidade em um único passo ou em vários passos. Recentemente, os modelos de consistência têm demonstrado um desempenho impressionante, alcançando resultados comparáveis aos modelos de difusão no espaço de pixels. No entanto, o sucesso em escalar o treinamento de consistência para conjuntos de dados em larga escala, especialmente para tarefas de geração de texto para imagem e vídeo, é determinado pelo desempenho no espaço latente. Neste trabalho, analisamos as diferenças estatísticas entre os espaços de pixels e latentes, descobrindo que os dados latentes frequentemente contêm outliers altamente impulsivos, que degradam significativamente o desempenho do iCT no espaço latente. Para lidar com isso, substituímos as perdas Pseudo-Huber por perdas de Cauchy, mitigando efetivamente o impacto dos outliers. Além disso, introduzimos uma perda de difusão nos primeiros passos e utilizamos acoplamento de transporte ótimo (OT) para melhorar ainda mais o desempenho. Por fim, introduzimos o programador de escalonamento adaptativo para gerenciar o processo de treinamento robusto e adotamos a LayerNorm não escalonável na arquitetura para capturar melhor as estatísticas das características e reduzir o impacto dos outliers. Com essas estratégias, treinamos com sucesso modelos de consistência latente capazes de amostragem de alta qualidade com um ou dois passos, reduzindo significativamente a lacuna de desempenho entre a consistência latente e os modelos de difusão. A implementação está disponível em: https://github.com/quandao10/sLCT/

English

Consistency models are a new family of generative models capable of producing high-quality samples in either a single step or multiple steps. Recently, consistency models have demonstrated impressive performance, achieving results on par with diffusion models in the pixel space. However, the success of scaling consistency training to large-scale datasets, particularly for text-to-image and video generation tasks, is determined by performance in the latent space. In this work, we analyze the statistical differences between pixel and latent spaces, discovering that latent data often contains highly impulsive outliers, which significantly degrade the performance of iCT in the latent space. To address this, we replace Pseudo-Huber losses with Cauchy losses, effectively mitigating the impact of outliers. Additionally, we introduce a diffusion loss at early timesteps and employ optimal transport (OT) coupling to further enhance performance. Lastly, we introduce the adaptive scaling-c scheduler to manage the robust training process and adopt Non-scaling LayerNorm in the architecture to better capture the statistics of the features and reduce outlier impact. With these strategies, we successfully train latent consistency models capable of high-quality sampling with one or two steps, significantly narrowing the performance gap between latent consistency and diffusion models. The implementation is released here: https://github.com/quandao10/sLCT/