An exact information theory of generalization phase transitions in Bayesian diffusion models

Titre : Une théorie exacte de l'information sur les transitions de phase de la généralisation dans les modèles de diffusion bayésiens

Énoncé du problème
Un mystère fondamental de l'IA générative est la manière dont les modèles de diffusion contournent la « malédiction de la dimensionnalité » pour apprendre des distributions complexes sur des espaces de haute dimension à partir d'ensembles de données finis, sans simplement les mémoriser. Naïvement, la transition de la mémorisation vers la généralisation pour des distributions de données arbitraires devrait nécessiter une quantité de données intraçable (exponentielle en dimension). Si les modèles de score empiriques (qui inversent la diffusion en utilisant le score empirique exact d'un ensemble de données fini) sont connus pour mémoriser, les modèles de diffusion neuronaux entraînés généralisent souvent de manière robuste. Les approches théoriques existantes ont peiné à fournir une théorie précise et prédictive de l'endroit et du moment où se produit la transition entre la mémorisation et la généralisation pour des modèles de diffusion non linéaires entraînés à travers diverses architectures et jeux de données.

Méthodologie
Les auteurs introduisent une nouvelle classe de modèles analytiquement traçables appelés modèles de Diffusion à Restriction d'Information Bayésienne (BIRD - Bayesian Information Restricted Diffusion).

• Concept central : Dans les modèles BIRD, chaque pixel $x$ agit comme un agent qui n'observe qu'une partie restreinte, potentiellement stochastique, de l'image bruitée $\phi_t$ (notée $C_{x,t}$), plutôt que l'image complète. Le pixel utilise ensuite une inférence bayésienne optimale pour deviner quel échantillon d'entraînement passé $\phi \in \mathcal{D}$ a produit cette observation restreinte.
• Mécanisme : Le modèle calcule une valeur de pixel débruitée comme étant la moyenne a posteriori sur l'ensemble des données d'entraînement, pondérée par la probabilité a posteriori $P_{\text{train}}(\phi | C_{x,t})$. En restreignant l'information disponible pour chaque pixel (par exemple, en observant uniquement un patch spatial local), le « jeu de devinette bayésien » devient plus difficile, empêchant théoriquement le modèle de s'effondrer sur un point d'entraînement unique (mémorisation).
• Cadre théorique : Les auteurs analysent la transition mémorisation-généralisation en utilisant la théorie de l'information. Ils dérivent une frontière de phase dans l'espace conjoint de :
  1. La quantité de données d'entraînement ($|\mathcal{D}|$).
  2. Le temps/le niveau de bruit dans le processus inverse ($t$ ou $\sigma_t$).
  3. La quantité de restriction d'information (par exemple, la taille du patch $L$).
• Critère clé : La transition se produit lorsque l'information mutuelle $I(\phi; C_{x,t})$ entre l'observation restreinte et la véritable distribution de données est égale au logarithme de la taille du jeu de données : $\ln |\mathcal{D}| = I(\phi; C_{x,t})$.
  • Phase de mémorisation : $I(\phi; C_{x,t}) > \ln |\mathcal{D}|$ (L'observation est trop informative ; le modèle devine parfaitement le point d'entraînement).
  • Phase de généralisation : $I(\phi; C_{x,t}) < \ln |\mathcal{D}|$ (L'observation est insuffisante pour identifier de manière unique un point d'entraînement ; le modèle effectue une moyenne sur les possibilités).

Principales contributions

1. Introduction des modèles BIRD : Un cadre général et analytiquement soluble qui englobe et généralise les modèles de score locaux (LS), les modèles de score équivariants (ES) et les modèles de score équivariants-locaux (ELS). Il ne nécessite aucune hypothèse architecturale explicite ou dynamique d'apprentissage, reposant uniquement sur l'inférence bayésienne avec restriction d'information.
2. Prédiction de la dynamique précoce de l'entraînement : Les auteurs démontrent que les modèles BIRD spatialement locaux (où $C_{x,t}$ est un patch d'image local) approximent étroitement le comportement des modèles de diffusion neuronaux entraînés (y compris les UNets et les Diffusion Transformers) pendant les premières étapes de l'entraînement. Ils atteignent une précision prédictive élevée (médiane $r^2 \sim 0,85-0,93$) sur les sorties d'images individuelles à travers diverses architectures et jeux de données (Celeba64, CIFAR10, MNIST, FashionMNIST).
3. Théorie de la frontière de phase exacte : Sous des hypothèses minimales sur les distributions de données, l'article dérive un critère informationnel précis pour la transition mémorisation-généralisation. Il prouve que la généralisation se produit si et seulement si l'information mutuelle entre l'observation restreinte et la distribution de données est inférieure au logarithme du nombre de points d'entraînement.
4. Échelle de patch critique et alignement spectral : Pour les modèles spatialement locaux, les auteurs identifient une échelle de patch critique dépendante du temps $L_c(t)$. Ils montrent que la génération optimale se produit près du bord de la mémorisation, là où l'échelle du patch $L$ suit l'échelle critique $L_c$. Ils démontrent également que pour les images naturelles ayant des densités spectrales de loi de puissance ($P(k) \sim k^{-2-\epsilon}$), l'échelle critique s'aligne avec l'« échelle spectrale » (l'échelle du débruitage linéaire optimal).
5. Contournement de la malédiction de la dimensionnalité : La théorie révèle que pour les images invariantes d'échelle ($\epsilon=0$), les modèles BIRD ne souffrent pas de la malédiction de la dimensionnalité ; la quantité de données requise n'a pas besoin de croître avec la taille de l'image. Pour les images quasi-invariantes d'échelle ($\epsilon \sim 0,1-0,3$), la quantité de données requise croît comme $\exp(L_I^\epsilon)$, ce qui est nettement sous-exponentiel par rapport à l'exigence naïve de $\exp(L_I^2)$.

