Resumo Técnico: Desenrolando o RoPE: Controle de Frequência para Atenção Compartilhada em DiTs

Declaração do Problema

Em Transformers de Difusão (DiTs), como a arquitetura Multi-modal Diffusion Transformer (MMDiT) utilizada em modelos como o Flux, a informação posicional é codificada via Rotary Positional Embeddings (RoPE). Embora o RoPE seja essencial para modelar a estrutura espacial na geração de imagens padrão, seu comportamento em configurações de atenção compartilhada — onde uma geração de imagem alvo atende a tokens de uma imagem de referência — permanece mal compreendido.

Os autores identificam um modo de falha crítico nesta configuração: a cópia de referência (reference copying). Em vez de transferir os atributos estilísticos de uma imagem de referência para um prompt alvo (por exemplo, gerar uma girafa no estilo de um touro), o modelo colapsa na reprodução do conteúdo da referência de forma quase literal. Isso ocorre porque o forte viés espacial introduzido pelo RoPE força as consultas (queries) da imagem alvo a atenderem desproporcionalmente a tokens espacialmente alinhados na imagem de referência, sobrepondo-se à similaridade semântica.

Tentativas existentes para mitigar isso, como restringir o compartilhamento de atenção para etapas posteriores de denoising ou usar índices posicionais deslocados (como no trabalho concorrente AlignedGen), oferecem apenas um controle grosseiro ou introduzem novos artefatos como o "inpainting cíclico fantasma" (ghost-cyclic inpainting, onde o modelo trata a referência como uma continuação do alvo).

Metodologia

O artigo propõe uma modulação consciente de frequência do RoPE para restaurar a atenção compartilhada significativa sem retreinar o modelo ou alterar sua arquitetura.

1. Análise da Estrutura de Frequência do RoPE

Os autores fornecem uma análise fundamentada mostrando que o RoPE se decompõe em componentes de frequência com distintas sensibilidades posicionais:

• Componentes de alta frequência: Correspondem às dimensões de índice inferior no embedding. Estes exibem mudanças angulares rápidas em relação ao deslocamento posicional ($\Delta$), impondo uma forte localidade. Eles fazem com que a atenção caia abruptamente conforme a distância espacial aumenta.
• Componentes de baixa frequência: Correspondem às dimensões de índice superior. Estes variam lentamente com o deslocamento, contribuindo com um viés posicional fraco e preservando a coerência espacial global.

Em cenários de atenção compartilhada, os componentes de alta frequência dominam o cálculo da atenção, fazendo com que o modelo se fixe em tokens de referência espacialmente alinhados e copie seu conteúdo.

2. Modulação Consciente de Frequência

Para combater isso, os autores introduzem um método para modular seletivamente as bandas de frequência das chaves da imagem de referência ($K_{ref}$) durante o cálculo da atenção.

• Mecanismo: Em vez de desativar o RoPE (o que causa gerações degeneradas) ou escalar a atenção uniformemente, o método aplica uma interpolação polinomial suave de fatores de escala através do espectro do RoPE.
• Estratégia de Escalonamento:
  • Atenção (Atenuação): Os componentes de alta frequência são reduzidos (fator $s_{hf} \in (0, 1)$) para reduzir o forte viés posicional que impulsiona a cópia.
  • Amplificação: Os componentes de baixa frequência são amplificados (fator $s_{lf} > 1$) para promover uma atenção global e semanticamente guiada.
• Agendamento de Timestep: Os escalonamentos de modulação ($s_{hf}$ e $s_{lf}$) aumentam linearmente ao longo dos passos de denoising. Os primeiros passos de tempo utilizam menor modulação para estabelecer a estrutura global e a orientação de estilo ampla, enquanto os passos posteriores aumentam o escalonamento para permitir a transferência de detalhes estilísticos finos e texturas.

A chave de atenção final para a atenção compartilhada é formulada como:
K = [K'_{img}^{tar} \oplus K'_{txt}^{tar}] \oplus s_{FA}^t \cdot K_{img}^{ref}
onde $s_{FA}^t$ é o fator de modulação de frequência consciente no timestep $t$, aplicado exclusivamente às chaves de referência.

Principais Contribuições

1. Análise Fundamentada do RoPE em Atenção Compartilhada: O artigo demonstra que o fenômeno de "cópia de referência" em DiTs não é uma falha geral de compartilhamento de atenção, mas uma consequência específica dos componentes de alta frequência do RoPE que impõem uma estrita localidade espacial.
2. Mecanismo de Controle de Frequência: Os autores introduzem um método zero-shot, livre de treinamento, para modular as bandas de frequência do RoPE. Isso permite um equilíbrio controlável entre alinhamento posicional e similaridade semântica.
3. Resolução de Artefatos: O método resolve os problemas de "inpainting cíclico fantasma" e vazamento de conteúdo observados em abordagens anteriores (como o RoPE deslocado), abordando diretamente a estrutura de frequência subjacente em vez de apenas deslocar coordenadas.
4. Restauração da Geração Alinhada ao Estilo: A abordagem permite que os DiTs gerem conteúdos semânticos diversos (ex: diferentes animais) enquanto mantêm um estilo visual consistente de uma referência, uma capacidade anteriormente limitada a modelos baseados em UNet ou que exigia ajuste fino pesado.

Resultados

Avaliado no modelo Flux 1.-dev, o método demonstra:

• Transferência de Estilo Eficaz: O modelo transfere com sucesso atributos estilísticos (cor, textura, pinceladas) de referências (incluindo obras de arte e fotos) para prompts alvo sem duplicar o conteúdo da referência.
• Controlabilidade: Ao ajustar os escalonamentos de modulação, os usuários podem controlar o grau de transferência de estilo versus preservação de conteúdo. Um $s_{hf}$ mais alto preserva a estrutura da referência (pose/forma) enquanto transfere o estilo; um $s_{hf}$ mais baixo permite maior variação estrutural.
• Superioridade sobre Baselines:
  • Comparado ao StyleAligned (baseado em UNet), o método evita vazamento de conteúdo e irregularidades estruturais em DiTs.
  • Comparado ao AlignedGen (RoPE deslocado), o método evita artefatos de cópia periódica e inpainting cíclico fantasma, proporcionando uma transferência de estilo mais suave mesmo em altas resoluções.
  • Comparado ao IP-Adapter e B-LoRA, o método não requer treinamento adicional ou encoders de imagem, oferecendo uma solução zero-shot mais generalizável.

Significância

O artigo afirma que as codificações posicionais não são componentes neutros, mas governam ativamente o equilíbrio entre localidade e associação semântica em modelos generativos. As descobertas explicam por que estratégias de compartilhamento de atenção desenvolvidas para modelos baseados em UNet (que dependem de vieses indutivos convolucionais) falham quando aplicadas diretamente a DiTs (que dependem exclusivamente de RoPE).

Ao revelar que os componentes de alta frequência do RoPE impulsionam a cópia de conteúdo, este trabalho fornece um controle direto e interpretável para moldar o fluxo de atenção. Isso sugere uma direção mais ampla para o controle de modelos generativos através de intervenções direcionadas em representações internas, permitindo a geração de estilo alinhado com precisão sem mudanças arquitetônicas ou retreinamento.