Resumo Técnico: DiverseDiT++

Declaração do Problema

Avanços recentes em Diffusion Transformers (DiTs) alcançaram um progresso notável na síntese visual devido à sua escalabilidade. No entanto, os mecanismos subjacentes que regem o aprendizado de representação dentro dos DiTs permanecem mal compreendidos. Embora trabalhos recentes como o REPA tentem melhorar o treinamento de DiTs alinhando representações internas com codificadores externos pré-treinados, esses métodos dependem de recursos externos massivos e não explicam totalmente por que funcionam. Além disso, abordagens existentes que buscam aprimorar representações sem orientação externa (por exemplo, perdas dispersivas) frequentemente carecem de uma compreensão fundamentada de como as representações internas evoluem. Questões fundamentais persistem: Como os modelos DiT aprendem representações significativas e por que técnicas de alinhamento externo são eficazes? Essa falta de compreensão fundamental dificulta o desenvolvimento de paradigmas de treinamento mais eficientes e fundamentados.

Metodologia

O artigo propõe o DiverseDiT++, um framework projetado para promover explicitamente o aprendizado de representação diversa dentro de DiTs sem depender de modelos de orientação externa. A metodologia é construída sobre uma análise sistemática da dinâmica de representação e consiste em dois componentes primários:

1. Análise Sistemática e a Pontuação de Diversidade Ponderada (WDS)

Antes de propor a solução, os autores conduziram uma análise sistemática da dinâmica de representação de DiT usando Alinhamento de Kernel Centrado (CKA) para medir similaridades entre diferentes blocos transformer. As observações principais incluíram:

• A discrepância de representação entre os blocos aumenta naturalmente durante o treinamento.
• Alinhar blocos específicos com modelos pré-treinados aumenta sua dissimilaridade em relação a outros blocos.
• O alinhamento excessivo (múltiplos blocos ou codificadores) não melhora necessariamente o desempenho e pode prejudicar a diversidade geral.
• Padrões de diversidade são intrínsecos e consistentes através dos passos de denoising (redução de ruído).

Com base nessas descobertas, os autores introduziram a Pontuação de Diversidade Ponderada (WDS), uma métrica escalar definida como o complemento ponderado pela distância das similaridades CKA par a par entre todos os blocos transformer:
$$WDS = 1 - \frac{\sum_{i<j} |j-i| \cdot CKA(H_i, H_j)}{\sum_{i<j} |j-i|}$$
Os autores encontraram uma forte correlação negativa entre WDS e log(FID) (Pearson's $r = -0.869$), estabelecendo o WDS como um indicador confiável do desempenho do modelo e da qualidade de geração.

2. Framework DiverseDiT++

Para promover ativamente a diversidade observada em modelos de alto desempenho, o DiverseDiT++ introduz dois mecanismos:

• Conexões Residuais Longas: Para evitar a homogeneização representacional causada por entradas homogêneas, o método injeta saídas de camadas iniciais ($i$) em camadas posteriores ($L-i$) via conexões de salto (skip connections) longas. Isso diversifica as representações de entrada para blocos mais profundos, incentivando o reuso de características e prevenindo o colapso.

• Perda de Diversidade de Representação: Uma função de perda explícita é introduzida para penalizar a similaridade entre características de diferentes blocos. Ela compreende três componentes:

  1. Perda de Ortogonalidade: Penaliza alta similaridade de cosseno entre as representações médias por bloco.
  2. Perda de Informação Mútua: Aproxima e minimiza a informação mútua entre representações de blocos para garantir independência estatística.
  3. Perda de Dispersão de Características: Maximiza a variância das ativações de características para incentivar o uso diverso de canais.

  A perda de diversidade total é ponderada por uma estratégia de ponderação adaptativa derivada da análise de conflito de gradiente. Este peso é limitado (clipped) e decai ao longo do treinamento para evitar conflitos de gradiente quando a diversidade satura e para permitir que o modelo foque em detalhes finos em estágios posteriores.

Principais Contribuições

1. Insight Mecanístico: O artigo é o primeiro a elucidar sistematicamente que a diversidade de representação entre blocos é um fator crítico para o aprendizado eficaz em DiTs e fornece uma explicação mecanística para o sucesso das técnicas de alinhamento existentes (ou seja, elas funcionam ao forçar a especialização).
2. Métrica WDS: Introdução da Pontuação de Diversidade Ponderada, uma métrica fundamentada que captura quantitativamente a diversidade inter-blocos e correlaciona-se fortemente com a qualidade de geração (FID) através de escalas, métodos e resoluções.
3. Framework DiverseDiT++: Um framework prático e agnóstico à arquitetura que aprimora o aprendizado de representação através de conexões residuais longas e uma perda de diversidade adaptativa, operando sem codificadores externos pré-treinados.
4. Generalização: Validação do método em diversos domínios, incluindo geração de imagens, dobramento inverso de proteínas e geração de moléculas 3D.

Resultados Experimentais

Os autores conduziram experimentos extensivos em ImageNet (256×256 e 512×512), dobramento inverso de proteínas e geração de moléculas 3D.

• Geração de Imagens: O DiverseDiT++ melhorou consistentemente o desempenho em vários backbones (SiT, REPA) e escalas de modelo (B, L, XL).
  • No ImageNet 256×256, aplicar o DiverseDiT++ ao REPA-B reduziu o FID de 22.99 para 17.22 (400k iterações), superando a linha de base SiT-L (18.77) com o mesmo custo de treinamento.
  • O método acelerou a convergência, atingindo um FID de 1.49 em 200 épocas, superando o SiT-XL/2 treinado por 1400 épocas (FID 2.06).
  • Alcançou um novo estado da arte de FID de 2.94 para geração de um passo (one-step) no ImageNet 256×256 usando MeanFlow-XL/2.
• Domínios Científicos:
  • Dobramento Inverso de Proteínas: Melhorou a taxa de recuperação de sequência para 38.30% e reduziu o scRMSD para 1.81 no conjunto de dados PDB, superando baselines como ProteinMPNN e REED.
  • Geração de Moléculas 3D: Aumentou a estabilidade e reduziu a energia/tensão no GEOM-Drug e QM9 em comparação com SemlaFlow e REED.
• Estudos de Ablação: Confirmaram que tanto as conexões residuais longas quanto a perda de diversidade completa (ortogonalidade, MI, dispersão) são necessárias para o desempenho ideal. O método também se mostrou complementar a técnicas existentes como DispLoss e SRA.

Significância e Alegações

O artigo afirma que a diversidade de representação entre blocos é um princípio universal para o treinamento eficaz de modelos de difusão, estendendo-se além da síntese visual para domínios científicos. Ao quantificar essa diversidade via WDS, os autores fornecem uma ferramenta de diagnóstico para entender a dinâmica de treinamento de DiT. O framework DiverseDiT++ proposto oferece uma estratégia prática para impulsionar o desempenho ao encorajar explicitamente os blocos a aprender características distintas e complementares sem o overhead computacional de modelos de alinhamento externos. O trabalho sugere que promover a diversidade interna é um viés indutivo generalizável que pode melhorar a modelagem generativa através de diversas modalidades de dados e arquiteturas.

Os autores observam limitações, incluindo o leve aumento de parâmetros devido às conexões residuais e a determinação empírica dos limiares de peso adaptativo, que podem exigir exploração sistemática adicional para novas tarefas. Eles também reconhecem que a generalização para texto-para-imagem e edição de imagem ainda precisa ser explorada.