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

Résumé technique de TOC-SR : Diffusion compacte optimale pour la tâche dédiée à la super-résolution d'images

Énoncé du problème
Bien que les modèles de diffusion aient démontré des performances supérieures dans les tâches de restauration d'images telles que la super-résolution d'une seule image (SR) en modélisant les statistiques des images naturelles et en fournissant des priors génératifs puissants, ils font face à des obstacles majeurs pour un déploiement pratique. Les pipelines de diffusion standard nécessitent de grands réseaux de base et des procédures d'échantillonnage itératif (souvent des dizaines d'étapes de débruitage), entraînant une latence élevée, une consommation mémoire importante et un coût computationnel élevé. Les tentatives existantes pour créer des modèles SR compacts en une seule étape souffrent souvent d'une optimisation instable et d'une qualité dégradée, car le modèle doit simultanément apprendre le prior génératif et la carte de restauration dans une capacité limitée.

Méthodologie
Les auteurs proposent TOC-SR, un cadre conçu pour construire des modèles SR efficaces en une seule étape en découvrant d'abord un backbone de diffusion compact qui préserve la fidélité générative. La méthodologie se déroule en trois étapes distinctes :

1. Expansion de la capacité latente (de 4 à 16 canaux) :
   Reconnaissant que la représentation latente standard à quatre canaux dans Stable Diffusion crée un goulot d'étranglement d'information pour les détails haute fréquence requis en SR, les auteurs étendent la dimensionnalité latente à seize canaux tout en maintenant le même facteur de compression spatiale. Étant donné qu'aucun modèle Stable Diffusion standard à 16 canaux n'existe, ils construisent un nouveau backbone en distillant la capacité de représentation d'un VAE haute capacité (de la famille de modèles FLUX) dans un VAE Stable Diffusion modifié. Cela implique l'entraînement d'un encodeur étudiant pour correspondre à la distribution latente à 16 canaux de l'enseignant tout en maintenant la fidélité de reconstruction. Le U-Net de diffusion est ensuite modifié pour accepter et produire des tenseurs à 16 canaux, suivi d'un fine-tuning conjoint pour adapter le prior de diffusion à cet espace étendu.

2. Découverte de backbone compact via une optimisation bayésienne contrainte par $\epsilon$ basée sur des substituts :
   Pour réduire le coût computationnel du modèle étendu à 16 canaux, les auteurs emploient une stratégie de recherche plutôt que l'entraînement direct de modèles compacts.

   • Construction de substituts : Le U-Net enseignant est décomposé en étapes séquentielles. Pour chaque étape, une bibliothèque de blocs substituts économes en paramètres est construite par élagage structuré et réduction de modules (par exemple, suppression de modules d'attention ou simplification de blocs résiduels) tout en préservant strictement la dimensionalité des tenseurs d'entrée et de sortie.
   • Distillation générative au niveau des caractéristiques (FwKD) : Chaque bloc substitut candidat est calibré indépendamment en utilisant une fonction de perte de distillation au niveau des caractéristiques par rapport aux caractéristiques de l'enseignant sur un sous-ensemble de calibration. Cela évite la nécessité d'un entraînement complet de diffusion pendant la phase de recherche.
   • Optimisation bayésienne contrainte par $\epsilon$ (BO) : La recherche d'architecture est formulée comme un problème d'optimisation contrainte. L'objectif est de minimiser le nombre de paramètres ($f_{param}$) soumis à une contrainte sur la fidélité générative ($f_{acc} \leq \epsilon$), où $f_{acc}$ mesure l'écart dans les prédictions de bruit entre le candidat et l'enseignant. Un modèle substitut de processus gaussien guide le processus BO pour explorer efficacement l'espace des architectures sur une grille de valeurs $\epsilon$, identifiant les architectures offrant le meilleur compromis entre compression et fidélité.

3. Spécialisation de la tâche et distillation en une étape :
   Une fois le backbone compact optimal identifié, il est adapté à la tâche SR par un entraînement de diffusion conditionnelle, où le modèle apprend à débruiter une représentation latente conditionnée par une entrée basse résolution (LR). Enfin, ce modèle SR de diffusion itératif est distillé en un générateur en une seule étape. Le réseau étudiant est entraîné à prédire directement le latent haute résolution propre final à partir d'un latent bruité et de la condition LR en une seule passe avant, éliminant ainsi le besoin de débruitage itératif.

Contributions clés

• Expansion latente pour la SR : L'article démontre que l'expansion de l'espace latent de 4 à 16 canaux atténue le goulot d'étranglement pour la modélisation des détails haute fréquence, fournissant un point de départ plus solide pour la compression que les architectures standard.
• Recherche d'architecture basée sur des substituts : Les auteurs introduisent un cadre novateur combinant la distillation générative au niveau des caractéristiques avec une optimisation bayésienne contrainte par $\epsilon$ pour découvrir des backbones de diffusion compacts. Cette approche réduit systématiquement les paramètres tout en garantissant que les architectures candidates restent fidèles aux dynamiques génératives de l'enseignant, sans nécessiter d'entraînement spécifique à la tâche pendant la recherche.
• SR efficace en une étape : Le cadre résultant, TOC-SR, intègre avec succès un prior de diffusion compact avec un mécanisme d'inférence en une seule étape, atteignant une reconstruction de haute qualité sans échantillonnage itératif.

Résultats expérimentaux
Des expériences ont été menées sur des benchmarks standards (DIV2K, RealSR, DRealSR) comparant TOC-SR aux méthodes SR basées sur la diffusion les plus avancées (par exemple, StableSR, DiffBIR, SeeSR) et aux modèles compacts/en une étape existants (par exemple, SinSR, OSEDiff, PocketSR).

• Efficacité : Par rapport au modèle de diffusion étendu à 16 canaux, le backbone TOC-SR découvert réalise une réduction de 6,6 fois des paramètres et une réduction de 2,8 fois des GMACs.
• Performance : Sur l'ensemble de validation DIV2K, TOC-SR atteint un PSNR de 24,26 dB et le meilleur score LPIPS (0,242) parmi les méthodes listées, tout en opérant en une seule étape.
• Généralisation au monde réel : Sur des ensembles de données réels (RealSR et DRealSR), TOC-SR surpasse les autres bases de référence compactes dans les métriques perceptuelles (MUSIQ) et maintient un PSNR compétitif, démontrant une robustesse face aux dégradations du monde réel.
• Analyse qualitative : Les comparaisons visuelles montrent que TOC-SR préserve mieux les textures fines et les structures minces que les autres modèles compacts, évitant les artefacts courants tels que le lissage excessif et les micro-textures non naturelles.

Signification et affirmations
L'article affirme que TOC-SR fournit une voie stable et efficace pour construire des systèmes de restauration de haute qualité et efficaces. En découplant la découverte d'un prior génératif compact de l'adaptation spécifique à la tâche, le cadre évite l'instabilité d'optimisation souvent observée lors de l'entraînement direct de modèles SR compacts. Les auteurs affirment que leur approche équilibre avec succès la puissance générative des modèles de diffusion avec l'efficacité requise pour un déploiement pratique, offrant une réduction significative du coût computationnel (de dizaines d'étapes à une seule) tout en maintenant une forte qualité de reconstruction. Le travail suggère que la découverte de priors de diffusion compacts via une optimisation contrainte est une stratégie viable pour la restauration d'images adaptée aux périphériques.