NNiT: Width-Agnostic Neural Network Generation with Structurally Aligned Weight Spaces

Ce travail présente NNiT, un modèle de diffusion qui génère des poids de réseaux de neurones de manière indépendante de la largeur en alignant structurellement l'espace des poids via des HyperRéseaux Graphiques et en traitant les poids comme des champs structurés, permettant ainsi une généralisation efficace à des architectures non vues lors de l'entraînement.

L'essentiel

🤖 NNiT : Le "Chef Cuisinier" qui invente des recettes pour n'importe quelle cuisine

Imaginez que vous voulez créer un robot capable de faire des tâches complexes (comme assembler des cubes). Pour cela, le robot a besoin d'un "cerveau" : un réseau de neurones artificiels.

Généralement, pour entraîner ce cerveau, il faut le construire, le nourrir de données et le laisser apprendre pendant des heures. C'est long et coûteux.

Les chercheurs de cet article ont une idée géniale : au lieu d'entraîner chaque cerveau individuellement, pourquoi ne pas créer un "générateur" capable de dessiner instantanément un cerveau parfait, même si on lui demande un cerveau qu'il n'a jamais vu auparavant ?

C'est là qu'intervient NNiT (Neural Network Diffusion Transformers). Voici comment ça marche, avec des analogies simples.

1. Le Problème : Le chaos des Lego 🧱

Le gros problème avec les réseaux de neurones, c'est qu'ils sont comme des boîtes de Lego géantes.

