Résumé Technique : Quand LeJEPA apprend-il un Modèle du Monde ?

Énoncé du Problème

L'article comble une lacune fondamentale dans l'apprentissage auto-supervisé (SSL) : bien que les Architectures Prédictives à Encodage Joint (JEPAs) aient démontré un succès empirique dans l'apprentissage de représentations pour la vision, la vidéo et la planification, il n'existe aucune garantie théorique que ces représentations constituent un « Modèle du Monde ». Plus précisément, il reste unclear dans quelles conditions une représentation apprise récupère fidèlement la structure latente sous-jacente du monde (par exemple, position, vitesse, couleur) plutôt que d'entremêler ces facteurs.

Une représentation qui brouille les vrais degrés de liberté (par exemple, en mélangeant la position d'un objet avec sa couleur) peut bien performer sur des tâches étroites, mais échoue lorsque le monde change ou lorsqu'elle est utilisée pour une généralisation compositionnelle. La question centrale est : Quand une représentation apprise récupère-t-elle linéairement les variables latentes du monde ? Les auteurs définissent cette propriété comme l'identifiabilité linéaire, une condition nécessaire pour un sondage linéaire fiable et une planification optimale dans l'espace latent.

Méthodologie et Cadre Théorique

Le Modèle du Monde

Les auteurs formalisent le « Monde » comme un système avec des variables latentes $z \in \mathbb{R}^n$ qui sont non observées. Les données observées $x$ sont générées par une fonction de mélange inconnue, potentiellement hautement non linéaire $g$, telle que $x = g(z)$. L'objectif d'apprentissage implique des paires positives $(z, z')$ générées par un processus stationnaire avec des transitions de bruit additif indépendant.

L'article se concentre sur une classe spécifique de mondes où :

1. Indépendance : Les variables latentes et leurs transitions sont indépendantes à travers les dimensions.
2. Stationnarité : La distribution marginale des variables latentes reste constante entre les vues.
3. Bruit Additif : Les transitions suivent $z' = m(z) + \eta$, où $\eta$ est un bruit indépendant.

Crucialement, les auteurs spécialisent cette classe au Monde Gaussien, où $z \sim \mathcal{N}(0, I_n)$. Sous ces hypothèses, les dynamiques de transition correspondent à un processus d'Ornstein–Uhlenbeck (OU) : $z' = \rho z + \sqrt{1-\rho^2}\eta$.

L'Apprenant : LeJEPA

L'architecture d'apprentissage est LeJEPA (JEPA de LeCun), qui entraîne un encodeur $f$ pour produire des plongements $h(z) = f(g(z))$ en minimisant une fonction de perte composée de deux termes :

1. Perte d'Alignement : Minimise la distance entre les plongements des paires positives, $E[\|h(z') - h(z)\|^2]$. Cela encourage la représentation à capturer des caractéristiques lentes et prévisibles.
2. SIGReg (Régularisation Isotrope Gaussienne Esquissée) : Régularise explicitement la distribution des plongements pour qu'elle soit une gaussienne isotrope, $h(z) \sim \mathcal{N}(0, I_n)$, empêchant l'effondrement de la représentation.

Le problème d'optimisation est :
$$\min_h E[\|h(z') - h(z)\|^2] \quad \text{sous contrainte} \quad h(z) \sim \mathcal{N}(0, I_n)$$

Approche Théorique

Les auteurs emploient une analyse spectrale de l'opérateur de transition $T$, défini par $(Th)(z) = E[h(z')|z]$.

• Développement de Hermite : Pour des latents gaussiens, les fonctions propres de l'opérateur de transition sont les polynômes de Hermite. Les valeurs propres décroissent géométriquement comme $\rho^d$, où $d$ est le degré du polynôme.
• Pénalité Spectrale : La perte d'alignement pénalise davantage les composantes de degré supérieur (non linéaires) que les composantes linéaires car elles se décorrèlent plus rapidement sous la transition.
• Théorie de Sturm-Liouville : Pour des distributions générales (non gaussiennes), les fonctions propres sont caractérisées par la théorie de Sturm-Liouville. La première fonction propre non constante est toujours monotone, mais la linéarité n'est garantie que si la distribution est gaussienne.

Contributions Clés

1. Sens Direct : Identifiabilité Linéaire pour les Mondes Gaussiens

Le Théorème 1 prouve que si les variables latentes sont gaussiennes et que la transition est un processus OU, tout minimiseur de l'objectif LeJEPA (sous contraintes de gaussianité) doit être une transformation orthogonale linéaire des vrais latents : $h(z) = Qz$ pour une certaine matrice orthogonale $Q$.

• Mécanisme : La preuve repose sur une décomposition spectrale où toute non-linéarité réduit strictement la corrélation entre les paires positives. Puisque la perte d'alignement cherche à maximiser cette corrélation, l'optimum global force la représentation à être purement linéaire (polynômes de Hermite de degré 1).
• Résultat : La représentation apprise récupère les variables latentes à une rotation globale près, et les dynamiques de transition sont récupérées exactement dans l'espace latent.

2. Sens Inverse : Unicité de la Gaussienne

Le Théorème 2 établit que la distribution gaussienne est la seule distribution latente pour laquelle l'identifiabilité linéaire vaut dans cette classe de mondes.

