Résumé technique : Échantillonnage de posterior inspiré de Langevin discret (∆LPS)

Énoncé du problème

L'article aborde le défi de la résolution de problèmes inverses dans des espaces d'états discrets. Les problèmes inverses consistent à retrouver un signal inconnu $x$ à partir de mesures indirectes, clairsemées et bruitées $y$ (modélisées comme $y = A(x) + n$). Alors que les modèles de diffusion continus sont devenus la norme pour résoudre de tels problèmes via l'échantillonnage de posterior, leurs équivalents discrets restent peu explorés.

Les méthodes existantes pour l'échantillonnage de posterior discret souffrent de trois limitations principales :

1. Relaxations continues : Des méthodes comme G2D2 relâchent les variables catégorielles dans des espaces continus pour utiliser une guidance par gradient, modifiant ainsi la géométrie de l'inférence et risquant que les tokens discrets projetés ne se décodent pas de manière cohérente.
2. Mises à jour séquentielles : Des méthodes comme SGDD reposent sur des mises à jour de type Gibbs, qui sont séquentielles et lentes à se mélanger dans des espaces de haute dimension.
3. Spécialisation : Des méthodes comme APS sont étroitement couplées à des processus de corruption spécifiques (par exemple, la diffusion masquée) et à des paradigmes d'entraînement, limitant leur généralité.

L'objectif est de développer un échantillonneur de posterior qui reste entièrement discret, met à jour toutes les dimensions de tokens en parallèle, et soit agnostique au paradigme d'entraînement spécifique du prior de diffusion discret.

Méthodologie : ∆LPS

Les auteurs proposent ∆LPS (Échantillonneur de posterior inspiré de Langevin discret), un algorithme sans entraînement qui utilise l'information du gradient pour guider les transitions discrètes sans quitter l'espace d'états discret.

Formulation de base

Le problème est formulé dans l'espace latent discret $z_0 \in \{1, \dots, K\}^L$ d'un modèle encodeur-décodeur. L'objectif est d'échantillonner à partir du posterior $p(z_0|y) \propto p(y|z_0)p_\theta(z_0)$, où $p_\theta$ est un prior de diffusion discret préentraîné. L'échantillonneur approxime la transition du posterior $p(z_s|z_t, y)$ en inférant d'abord un latent propre $z_0$ cohérent avec la mesure et le prior, puis en le ré-ajoutant de bruit jusqu'à l'étape de temps suivante.

Proposition de Langevin discret

L'innovation clé consiste à adapter la dynamique de Langevin aux espaces discrets. Dans des contextes continus, une mise à jour de Langevin déplace un échantillon $z_0$ dans la direction du gradient du potentiel de log-postérieur $U(z_0) = \log p(y|z_0) + \log p_\theta(z_0; z_t)$. Puisque $z_0$ est discret, les auteurs n'ajoutent pas directement un vecteur de gradient. À la place, ils utilisent la distribution de proposition gaussienne de la dynamique de Langevin continue pour définir des probabilités de transition sur les candidats discrets valides $z'_0$ :

$$q(z'_0|z_0; z_t, y) \propto \exp\left( -\frac{1}{4\eta} \|z'_0 - z_0 - \eta \nabla_{z_0} U(z_0; z_t, y)\|_2^2 \right)$$

