Semantic-Aware Prefix Learning for Token-Efficient Image Generation

1. O Problema

Os tokenizadores visuais são componentes centrais na geração de imagens latentes, atuando como uma ponte entre imagens de alta dimensão (pixels) e modelos generativos tratáveis. No entanto, a maioria dos tokenizadores existentes é treinada com objetivos dominados pela reconstrução (focados em minimizar o erro de pixel). Isso resulta em representações latentes que possuem um alinhamento fraco com semânticas de alto nível (conceitos, categorias, estrutura global).

Embora abordagens recentes tentem melhorar o alinhamento semântico, elas geralmente tratam os sinais semânticos apenas como regularização auxiliar, em vez de torná-los funcionalmente necessários para o aprendizado da representação. Consequentemente, o espaço latente muitas vezes falha em codificar eficientemente a estrutura global e os conceitos de alto nível, limitando a eficiência e o desempenho dos modelos generativos downstream, especialmente sob orçamentos de tokens compactos.

2. Metodologia

Os autores propõem duas contribuições principais: o tokenizador SMAP e o gerador híbrido CARD.

A. SMAP (SeMantic-Aware Prefix Tokenizer)

O SMAP é um tokenizador 1D baseado em consultas (query-based) projetado para injetar condições semânticas de nível de classe diretamente no processo de formação de tokens e reconstrução.

• Injeção Semântica Explícita: Diferente de métodos anteriores que usam alinhamento solto, o SMAP injeta embeddings condicionais (derivados de rótulos de classe) como elementos explícitos na sequência de tokens. Esses embeddings são inseridos entre os patches visuais e as consultas latentes aprendidas no codificador, e também entre os tokens de máscara e os tokens latentes processados no decodificador. Isso força a semântica a participar ativamente da formação da representação.
• Estratégia de Descarte de Tokens de Cauda (Tail Token Dropping): Para garantir que a semântica seja indispensável, o SMAP utiliza uma estratégia de treinamento onde, a cada iteração, uma parte dos tokens latentes finais (a "cauda") é descartada aleatoriamente.
  • Isso força os embeddings condicionais e os prefixos latentes iniciais a assumirem uma responsabilidade crescente pela reconstrução da imagem à medida que o orçamento de tokens diminui.
  • O caso extremo ($k=0$) exige que o decodificador reconstrua a imagem usando apenas as condições semânticas e os tokens de máscara, forçando a codificação de identidade de categoria e estrutura global no prefixo.
• Aprendizado de Prefixo Consciente de Semântica: O resultado é uma sequência de tokens ordenada por informação: o prefixo (condição + primeiros tokens) codifica a identidade e estrutura global, enquanto os tokens subsequentes refinam detalhes de nível de instância.

B. CARD (Hybrid Causal AutoRegressive–Diffusion Generator)

Baseado no espaço latente do SMAP, os autores propõem o CARD, um gerador híbrido que combina modelagem autoregressiva e difusão.

• Arquitetura Híbrida:
  1. Módulo Autoregressivo: Um Transformer causal modela as dependências estruturais de alto nível entre os tokens latentes, capturando a estrutura global e interações de longo alcance.
  2. Refinamento Contínuo: Um modelo de densidade contínua (baseado em Flow Matching) refina as previsões do módulo autoregressivo, adicionando detalhes e melhorando a fidelidade.
• Compartilhamento de Condição Semântica: O CARD reutiliza diretamente os embeddings semânticos aprendidos durante o pré-treinamento do SMAP como entrada de condição para a geração. Isso elimina a necessidade de um codificador de condição separado e garante consistência semântica entre a tokenização e a geração.

3. Principais Contribuições

1. Identificação de uma Limitação Central: O artigo argumenta que o alinhamento semântico solto é insuficiente; a semântica deve ser uma componente funcional e necessária do aprendizado de representação, não apenas um sinal de regularização.
2. SMAP: Um tokenizador 1D que utiliza injeção de condições e descarte progressivo de tokens para aprender sequências latentes com base semântica forte, ordenadas por informação e adaptáveis em comprimento.
3. CARD: Um gerador híbrido autoregressivo-difusão que explora a estrutura ordenada do SMAP, demonstrando que a tokenização baseada em semântica melhora tanto a reconstrução quanto a geração downstream.
4. Eficiência: Demonstração de que é possível obter alta qualidade de geração com orçamentos de tokens muito compactos (ex: 128 tokens) graças à eficiência do espaço latente semântico.

4. Resultados Experimentais

Os experimentos foram conduzidos no conjunto de dados ImageNet-1K (256x256).

• Reconstrução: O SMAP superou consistentemente os tokenizadores existentes (como TiTok e TA-TiTok) em métricas de FID de reconstrução (rFID) e Pontuação de Inception (IS), tanto em configurações discretas (VQ) quanto contínuas (KL, SoftVQ). O ganho foi particularmente notável sob orçamentos de tokens limitados.
• Geração: O sistema combinado SMAP + CARD alcançou desempenho competitivo contra modelos autoregressivos e baseados em difusão de última geração.
  • Com um orçamento de apenas 128 tokens, o modelo SMAP(KL) + CARD (1.1B parâmetros) alcançou um gFID de 1.85 (com CFG) e 325.1 de IS.
  • O modelo SMAP(SoftVQ) + CARD (1.1B) alcançou um gFID de 1.79 e 328.9 de IS.
  • Esses resultados foram obtidos com significativamente menos tokens do que métodos que utilizam 256 a 4096 tokens, evidenciando a eficiência da representação.
• Ablação: Estudos mostraram que a estratégia de descarte de tokens de cauda é crucial para forçar a dependência semântica. Além disso, o uso de embeddings semânticos compartilhados entre tokenizador e gerador trouxe melhorias adicionais, embora a principal melhoria venha do próprio pré-treinamento semântico do tokenizador.

5. Significado e Impacto

Este trabalho redefine a abordagem para o aprendizado de representações latentes em geração de imagens. Ao tornar a semântica um componente funcional e indispensável do processo de tokenização (e não apenas um sinal de regularização), os autores conseguem criar um espaço latente que é:

1. Mais Eficiente: Permite alta qualidade de geração com menos tokens.
2. Mais Estruturado: Separa claramente a identidade global (prefixo) dos detalhes locais (tokens subsequentes).
3. Mais Transferível: O espaço latente aprendido facilita o treinamento de modelos generativos downstream.

A proposta de SMAP e CARD sugere um novo paradigma onde a compreensão semântica é integrada desde a fase de compressão de dados, melhorando todo o pipeline de geração de imagens latentes.