Résumé Technique : Apprentissage Variationnel pour la Génération par Insertion

1. Énoncé du Problème

Les modèles de séquences autorégressifs génèrent typiquement des jetons selon un ordre fixe de gauche à droite. Bien que cette factorisation soit efficace pour de nombreuses tâches, elle n'est pas inhérente à toutes les structures de données et peut être mal alignée avec des problèmes de génération du monde réel tels que la conception de séquences biologiques, la synthèse de programmes et la planification structurée, où les parties de la sortie peuvent dépendre d'un contexte futur.

Les alternatives existantes, telles que les modèles de diffusion masqués (génération non monotone), offrent de la flexibilité en mettant à jour des sous-ensembles de positions masquées sur une grille de longueur fixe. Cependant, ces approches souffrent de deux limitations principales :

1. Toile Statique : Elles opèrent sur des grilles de longueurs prescrites et fixes, nécessitant des spécifications de longueur externes et rendant la génération de longueur variable indirecte.
2. Agnosticisme de l'Ordre : Leurs objectifs d'entraînement marginalisent souvent ou supposent implicitement de nombreux schémas de mise à jour possibles. Cela force le modèle à approximer des conditionnelles pour un nombre factoriel d'ordres de génération, même si seul un petit sous-ensemble s'aligne avec la structure sous-jacente des données.

Des travaux récents ont traité ces problèmes de manière isolée : les modèles basés sur l'insertion permettent des séquences de longueur variable via des insertions arbitraires, tandis que les méthodes d'apprentissage de l'ordre introduisent des distributions sur les permutations. Cependant, combiner ces approches est non trivial car, dans les modèles d'insertion, les emplacements d'insertion ne sont pas des coordonnées fixes mais évoluent avec la séquence partielle, couplant étroitement l'ordre de génération et les états intermédiaires.

2. Méthodologie : Le Processus d'Insertion (IP)

Les auteurs proposent l'Insertion Process (IP), un cadre probabiliste qui apprend conjointement :

1. Où insérer : Une préférence basée sur les données concernant les décisions d'insertion (sélection de créneaux/slots).
2. Quoi insérer : Une distribution de jetons.
3. Quand terminer : Une règle d'arrêt (basée sur une politique ou un classifieur).

2.1. Reparamétrage via les Permutations

Une contribution théorique centrale est la formalisation d'une correspondance bijective entre les trajectoires d'insertion et les permutations des indices de la séquence cible.

• Le Défi : L'inférence variationnelle directe sur les trajectoires d'insertion est difficile car les indices de créneaux relatifs ($z_i$) dépendent de la trajectoire, ce qui rend difficile la définition d'une distribution postérieure stable sur l'ordre dans l'espace des actions relatives.
• La Solution : Les auteurs associent chaque trajectoire d'insertion valide produisant une séquence $y_L$ à une permutation unique $\sigma \in S_L$, où $\sigma_i$ est l'indice de l'élément rempli à l'étape $i$.
• Correspondance Déterministe : Le créneau d'insertion relatif à l'étape $i$ devient une fonction déterministe du préfixe de la permutation ($\sigma_{\le i}$). Spécifiquement, $z_i = 1 + |\{j \in \sigma_{<i} : j < \sigma_i\}|$.
• Factorisation de la Vraisemblance : Ce changement de variable convertit la marginalisation sur les actions dépendantes de la trajectoire en une somme exacte sur les permutations, permettant un reparamétrage exact de la vraisemblance des données :
  $$p(y_L) = \sum_{\sigma \in S_L} p(y_L, \sigma)$$

2.2. Inférence Variationnelle et Entraînement

Pour entraîner le modèle, les auteurs maximisent une borne inférieure de l'évidence (ELBO) sur la log-vraisemblance des données, traitant l'ordre de génération (la permutation $\sigma$) comme une variable latente.

• Postérior de l'Inférence : Ils introduisent un postérior variationnel amorti $q_\theta(\sigma | y_L)$, paramétré par un modèle de Plackett-Luce (PL).
• Optimisation : Puisque l'ordre latent est discret, l'objectif est non différentiable. Les auteurs optimisent l'ELBO en utilisant :
  • Astuce Gumbel-Top-k : Pour échantillonner efficacement des permutations à partir de la distribution PL.
  • REINFORCE avec Leave-One-Out (RLOO) : Pour estimer les gradients.
  • Réduction de la Variance : En exploitant la structure permutation-trajectoire, l'espérance interne sur le prochain indice $\sigma_i$ (conditionnée sur un préfixe fixe) est évaluée exactement en forme fermée, réduisant la variance de l'estimateur de Monte Carlo.

2.3. Architecture du Réseau

• Décodeur Génératif : Un décodeur de type Transformer traite la séquence partielle (augmentée de jetons de bordure) pour produire des plongements contextuels (embeddings).
  • Une tête de localisation évalue les créneaux d'insertion candidats (incluant une action de terminaison) basés sur les embeddings.
  • Une tête de contenu prédit la distribution des jetons conditionnée par le créneau sélectionné.