En développant l'exposant et en normalisant sur l'ensemble fini des tokens valides, la proposition se factorise à travers les coordonnées des tokens. Cela permet à l'échantillonneur de calculer une distribution catégorielle pour chaque position de token $\ell$ indépendamment :
$$q_\ell(z'_0[\ell]) = \text{Cat}\left( \sigma\left( \frac{1}{2} \nabla_{z_0[\ell]} U \cdot (k - z_0[\ell]) - \frac{(k - z_0[\ell])^2}{4\eta} \right) \right)$$

Ce mécanisme permet des mises à jour entièrement parallèles sur toutes les $L$ dimensions de tokens, évitant le goulot d'étranglement séquentiel des échantillonneurs de Gibbs.

Généralisation et préconditionnement

• Agnosticisme vis-à-vis du prior : La méthode ne nécessite que des probabilités prédictives sur les états catégoriels, ce qui la rend compatible avec les priors de diffusion masquée (MDLM) et de diffusion à état uniforme (Duo).
• Problèmes inverses aveugles : Pour des opérateurs directs non différentiables ou inconnus, le gradient de vraisemblance est remplacé par un surrogate contrastif appris (utilisant des encodeurs pour les données et les mesures) afin de fournir un signal de guidance lisse.
• Préconditionnement Adam : Pour gérer des gradients mal mis à l'échelle dans des espaces latents de haute dimension, le gradient brut est remplacé par une direction préconditionnée par Adam. Cela redimensionne adaptativement les mises à jour, atténuant les grands gradients bruyants tout en préservant ceux qui sont constamment utiles, améliorant ainsi le mélange et la convergence.

Contributions clés

1. Échantillonnage de posterior entièrement discret : ∆LPS est un échantillonneur sans entraînement qui maintient chaque itéré dans l'espace d'états discret, évitant l'optimisation de tokens mous et les erreurs de projection.
2. Mises à jour parallèles efficaces : En factorisant la proposition à travers les coordonnées, la méthode met à jour toutes les dimensions de tokens en parallèle, améliorant considérablement l'efficacité par rapport aux échantillonneurs de type Gibbs séquentiels.
3. Conception agnostique du prior : L'échantillonneur n'est pas lié à un processus de corruption spécifique (masqué vs uniforme) et fonctionne avec n'importe quel prior de diffusion discret préentraîné fournissant des probabilités prédictives.

Résultats expérimentaux

Les auteurs évaluent ∆LPS sur un ensemble diversifié de problèmes inverses couvrant MNIST, CIFAR-10 et FFHQ, englobant des configurations linéaires, non linéaires et aveugles.

• Restauration d'images (FFHQ) : Sur des tâches incluant l'inpainting, le défloutage gaussien, la super-résolution, le HDR et le défloutage de mouvement, ∆LPS atteint des performances comparables ou supérieures aux bases de référence continues (par exemple, DPS, PSLD) et surpasse les bases de référence discrètes existantes (SGDD, G2D2, APS) en qualité perceptuelle (LPIPS).
• Tâches discrètes (MNIST/CIFAR) : La méthode démontre une robustesse à travers des niveaux de difficulté variables (Facile, Moyen, Difficile). Elle correspond systématiquement ou dépasse SGDD et G2D2 en précision des tokens et en PSNR, même avec des budgets de calcul significativement inférieurs (moins d'évaluations de fonctions) que la variante "complète" de SGDD.
• Problèmes inverses aveugles : Dans une tâche de cartographie spatiale (inférer des plans d'étage à partir de trajectoires de marche où l'opérateur direct est inconnu), ∆LPS surpasse les bases de référence continues (CoGuide, variantes DPS) en score F1 et en Intersection sur Union (IoU).

Importance et affirmations

L'article affirme que les échantillonneurs de posterior entièrement discrets et informés par le gradient offrent une voie évolutive et générale pour résoudre des problèmes inverses sur des représentations discrètes.

• Évolutivité : Le mécanisme de mise à jour parallèle répond aux limitations de vitesse de mélange des échantillonneurs discrets dans les espaces de haute dimension.
• Généralité : La méthode découple l'algorithme d'inférence du paradigme d'entraînement spécifique du prior, lui permettant de tirer parti de l'écosystème croissant des modèles de diffusion discrets (texte, code, images, audio).
• Praticité : En évitant les relaxations continues, ∆LPS garantit que les échantillons finaux sont des tokens discrets valides, empêchant l'"écart de projection" où une approximation continue échoue à se décoder en un signal discret cohérent.

Les auteurs concluent que, à mesure que les modèles génératifs discrets s'échelle, des méthodes d'inférence fondées sur des principes comme ∆LPS deviendront de plus en plus critiques pour des applications dans la conception biomoléculaire, la musique symbolique, la complétion de code et d'autres domaines où les données sont intrinsèquement discrètes.