CoLoGen: Progressive Learning of Concept-Localization Duality for Unified Image Generation

CoLoGen est un cadre de diffusion unifié qui résout le conflit de représentation entre la compréhension conceptuelle et la précision spatiale en apprenant progressivement leur dualité grâce à un curriculum en étapes et un module d'entrelacement de représentations.

L'essentiel

Imaginez que vous essayez d'enseigner à un seul artiste tout l'art du monde. D'un côté, vous voulez qu'il soit un poète capable de comprendre des concepts abstraits (comme "une ambiance triste" ou "un dragon mythique"). De l'autre, vous voulez qu'il soit un architecte capable de placer des objets avec une précision chirurgicale (comme "un chat assis exactement sur ce coussin rouge").

Le problème, c'est que si vous forcez cet artiste à faire les deux en même temps, il se perd. Il commence à confondre les idées avec les lieux, et le résultat est souvent décevant : soit le chat est flou, soit l'ambiance est absente. C'est ce que les chercheurs appellent le conflit entre le concept et la localisation.

Voici comment CoLoGen résout ce problème, expliqué simplement :

1. Le Problème : Le "Cerveau" qui se dispute

Dans les modèles d'IA actuels, tout est mélangé dans un seul gros cerveau. Quand on demande à l'IA de générer une image, elle doit à la fois comprendre ce qu'elle dessine (le concept) et où elle le dessine (la localisation).

• L'analogie : C'est comme si vous demandiez à un chef cuisinier de préparer un plat complexe tout en tenant un marteau pour construire la maison en même temps. Ses mains ne savent pas si elles doivent tenir une cuillère ou un marteau, et le résultat est un désastre.

2. La Solution : L'École Progressive (Le "Curriculum")

Au lieu de tout apprendre d'un coup, CoLoGen utilise une méthode d'apprentissage progressive, comme un système scolaire bien structuré :