• La symétrie : Si vous avez un mur de Lego rouge, vous pouvez permuter (changer l'ordre) des briques rouges entre elles. Le mur reste identique visuellement, mais l'ordre des briques est différent.
• Le chaos : Pour un ordinateur, ces deux murs sont totalement différents, même s'ils font la même chose. C'est comme si vous essayiez d'apprendre à un enfant à dessiner un chat, mais que chaque fois qu'il dessine un chat, vous lui disiez : "Non, la queue est à gauche, puis à droite, puis en haut !". C'est impossible à apprendre car il n'y a pas de règle fixe.

Les méthodes actuelles échouent souvent quand on change la taille du réseau (le "largeur"). C'est comme si un modèle apprenait à dessiner un chat avec 100 pixels, mais qu'il paniquait dès qu'on lui demandait un chat avec 150 pixels.

2. La Solution Magique : L'Alignement par le GHN 🧭

Pour résoudre ce chaos, les chercheurs utilisent un outil appelé GHN (Graph HyperNetwork).

• L'analogie du Chef de Chantier : Imaginez un chef de chantier (le GHN) qui supervise la construction. Au lieu de laisser les ouvriers (les neurones) placer les briques au hasard, le chef impose un ordre strict.
• L'alignement : Grâce à ce chef, chaque fois qu'on construit un mur, les briques rouges sont toujours à la même place relative, les bleues à une autre. Même si le mur change de taille, la "logique" de construction reste la même.
• Le résultat : Le cerveau du robot n'est plus un tas de Lego chaotique, mais un champ structuré, comme une image bien rangée où les couleurs voisines ont des liens logiques.

3. La Révolution : NNiT et les "Patches" (Timbres-poste) 🎨

Une fois que le cerveau est bien rangé, les chercheurs utilisent NNiT.

• L'analogie du Puzzle : Au lieu de regarder tout le cerveau comme une seule grande liste de chiffres (ce qui est rigide), NNiT le découpe en petits carrés, comme des timbres-poste ou des pièces de puzzle.
• La magie de la taille :
  • Si vous voulez un cerveau plus large, vous n'avez pas besoin d'apprendre une nouvelle règle. Vous ajoutez simplement plus de timbres à la suite.
  • C'est comme si vous dessiniez une image : si vous voulez une image plus large, vous ajoutez des pixels à droite. Le style de dessin reste le même, vous ajoutez juste de la matière.
• Génération "Agnostique" : Cela signifie que NNiT peut générer un cerveau pour une tâche donnée, que ce cerveau soit petit, grand, ou d'une forme bizarre qu'il n'a jamais vue dans sa formation.

4. Le Résultat : Un Robot qui s'adapte instantanément 🤖✨

Les chercheurs ont testé cela sur des robots (dans un simulateur appelé ManiSkill3) qui doivent manipuler des objets.

• Avant (Les anciennes méthodes) : Si on changeait la taille du cerveau du robot, il devenait stupide et échouait à sa tâche. Il fallait tout réapprendre.
• Avec NNiT : On demande au générateur : "Fais-moi un cerveau pour ce robot, mais avec une largeur de 32 neurones ici et 64 là-bas".
  • Même si cette configuration exacte n'existait jamais pendant l'entraînement, NNiT réussit plus de 85% du temps.
  • Il comprend la "logique" du cerveau, pas juste la taille.

En résumé 🌟

Imaginez que vous avez un magicien (NNiT) qui sait comment construire des cerveaux de robots.

1. Il utilise un chef d'orchestre (GHN) pour s'assurer que toutes les notes sont bien alignées, même si l'orchestre change de taille.
2. Il compose la musique par petites phrases (les patches) plutôt que par une seule longue mélodie rigide.
3. Résultat : Il peut improviser une symphonie parfaite pour un orchestre de 10 musiciens, ou pour un orchestre de 1000 musiciens, sans jamais avoir répété cette configuration spécifique.

C'est une avancée majeure pour l'Intelligence Artificielle, car cela permet de créer des robots et des logiciels qui s'adaptent instantanément à n'importe quel matériel ou contrainte, sans avoir besoin de réapprendre tout depuis zéro.

Résumé technique

1. Problématique

La génération de paramètres pour des réseaux de neurones (synthèse de poids) se heurte à deux obstacles majeurs :

• La dépendance aux dimensions fixes : Les méthodes existantes flattent généralement les matrices de poids en vecteurs de dimension fixe. Cela lie le modèle génératif à des architectures spécifiques. Dès que la largeur (nombre de neurones) d'une couche change, la dimensionnalité des tokens change, brisant les correspondances apprises et empêchant la généralisation à des architectures non vues lors de l'entraînement.
• La symétrie de permutation : Dans les réseaux de neurones (notamment les MLP), l'ordre des neurones est arbitraire. De nombreuses configurations de poids distinctes produisent la même fonction entrée-sortie. Cela crée un espace de paramètres non aligné où les poids adjacents dans une matrice ne présentent aucune corrélation spatiale, rendant difficile l'apprentissage de structures locales par des modèles génératifs (comme les modèles de diffusion).

2. Méthodologie : NNiT

Les auteurs proposent NNiT (Neural Network Diffusion Transformers), un cadre unifié qui traite la synthèse de réseaux de neurones comme une tâche de modélisation de séquence multimodale.

A. Alignement Structurel via Graph HyperNetworks (GHN)

Pour résoudre le problème de la symétrie de permutation et créer une structure spatiale exploitable, les auteurs utilisent des Graph HyperNetworks (GHN) avec un décodeur CNN (Convolutional Neural Network) :

• Le GHN génère les paramètres du réseau en propageant des informations sur un graphe d'architecture.
• L'utilisation d'un décodeur CNN impose un biais de localité explicite dans l'espace des poids. Contrairement à l'entraînement par descente de gradient stochastique (SGD) qui produit des matrices désordonnées, le GHN génère des tenseurs de poids avec des corrélations spatiales locales cohérentes (ex: structures en bandes verticales).
• Cela transforme les poids en un champ spatial continu et aligné, prêt pour une tokenisation basée sur des patchs.

B. Tokenisation par Patchs (Patch Tokenization)

Au lieu de vectoriser les poids globalement, NNiT découpe les tenseurs de poids alignés en patchs de taille $p \times p$.

• Agnosticisme à la largeur : Élargir une couche ne change pas le schéma de tokenisation, mais correspond simplement à générer un nombre supplémentaire de patchs. Cela permet au modèle de généraliser à des largeurs de couches jamais vues.
• Représentation Unifiée : Le modèle traite conjointement :
  1. Des tokens discrets représentant l'architecture (largeurs des couches).
  2. Des patchs continus représentant les poids.
     Ces deux modalités sont unifiées dans une seule séquence pour le modèle Transformer.

C. Architecture du Modèle (Diffusion Transformer)

NNiT utilise un Diffusion Transformer (DiT) comme colonne vertébrale :

• Conditionnement Multi-modal : Il utilise des embeddings d'étape de temps (timestep) distincts pour l'architecture et les poids, permettant de réguler dynamiquement le processus de débruitage.
• Mixture of Noise Levels (MoNL) : Le modèle est entraîné avec deux modes :
  • Génération conjointe : L'architecture et les poids sont tous deux bruités ($p(a, w)$).
  • Synthèse conditionnelle : L'architecture est fixe (non bruitée) et seuls les poids sont générés ($p(w|a)$).

3. Contributions Clés

1. Alignement de l'espace de poids : Démonstration que les GHNs avec décodeur CNN réduisent la variabilité induite par la permutation et créent un champ de paramètres structuré, essentiel pour la tokenisation par patchs.
2. Tokenisation par patchs agnostique à la largeur : Une nouvelle représentation des poids qui permet la synthèse "zero-shot" (sans entraînement préalable) pour des topologies d'architectures totalement nouvelles.
3. NNiT : Un modèle de diffusion multimodal capable de générer conjointement des architectures et des poids, ou de synthétiser des poids pour une architecture donnée, en traitant les deux comme une séquence unique.

4. Résultats Expérimentaux

Les expériences ont été menées sur la tâche de contrôle robotique ManiSkill3 (tâches de manipulation comme PickCube, PushCube, StackCube) utilisant des politiques MLP.

• Génération Conditionnelle (Zero-Shot) :
  • Sur des architectures vues lors de l'entraînement, NNiT et les méthodes de base (D2NWG, SANE) obtiennent de bonnes performances.
  • Sur des architectures non vues (largeurs différentes), les méthodes de base échouent drastiquement (taux de succès chutant à 0-59%).
  • NNiT maintient un taux de succès > 85% sur des topologies jamais vues, prouvant sa capacité de généralisation structurelle.
• Génération Conjointe :
  • NNiT réussit à générer spontanément des politiques complètes (architecture + poids) avec des taux de succès proches de 100% (99-100% sur PickCube et PushCube), sans avoir besoin d'une architecture d'entrée fixe.
• Validation de l'alignement :
  • L'analyse visuelle (Figure 3) confirme que les poids générés par GHN présentent des structures spatiales cohérentes (bandes), contrairement au bruit non structuré des poids SGD.
  • L'analyse de diversité (Tableau 1) montre que l'alignement par GHN n'entraîne pas d'effondrement de mode (mode collapse) ; les politiques restent diversifiées.

5. Signification et Impact

• Avancée Théorique : L'article résout le problème fondamental de la symétrie de permutation en transformant les poids en un champ spatial structuré, rendant possible l'application des modèles de diffusion (initialement conçus pour les images/vidéos) à la génération de réseaux de neurones.
• Flexibilité Architecturale : NNiT permet de concevoir des réseaux de neurones pour des contraintes matérielles spécifiques (taille de mémoire, puissance de calcul) ou des tâches nouvelles sans réentraînement, simplement en spécifiant les tokens d'architecture.
• Potentiel pour l'IA Embodied : Cette approche ouvre la voie à l'apprentissage méta et à l'adaptation rapide de politiques robotiques pour différents corps (morphologies) ou environnements, en générant instantanément des contrôleurs optimaux adaptés à la structure requise.

En résumé, NNiT représente un changement de paradigme : au lieu d'apprendre à prédire des vecteurs de poids rigides, il apprend à générer des champs de poids structurés et flexibles, découplant la logique fonctionnelle des dimensions fixes de la matrice.