Consist-Retinex: One-Step Noise-Emphasized Consistency Training Accelerates High-Quality Retinex Enhancement

1. Énoncé du problème

L'amélioration d'images en faible luminosité (LLIE) est cruciale pour des applications allant de la photographie mobile à la navigation autonome. Bien que la théorie de Retinex (décomposant une image en Réflexance $R$ et Éclairage $L$) fournisse un cadre robuste pour séparer la structure de l'éclairage, les approches génératives modernes de haute qualité (comme les modèles de Diffusion) reposent généralement sur un échantillonnage itératif (par exemple, 1000 étapes). Cela crée un goulot d'étranglement pour le déploiement en temps réel sous des budgets de latence stricts.

Les Modèles de Cohérence (CM) offrent une solution en permettant une génération en une étape grâce à l'auto-cohérence le long des trajectoires de flux de probabilité. Cependant, l'application directe des Modèles de Cohérence standards à l'amélioration basée sur Retinex se heurte à deux défis critiques :

1. Inadéquation des extrémités : L'inférence en une étape opère au niveau de bruit maximal ($\sigma_{max}$), mais les programmes d'entraînement standards (échantillonnage log-uniforme) visitent rarement cette région de bruit élevé (seulement ~0,5 % des échantillons), conduisant à des prédictions instables au point d'inférence.
2. Ambiguïté conditionnelle : Les CM standards imposent l'auto-cohérence mais ne garantissent pas que la sortie s'aligne avec la cible conditionnelle correcte (l'image améliorée) lorsque des données de vérité terrain appariées sont disponibles. Un modèle peut être auto-cohérent mais ancré à une valeur incorrecte.

2. Méthodologie : Consist-Retinex

Les auteurs proposent Consist-Retinex, un cadre qui intègre la décomposition Retinex avec une stratégie d'entraînement de cohérence spécialisée conçue pour une inférence en une étape.

A. Architecture et Décomposition

• Réseau de Décomposition Transformer Retinex (TDN) : Le pipeline commence par un module TDN qui décompose l'image d'entrée en faible luminosité ($I_l$) en cartes de réflexance en faible luminosité ($R_l$) et d'éclairage ($L_l$). De même, il traite les images en lumière normale pour générer les cibles ($R_n, L_n$).
• Modèles de Cohérence Conditionnels Doubles : Deux modèles de cohérence distincts, $f^R_\theta$ et $f^L_\theta$, sont entraînés. Ils mappent les composantes en lumière normale bruitées ($R_n + \sigma\epsilon, L_n + \sigma\epsilon$) vers des composantes améliorées, conditionnées par les cartes correspondantes en faible luminosité ($R_l, L_n$).
• Reconstruction : L'image améliorée finale est reconstruite par multiplication élément par élément : $\hat{I} = \hat{R} \odot \hat{L}$.

B. Innovations principales

La stratégie d'entraînement adresse l'inadéquation des extrémités et l'ambiguïté conditionnelle grâce à deux mécanismes clés :

1. Échantillonnage Adaptatif avec Accentuation du Bruit :

   • L'échantillonnage log-uniforme standard est inefficace pour une inférence en une étape.
   • Les auteurs introduisent une stratégie d'échantillonnage mixte : 95 % des échantillons d'entraînement sont tirés d'un intervalle de bruit élevé $[0.95\sigma_{max}, \sigma_{max}]$, tandis que 5 % sont tirés de la plage complète.
   • Cela garantit que le modèle reçoit une supervision dense spécifiquement au niveau de bruit utilisé lors de l'inférence ($\sigma_{max}$).

2. Fonction de Perte à Double Objectif :
   La perte totale combine deux termes complémentaires :

   • Perte de Cohérence Temporelle ($L_{consist}$) : Impose l'auto-cohérence le long d'une grille discrète de Karras (entraînement de cohérence standard) pour assurer la cohérence de la trajectoire.
   • Perte d'Alignement du Point Fixe de Vérité Terrain ($L_{fixed}$) : Ancre la prédiction du modèle à l'extrémité de bruit élevé directement vers les composantes de vérité terrain appariées ($R_n, L_n$). Cela empêche le modèle de converger vers une valeur auto-cohérente mais incorrecte.
   • Formule : $L_{total} = \lambda_{consist} L_{consist} + \lambda_{fixed} L_{fixed}$.

C. Analyse Théorique

L'article fournit une justification théorique de la conception :

• Limite d'Erreur d'Extrémité : Prouve que l'erreur d'inférence en une étape est contrôlée par l'erreur du point fixe à bruit élevé plus un terme d'inadéquation.
• Propagation de l'Ancrage : Démontre que si l'extrémité de bruit élevé est correctement ancrée à la vérité terrain, la cohérence temporelle propage cette correctionnalité en arrière le long de toute la trajectoire.
• Allocation des Échantillons : Montre mathématiquement que l'accentuation de l'échantillonnage à bruit élevé réduit le nombre total d'échantillons d'entraînement nécessaires pour couvrir le régime d'inférence d'un facteur près de 200 par rapport à un échantillonnage uniforme.

3. Contributions Clés

1. Formulation : Un cadre novateur d'entraînement conditionnel à double objectif spécifiquement pour les composantes décomposées Retinex, combinant la cohérence de trajectoire avec l'alignement sur la vérité terrain appariée.
2. Stratégie d'Échantillonnage : Développement d'une règle d'échantillonnage adaptatif avec accentuation du bruit qui concentre la supervision près de l'extrémité d'inférence en une étape, validée par des études d'ablation extensives.
3. Insights Théoriques : Fourniture de limites d'erreur d'extrémité et d'analyses d'allocation d'échantillons expliquant pourquoi l'ancrage à bruit élevé et la cohérence temporelle sont complémentaires pour l'amélioration en une étape.
4. Efficacité : Atteinte d'une performance de pointe avec une inférence en une étape, réduisant drastiquement à la fois les étapes d'échantillonnage (1 vs 1000) et les coûts d'entraînement (1/8ème des échantillons effectifs des bases de diffusion).

4. Résultats Expérimentaux

La méthode a été évaluée sur des benchmarks appariés (LOL, VE-LOL-L) et des jeux de données non appariés (DICM, VV).

• Performance sur VE-LOL-L : Consist-Retinex a obtenu le meilleur PSNR (25,51 dB) et le meilleur SSIM (0,880) parmi toutes les méthodes comparées sous inférence en une étape, surpassant les modèles de diffusion multi-étapes (par exemple, Diff-Retinex++) et d'autres approches basées sur la cohérence.
• Performance sur LOL : Elle est restée compétitive (22,21 dB PSNR), surpassant significativement les autres méthodes en une étape et s'approchant de la qualité des bases de diffusion à 1000 étapes.
• Efficacité :
  • Inférence : ~12,5 FPS (0,08s par image) sur un seul GPU, représentant une réduction de 1000x des étapes d'échantillonnage par rapport aux bases de diffusion.
  • Entraînement : Nécessite seulement 200K itérations (1,6M d'échantillons effectifs) contre 800K itérations (12,8M d'échantillons) pour les bases de diffusion, réalisant une réduction de 8x du coût d'entraînement pour l'étape de restauration.
• Études d'Ablation : Ont confirmé que le retrait soit de la cohérence temporelle soit de l'alignement sur la vérité terrain entraîne un effondrement des performances, et que l'échantillonnage avec accentuation du bruit est critique pour un PSNR élevé.

5. Importance

Consist-Retinex comble le fossé entre la restauration générative de haute qualité et le déploiement en temps réel. En résolvant les défis spécifiques de l'application des modèles de cohérence à l'amélioration d'images factorisée (Retinex), il démontre que l'inférence en une étape peut atteindre une qualité comparable aux modèles de diffusion itératifs sans la latence associée. La stratégie d'échantillonnage avec accentuation du bruit et la fonction de perte à double objectif proposées offrent une stratégie généralisable pour l'entraînement de modèles de cohérence dans des tâches de restauration conditionnelle où l'inférence se produit à l'extrémité extrême du programme de bruit. Ce travail ouvre la voie au déploiement d'améliorations en faible luminosité haute fidélité sur des appareils contraints par la latence tels que les téléphones mobiles et les véhicules autonomes.