Résumé Technique : Inversion de Transformations de Données via l'Échantillonnage par Diffusion

Énoncé du Problème

Le papier traite du problème de l'inversion de transformation sur des groupes de Lie généraux. Dans de nombreux scénarios d'apprentissage automatique et de modélisation scientifique, une donnée observée $\tilde{x}$ est générée en appliquant une transformation inconnue $g$ issue d'un groupe de Lie $G$ à un échantillon conforme à la distribution $\textit{in-distribution}$ $x$ (c'est-à-dire $\tilde{x} = g \cdot x$). L'objectif est de récupérer une transformation inverse $g^{-1}$ qui ramène l'observation hors-distribution vers la distribution de données originale.

Ceci est formulé comme un problème inverse aveugle car la transformation nuisance est inconnue et la solution est généralement non unique. Alors que les méthodes existantes opèrent souvent dans l'espace des données ou sont spécialisées pour des familles de transformations spécifiques (par exemple, les groupes affines ou de rotation), ce travail cible les actions de groupes de Lie généraux, incluant des groupes non compacts, non abéliens et sans métriques bi-invariantes (par exemple, les homographies d'image $PGL(3, \mathbb{R})$ et les symétries de points de Lie pour les EDP).

Une application clé mise en évidence est l'équivariance au moment du test (test-time equivariance). Les réseaux de neurones pré-entraînés manquent souvent de robustesse face aux transformations géométriques. Les auteurs proposent d'inverser ces transformations au moment du test pour ramener les entrées dans la distribution d'entraînement du modèle, améliorant ainsi la robustesse sans nécessiter d'architectures équivariantes spécialisées ou de ré-entraînement.

Méthodologie : Diffusion d'Énergie pour l'Inversion de Transformation (TIED)

La méthode proposée, TIED, modélise la distribution postérieure sur les transformations $p(g | \tilde{x})$ comme une distribution de Boltzmann définie par une fonction d'énergie $E$ sur l'espace des données. L'énergie est dérivée d'un a priori de données $p_x(x)$, où $E(x) = -\log p_x(x)$. La densité postérieure est donnée par :
$$p(g | \tilde{x}) \propto p_x(g^{-1} \cdot \tilde{x}) |\det J_{g^{-1}}(\tilde{x})|$$

Pour échantillonner de cette distribution postérieure, TIED introduit un processus de diffusion directement sur le groupe de Lie. La méthodologie surmonte deux défis majeurs : la nature accidentée et multimodale des paysages d'énergie (souvent dérivés de réseaux de neurones) et les contraintes géométriques de l'échantillonnage sur les variétés.

1. Diffusion Trivialisée sur les Groupes de Lie

Au lieu d'effectuer des mises à jour dans les espaces tangents d'arbitraires éléments de groupe, TIED utilise la trivialisation. Les mises à jour sont effectuées dans l'algèbre de Lie $\mathfrak{g}$ (l'espace tangent à l'identité) et transportées vers l'élément de groupe actuel via l'application exponentielle.

• Processus Direct (Forward Process) : Une équation différentielle stochastique (SDE) à variance croissante est définie sur le groupe en tirant des perturbations infinitésimales dans l'algèbre de Lie et en les transportant par multiplication à droite :
  $$dg_t = d(L_{g_t})_e [\gamma(t) dw^{\mathfrak{g}}_t]$$
  où $w^{\mathfrak{g}}_t$ est un mouvement brownien dans l'algèbre de Lie.
• Processus Inverse (Reverse Process) : La SDE de temps inverse est dérivée pour récupérer des échantillons de la distribution postérieure. Crucialement, le terme de dérive implique le score trivialisé $\nabla_g \log p_t(g_t)$, qui est un élément de l'algèbre de Lie.

2. Identité du Score Cible Trivialisé

Une contribution théorique centrale est une nouvelle identité de score cible trivialisé. Le "score matching" standard nécessite l'accès à des échantillons propres, qui sont indisponibles dans ce contexte d'inversion. Au lieu de cela, TIED exprime le score de la distribution bruitée en termes du gradient de la fonction d'énergie propre.
Les auteurs prouvent que pour les groupes de Lie généraux (sans nécessiter de métrique bi-invariante), le score trivialisé peut être estimé via une espérance de Monte Carlo sur le noyau de bruit $k_t$ :
$$\nabla_g \log p_t(g_t) = \nabla_g \log \int_G k_t(w) \exp\left(-E(g_t w^{-1}) - \log \lambda(w)\right) d\mu(w)$$
où $\lambda(w)$ rend compte du changement de la mesure de Haar sous inversion. Cela permet d'estimer le score en utilisant uniquement la fonction d'énergie $E$ et ses gradients, calculés via la différenciation automatique.

3. Implémentation Pratique

L'algorithme d'échantillonnage (Algorithme 1) discrétise la SDE inverse. À chaque étape, il :