• Encodeur d'Inférence : Un encodeur Transformer bidirectionnel traite la séquence observée complète pour produire des scores pour chaque indice global, qui paramètrent la distribution de Plackett-Luce pour l'échantillonnage de l'ordre de génération.

3. Contributions Clés

1. Cadre Probabiliste pour l'Insertion de Longueur Variable : Le premier cadre apprenant l'ordre d'insertion pour des séquences de longueur variable en utilisant une correspondance bijective avec les permutations, permettant un reparamétrage exact de la vraisemblance des données.
2. Vraisemblance Marginalisée par Permutation : Une dérivation théorique montrant que la vraisemblance des données peut être exprimée comme une somme sur les permutations, facilitant l'inférence variationnelle avec une variable d'ordre latente.
3. Algorithme d'Entraînement Efficace : Une méthode d'entraînement à variance réduite utilisant des modèles de Plackett-Luce, l'échantillonnage Gumbel-Top-k et des gradients RLOO pour gérer la nature discrète et non différentiable de l'ordre latent.
4. Support Natif de Longueur Variable : Contrairement aux modèles de diffusion à toile fixe, IP supporte nativement la génération de longueur variable et apprend les préférences basées sur les données pour les ordres d'insertion.

4. Résultats Expérimentaux

Les auteurs ont évalué IP sur deux domaines : la planification conditionnée par un objectif et la génération de chaînes moléculaires.

4.1. Planification Conditionnée par un Objectif

• Benchmarks : Planification de labyrinthes (Braided, Imperfect, Perfect) et planification de graphes en étoile (Star Graph).
• Baselines : Modèles AR de type ordre fixe (FO-ARM), modèles AR de type n'importe quel ordre (AO-ARM), modèles de diffusion masqués (MDM) et modèles AR d'apprentissage d'ordre (LO-ARM).
• Résultats : IP a atteint des performances de pointe (state-of-the-art) sur toutes les familles de labyrinthes et les benchmarks de graphes.
  • Sur les Labyrinthes Parfaits (hautement contraints), les FO-ARM se sont considérablement dégradés, tandis qu'IP a maintenu une précision quasi parfaite (ex: 99,9 % sur les Hard Perfect Mazes contre 27,2 % pour FO-ARM).
  • IP a considérablement amélioré la précision des séquences et des jetons, et a réduit les distances d'édition par rapport à toutes les baselines, démontant que l'apprentissage d'ordres d'insertion spécifiques à l'instance est crucial pour les tâches avec des contraintes à longue portée.

4.2. Génération de Chaînes Moléculaires (SMILES)

• Benchmark : Génération inconditionnelle et conditionnelle GuacaMol.
• Résultats :
  • Validité et Nouveauté : IP a atteint une validité élevée (97,2–97,4 %) et la plus haute nouveauté (~95 %) parmi tous les modèles, surpassant significativement l'insertion d'ordre aléatoire (AO-IP) qui chutait à 88,2 % de validité. Cela indique qu'IP intériorise avec succès les règles de syntaxe chimique.
  • Métriques Distributionnelles : Bien que les FO-ARM aient légèrement surpassé IP sur la similitude distributionnelle (divergence KL, FCD), IP offre un meilleur compromis en explorant l'espace chimique plus efficacement.
  • Génération Conditionnée : Dans les tâches nécessitant l'insertion entre des fragments fixes (ex: conception de lieurs), IP a maintenu environ 100 % de validité, là où les modèles à ordre fixe peinent face à de telles contraintes.
  • Interprétabilité : L'analyse des calendriers d'insertion appris a révélé un ordre de génération cohérent et interprétable : (1) construction du squelette moléculaire (cycles/parenthèses), (2) insertion des atomes, et (3) terminaison. Cela s'est produit sans priors grammaticaux explicites.

5. Signification et Revendications

L'article affirme que l'apprentissage d'ordres d'insertion dépendants des données améliore considérablement la qualité de la modélisation et la généralisation dans les domaines dépourvus de structure canonique de gauche à droite. En formalisant la connexion entre les trajectoires d'insertion et les permutations, les auteurs fournissent une méthode variationnelle traitable et non biaisée pour entraîner des modèles basés sur l'insertion.

Les auteurs positionnent leur travail comme une étape vers la modélisation générative dans des domaines où l'ordre de génération "naturel" est ambigu ou non linéaire. Ils déclarent explicitement que des travaux futurs pourraient appliquer cette méthode à d'autres domaines tels que la conception de protéines. Le papier maintient un ton modeste, notant que bien qu'IP améliore la généralisation et la validité dans des tâches structurelles complexes, il ne surpasse pas nécessairement les modèles à ordre fixe sur les métriques simples de correspondance distributionnelle, soulignant un compromis entre flexibilité et fidélité distributionnelle stricte.