Cross-Space Distillation: Teaching One-Step Students with Modern Diffusion Teachers

Resumo Técnico: Destilação de Espaço Cruzado (Cross-Space Distillation)

Definição do Problema
Modelos de difusão de um único passo modernos alcançam geração de alta qualidade por meio da destilação de passo de tempo baseada em distribuição. No entanto, esses métodos dependem de uma restritiva restrição de Espaço Compartilhado: o Professor e o Aluno devem habitar o mesmo espaço latente, compartilhando resolução latente e parametrização de VAE idênticas. Essa premissa impede a transferência de conhecimento de Professores modernos e de alta capacidade (ex: Stable Diffusion 3.5, Flux, Kolors) que operam em altas resoluções (ex: 1024×1024) e utilizam VAEs avançados, para Alunos compactos e amigáveis para implantação (ex: Stable Diffusion 1.5) que operam em resoluções menores (ex: 512×512) com diferentes VAEs.

Os objetivos de destilação padrão falham nesse cenário porque os tensores latentes ($z_S$ e $z_T$) são inerentemente incompatíveis devido a:

1. Resolução Cruzada: Formatos de grade espacial discordantes (ex: $h \times w$ vs. $H \times W$).
2. VAE Cruzado: Diferentes parametrizações de autoencoder resultando em espaços latentes desalinhados.
3. Arquitetura/Mecanismo Cruzado: Diferenças entre backbones UNet e DiT, ou entre formulações de difusão padrão e de correspondência de fluxo (flow matching).

Alternativas existentes, como poda de rede (network pruning) ou a construção de novos backbones compactos do zero, enfrentam desafios que incluem limiares de esparsidade rígidos, requisitos de retreinamento extensos e a perda de compatibilidade com o ecossistema.

Metodologia: Destilação de Espaço Cruzado e a Ponte (The Bridge)
Os autores formalizam este cenário como Destilação de Espaço Cruzado e introduzem a Ponte ($B_\phi$), um módulo leve projetado para mapear os latentes do Aluno para o espaço latente do Professor sem modificar o backbone do Aluno.

• Arquitetura: A Ponte é fatorizada em duas etapas para lidar com o descompasso de forma eficiente:

  1. Prior Espacial (Congelado): Os primeiros $n$ blocos de decodificação do decodificador do VAE do Aluno ($D_S^{(n)}$) são mantidos congelados. Esses blocos atuam como um upsampler espacial pré-treinado, expandindo o latente do Aluno $z_S$ de baixa resolução para a alta resolução do Professor ($H \times W$). Isso aproveita o conhecimento espacial existente do Aluno sem adicionar parâmetros treináveis para o upsampling.
  2. Projetor Aprendível ($g_\phi$): Uma rede compacta e aprendível (implementada como SwinIR) mapeia as características upsampled do espaço de características do Aluno para o espaço de representação latente do Professor, produzindo um latente compatível com o Professor $\hat{z}_T$.

• Objetivos de Treinamento: A Ponte é treinada usando um conjunto de dados sintético de latentes pareados $(z_S, z_T)$ derivados da mesma imagem codificada por ambos os VAEs. A otimização minimiza uma soma ponderada de duas perdas:

  1. Reconstrução de Latente ($L_{rec}$): Uma perda $\ell_1$ entre o latente previsto $\hat{z}_T$ e o latente real do Professor $z_T$. Os autores observam que a $\ell_1$ proporciona melhor estabilidade e preservação de estrutura do que a $\ell_2$ neste contexto.
  2. Fidelidade de Atenção ($L_{attn}$): Uma perda de divergência KL reversa que alinha as distribuições de autoatenção do denoiser congelado do Professor quando alimentado com $\hat{z}_T$ versus $z_T$. Isso garante que o latente mapeado induza as mesmas respostas estruturais internas (dependências de longo alcance) que o verdadeiro latente do Professor, preservando detalhes finos que perdas de nível de pixel poderiam perder.

