OmniAlpha: Aligning Transparency-Aware Generation via Multi-Task Unified Reinforcement Learning

1. Énoncé du problème

Les modèles d'IA générative actuels, en particulier les systèmes basés sur la diffusion, excellent dans la création d'images RVB haute fidélité. Cependant, de nombreux flux de travail visuels réels (par exemple, les effets visuels, le design graphique, l'édition d'images) nécessitent une génération consciente de la transparence (RGBA), qui implique de modéliser non seulement la couleur, mais aussi l'opacité par pixel (canaux alpha) et la composition inter-couches.

Les solutions existantes souffrent de deux limitations principales :

• Fragmentation : Les méthodes actuelles sont des silos spécifiques à une tâche. Des pipelines séparés existent pour le détourage d'images, la suppression d'objets, la décomposition de couches et la génération d'images à partir de texte. Il n'existe aucun modèle unifié capable de gérer divers flux de travail RGBA.
• Limites du fine-tuning supervisé (SFT) : Bien que le passage à l'échelle du SFT à travers les tâches soit une étape logique vers l'unification, il repose sur des objectifs de régression localisés (par exemple, perte L1/L2 au niveau des pixels). Ces objectifs échouent à optimiser directement les propriétés critiques de la génération consciente de la transparence, telles que :
  • La cohérence et la physique inter-couches.
  • La précision aux limites alpha (par exemple, cheveux, fourrure).
  • La fidélité compositionnelle lorsque les couches sont mélangées.
  • Par conséquent, les modèles basés uniquement sur le SFT atteignent souvent un plateau, peinant à capturer les détails structurels fins requis pour une transparence fluide.

2. Méthodologie : OMNIALPHA

Les auteurs proposent OMNIALPHA, un cadre unifié multi-tâches qui passe d'un apprentissage purement supervisé à un alignement par Apprentissage par Renforcement (RL). L'architecture se compose de trois composants principaux :

A. Architecture unifiée

• VAE conscient de l'alpha : Un Autoencodeur Variationnel (VAE) de bout en bout initialisé à partir d'un VAE RVB pré-entraîné en utilisant une stratégie d'"initialisation opaque". L'encodeur met les poids alpha à zéro initialement pour se concentrer sur le RVB, tandis que le décodeur est initialisé pour prédire un canal alpha entièrement opaque ($\alpha=1.0$). Cela permet une adaptation efficace à l'espace RGBA à 4 canaux.
• Transformateur de diffusion Séquence-à-Séquence (DiT) : Le backbone de débruitage traite plusieurs entrées et sorties simultanément.
  • Axe de couche bidirectionnel : Pour gérer plusieurs images (entrées et cibles) en une seule passe avant, les auteurs étendent les embeddings de position rotative 2D standards (RoPE) avec un axe z bidirectionnel.
  • Les latents d'entrée reçoivent des indices non négatifs ($0$à$n-1$), les latents cibles reçoivent des indices négatifs ($-1$à$-m$), et les embeddings de texte reçoivent des indices positifs distincts. Cela permet au modèle de différencier les couches spatialement au sein de la séquence.

B. Paradigme d'entraînement

L'entraînement suit un processus en deux étapes :

1. Démarrage à froid (SFT multi-tâches) : Le modèle est d'abord entraîné sur un jeu de données unifié couvrant cinq catégories de tâches en utilisant des objectifs de Flow Matching. Cela établit une capacité de base à travers des tâches diverses.
2. Post-entraînement (Optimisation de la politique relative par groupe - GRPO) : Pour surmonter les limites du SFT, les auteurs formulent le processus de génération comme un Processus de Décision Markovien (MDP).
   • État : Conditions multimodales, étape de temps et séquence de latents intermédiaires.
   • Action : L'étape de débruitage guidée par le champ de vitesse prédit.
   • Récompense : Une récompense consciente de la couche est calculée uniquement à la fin de la trajectoire (basée sur le résultat).

C. Conception de la récompense consciente de la couche