1. Échantillonne le bruit $w$ à partir du noyau de bruit $k_t$.
2. Calcule l'énergie $E$ pour les candidats transformés.
3. Estime le score trivialisé en utilisant l'identité de Monte Carlo (en employant l'astuce log-sum-exp pour la stabilité).
4. Met à jour l'élément de groupe en utilisant l'application exponentielle.

La méthode supporte diverses définitions d'énergie, y compris celles dérivées de Variational Autoencoders (VAE) ou directement de la confiance de prédiction d'un classificateur pré-entraîné.

Contributions Clés

1. Échantillonneur TIED : Un nouvel échantillonneur de diffusion pour inverser les transformations de données sur des groupes de Lie généraux. Il garantit que les mises à jour restent sur la variété et gère les actions non compactes et non linéaires en utilisant uniquement la connaissance de l'algèbre de Lie.
2. Cadre Théorique : Une preuve que l'inversion de transformation peut être utilisée pour l'équivariance au moment du test de modèles pré-entraînés. Les auteurs démontrent que la distribution postérieure de l'inversion de transformation est $G$-equivariante, assurant la validité de l'approche.
3. Nouvelle Identité de Score : Une généralisation de l'identité du score cible aux groupes de Lie qui fonctionne sans métriques bi-invariantes, permettant un échantillonnage basé sur le score efficace en utilisant les gradients d'énergie plutôt que des échantillons de données propres.
4. Validation Empirique : Démonstrations sur les homographies d'images ($PGL(3, \mathbb{R})$) et les symétries d'EDP ($SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$), montrant des améliorations de performance constantes par rapport aux bases solides.

Résultats Expérimentaux

Le papier évalue TIED sur trois tâches principales :

1. Échantillonnage Synthétique sur $SO(10)$ : TIED échantillonne avec succès une distribution d'énergie multimodale sur un groupe de Lie de haute dimension ($dim=45$), convergeant de manière significativement plus rapide qu'un échantillonneur de Langevin cinétique trivialisé.
2. Classification d'Images (MNIST) :
   • Configuration : Un ResNet18 pré-entraîné est testé sur des images MNIST soumises à des transformations inconnues affines ($Aff(2, \mathbb{R})$) et d'homographie ($PGL(3, \mathbb{R})$).
   • Baselines : Comparé contre des méthodes basées sur l'optimisation (LieLAC, FoCal, ITS) et des méthodes d'échantillonnage (Langevin cinétique).
   • Résultats : TIED surpasse significativement toutes les baselines. Par exemple, sous des transformations affines avec une énergie basée sur le classificateur, TIED atteint une précision de 82,64 % contre 55,48 % pour le ResNet18 brut et 74,72 % pour le Langevin cinétique. Il surpasse également les réseaux équivariants spécialisés (affConv) dans certains contextes.
3. Résolution d'EDP :
   • Configuration : Un opérateur neuronal DeepONet est testé sur les équations de la chaleur et de Burgers en 1D sous des transformations de symétrie de points de Lie inconnues.
   • Résultats : TIED obtient l'erreur relative $L_2$ la plus faible sur tous les cas de test, surpassant LieLAC et FoCal. Pour l'équation de la chaleur en 1D avec les symétries $SL(2, \mathbb{R}) \ltimes H(1, \mathbb{R})$, TIED réduit l'erreur à 0,043, contre 0,066 pour LieLAC et 0,690 pour le DeepONet non corrigé.

Signification et Revendications

Les auteurs affirment que TIED fournit un pipeline sans entraînement pour rendre n'importe quel modèle pré-entraîné équivariant aux transformations de groupes de Lie au moment du test. En exploitant les modèles basés sur l'énergie (qui peuvent être approximés par le modèle pré-entraîné lui-même), la méthode évite le besoin d'architectures équivariantes spécialisées ou de ré-entraînements coûteux.

La signification réside dans :

• Généralité : Elle s'applique à une large classe de groupes de Lie, incluant ceux qui sont non compacts et dépourvus de métriques bi-invariantes, ce qui a posé des défis pour les méthodes de diffusion ou d'optimisation précédentes.
• Robustesse : Elle restaure efficacement les entrées transformées vers la distribution d'entraînement, améliorant considérablement la généralisation des réseaux pré-entraînés face aux transformations géométriques hors-distribution.
• Efficacité : L'approche par diffusion, combinée à l'identité du score trivialisé, offre un processus d'échantillonnage plus efficace et plus stable que le MCMC traditionnel ou l'optimisation itérative sur des paysages d'énergie accidentés.

Le papier conclut que cette approche représente une forme de mise à l'échelle de l'équivariance au moment de l'inférence, offrant une solution pratique pour la robustesse à l'ère des grands modèles de fondation où la gestion des transformations hors-distribution est critique.