Resumo Técnico: MergeTok

Declaração do Problema

Os atuais tokenizadores visuais para geração de imagens estão bifurcados em duas famílias distintas, cada uma sofrendo de limitações estruturais complementares:

1. Tokenizadores baseados em VAE contínuo: Embora ofereçam reconstrução de alta fidelidade, dependem de espaços latentes que carecem de organização intrínseca. Sob compressão forte, os fatores semânticos tornam-se emaranhados dentro de representações latentes densas, resultando em códigos fracamente desentranhados que dificultam o controle semântico.
2. Tokenizadores baseados em VQ discreto: Embora permitam a geração autorregressiva, enfrentam gargalos de otimização. A natureza não diferenciável da quantização leva a atualizações de gradiente esparsas e desiguais, fazendo com que muitas entradas do código (codebook) recebam supervisão insuficiente. Isso aumenta o risco de colapso do codebook e reduz a eficiência da representação.

Essas limitações criam uma tensão persistente entre fidelidade de reconstrução, estrutura semântica, eficiência e estabilidade de treinamento.

Metodologia: Estrutura MergeTok

O MergeTok propõe um tokenizador de braço duplo unificado que otimiza conjuntamente a tokenização contínua (VAE) e discreta (VQ) dentro de uma única arquitetura encoder-decoder. A inovação central é o uso de Token Merging (ToMe) como uma ponte semântica entre os dois braços.

Visão Geral da Arquitetura

Dado um input de imagem $X$, um codificador CNN compartilhado $E_c$ produz um mapa de características $Z$. Este mapa é processado por dois braços paralelos:

1. Braço VAE (Contínuo):

   • Token Merging: Um codificador baseado em atenção equipado com módulos ToMe comprime a sequência de tokens de $L$ tokens para $K$ tokens ($K < L$) ao fundir tokens semelhantes. Este processo gera uma representação condensada $Z^{(vae)}_K$ e um mapa de origem $S$ que registra a atribuição dos tokens originais aos grupos fundidos.
   • Alinhamento Semântico: Os tokens fundidos são alinhados com um modelo professor do tipo DINO congelado (também equipado com ToMe) usando uma perda de distância de cosseno ($L_{align}$). Isso regulariza o espaço latente contínuo em direção a representações semanticamente conscientes e desentranhadas.
   • Reconstrução: Um decoder híbrido utiliza tanto os tokens condensados quanto o mapa de origem $S$ para "desfundir" os grupos e reconstruir a imagem no espaço de pixels.

2. Braço VQ (Discreto):

   • Quantização Consciente de Grupo: Este braço ignora a operação ToMe durante o passe direto (forward pass) para preservar o comprimento total da sequência ($L$ tokens). No entanto, ele utiliza o mapa de origem $S$ gerado pelo braço VAE como um prior estrutural.
   • Regularização: O mapa de origem particiona os $L$ tokens em $K$ grupos semânticos. Duas perdas conscientes de grupo são aplicadas às atribuições do codebook:
     • Diversidade Intra-grupo ($L_{div}$): Maximiza a entropia dentro de cada grupo para evitar o colapso para um único código.
     • Exclusividade Inter-grupo ($L_{cons}$): Minimiza a sobreposição entre grupos para garantir que diferentes grupos semânticos utilizem subconjuntos de códigos distintos.
   • Reconstrução: Um decoder VQ padrão reconstrói a imagem a partir dos tokens quantizados.

Estratégia de Treinamento

• Objetivo Unificado: O modelo é treinado de ponta a ponta com uma perda total combinando regularizadores de reconstrução, adversariais e de fusão consciente:
  $$L_{total} = \lambda_{vae}L_{VAE} + \lambda_{vq}L_{VQ} + \lambda_{ali}L_{align} + \lambda_{div}L_{div} + \lambda_{cons}L_{cons}$$