Au lieu de récompenses spécifiques à une tâche, OMNIALPHA utilise quatre objectifs universels définis sur la séquence RGBA décodée pour guider la politique RL :

1. Fidélité de la couche ($r_{layer}$) : Mesure la similarité RVB (SSIM/PSNR) sur toutes les couches valides.
2. Fidélité de la composition ($r_{comp}$) : Évalue l'image composée globalement par rapport à la vérité terrain.
3. Structure d'arrière-plan ($r_{bg}$) : Pénalise les déviations structurelles dans les couches d'arrière-plan en utilisant LPIPS.
4. Limite de premier plan ($r_{fg}$) : Impose la précision et le rappel aux limites alpha (crucial pour les cheveux/la fourrure) en utilisant des masques de limites souples.

Ces récompenses sont agrégées dynamiquement en fonction de la tâche spécifique (par exemple, la génération d'images à partir de texte saute les récompenses alignées sur les pixels car la vérité terrain n'est pas disponible, tandis que les tâches de détourage accordent un poids important aux récompenses de limite).

3. Contributions clés

1. Premier cadre RL unifié pour RGBA : OMNIALPHA est le premier cadre à appliquer l'apprentissage par renforcement multi-tâches pour unifier la génération et la manipulation conscientes de la transparence, dépassant les pipelines fragmentés et spécifiques à une tâche.
2. Paradigme de récompense structurelle : Les auteurs introduisent des récompenses ancrées structurellement qui optimisent explicitement la cohérence inter-couches et la précision des limites alpha, transcendant avec succès le plafond de performance du SFT standard.
3. Encodage de couche bidirectionnel : Une extension novatrice des embeddings de position permettant à un seul DiT de modéliser conjointement plusieurs couches RGBA d'entrée et de sortie en une seule passe avant.

4. Résultats expérimentaux

Le modèle a été évalué sur cinq catégories de tâches : Image à partir de texte, Suppression d'objets, Détourage automatique, Détourage par référence et Décomposition de couches.

• Performance globale : OMNIALPHA surpasse constamment sa propre base de référence SFT et les modèles experts spécialisés sur divers benchmarks.
• Métriques clés :
  • Décomposition de couches : A réalisé une réduction relative de 9,07 % du L1 RVB par rapport à la base de référence SFT.
  • Détourage automatique : A surpassé les outils conventionnels avec une amélioration de 74 % du SAD (Somme des différences absolues) et une amélioration de 68 % du Grad (Différence de gradient).
  • Détourage par référence : A obtenu des résultats de l'état de l'art sur le benchmark RefMatte-RW100 (SAD : 14,768), battant nettement des modèles spécialisés comme CLIPMat.
  • Suppression d'objets : A démontré une forte généralisation sur des jeux de données hors domaine (OBER-Wild), obtenant des scores ReMOVE comparables au modèle ObjectClear hautement spécialisé.
• Études d'ablation : La suppression de n'importe quel composant de récompense unique ($r_{layer}, r_{comp}, r_{bg}, r_{fg}$) a entraîné une dégradation des performances, confirmant que l'ensemble complet des récompenses est nécessaire pour un alignement multi-tâches robuste.

5. Signification

OMNIALPHA représente un changement de paradigme significatif dans l'IA générative pour l'édition visuelle. En démontrant que l'Apprentissage par Renforcement peut aligner efficacement les modèles de diffusion pour des tâches RGBA complexes et multi-couches, l'article comble le fossé entre la génération RVB opaque et les besoins nuancés des flux de travail visuels professionnels.

Cela suggère que les futurs modèles de base pour l'édition d'images devraient s'éloigner du fine-tuning isolé et spécifique à une tâche vers des politiques unifiées optimisées via des récompenses basées sur les résultats. Cette approche améliore non seulement la fidélité des canaux alpha et de la composition des couches, mais permet également à un seul modèle de gérer un large spectre de tâches créatives, de la génération d'objets transparents à la décomposition de scènes complexes.