TOC-SR: Task-Optimal Compact diffusion for Image Super Resolution

Resumo Técnico do TOC-SR: Difusão Compacta Otimizada para Tarefa para Super-Resolução de Imagem

Declaração do Problema
Embora os modelos de difusão tenham demonstrado desempenho superior em tarefas de restauração de imagem, como Super-Resolução de Imagem Única (SR), ao modelar estatísticas de imagens naturais e fornecer priores generativos robustos, eles enfrentam barreiras significativas para implantação prática. Pipelines de difusão padrão exigem grandes redes de base e procedimentos de amostragem iterativa (muitas vezes dezenas de etapas de remoção de ruído), resultando em alta latência, uso de memória e custo computacional. Tentativas existentes de criar modelos SR compactos de uma única etapa frequentemente sofrem com otimização instável e qualidade degradada, pois o modelo deve aprender simultaneamente o prior generativo e o mapeamento de restauração dentro de uma capacidade limitada.

Metodologia
Os autores propõem o TOC-SR, um framework projetado para construir modelos SR eficientes de uma única etapa, descobrindo primeiro uma base de difusão compacta que preserva a fidelidade generativa. A metodologia prossegue em três etapas distintas:

1. Expansão de Capacidade Latente (de 4 para 16 Canais):
   Reconhecendo que a representação latente padrão de quatro canais no Stable Diffusion cria um gargalo de informação para detalhes de alta frequência necessários na SR, os autores expandem a dimensionalidade latente para dezesseis canais, mantendo o mesmo fator de compressão espacial. Como não existe um modelo padrão Stable Diffusion de 16 canais, eles constroem uma nova base destilando a capacidade representacional de um VAE de alta capacidade (da família de modelos FLUX) em um VAE modificado do Stable Diffusion. Isso envolve treinar um codificador estudante para corresponder à distribuição latente de 16 canais do professor, mantendo a fidelidade de reconstrução. O U-Net de difusão é subsequentemente modificado para aceitar e produzir tensores de 16 canais, seguido de ajuste fino conjunto para adaptar o prior de difusão a esse espaço expandido.

2. Descoberta de Base Compacta via Otimização Bayesiana $\epsilon$-Constrita Baseada em Surrogato:
   Para reduzir o custo computacional do modelo expandido de 16 canais, os autores empregam uma estratégia de busca em vez do treinamento direto de modelos compactos.

   • Construção de Surrogato: O U-Net professor é decomposto em etapas sequenciais. Para cada etapa, uma biblioteca de blocos surrogados eficientes em parâmetros é construída através de poda estruturada e redução de módulos (por exemplo, removendo módulos de atenção ou simplificando blocos residuais), enquanto a dimensionalidade dos tensores de entrada e saída é estritamente preservada.
   • Destilação Generativa por Características (FwKD): Cada bloco surrogado candidato é calibrado independentemente usando uma função de perda de destilação por características contra as características do professor em um subconjunto de calibração. Isso evita a necessidade de treinamento completo de difusão durante a fase de busca.
   • Otimização Bayesiana ($\epsilon$-Constrita - BO): A busca de arquitetura é formulada como um problema de otimização com restrições. O objetivo é minimizar a contagem de parâmetros ($f_{param}$) sujeito a uma restrição na fidelidade generativa ($f_{acc} \leq \epsilon$), onde $f_{acc}$ mede a discrepância nas previsões de ruído entre o candidato e o professor. Um modelo surrogado de Processo Gaussiano guia o processo de BO para explorar eficientemente o espaço de arquiteturas em uma grade de valores de $\epsilon$, identificando arquiteturas que oferecem o melhor compromisso entre compressão e fidelidade.

3. Especialização de Tarefa e Destilação de Uma Etapa:
   Uma vez identificada a base compacta ótima, ela é adaptada para a tarefa de SR através de treinamento de difusão condicional, onde o modelo aprende a remover ruído de uma representação latente condicionada a uma entrada de baixa resolução (LR). Finalmente, este modelo iterativo de difusão SR é destilado em um gerador de uma única etapa. A rede estudante é treinada para prever o latente limpo de alta resolução final diretamente a partir de um latente ruidoso e da condição LR em uma única passagem forward, eliminando a necessidade de remoção de ruído iterativa.

Principais Contribuições

• Expansão Latente para SR: O artigo demonstra que expandir o espaço latente de 4 para 16 canais alivia o gargalo para modelagem de detalhes de alta frequência, fornecendo um ponto de partida mais forte para compressão do que arquiteturas padrão.
• Busca de Arquitetura Baseada em Surrogato: Os autores introduzem um framework inovador que combina destilação generativa por características com Otimização Bayesiana $\epsilon$-constrita para descobrir bases de difusão compactas. Esta abordagem reduz sistematicamente os parâmetros, garantindo que as arquiteturas candidatas permaneçam fiéis à dinâmica generativa do professor, sem exigir treinamento específico da tarefa durante a busca.
• SR Eficiente de Uma Etapa: O framework resultante, TOC-SR, integra com sucesso um prior de difusão compacto com um mecanismo de inferência de uma única etapa, alcançando reconstrução de alta qualidade sem amostragem iterativa.

Resultados Experimentais
Os experimentos foram conduzidos em benchmarks padrão (DIV2K, RealSR, DRealSR), comparando o TOC-SR com métodos SR baseados em difusão de última geração (por exemplo, StableSR, DiffBIR, SeeSR) e modelos compactos/de uma etapa existentes (por exemplo, SinSR, OSEDiff, PocketSR).

• Eficiência: Comparado ao modelo de difusão expandido de 16 canais, a base TOC-SR descoberta alcança uma redução de 6,6× nos parâmetros e uma redução de 2,8× nos GMACs.
• Desempenho: No conjunto de validação DIV2K, o TOC-SR alcança um PSNR de 24,26 dB e a melhor pontuação LPIPS (0,242) entre os métodos listados, operando em uma única etapa.
• Generalização para o Mundo Real: Em conjuntos de dados do mundo real (RealSR e DRealSR), o TOC-SR supera outras bases compactas em métricas perceptuais (MUSIQ) e mantém um PSNR competitivo, demonstrando robustez a degradações do mundo real.
• Análise Qualitativa: Comparações visuais mostram que o TOC-SR preserva melhor texturas finas e estruturas finas do que outros modelos compactos, evitando artefatos comuns como suavização excessiva e micro-texturas não naturais.

Significado e Alegações
O artigo afirma que o TOC-SR fornece um caminho estável e eficaz para construir sistemas de restauração de alta qualidade e eficientes. Ao desacoplar a descoberta de um prior generativo compacto da adaptação específica da tarefa, o framework evita a instabilidade de otimização frequentemente observada no treinamento direto de modelos SR compactos. Os autores afirmam que sua abordagem equilibra com sucesso a força generativa dos modelos de difusão com a eficiência necessária para implantação prática, oferecendo uma redução significativa no custo computacional (de dezenas de etapas para uma) enquanto mantém forte qualidade de reconstrução. O trabalho sugere que descobrir priores de difusão compactos via otimização com restrições é uma estratégia viável para restauração de imagem amigável a dispositivos de borda.