Concept-Centric Token Interpretation for Vector-Quantized Generative Models

1. Problema e Motivação

Os Modelos Gerativos Vetoriais Quantizados (VQGMs), como VQGAN e DALL-E, tornaram-se ferramentas poderosas para geração de imagens de alta qualidade, utilizando representações de espaço latente discreto. No entanto, um componente fundamental desses modelos — o livro de códigos (codebook) de tokens discretos — permanece uma "caixa preta".

O problema central abordado pelo artigo é a falta de compreensão sobre quais tokens específicos são críticos para gerar um determinado conceito visual (ex: "médico", "gato", "cor vermelha").

• Limitações de abordagens anteriores: Métodos simples que selecionam tokens frequentes em imagens geradas para um conceito tendem a capturar elementos de contexto ou fundo (ruído), em vez das características essenciais do objeto.
• Consequências: A incapacidade de distinguir tokens essenciais de irrelevantes dificulta a identificação de vieses (ex: viés racial ou de gênero na representação profissional) e impede a edição de imagens precisa e controlada.

2. Metodologia: Framework CORTEX

O artigo propõe o CORTEX (Concept-Oriented Token Explanation), um framework que utiliza o Princípio do Gargalo de Informação (Information Bottleneck) para interpretar VQGMs. A ideia central é treinar um módulo que "inverte" o fluxo de informação do gerador: em vez de mapear conceitos para imagens, o modelo mapeia tokens de volta para conceitos semânticos, forçando a extração apenas das informações mais relevantes.

O framework consiste em três componentes principais:

A. Extrator de Informação (Information Extractor - IEM)

É um módulo treinado supervisionadamente que recebe as embeddings baseadas em tokens de uma imagem e prevê a probabilidade de presença de conceitos semânticos.

• Função: Atua como um classificador que aprende a reconhecer padrões de tokens específicos para cada conceito, descartando informações irrelevantes (ruído de fundo).
• Treinamento: Utiliza imagens geradas pelo VQGM e seus tokens correspondentes para treinar o IEM a prever rótulos (ex: classes do ImageNet).

B. Método de Explicação em Nível de Amostra (Sample-level Explanation)

Analisa a importância de tokens em imagens individuais geradas.

1. Calcula um Score de Importância de Token (TIS) usando gradientes (saliency maps) em relação ao IEM treinado.
2. Identifica os top-$n$ tokens com maior TIS para uma imagem específica.
3. Agrega esses resultados em múltiplas imagens para encontrar os tokens mais frequentes e consistentes para um conceito.

• Objetivo: Revelar quais tokens formam o conceito em imagens reais geradas.

C. Método de Explicação em Nível de Código (Codebook-level Explanation)

Explora todo o espaço do codebook para encontrar combinações de tokens que melhor representam um conceito, sem depender de imagens geradas previamente.

1. Define uma matriz de seleção de tokens $P$ sobre todo o codebook.
2. Utiliza Gumbel-Softmax para tornar a seleção de tokens discretos diferenciável, permitindo otimização baseada em gradiente.
3. Otimiza a matriz $P$ para maximizar a probabilidade de um conceito alvo no IEM, enquanto aplica uma máscara para focar em regiões específicas.

• Objetivo: Descobrir as combinações fundamentais de tokens que definem um conceito no espaço latente do modelo.

3. Contribuições Principais

1. Método de Explicação em Nível de Amostra: Identifica combinações de tokens relevantes para conceitos em imagens geradas individuais, fornecendo insights iniciais sobre a representação visual.
2. Método de Explicação em Nível de Codebook: Estende a análise para todo o espaço do codebook, descobrindo combinações de tokens fundamentais que caracterizam conceitos, utilizando otimização direta.
3. Validação e Aplicações Práticas:
   • Demonstra que os tokens selecionados pelo CORTEX são realmente críticos (ao mascará-los, a precisão de classificadores pré-treinados cai drasticamente).
   • Habilita detecção de vieses (ex: mostrar que tokens associados a "médicos brancos" aparecem 4x mais frequentemente que "médicos negros" em prompts neutros).
   • Permite edição de imagem direcionada, modificando regiões específicas de uma imagem alterando apenas os tokens relevantes para um novo conceito.

4. Resultados Experimentais

Os autores validaram o CORTEX em diversos modelos VQGM pré-treinados e utilizando classificadores de estado da arte (ResNet, Vision Transformer) como benchmarks.

• Avaliação de Mascaramento (Counterfactual):
  • Quando os tokens identificados pelo CORTEX foram mascarados (removidos) das imagens, a probabilidade de classificação do conceito alvo caiu significativamente mais do que quando tokens aleatórios foram mascarados.
  • Exemplo: Ao mascarar 50 tokens de alta importância, modelos como ResNet50 tiveram quedas de precisão de ~46%, enquanto o baseline aleatório teve quedas menores.
• Detecção de Viés:
  • Em prompts neutros como "um médico no hospital", o modelo gerou imagens onde tokens associados a "médicos brancos" apareceram em média 8,55 vezes por imagem, contra 2,32 vezes para "médicos negros". O CORTEX quantificou esse viés sistêmico.
• Edição de Imagem:
  • O método de nível de codebook permitiu transformar gradualmente uma espécie de pássaro em outra (ex: de um "Indigo Bird" para outro) alterando apenas os tokens na região da cabeça, mantendo o resto da imagem intacto.

5. Significado e Impacto

O trabalho CORTEX é significativo por várias razões:

• Transparência: Oferece uma janela para o processo interno de modelos gerativos baseados em tokens, que eram anteriormente difíceis de interpretar em nível semântico.
• Segurança e Ética: Fornece uma ferramenta robusta para detectar e quantificar vieses demográficos e sociais em modelos de IA generativa, permitindo intervenções mais justas.
• Controle Criativo: Permite a edição precisa de imagens geradas, dando aos usuários controle sobre conceitos específicos sem re-treinar o modelo.
• Avanço Teórico: Aplica o Princípio do Gargalo de Informação ao domínio de modelos generativos discretos, estabelecendo uma nova linha de pesquisa para interpretabilidade em VQGMs.

Em suma, o CORTEX transforma a interpretação de modelos generativos de uma análise visual superficial para uma análise estrutural baseada em tokens, permitindo entender, auditar e controlar como esses modelos "pensam" e geram conceitos visuais.