Résumé Technique : De la Composition d'Experts Globale à l'Échelle des Facteurs dans les Modèles de Diffusion Discrète

Problématique

Les modèles de diffusion discrète sont apparus comme un cadre puissant pour les tâches de raisonnement complexe, notamment grâce à la génération compositionnelle, qui combine plusieurs modèles « experts » pré-entraînés pour généraliser au-delà de leurs distributions d'entraînement individuelles. Les avancées théoriques récentes (ex. : SuperDiff, RNE, correcteurs de Feynman-Kac) ont amélioré l'échantillonnage compositionnel en introduisant des poids de mélange dépendant du temps pour aligner la dynamique composée avec les distributions cibles.

Cependant, ces méthodes de pointe partagent une limitation fondamentale : elles opèrent sur une base par échantillon, en assignant un poids scalaire global unique à chaque expert à chaque étape temporelle. Cette hypothèse « monolithique » traite implicitement chaque expert comme étant également compétent sur l'ensemble de l'espace d'états. Dans les tâches génératives complexes — comme le raisonnement spatial où un expert peut se spécialiser dans les bordures tandis qu'un autre gère les textures — ce pondération globale force un compromis. Elle dilue le signal de l'expert correct dans sa région spécifique de compétence, échouant à capturer l'indépendance spatiale et fonctionnelle des différents experts.

Méthodologie : FactorDiff

Les auteurs proposent FactorDiff, un cadre de composition par facteurs qui décompose le processus de génération du niveau de l'échantillon jusqu'au niveau du facteur (par exemple, les pixels spatiaux ou les cellules de grille). Au lieu d'appliquer un poids unique à un état entier généré, FactorDiff route dynamiquement chaque facteur vers l'expert le plus pertinent.

Composants Clés

1. Routage par Facteur : L'espace d'états est décomposé en facteurs définis par l'utilisateur. Pendant le processus d'échantillonnage, une carte de routage assigne chaque facteur (par exemple, un pixel dans une grille) à un expert spécifique. Cela permet des contributions d'experts hétérogènes au sein d'un même état généré.
2. Sélection Basée sur la Confiance : Le cadre instancie le routage en sélectionnant dynamiquement l'expert présentant la plus haute confiance à chaque position lors de l'inférence. La variable de routage $z_v$ pour une variable $v$ est déterminée par :
   $$z_v = \arg\max_{k} C(p_k(x_v | x_t))$$
   où $C$ est une métrique de confiance (la marge entre les deux premières probabilités de prédiction). L'étape de retour composée sélectionne ensuite la prédiction de l'expert choisi pour cette position spécifique.
3. Architecture de Diffusion Masquée 2D Native : Pour soutenir cela, les auteurs introduisent une architecture de diffusion masquée 2D qui préserve la topologie de la grille. Contra contrairement aux travaux antérieurs qui aplatissent les grilles en séquences 1D (introduisant des décalages topologiques), cette approche traite la grille comme un graphe de réseau.
   • Experts : Trois types de modèles complémentaires sont entraînés : un modèle Complet (toutes les couleurs), un modèle d'Occupation (arrière-plan vs premier plan), et un modèle de Couleur (couleurs du premier plan uniquement).
   • Backbone : Le réseau de débruitage utilise une architecture de raffinement itératif (soit convolutionnelle, soit basée sur les Transformers) avec des plongements de position 2D apprenables pour assurer la localité spatiale.

Justification Théorique

Les méthodes existantes définissent souvent une distribution cible comme un produit pondéré d'experts ($p(x) \propto \prod p_m(x)^{\alpha_m}$). Les auteurs soutiennent que pour des spécialistes complémentaires, un exposant global $\alpha_m$ est mal spécifié car il force les prédictions d'un expert (même erronées en dehors de son domaine) à influencer l'état entier. FactorDiff remplace cela par un produit par facteur :
$$p_{comp}(x) \propto \prod_{f \in \mathcal{F}} \prod_{m=1}^M (p_m(x_f))^{\alpha_{m,f}}$$
Cette formulation permet à l'expertise de varier selon les facteurs. Dans le cas du routage strict, $\alpha_{m,f}$ est un vecteur "one-hot", sélectionnant un seul expert par facteur sans nécessiter de normalisation supplémentaire ou de correction d'importance, car le produit de conditionnelles valides définit une transition jointe valide.

