Variable-Length Tokenization via Learnable Global Merging for Diffusion Transformers

Resumo Técnico: Tokenização de Comprimento Variável via Fusão Global Aprendível para Transformers de Difusão

1. Declaração do Problema

Modelos de Difusão Latente (LDMs) dependem de tokenizadores para comprimir imagens em espaços latentes, criando um compromisso entre qualidade de geração e custo computacional. Tokenizadores convencionais utilizam uma razão de compressão fixa, o que exige múltiplas variantes de modelos para acomodar diferentes orçamentos computacionais ou requisitos de qualidade.

Tokenizadores de Comprimento Variável (VLTs) foram propostos para abordar isso, ajustando adaptativamente o número de tokens com base em condições específicas (ex: orçamento computacional). O método predominante, dropout aninhado (nested dropout), treina modelos para descartar tokens de cauda de uma sequência ordenada, posicionando semânticas de alto nível nos primeiros tokens e detalhes de baixo nível nos tokens posteriores.

No entanto, os autores identificam uma falha crítica no dropout aninhado quando aplicado a Transformers de Difusão (DiTs):

• Deslocamento de Comprimento Cruzado (Cross-Length Shift): Truncar uma sequência ordenada altera a estrutura de similaridade pareada (alinhamento representacional) do espaço latente. Sequências mais curtas focam em semânticas centrais, enquanto sequências mais longas incluem detalhes residuais.
• Descompasso Distribucional: Esse deslocamento faz com que a distribuição latente e a respectiva função de score variem significamente entre comprimentos de tokens. Consequentemente, um único modelo de difusão treinado conjuntamente em vários comprimentos falha em generalizar efetivamente, muitas vezes performando pior do que modelos treinados para comprimentos específicos.
• Limitação de Soluções Existentes: Perdas de alinhamento de representação padrão exigem correspondência de dimensões de tokens (N = M), o que é impossível em configurações de comprimento variável onde a cardinalidade difere.

2. Metodologia

O artigo propõe um novo tokenizador de comprimento variável baseado em fusão (merging-based) que modula o comprimento combinando tokens em vez de truncá-los. A inovação central é a Fusão Global Aprendível (LGM - Learnable Global Merging), projetada para garantir o alinhamento representacional entre diferentes contagens de tokens, mantendo a compatibilidade com modelos generativos.

A. Modulação de Comprimento Baseada em Fusão
Em vez de descartar tokens, o método funde tokens com base em uma matriz de atribuição $\Gamma$.

• Atenção Proporcional: Para lidar com tokens fundidos, o decodificador e o transformer de difusão utilizam atenção proporcional, onde o escore de atenção é ajustado pelo tamanho efetivo do cluster ($m_j$) de cada token fundido.
• Representação Equivalente de Comprimento Total: Os autores demonstram que processar latentes modulados com atenção proporcional é matematicamente equivalente a processar latentes de comprimento total projetados ($Wz$) com atenção padrão. Essa equivalência permite medir e minimizar diretamente a diferença entre os latentes originais ($z$) e os latentes fundidos ($Wz$), efetivamente minimizando o deslocamento representacional ao incentivar a fusão de tokens similares.

B. Fusão Global Aprendível (LGM)
Surge um desafio porque a fusão padrão (ex: agrupamento aglomerativo) é dependente de dados, derivando padrões de fusão a partir da imagem de entrada. Isso é incompatencial com a geração por difusão, onde a imagem alvo não está disponível.

• Independência de Dados: Os autores introduzem Embeddings Globais Aprendíveis ($e$) que são independentes dos dados de entrada. Esses embeddings determinam o padrão de fusão para um determinado alvo de redução.
• Otimização: O padrão de fusão é derivado via agrupamento aglomerativo nos embeddings aprendíveis. Como o agrupamento envolve operações discretas, um truque straight-through é aplicado para permitir a otimização baseada em gradiente dos embeddings.
• Perda de Alinhamento: Uma perda de alinhamento ($L_{align}$) é introduzida para garantir que a estrutura de similaridade dos embeddings globais aprendíveis ($e$) corresponda à estrutura de similaridade dos tokens latentes reais ($z$). Isso guia o padrão de fusão global para agrupar tokens similares, minimizando assim o deslocamento representacional definido no objetivo de fusão.