• Amostragem Consciente de Granularidade: Durante o treinamento, a razão de fusão (e, portanto, o número de tokens mantidos $K$) é amostrada dinamicamente de uma distribuição discreta centrada em torno de um hiperparâmetro $K^*$. Isso expõe o encoder a variações na largura do gargalo, incentivando a robustez através de diferentes granularidades semânticas.
• Inferência: Na inferência, o tokenizador sempre produz uma sequência fixa de 256 tokens. Geradores downstream (tanto autorregressivos quanto de difusão) operam nesta sequência fixa e são agnósticos à estrutura interna de fusão.

Principais Contribuições

1. Tokenizador de Braço Duplo Unificado: O MergeTok é a primeira arquitetura a otimizar conjuntamente latentes contínuos e discretos sob um único objetivo, usando o token merging como uma interface semântica para reforçar a estabilidade do treinamento VAE e as vantagens estruturais do VQ.
2. Restrições de Treinamento Conscientes de Fusão: Os autores introduzem dois novos objetivos derivados do processo de fusão:
   • Alinhamento de token fundido: Melhora a estrutura semântica no espaço latente contínuo.
   • Quantização consciente de grupo: Estabiliza o treinamento VQ e melhora a utilização do codebook via diversidade intra-grupo e exclusividade inter-grupo.
3. Amostragem de Razão de Fusão Consciente de Granularidade: Uma estratégia de tempo de treinamento que amostra diferentes granularidades de tokens, permitindo que o modelo aprenda uma tokenização unificada que é robusta através de diferentes taxas de compressão.

Resultados Experimentais

Experimentos foram conduzidos no ImageNet-256 (resolução 256×256) com um orçamento fixo de 256 tokens.

• Fidelidade de Reconstrução:
  • Discreto (VQ): O MergeTok-BL alcançou um rFID de 0.78, superando fortes baselines como GigaTok-BL (0.81) e MergeVQ-GR (1.12). Com um discriminador DINO congelado, o MergeTok-BL* atingiu 0.48 de rFID.
  • Contínuo (VAE): O MergeTok-BL alcançou um rFID de 0.47, superando marginalmente o RAE (0.49) e vencendo significativamente o SoftVQ-L (0.61) e o DC-AE-f32 (0.69).
• Qualidade de Representação:
  • O MergeTok-BL obteve 78.3% de acurácia em linear probing, a mais alta entre os tokenizadores discretos listados, superando o UniTok (70.8%) e o MergeVQ-GR (77.9%). Isso indica uma organização semântica superior.
• Geração Downstream:
  • Autorregressiva (Braço VQ): Emparelhado com LlamaGen-XXL, o MergeTok-BL alcançou um gFID de 1.93 (sem CFG) e 2.13 (com CFG), superando o GigaTok-BL (2.03).
  • Difusão/Flow (Braço VAE): Emparelhado com SiT-XL/1, o MergeTok-BL alcançou um gFID de 1.79 (sem CFG) e 1.29 (com CFG), igualando os melhores registros da tabela.

Significância e Alegações

O artigo afirma que o MergeTok demonstra que uma única arquitetura pode dotar os tokenizadores visuais tanto de uma organização semântica robusta quanto de uma discreticidade amigável ao gerador.

• Unificação: Ele resolve o trade-off entre a alta fidelidade de reconstrução dos VAEs e a estrutura discreta necessária para a geração autorregressiva.
• Ponte Semântica: Ao aproveitar o token merging como um prior estrutural, o modelo transfere com sucesso informações de agrupamento semântico do braço contínuo para estabilizar o braço discreto, e vice-versa.
• Compatibilidade: O tokenizador produz sequências de 256 tokens que são compatíveis tanto com geradores autorregressivos (ex: LlamaGen) quanto com geradores de difusão/flow (ex: DiT, SiT) sem exigir modificações nos próprios geradores.

Os autores posicionam o MergeTok não meramente como uma ferramenta de reconstrução, mas como uma interface unificada que suporta simultaneamente a geração autorregressiva discreta, a geração de difusão contínua e o aprendizado de representação semanticamente estruturado.