VibeToken: Scaling 1D Image Tokenizers and Autoregressive Models for Dynamic Resolution Generations

1. Declaração do Problema

Modelos de geração de imagem autoregressivos (AR) (por exemplo, LlamaGen) demonstraram desempenho competitivo, mas enfrentam barreiras significativas para adoção em produção em comparação com modelos de difusão. O principal gargalo é a escalabilidade em relação à resolução e às proporções:

• Explosão de Tokens: Tokenizadores 2D convencionais (por exemplo, VQGAN) geram sequências de tokens que crescem linearmente com a resolução da imagem. Para uma imagem de 1024×1024, isso resulta em milhares de tokens (por exemplo, 4.096 tokens).
• Custo Computacional Quadrático: Modelos AR dependem de autoatenção, que possui complexidade $O(T^2)$, onde $T$ é o comprimento da sequência. À medida que a resolução aumenta, os FLOPs de inferência crescem quadraticamente, tornando a geração de alta resolução computacionalmente proibitiva.
• Rigidez: A maioria dos modelos AR é treinada em resoluções fixas (por exemplo, 256×256) e luta para generalizar para proporções arbitrárias ou resoluções mais altas sem retreinamento ou pipelines complexos de super-resolução.
• Domínio da Difusão: Modelos de difusão lidam naturalmente com resoluções arbitrárias, deixando os modelos AR para trás em casos de uso flexíveis e de nível de produção.

2. Metodologia

Os autores propõem uma solução em duas partes: VibeToken (um tokenizador inovador) e VibeToken-Gen (um gerador AR).

A. VibeToken: Tokenizador 1D Agnóstico à Resolução

O VibeToken é um tokenizador de imagem baseado em Transformer 1D projetado para codificar imagens de qualquer resolução em um número fixo e controlável pelo usuário de tokens (por exemplo, 32–256), desacoplando a contagem de tokens da resolução espacial.

Inovações arquitetônicas chave incluem:

1. Embeddings Posicionais de Grade Dinâmica: Em vez de embeddings posicionais absolutos fixos, o VibeToken usa uma grade aprendida de embeddings que é redimensionada diferencialmente (via interpolação bilinear/bicúbica) para corresponder à grade de patches de entrada ($T_H \times T_W$). Isso permite que o modelo lide com proporções arbitrárias sem degradação de qualidade.
2. Embeddings de Patch Adaptativos: Inspirado no FlexiViT, o modelo suporta tamanhos de patch variáveis ($k$) durante o treinamento e a inferência. Um operador de redimensionamento de pesos mapeia os pesos de projeção base para o tamanho de patch atual, permitindo que o modelo se adapte a diferentes resoluções sem reparametrização.
3. Resolução Adaptativa do Decodificador: O decodificador prevê pixels em uma resolução alta fixa (4× o tamanho do patch) e usa uma camada de subamostragem convolucional 2D aprendível para mapear para a resolução alvo. Isso habilita super-resolução nativa sem um módulo de upscaler separado.
4. Tokenização de Comprimento Dinâmico: O modelo é treinado com uma amostragem uniforme de comprimentos latentes $L \in [32, 256]$. O codificador produz exatamente $L$ tokens, e o decodificador consome exatamente $L$, eliminando a necessidade de preenchimento ou restrições de comprimento fixo.
5. Quantização: Utiliza Quantização de Vetor Múltiplo (MVQ) com 8 livros de códigos para alcançar um tamanho de vocabulário de 32.768, equilibrando compressão e qualidade de reconstrução.

B. VibeToken-Gen: Gerador AR Eficiente

O VibeToken-Gen é um modelo AR condicionado a classes (baseado na arquitetura LlamaGen) que opera nas sequências de tokens curtas e agnósticas à resolução produzidas pelo VibeToken.

