Alinhando Codificadores Visuais de Base a Tokenizadores para Modelos de Difusão

Resumo

Neste trabalho, propomos alinhar codificadores visuais pré-treinados para atuarem como tokenizadores em modelos de difusão latente para geração de imagens. Diferente de treinar um autoencoder variacional (VAE) do zero, que prioriza principalmente detalhes de baixo nível, nossa abordagem aproveita a rica estrutura semântica de codificadores de base. Introduzimos uma estratégia de alinhamento em três estágios: (1) congelar o codificador e treinar um adaptador e um decodificador para estabelecer um espaço latente semântico; (2) otimizar conjuntamente todos os componentes com uma função de perda adicional de preservação semântica, permitindo que o codificador capture detalhes perceptuais enquanto mantém a semântica de alto nível; e (3) refinar o decodificador para melhorar a qualidade da reconstrução. Esse alinhamento resulta em tokenizadores de imagens semanticamente ricos que beneficiam modelos de difusão. No ImageNet 256x256, nosso tokenizador acelera a convergência dos modelos de difusão, atingindo um gFID de 1,90 em apenas 64 épocas, e melhora a geração tanto com quanto sem orientação livre de classificador. Ao escalar para o LAION, um modelo de texto para imagem com 2 bilhões de parâmetros treinado com nosso tokenizador supera consistentemente o FLUX VAE sob o mesmo número de passos de treinamento. No geral, nosso método é simples, escalável e estabelece um paradigma semanticamente fundamentado para o design contínuo de tokenizadores.

English

In this work, we propose aligning pretrained visual encoders to serve as tokenizers for latent diffusion models in image generation. Unlike training a variational autoencoder (VAE) from scratch, which primarily emphasizes low-level details, our approach leverages the rich semantic structure of foundation encoders. We introduce a three-stage alignment strategy: (1) freeze the encoder and train an adapter and a decoder to establish a semantic latent space; (2) jointly optimize all components with an additional semantic preservation loss, enabling the encoder to capture perceptual details while retaining high-level semantics; and (3) refine the decoder for improved reconstruction quality. This alignment yields semantically rich image tokenizers that benefit diffusion models. On ImageNet 256times256, our tokenizer accelerates the convergence of diffusion models, reaching a gFID of 1.90 within just 64 epochs, and improves generation both with and without classifier-free guidance. Scaling to LAION, a 2B-parameter text-to-image model trained with our tokenizer consistently outperforms FLUX VAE under the same training steps. Overall, our method is simple, scalable, and establishes a semantically grounded paradigm for continuous tokenizer design.

Alinhando Codificadores Visuais de Base a Tokenizadores para Modelos de Difusão

Aligning Visual Foundation Encoders to Tokenizers for Diffusion Models

Resumo

Support