• Mécanisme : En utilisant la théorie de Sturm-Liouville, les auteurs montrent que l'exigence que la première fonction propre de l'opérateur de transition soit affine (linéaire) force la fonction de score de la distribution latente à être linéaire. Cette condition caractérise de manière unique la distribution gaussienne.
• Implication : Si les variables latentes sont non gaussiennes, LeJEPA ne peut garantir l'identifiabilité linéaire ; la représentation apprise sera nécessairement une transformation non linéaire des latents.

3. Identifiabilité Approximative

Le Théorème 3 fournit une borne quantitative sur l'erreur de récupération lorsque les objectifs ne sont satisfaits qu'approximativement (c'est-à-dire un alignement imparfait et un blanchiment imparfait).

• L'erreur évolue gracieusement avec l'écart d'alignement $\delta$ et l'erreur de blanchiment $\epsilon$.
• La borne met en évidence que l'écart d'alignement est le facteur dominant dans l'erreur de récupération, tandis que le blanchiment approximatif est relativement « gratuit » à obtenir.

4. Planification Latente Optimale

Le Théorème 4 relie l'identifiabilité linéaire à la théorie du contrôle. Il prouve que si l'encodeur est identifiable linéairement ($h(z) = Qz$), alors tout problème de contrôle optimal avec des coûts invariants par transformations orthogonales (par exemple, la distance euclidienne à un objectif) peut être résolu de manière optimale dans l'espace latent appris. Les actions optimales et les fonctions de valeur dans l'espace latent se mappent directement au monde réel, ne nécessitant aucune modification de la fonction de coût.

Résultats Expérimentaux

Les auteurs valident leur théorie à travers plusieurs configurations expérimentales :

1. Simulations 2D : En utilisant diverses fonctions de mélange non linéaires (spirale, cisaillement, RealNVP), LeJEPA récupère avec succès les variables latentes à une rotation près lorsque les latents sont gaussiens, confirmant le Théorème 1.
2. Passage à l'Échelle Haute-Dimensionnelle : Des expériences jusqu'à $N=1024$ dimensions montrent que les régularisateurs basés sur les statistiques de lot (SIGReg, VICReg) maintiennent une identifiabilité linéaire quasi parfaite ($R^2 > 0,999$), tandis que les méthodes de contraste (InfoNCE) se dégradent à grande échelle en raison de décalages de largeur de noyau.
3. Ablations Distributionnelles : Balayer la forme de la distribution latente (de lourde queue à uniforme) démontre que la récupération linéaire atteint un pic aigu à la gaussienne ($\alpha=2$) et se dégrade pour les distributions non gaussiennes, validant le Théorème 2.
4. Contrôle Robotique Basé sur les Pixels : En utilisant le DeepMind Control Suite Reacher, les auteurs comparent l'entraînement sur des transitions gaussiennes synthétiques (OU) versus de vraies trajectoires RL.
   • Condition OU : Atteint une haute identifiabilité linéaire ($R^2 \approx 0,95$).
   • Condition Trajectoire : Les politiques du monde réel induisent des marginales non gaussiennes et anisotropes, résultant en une identifiabilité significativement plus faible ($R^2 < 0,5$).
5. Validation de la Planification : Dans la tâche Reacher, l'interpolation en ligne droite dans l'espace latent appris de l'encodeur entraîné sur la gaussienne se décode en chemins quasi droits dans l'espace des articulations (planification optimale). En revanche, l'encodeur entraîné sur les trajectoires produit des chemins courbes et sous-optimaux, confirmant le Théorème 4.

Signification et Revendications

L'article revendique fournir le premier résultat d'identifiabilité pour les JEPAs, transformant une recette empiriquement réussie en une garantie mathématique.

• Garantie Fondamentale : Le travail établit que LeJEPA récupère prouvablement la structure du monde (variables latentes et dynamiques) sous des conditions spécifiques et bien définies (latents gaussiens, bruit additif stationnaire).
• Rôle de la Gaussianité : L'article inverse le récit classique de l'ICA linéaire. Dans l'ICA linéaire, la gaussianité est la distribution où la séparation de sources échoue. Dans ce cadre non linéaire et auto-supervisé, la gaussianité est la seule distribution qui permet l'identifiabilité linéaire.
• Implications Pratiques : La théorie explique pourquoi les stratégies d'exploration qui approximent des marches aléatoires isotropes (maintenant les données dans le régime gaussien) sont cruciales pour apprendre des modèles du monde utiles, tandis que les politiques orientées vers un but (qui effondrent les données dans des régions de faible entropie) peuvent briser l'identifiabilité.
• Fondement de la Planification : En prouvant que l'identifiabilité linéaire permet une planification optimale avec des coûts invariants par rotation, l'article fournit une justification théorique pour l'utilisation d'espaces latents appris pour le contrôle sans avoir besoin de mapper vers les vraies coordonnées physiques.

Les auteurs notent des limites : la théorie suppose que la dimension de l'encodeur correspond à la dimension latente et repose sur l'optimalité au niveau de la population. Cependant, ils démontrent que les garanties se dégradent gracieusement avec des échantillons finis et des erreurs d'optimisation, et ils vérifient formellement tous les théorèmes principaux en utilisant le prouveur de théorèmes Lean 4.