C. Integração com Transformer de Difusão

• Embeddings Posicionais: O método utiliza embeddings posicionais fundidos, computados da mesma forma que os latentes, para lidar com a informação posicional dos tokens agregados.
• Pós-treinamento: Um passo de pós-treinamento leve de LoRA (Low-Rank Adaptation) é introduzido. Após o treinamento conjunto em vários comprimentos, parâmetros LoRA específicos de comprimento são treinados por algumas épocas para adaptar ainda mais o modelo a contagens de tokens específicas com overhead negligenciável.

3. Principais Contribuições

1. Tokenização Baseada em Fusão: Uma nova arquitetura VLT que modula o comprimento via fusão de tokens, permitindo o alinhamento direto da representação entre diferentes contagens de tokens ao minimizar o deslocamento entre latentes originais e fundidos.
2. Fusão Global Aprendível: Um mecanismo de fusão independente de dados que resolve a incompatibilidade entre a fusão dependente de dados e a geração por difusão, garantindo que o padrão de fusão seja acessível durante a inferência.
3. Melhoria na Relação Qualidade-Computação: A integração deste tokenizador com um Transformer de Difusão alcança um desempenho superior no equilíbrio entre qualidade de geração (gFID) e custo computacional em comparação com métodos VLT existentes.

4. Resultados Experimentais

Experimentos foram conduzidos na geração condicional de classe ImageNet 256×256.

• Alinhamento Representacional: O método proposto alcança pontuações de Centered Kernel Nearest-Neighbor Alignment (CKNNA) significativamente mais altas em diferentes comprimentos de tokens comparado ao dropout aninhado. Notavelmente, o alinhamento é frequentemente superior ao de tokenizadores de comprimento fixo treinados de forma independente.
• Desempenho de Treinamento Conjunto:
  • Dropout Aninhado: Apresenta uma grande lacuna de desempenho (0.8–1.7 gFID) entre modelos treinados conjuntamente em vários comprimentos e aqueles treinados separadamente.
  • Método Proposto: Reduz essa lacuna para menos de 0.3 gFID. Com o pós-treinamento LoRA, o modelo conjunto alcança um desempenho dentro de 0.05 gFID dos modelos treinados separadamente.
• Comparação de Sistema:
  • O método supera os VLTs existentes (Semanticist, FlexTok) e baselines de comprimento fixo (SoftVQ, MAETok) em termos de relação gFID–computação.
  • Alcança qualidade de geração competitiva em altas taxas de compressão (ex: 32 tokens) sem exigir um decodificador de tokenizador baseado em difusão, ao contrário de alguns baselines.
• Qualidade de Reconstrução: O tokenizador mantém métricas de reconstrução fortes (rFID, PSNR, SSIM) através de vários comprimentos de tokens, particularmente em altas taxas de compressão.

5. Significância e Alegações

O artigo alega que a principal barreira para modelos de difusão de comprimento variável eficazes é o deslocamento de distribuição latente em comprimento cruzado causado pelos métodos de truncamento convencionais. Ao mudar para uma abordagem baseada em fusão com padrões aprendíveis e independentes de dados, os autores conseguem manter o alinhamento representacional.

A significância reside em permitir que um único modelo de difusão opere efetivamente através de diversas contagens de tokens, permitindo o controle flexível de qualidade-computação sem a necessidade de múltiplas variantes de modelos. Os autores posicionam isso como um passo em direção ao deployment mais eficiente de modelos generativos em cenários com diversas restrições computacionais, observando que, embora permaneça uma lacuna entre tokenizadores 1D e 2D, sua abordagem oferece uma direção promissora para arquiteturas híbridas. O trabalho é enquadrado como uma melhoria de eficiência, e não como uma nova capacidade generativa, sem alegação de preocupações éticas.