Purrception: Variational Flow Matching for Vector-Quantized Image Generation

Resumo Técnico: PURRCEPTION

1. O Problema

A geração de imagens de alta resolução em espaços latentes vetor-quantizados (VQ) apresenta um dilema fundamental de modelagem devido à natureza dual desses latentes: eles são simultaneamente índices discretos (escolhidos de um código-book finito) e vetores de incorporação contínuos (que capturam relações geométricas como proximidade e direção).

As abordagens existentes forçam uma escolha entre dois extremos degenerados:

• Métodos Contínuos (ex: Diffusion Latente, Flow Matching - CFM): Operam no espaço de incorporação contínuo ($\mathbb{R}^D$). Embora preservem a geometria e permitam transporte suave, eles ignoram a estrutura categórica. O modelo nunca recebe um sinal de aprendizado categórico, não consegue expressar incerteza sobre múltiplos códigos plausíveis e carece de logits para controle via temperatura.
• Métodos Discretos (ex: VQ-Diffusion, Discrete Flow Matching - DFM): Predizem índices diretamente. Embora alinhados com a estrutura quantizada, eles colapsam a geometria. Códigos semanticamente relacionados são tratados como tokens não relacionados, resultando em "teletransportes" discretos entre índices, eliminando a interpolação e tornando o controle de temperatura e a modelagem de incerteza ineficazes.

O objetivo é criar um modelo que aproveite a geometria suave dos vetores de incorporação enquanto mantém a supervisão categórica sobre os índices.

2. Metodologia: Purrception

O Purrception propõe uma adaptação do Variational Flow Matching (VFM) para latentes vetor-quantizados. A abordagem híbrida combina a dinâmica de transporte contínuo com a supervisão categórica.

• Posterior Categórica Variacional: Em vez de prever diretamente o campo de velocidade ou o índice final, o modelo aprende uma distribuição posterior variacional $q_\theta(c | z_t)$ sobre os índices do código-book ($c$) dado o ponto de interpolação atual ($z_t$).
• Campo de Velocidade Híbrido: O campo de velocidade é calculado como a esperança dessa distribuição categórica sobre os vetores de incorporação do código-book:
  $$v_\theta(z_t) = \sum_{k=1}^K \pi_{\theta, k}(z_t) \left( \frac{e_k - z_t}{1-t} \right)$$
  Onde $\pi_{\theta, k}$ é a probabilidade de escolher o vetor $e_k$. Isso permite que a incerteza sobre múltiplos códigos plausíveis se traduza em um movimento suave e consciente da geometria, em vez de saltos discretos.
• Objetivo de Treinamento: O treinamento minimiza a perda de Entropia Cruzada entre a distribuição categórica predita e os índices de código verdadeiros (ground-truth). Isso fornece um sinal de aprendizado categórico direto.
• Controle por Temperatura: Como o modelo opera com logits antes da softmax, ele herda naturalmente o controle de temperatura ($\tau$) no momento da inferência:
  • Baixa Temperatura ($\tau \downarrow$): A distribuição colapsa para o índice mais provável, gerando imagens mais nítidas e fiéis, mas potencialmente menos variadas.
  • Alta Temperatura ($\tau \uparrow$): A distribuição se alarga, atribuindo peso a códigos vizinhos, o que adiciona detalhes e variabilidade, mas pode reduzir a fidelidade global.

3. Contribuições Principais

1. Arquitetura Híbrida: Introdução do Purrception, que une a eficiência do Flow Matching contínuo com a supervisão categórica necessária para latentes VQ, resolvendo o trade-off entre preservação geométrica e estrutura discreta.
2. Eficiência de Treinamento: Demonstração de que a supervisão categórica direta acelera a convergência do modelo em comparação com métodos puramente contínuos ou discretos.
3. Controle de Inferência: Estabelecimento de que a temperatura de softmax pode ser usada como um "botão" controlável para equilibrar nitidez e diversidade durante a geração, uma característica ausente em métodos puramente contínuos e ineficaz em métodos puramente discretos.
4. Desempenho SOTA: Alcançamento de resultados competitivos no benchmark ImageNet-1k (256x256), superando modelos de difusão discreta e mascarada, e competindo com modelos de difusão contínua de ponta, com menor custo computacional de treinamento.

4. Resultados Experimentais

Os experimentos foram conduzidos no ImageNet-1k (256x256) utilizando backbones DiT-L/2 e DiT-XL/2.

• Velocidade de Convergência:
  • O Purrception convergiu significativamente mais rápido que as bases de referência (CFM e DFM).
  • Com o backbone DiT-L/2, atingiu o desempenho final do DFM em apenas 325k iterações (3.0x mais rápido) e o do CFM em 1.2M iterações (1.65x mais rápido).
  • Com o backbone DiT-XL/2, foi 2.3x mais rápido que o CFM e 3.5x mais rápido que o DFM.
• Qualidade da Amostra (FID):
  • O modelo alcançou um FID de 3.88 (com $\tau=0.9$), superando todos os modelos de difusão discreta e geradores mascarados comparados na tabela.
  • Embora fique ligeiramente atrás dos modelos de difusão contínua de ponta (como SiT-XL/2 com FID ~2.06), o Purrception se destaca por ser um modelo VQ nativo e por sua eficiência de treinamento.
• Efeito da Temperatura:
  • Uma análise de ablação mostrou uma relação em forma de "U" entre a temperatura e o FID. O desempenho ótimo ocorreu em $\tau \approx 0.8 - 0.9$.
  • Temperaturas muito baixas geraram imagens simplificadas demais; temperaturas muito altas geraram ruído e incoerência.

5. Significado e Conclusão

O Purrception representa um avanço significativo na geração de imagens latentes, demonstrando que é possível pontear o transporte contínuo e a supervisão discreta. Ao tratar o problema como uma inferência variacional sobre índices discretos dentro de um espaço contínuo, o modelo aprende "o que escolher" (supervisão categórica) e "para onde ir" (geometria contínua).

Isso permite:

• Quantificação de incerteza sobre códigos plausíveis.
• Treinamento mais eficiente e rápido.
• Controle flexível na inferência via temperatura.

Embora limitado atualmente pela dependência de um autoencoder VQ pré-treinado fixo, o trabalho abre caminho para futuras pesquisas em modelos híbridos para áudio, vídeo e formas 3D, além de sugerir a possibilidade de destilação para amostradores de poucos passos, mantendo a eficiência dos fluxos contínuos.