Circuit Mechanisms for Spatial Relation Generation in Diffusion Transformers

Título: Mecanismos de Circuito para Geração de Relações Espaciais em Diffusion Transformers

1. Problema Investigado

Os modelos de geração de imagens baseados em Diffusion Transformers (DiT) têm avançado significativamente na síntese de imagens de alta fidelidade a partir de prompts de texto. No entanto, eles ainda enfrentam dificuldades crônicas ao gerar relações espaciais corretas entre múltiplos objetos (ex: "um quadrado vermelho à esquerda de um círculo azul"). Embora a geração de atributos individuais (cor, forma) tenha melhorado, a composição espacial permanece um gargalo.

A questão central do trabalho é: Como os DiTs implementam internamente o raciocínio espacial? A literatura anterior atribuiu falhas a estágios específicos (como atenção cruzada ou codificação de texto), mas não havia uma compreensão mecanicista de como as redes neurais codificam e utilizam relações não comutativas entre objetos.

2. Metodologia

Os autores adotaram uma abordagem de interpretabilidade mecanicista, projetando um ambiente controlado para isolar e analisar os circuitos internos do modelo.

• Dataset Sintético: Criaram um conjunto de dados minimalista onde o modelo deve gerar duas formas geométricas (círculo, quadrado, triângulo) com cores específicas (vermelho, azul) e uma das oito relações espaciais possíveis (ex: "acima", "esquerda", "diagonal inferior direita").
• Arquitetura do Modelo: Utilizaram a arquitetura PixArt (baseada em DiT) treinada do zero em várias escalas (Nano a B).
• Variável Experimental (Codificadores de Texto): O ponto crucial do estudo foi comparar dois tipos de codificadores de texto:
  1. RTE (Random Token Embeddings): Embeddings aleatórios com codificação posicional sinusoidal, sem estrutura semântica prévia.
  2. T5 (Pretrained): Um codificador de linguagem natural pré-treinado (T5-XXL) que gera embeddings contextuais e semanticamente misturados.
• Técnicas de Análise:
  • Resumo de Atenção (Attention Synopsis): Agregação de mapas de atenção cruzada através de etapas de amostragem e categorias de tokens (objeto vs. relação) para identificar padrões estáveis.
  • Rastreamento de Circuitos: Identificação de cabeças de atenção específicas responsáveis por transmitir informações.
  • Manipulação Causal: Ablação de cabeças de atenção e manipulação vetorial de embeddings para testar a causalidade.
  • Particionamento de Variância: Análise estatística para decompor como a informação espacial e de objeto está codificada nos vetores de embedding do T5.

3. Contribuições e Descobertas Principais

O estudo revela que, embora ambos os modelos alcancem alta precisão na tarefa, eles utilizam mecanismos internos drasticamente diferentes dependendo do codificador de texto:

A. Mecanismo no DiT com Embeddings Aleatórios (RTE-DiT)

O modelo com RTE desenvolve um circuito modular de dois estágios:

1. Cabeça de Relação Espacial (Spatial Relation Head): Uma cabeça de atenção cruzada especializada (ex: L2H8) lê os tokens de relação ("acima", "esquerda") e interage com as codificações posicionais dos tokens de imagem. Isso gera gradientes espaciais (mapas de calor) que marcam onde o objeto deve aparecer.
2. Cabeça de Geração de Objeto (Object Generation Head): Uma segunda cabeça (ex: L4H3) lê os tokens de forma ("quadrado", "círculo") e, crucialmente, lê as "tags" posicionais geradas pela primeira cabeça. Ela atribui a forma correta à localização marcada.

• Conclusão: O processo é disjunto: primeiro define-se o local, depois preenche-se o objeto.

B. Mecanismo no DiT com Codificador T5 (T5-DiT)

O modelo com T5 utiliza um mecanismo baseado no objeto (object-centric):

• Fusão de Informação: Devido à auto-atenção do T5, a informação espacial não fica isolada em tokens de relação. Em vez disso, o contexto da frase inteira é fundido nos embeddings dos tokens de objeto (especificamente o segundo objeto, shape2).
• Decodificação Única: O DiT não precisa de cabeças especializadas para ler palavras de relação separadamente. Ele decodifica tanto a identidade do objeto quanto a relação espacial diretamente do embedding contextual do token do objeto.
• Evidência: A ablação de palavras de relação no prompt não afeta a geração, mas a manipulação vetorial do embedding do token do objeto altera diretamente a posição gerada.

4. Resultados e Robustez

• Desempenho In-Domain: Ambos os modelos (RTE e T5) alcançam alta precisão na tarefa sintética quando o prompt segue o formato de treinamento.
• Robustez a Perturbações (Out-of-Domain):
  • O RTE-DiT é robusto a pequenas variações no prompt (ex: adicionar a palavra "o" ou "the").
  • O T5-DiT é extremamente frágil. A adição de uma palavra de preenchimento (ex: "the") ou inversão da ordem da frase causa uma queda drástica na precisão espacial (colapso de ~80% para ~40%).
  • Causa: No T5, a adição de uma palavra altera o embedding contextual do token do objeto, deslocando-o fora da direção vetorial aprendida para a relação espacial correta, confundindo o mecanismo de geração.

5. Significado e Implicações

1. Reinterpretação de Falhas: O trabalho sugere que a dificuldade dos modelos atuais em gerar relações espaciais pode não ser um defeito do DiT em si, mas sim uma limitação do modelo de embedding de texto (como CLIP ou T5) em preservar informações espaciais de forma robusta e desacoplada.
2. Design de Arquitetura: Para aplicações que exigem controle espacial rigoroso, embeddings explícitos e modulares (como os aprendidos no RTE) podem oferecer mecanismos de controle mais claros e robustos do que embeddings contextuais densos.
3. Ferramentas de Análise: O artigo introduz métodos sistemáticos (como Attention Synopsis e Weight-Space Screening) para mapear circuitos em modelos de difusão, superando a complexidade da análise de mapas de atenção iterativos.
4. Generalização: A análise foi validada em modelos pré-treinados (PixArt-Sigma) e com outros codificadores (CLIP), mostrando que os circuitos espaciais existem mesmo em modelos complexos, embora sejam menos disjuntos no T5/CLIP.

Em resumo, o paper demonstra que a "inteligência" espacial em DiTs não é monolítica; ela emerge de circuitos distintos dependendo de como a linguagem é codificada, revelando um trade-off entre a eficiência da integração contextual (T5) e a robustez/interpretabilidade do raciocínio modular (RTE).