Résultats

• Validation empirique : Les expériences confirment l'emplacement prédit par la théorie de la transition de phase. Les courbes d'information mutuelle dérivées de la théorie suivent de près le déficit d'entropie observé dans les expériences numériques sur des jeux de données de jouets et sur des jeux de données d'images réelles.
• Consistance : Tant les modèles BIRD que les réseaux neuronaux précocement entraînés présentent une « généralisation consistante », où les modèles entraînés sur des sous-ensembles de données disjoints produisent des sorties presque identiques lorsqu'ils reçoivent le même bruit en entrée.
• Lois d'échelle : L'échelle de patch critique $L_c(\sigma_t)$ est montrée comme décroissant de manière monotone à mesure que le bruit diminue (le temps progresse dans le processus inverse), s'alignant sur le processus de génération du grossier au fin. L'échelle critique est trouvée hautement corrélée à l'échelle spectrale $L_{\text{spec}}$ dérivée de la densité spectrale de puissance des données.
• Modes de défaillance : L'article illustre que si l'échelle du patch est trop petite ($L \ll L_c$), le modèle généralise mais performe mal (erreur de débruitage élevée). Si l'échelle du patch est trop grande ($L > L_c$), le modèle mémorise les données d'entraînement et échoue à généraliser aux données de test.

Signification
L'article prétend fournir un aperçu fondamental de la manière dont l'IA générative contourne la malédiction de la dimensionnalité. En formulant le problème à travers le prisme de la restriction d'information, les auteurs montrent que la généralisation réussie dans les modèles de diffusion n'est pas seulement le résultat de biais inductifs architecturaux (comme la lissité ou la faible dimensionnalité), mais est fondamentalement pilotée par la limitation de l'information disponible pour le processus génératif à chaque étape. Les résultats suggèrent que la génération se produit « au bord de la mémorisation », où la restriction d'information est soigneusement ajustée pour empêcher le modèle de s'effondrer vers des exemples d'entraînement spécifiques tout en conservant suffisamment d'informations pour reconstruire la distribution de données sous-jacente. Ce travail offre un nouveau laboratoire théorique pour étudier la transition mémorisation-généralisation à travers une grande variété de distributions de données et de classes de modèles.