Contributions Clés

1. Cadre FactorDiff : Une nouvelle méthode de composition qui traite la combinaison d'experts comme un routage sur des facteurs, spécifiquement instanciée avec une sélection d'experts au niveau de la position pour les modèles de diffusion structurés en grille.
2. Architecture de Diffusion Masquée 2D : Une architecture native qui maintient la topologie de la grille, permettant l'entraînement d'experts complémentaires (Occupation, Couleur, Complet) dont les prédictions peuvent être composées lors de l'inférence.
3. Validation Empirique sur ARC-AGI : Une évaluation étendue sur un sous-ensemble de 120 tâches ARC-AGI utilisant le générateur RE-ARC pour l'évaluation en distribution (IID) démontre que le routage simple par facteur surpasse systématiquement les schémas de pondération scalaire globale complexes, particulièrement sur les tâches exigeant une cohérence logique et un désenchevêtrement spatial.

Résultats Expérimentaux

Les auteurs ont évalué FactorDiff sur un sous-ensemble de 120 tâches ARC-AGI en utilisant le générateur RE-ARC pour l'évaluation IID.

• Régime de Spécialisation (RQ1) : Lors de la combinaison d'experts hautement spécialisés (ex. : Occupation + Complet, ou Couleur + Occupation), FactorDiff a nettement surpassé les bases de comparaison par échantillon (SuperDiff, RNE, FKC).
  • Couleur + Occupation : Les spécialistes individuels ont obtenu un score proche de zéro sur la tâche complète (3,3 % et 0,2 % de précision exacte). FactorDiff les a composés pour atteindre 90,5 % de précision exacte, une amélioration massive par rapport au meilleur baseline par échantillon (78,4 %).
  • Occupation + Complet : FactorDiff a atteint 94,0 % de précision exacte, surpassant le meilleur expert individuel (93,0 %) et toutes les méthodes de composition par échantillon.
• Régime de Généraliste (RQ2) : Lors de la combinaison de deux modèles généralistes forts (Complet + Complet), FactorDiff s'est montré "sûr", égalant le meilleur baseline par échantillon (95,2 %) sans dégrader les performances. Cela indique que la sélection spatialement variable n'introduit pas d'instabilité lorsque les experts sont largement en accord.
• Ablation : Même en composant deux modèles complets ayant des objectifs identiques mais des architectures différentes (Trans-IR vs. Conv-IR), FactorDiff a atteint une précision plus élevée (94,8 %) que l'un ou l'autre des modèles individuels, réussissant à assembler les régions correctes des deux.

Signification et Revendications

L'article affirme que FactorDiff fournit un mécanisme simple et efficace pour combiner des spécialistes arbitraires au sein d'une seule trajectoire de diffusion. Sa principale importance réside dans la démonstration que le pondage global par échantillon est insuffisant pour les tâches où les experts ont des domaines de spécialisation spatiale ou fonctionnelle. En passant à une composition par facteur, la méthode capture une expertise complémentaire que les poids globaux ne peuvent représenter, atteignant une généralisation supérieure sur les tâches exigeant une stricte cohérence logique.

Les auteurs notent que bien que leur implémentation actuelle repose sur un routage basé sur la confiance (supposant que la confiance est corrélée à la justesse), le cadre s'étend naturellement à un routage appris ou conditionné par la tâche. Ils positionnent ce travail comme une étape vers une composition plus granulaire dans la diffusion discrète, allant au-delà des hypothèses "monolithiques" des méthodes actuelles de pointe.