Resumo Técnico: SEGA – Atenção Guiada por Energia Espectral para Extrapolacão de Resolução em Transformers de Difusão

Declaração do Problema

Os Transformers de Difusão (DiTs) estabeleceram-se como a arquitetura dominante para geração de imagem a partir de texto (T2I). No entanto, seu desempenho degrada-se significativamente ao gerar imagens em resoluções além de sua faixa de treinamento (tipicamente 1024² a 2048²). A extrapolação para resoluções mais altas frequentemente resulta em texturas borradas, padrões repetitivos e colapsos estruturais.

Soluções existentes sem treinamento tentam mitigar isso modificando o comportamento da atenção no momento da inferência, principalmente através da extrapolação de Embeddings de Posição Rotativos (RoPE) combinada com escalonamento da atenção. Métodos atuais, como o YaRN, aplicam um escalonamento uniforme e agnóstico ao conteúdo em todos os componentes do RoPE. Essa abordagem trata componentes de comprimento de onda curto (que governam texturas de alta granularidade) de forma idêntica aos componentes de comprimento de onda longo (que moldam a estrutura global). O artigo argumenta que essa uniformidade estática induz uma compensação fundamental: estratégias que preservam a estrutura global frequentemente falham em recuperar detalhes finos, e vice-versa. Além disso, a distribuição espectral da representação latente evolui dinamicamente durante o processo de remoção de ruído e varia entre diferentes conteúdos de imagem, tornando o escalonamento estático subótimo.

Metodologia: SEGA

Os autores introduzem o SEGA (Atenção Guiada por Energia Espectral), um método sem treinamento que escala dinamicamente a atenção através dos componentes do RoPE com base na estrutura de frequência espacial do latente em cada etapa de remoção de ruído.

Insight Central

O SEGA é motivado pela observação de que os componentes do RoPE estão acoplados a frequências espaciais específicas. Em vez de aplicar um fator de escalonamento fixo, o SEGA deriva magnitudes de escalonamento por dimensão analisando a distribuição de energia do latente intermediário atual.

Fluxo de Trabalho Algorítmico

1. Análise Espectral: Em cada etapa de remoção de ruído, o latente intermediário $Z$ é transformado do domínio espacial para o domínio da frequência usando uma Transformada Rápida de Fourier 2D (FFT).

   • Perfis por Eixo: O espectro de potência 2D é marginalizado sobre eixos ortogonais para obter perfis de energia 1D ($E_H, E_W$) para frequências horizontais e verticais.
   • Perfil Radial: O espectro de potência é médio dentro de anéis concêntricos para obter um perfil invariante à rotação ($E_{iso}$) representando a distribuição de energia através das escalas espaciais.

2. Correção por Dimensão ($s_d$):

   • Para cada dimensão do RoPE $d$, o comprimento de onda correspondente é identificado para mapeá-lo a uma banda de frequência específica no perfil por eixo.
   • A energia logarítmica dessa banda é padronizada.
   • Uma não-linearidade (tanh) é aplicada, e o resultado é centralizado em zero para garantir uma redistribuição estrita do escalonamento. Dimensões associadas a bandas de baixa energia recebem escalonamento mais forte para preservar a discriminação posicional, enquanto bandas de alta energia recebem escalonamento mais fraco para evitar superamplificação.

3. Fator de Amplitude Global ($\sigma$):

   • O perfil radial é reduzido a um único escalar usando planicidade espectral (entropia de Wiener).
   • Esse fator controla a força do escalonamento. Se o espectro do latente for plano (dominado por ruído), $\sigma \to 0$, suprimindo a correção. À medida que o conteúdo estrutural emerge e o espectro se torna menos plano, $\sigma \to 1$, aplicando a correção com força total.

4. Fórmula de Escalonamento Final:
   O fator de escalonamento final $m_d$ para uma dimensão $d$ é calculado como:
   $$m_d = m_{ref} \cdot (1 - \sigma \cdot s_d)$$
   Onde $m_{ref}$ é uma escala de referência determinada pela razão de resolução (usando uma formulação de lei de potência). Isso permite que o SEGA se adapte continuamente ao conteúdo espectral do latente, afiando a atenção em frequências sub-resolvidas e suavizando-a em frequências superenfatizadas.

Contribuições Principais

• Escalonamento Dinâmico e Consciente do Conteúdo: Diferentemente de métodos anteriores de escalonamento uniforme, o SEGA adapta o escalonamento da atenção à estrutura espectral específica do latente em cada etapa de remoção de ruído.
• Resolução da Compensação Estrutura-Detalhe: Ao diferenciar componentes do RoPE com base em suas bandas de frequência associadas e níveis de energia, o SEGA melhora simultaneamente a coerência estrutural global e a fidelidade de detalhes finos.
• Sem Treinamento e Agnóstico à Arquitetura: O método não requer ajuste fino, não possui parâmetros aprendíveis e não exige alterações arquiteturais. Integra-se diretamente em pipelines padrão baseados em RoPE (demonstrado no Flux e no Qwen).
• Desempenho Superior em Resoluções Ultra-altas: O SEGA permite geração estável em resoluções superiores a 36 milhões de pixels (por exemplo, 6144²), um regime onde outros métodos frequentemente falham completamente.

Resultados Experimentais

Os autores avaliaram o SEGA nas arquiteturas Flux e Qwen em múltiplas configurações de alta resolução (até 6144²) usando o conjunto de dados Aesthetic-4K e um benchmark personalizado zero-shot.

• Desempenho Quantitativo: O SEGA superou consistentemente as linhas de base mais avançadas (incluindo NTK, YaRN, DyPE, UltraImage e métodos de orientação multi-estágio como HiFlow e I-Max) em todas as métricas: ImageReward (IR), PickScore (PS), CLIP Score (CS), MUSIQ (MSQ) e FID. Notavelmente, em 4096², o SEGA alcançou as melhores pontuações em todas as métricas avaliadas.
• Melhorias Qualitativas: Comparações visuais demonstraram que o SEGA previne efetivamente modos de falha comuns, como artefatos de telhagem repetitiva, colapso estrutural e perda de clareza semântica. Ele preserva geometrias complexas e texturas finas (por exemplo, pêlos, detalhes de tecido) que outros métodos borram ou distorcem.
• Robustez: O método manteve estabilidade em proporções de aspecto não quadradas e fatores de extrapolação extremos (até 36× a área de treinamento), enquanto métodos concorrentes mostraram degradação significativa de desempenho à medida que a resolução aumentava.

Significado e Alegações

O artigo afirma que o SEGA representa uma solução mínima, porém eficaz, para desbloquear as capacidades de geração em alta resolução de DiTs pré-treinados. Ao mudar de um escalonamento estático baseado em heurísticas para uma abordagem dinâmica guiada por energia espectral, os autores abordam uma limitação fundamental nas estratégias atuais de extrapolação.

Os autores enfatizam que seu trabalho fornece uma "perspectiva espectral" sobre o escalonamento da atenção, revelando um acoplamento entre as dimensões do RoPE e as frequências espaciais que pode ser explorado para guiar a adaptação no momento da inferência. Eles posicionam o SEGA não como uma substituição para abordagens baseadas em treinamento, mas como um aprimoramento poderoso, plug-and-play, que permite que modelos existentes generalizem muito além de suas resoluções nativas de treinamento, sem o custo computacional de retreinamento ou a complexidade de pipelines multi-estágio. O método é apresentado como um passo rumo a uma melhor compreensão e controle do comportamento dos mecanismos de atenção em tarefas de extrapolação de resolução.