• Destilação e Inferência: Uma vez treinada, a Ponte é anexada ao Aluno. Durante a destilação, a saída do Aluno é passada pela Ponte para gerar $\hat{z}_T$, permitindo que objetivos de destilação baseados em distribuição (ex: Variational Score Distillation, Adversarial Distillation) sejam aplicados diretamente contra o Professor. Crucialmente, o backbone do Aluno permanece inalterado.

Principais Contribuições

1. Formalização da Destilação de Espaço Cruzado: O artigo define um novo regime onde a destilação ocorre através de resoluções latentes e espaços de VAE incompatíveis, abordando uma lacuna crítica na transferência de conhecimento de modelos de alta resolução modernos para modelos compactos legados.
2. O Módulo Ponte (The Bridge Module): Uma interface modular e leve que permite objetivos de destilação padrão sem alterar a arquitetura do Aluno. Combina um prior espacial congelado com um projetor aprendível e utiliza um objetivo de Fidelidade de Atenção para um alinhamento estável.
3. Efetividade Prática: O método permite que Alunos compactos (ex: SD 1.5) herdem os priors visuais de diversos Professores modernos (SDXL, Kolors, PixArt-$\Sigma$, Flux, SD 3.5) enquanto mantêm inferência de um único passo, baixa latência e ampla compatibilidade de ecossistema.
4. Capacidade de Fusão de Modelos (Model Merging): Como o método destila diversos Professores para Alunos com o mesmo backbone e espaço latente, os checkpoints resultantes podem ser trivialmente fundidos via média de parâmetros para criar um único modelo que agrega o conhecimento de múltiplas fontes de Professor.

Resultados
Experimentos demonstram melhorias substanciais na qualidade de geração para Alunos compactos de um único passo:

• Ganhos Quantitativos: Destilar o SD 3.5 Medium (Professor) para o SD 1.5 (Aluno) via Ponte melhora a pontuação HPSv3 de 5.37 para 9.42 e o ImageReward de -0.29 para 0.62. A fusão de checkpoints de cinco Professores diferentes eleva ainda mais o HPSv3 do SD 1.5 para 10.53.
• Melhorias Qualitativas: Comparações visuais mostram que os alunos habilitados pela Ponte produzem imagens com texturas mais claras, consistência estrutural mais forte e menos artefatos em comparação com baselines inicializados ou modelos podados.
• Eficiência: A Ponte adiciona um overhead negligenciável (aprox. 5M de parâmetros) e permite a inferência de um único passo. Ela também possibilita o upgrade de resolução em tempo de inferência, onde uma amostra de 512×512 do Aluno pode ser mapeada para 1024×1024 via Ponte e decodificada com o decodificador do Professor sem retreinar o Aluno.
• Comparação com Poda (Pruning): A destilação habilitada pela Ponte supera a poda de rede (ex: podar SD 3.5 para 30% de esparsidade), que tende a produzir desfoque e artefatos.

Significância
O artigo afirma que a Destilação de Espaço Cruzado desbloqueia o potencial de transferir as capacidades generativas de Professores massivos e modernos para backbones eficientes e prontos para implantação, sem o custo de treinar novas arquiteturas do zero. Ao tratar o alinhamento de representações latentes heterogêneas como um problema de interface reutilizável, o trabalho permite:

• Compatibilidade de Ecossistema: Modelos compactos como o SD 1.5 podem ser atualizados com qualidade de ponta enquanto mantêm compatibilidade com ferramentas e hardware existentes.
• Unificação de Conhecimento: A capacidade de fundir o conhecimento de múltiplos Professores, arquitetonicamente distintos (ex: Flux e SD 3.5), em um único modelo leve (Aluno).
• Modularidade: A Ponte serve como um componente plug-and-play que pode ser reutilizado tanto para destilação quanto para o aumento de resolução em tempo de inferência, desacoplando a capacidade de geração do Aluno de suas restrições de representação latente.

Os autores posicionam este trabalho como um passo em direção à exploração mais ampla da destilação através de espaços latentes heterogêneos, aplicável além de texto-para-imagem para geração de vídeo e multimodal.