• Orçamento Computacional Constante: Como o comprimento da sequência de tokens $L$ é fixo (por exemplo, 64 tokens) independentemente da resolução de saída, os FLOPs de inferência permanecem constantes.
• Condicionamento: O modelo condiciona-se à resolução alvo $(H, W)$ e à proporção, juntamente com o rótulo da classe, para evitar artefatos de distorção geométrica.
• Treinamento: Treinado no ImageNet-1k com resoluções mistas (256×256 a 512×512) e proporções, usando um orçamento de tokens comparável a baselines de resolução fixa.

3. Contribuições Principais

1. Tokenização Agnóstica à Resolução: O primeiro tokenizador Transformer 1D que suporta nativamente resoluções e proporções arbitrárias, mantendo um orçamento de tokens curto e dinâmico (32–256 tokens).
2. Eficiência Computacional: Desacopla o custo de inferência AR da resolução da imagem. O VibeToken-Gen mantém 179 GFLOPs constantes para qualquer resolução, enquanto baselines como o LlamaGen escalam para ~11 TFLOPs para imagens de 1024×1024 (um ganho de eficiência de 63,4×).
3. Super-Resolução Nativa: A arquitetura suporta a geração de imagens de alta resolução a partir de entradas de baixa resolução (ou alvos arbitrários) sem módulos de upscaling separados.
4. Desempenho de Estado da Arte: Alcanza qualidade de reconstrução (rFID) e geração (gFID) competitiva em comparação com tokenizadores 2D e modelos de difusão, apesar de usar significativamente menos tokens.

4. Resultados Experimentais

O artigo avalia o VibeToken e o VibeToken-Gen nos conjuntos de dados ImageNet-1k e FFHQ em várias resoluções (256×256 a 1024×1024) e proporções.

• Qualidade de Reconstrução:
  • O VibeToken-LL alcança 0,40 rFID em 256×256 e 2,40 rFID em 1024×1024.
  • Supera outros tokenizadores 1D (por exemplo, TiTok, One-D-Piece) na generalização de resolução e iguala ou excede tokenizadores 2D (por exemplo, IBQ, MAGViT) em eficiência de compressão.
• Qualidade de Geração:
  • O VibeToken-Gen (XXL) sintetiza imagens de 1024×1024 usando apenas 64 tokens com um gFID de 3,94.
  • Comparação:
    • Vs. Difusão (NiT): O VibeToken-Gen é 2,35× mais rápido (0,46s vs 1,08s) e alcança melhor qualidade (3,94 gFID vs 5,87 gFID) em 1024×1024.
    • Vs. Baseline AR (LlamaGen): O LlamaGen requer ~11 TFLOPs para 1024×1024, enquanto o VibeToken-Gen usa 179 GFLOPs constantes.
• Eficiência:
  • A contagem de tokens é reduzida de ~4.096 (em tokenizadores 2D) para 64 tokens para imagens de 1024×1024.
  • A velocidade de inferência é subsegundo (0,46s) para geração de alta resolução, comparada a 32+ segundos para modelos AR de baseline em resoluções similares.

5. Significado

Este trabalho preenche a lacuna entre modelos Autoregressivos e modelos de Difusão em termos de flexibilidade e eficiência.

• Viabilidade de Produção: Ao permitir que modelos AR lidem com resoluções arbitrárias com computação constante, o VibeToken torna a geração AR viável para aplicações do mundo real onde as dimensões de entrada/saída variam (por exemplo, aplicativos móveis, formatos de exibição diversos).
• Escalabilidade: Demonstra que modelos AR podem escalar para altas resoluções sem a explosão computacional quadrática, desafiando a noção de que a difusão é o único caminho para IA generativa de alta resolução.
• Direção Futura: A abordagem abre caminho para modelos multimodais unificados (texto, imagem, vídeo) que podem operar sob um orçamento computacional fixo e controlável pelo usuário, independentemente da complexidade espacial do conteúdo gerado.

Em resumo, o VibeToken resolve o problema de "escalabilidade de resolução" na geração AR ao transferir o fardo da adaptação de resolução do modelo generativo para um tokenizador flexível e agnóstico à resolução, resultando em um sistema mais rápido, barato e flexível do que as alternativas atuais de estado da arte.