• Étape 1 : Les bases (L'Apprentissage Intégré)
  L'IA commence par apprendre séparément. D'abord, elle s'entraîne à dessiner des objets à partir de mots (le poète). Ensuite, elle s'entraîne à placer des objets précis dans une image (l'architecte). Elle ne mélange pas encore les deux.

  • Analogie : C'est comme apprendre à jouer de la guitare et à chanter séparément avant de former un groupe de rock.

• Étape 2 : L'Adaptation
  Une fois qu'elle maîtrise les bases, on lui apprend à utiliser ces compétences avec différents outils (des lignes de contour, des cartes de profondeur, etc.).

• Étape 3 : La Synergie (Le Grand Final)
  Enfin, on lui donne des instructions complexes du type : "Change le chien en robot, mais garde le jardin en arrière-plan". À ce stade, l'IA sait exactement quand faire appel à son côté "poète" et quand faire appel à son côté "architecte".

3. La Magie : Le Tissage Progressif (PRW)

C'est le cœur de l'invention. CoLoGen ne force pas le cerveau à tout faire en même temps. Il utilise un système intelligent appelé PRW (Progressive Representation Weaving).

• L'analogie du Chef d'Orchestre :
  Imaginez un chef d'orchestre (le modèle) qui dirige un groupe de musiciens spécialisés (les "experts").
  • Il y a un expert pour les concepts (les violons).
  • Il y a un expert pour la localisation (les percussions).
  • Le chef d'orchestre possède un tissu dynamique. Selon la musique qu'il faut jouer (la tâche demandée), il "tisse" ensemble les bons musiciens au bon moment.
  • Si la tâche demande de dessiner un paysage, il active les violons. Si elle demande de placer un objet précis, il active les percussions. Si les deux sont nécessaires, il les fait jouer en harmonie sans qu'ils ne se marchent dessus.

Ce système utilise aussi un "Portail de Vétéran" (Veteran Gate). C'est un petit garde qui s'assure que les musiciens qui ont déjà appris des choses lors des étapes précédentes ne sont pas oubliés, mais qu'ils continuent de jouer leur rôle au bon moment.

4. Les Résultats

Grâce à cette méthode, CoLoGen est capable de :

• Modifier des images avec des instructions précises (ex: "Enlève le chien, mets un chat à sa place").
• Créer des images personnalisées (ex: "Dessine mon chien dans un style Van Gogh").
• Contrôler la forme (ex: "Dessine quelque chose qui suit cette ligne rouge").

En résumé :
CoLoGen ne force pas l'IA à être un génie universel instantané. Il lui apprend d'abord à être un expert en idées, puis un expert en placement, et enfin, il lui apprend à orchestrer ces deux talents ensemble. C'est comme passer d'un brouillard de confusion à une symphonie parfaitement coordonnée, où chaque note (concept ou localisation) est jouée au bon moment pour créer une image magnifique et précise.

Résumé technique

1. Problématique : La Dualité Concept-Localisation

L'article identifie un défi fondamental dans la génération d'images conditionnelle unifiée : le conflit entre deux types de représentations internes nécessaires à des tâches différentes.

• Représentation Conceptuelle ($R_c$) : Nécessaire pour la synthèse sémantique (ex: génération personnalisée, compréhension d'instructions textuelles). Elle encode la cohérence sémantique et la compréhension des objets.
• Représentation de Localisation ($R_l$) : Nécessaire pour la précision spatiale (ex: inpainting, ancrage d'images, génération contrôlée par des bords ou des masques). Elle encode l'alignement spatial, la géométrie et la cohérence structurelle.

Le Conflit : Les modèles unifiés actuels tentent d'apprendre ces deux capacités simultanément dans un espace latent unique. Cela crée une interférence représentative : l'optimisation pour la compréhension conceptuelle dégrade souvent la précision spatiale, et vice-versa. Les approches d'apprentissage multi-tâches classiques (co-training) échouent souvent à équilibrer ces besoins contradictoires, menant à une instabilité de l'entraînement et à des performances médiocres sur certaines tâches.

2. Méthodologie : CoLoGen

Pour résoudre ce problème, les auteurs proposent CoLoGen, un cadre de diffusion unifié qui apprend et réconcilie progressivement cette dualité via une stratégie de formation en plusieurs étapes et une architecture spécifique.

A. Stratégie d'Apprentissage Progressif (Curriculum)

Au lieu d'entraîner le modèle sur toutes les tâches simultanément, CoLoGen utilise un curriculum "du facile au difficile" structuré autour des deux représentations :

1. Pré-entraînement Endogène : Le modèle apprend d'abord les capacités fondamentales de manière séparée :
   • Remplissage de masques (Mask Inpainting) pour renforcer la génération de concepts visuels.
   • Ancrage d'images (Image Grounding) pour renforcer la localisation précise.
2. Injection Conditionnelle : Adaptation des capacités acquises à divers signaux conditionnels (segmentation, profondeur, bords Canny).
3. Alignement Instruction-Image : Raffinement de la synergie entre les deux représentations pour des tâches complexes d'édition par instruction et de génération personnalisée.

B. Architecture : Progressive Representation Weaving (PRW)

Le cœur technique de CoLoGen est le module PRW, intégré dans les blocs d'attention du modèle de diffusion (basé sur FLUX.1).

• Experts Spécialisés : Le modèle utilise un pool d'experts paramétriques efficaces (basés sur des projections Key-Value) qui apprennent séparément les compétences de concept et de localisation lors des premières étapes.
• Routage Dynamique (Veteran Gate) : Un routeur dynamique sélectionne l'expert le plus pertinent pour chaque entrée. Contrairement aux méthodes statiques, ce routeur est guidé par une supervision de "porte vétérane" (Veteran Gate Routing).
• Supervision de la Porte Vétérane : Une fonction de perte auxiliaire ($L_{veteran}$) force le routeur à maintenir une densité d'utilisation spécifique pour les experts appris aux étapes précédentes. Cela empêche l'oubli catastrophique et assure que les connaissances acquises (concepts ou localisation) sont préservées et intégrées de manière stable dans les étapes ultérieures.

3. Contributions Clés

1. Identification de la Dualité Concept-Localisation : Formalisation du conflit représentatif comme le principal obstacle à la génération d'images unifiée fiable.
2. Cadre CoLoGen : Proposition d'un cadre unifié qui structure explicitement les tâches autour de ces deux représentations distinctes plutôt que de les fusionner aveuglément.
3. Stratégie d'Apprentissage et Architecture PRW : Introduction d'une formation progressive combinée à un mécanisme de "tissage" de représentations (PRW) avec un routage dynamique supervisé pour gérer les conflits et éviter l'oubli.
4. Performance Supérieure : Démonstration que cette approche permet d'atteindre des performances compétitives ou supérieures par rapport aux modèles spécialisés et aux modèles généralistes existants.

4. Résultats Expérimentaux

Les expériences ont été menées sur trois catégories de tâches principales, comparant CoLoGen à des modèles spécialisés (SOTA) et des modèles généralistes (comme OmniGen, UniReal, PixWizard).

• Édition par Instruction (Instruction Editing) : Sur les benchmarks MagicBrush et Emu Edit, CoLoGen obtient les meilleurs scores de conformité aux instructions (CLIP-T) et de similarité sémantique (DINO), surpassant les modèles de pointe comme Emu Edit et OmniGen.
• Génération Contrôlée (Controllable Generation) : Sur des tâches de contrôle par Canny, profondeur, et segmentation (benchmarks MultiGen-20M, ADE20K), CoLoGen atteint des performances comparables aux méthodes spécialisées (ControlNet, Uni-ControlNet) tout en restant un modèle unifié.
• Génération Personnalisée (Customized Generation) : Sur DreamBench, CoLoGen montre une capacité supérieure à préserver l'identité du sujet (scores DINO et CLIP-I élevés) tout en suivant les prompts textuels, malgré un jeu de données d'entraînement plus petit que certains concurrents.
• Études d'Ablation : Les résultats confirment que l'ajout séparé des représentations $R_c$ et $R_l$ améliore les performances, et que la formation progressive (CoLoGen) est supérieure au co-training simultané (qui dégrade souvent les performances sur la génération personnalisée).

5. Signification et Impact

Ce travail offre une perspective représentative fondamentale pour la génération d'images unifiée. Il démontre que la clé pour créer des modèles "généralistes" performants ne réside pas simplement dans l'augmentation de la taille des données ou des paramètres, mais dans la gestion structurelle des conflits internes entre différentes exigences de tâche.

En séparant et en réconciliant progressivement les besoins de sémantique (concept) et de géométrie (localisation), CoLoGen établit une nouvelle voie pour développer des modèles de génération d'images plus robustes, polyvalents et capables de gérer une large gamme de tâches complexes sans compromis de performance. Cela ouvre la porte à des systèmes d'IA générative plus fiables pour des applications pratiques allant de l'édition créative à la conception assistée par ordinateur.