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

Each language version is independently generated for its own context, not a direct translation.

Stel je voor dat je een meesterchef bent die altijd dezelfde soeprecepten maakt, maar dan met een heel specifiek probleem: elke keer als je een nieuw recept probeert te maken, moet je precies weten hoeveel wortels en aardappels erin gaan. Als je het recept van iemand anders wilt kopiëren, maar die persoon heeft net iets meer aardappels gebruikt dan jij gewend bent, dan werkt je kopie niet meer. Je moet het hele recept opnieuw uitvinden.

In de wereld van kunstmatige intelligentie (AI) gebeurt precies dit. AI-modellen zijn als die soeprecepten. Ze bestaan uit miljoenen getallen (de "gewichten") die bepalen hoe slim ze zijn. Het grote probleem is dat deze getallen vaak in een willekeurige volgorde staan. Als je een AI wilt "leren" om nieuwe modellen te bedenken, raakt hij de draad kwijt zodra de grootte van het model verandert.

Hier komt NNiT (Neural Network Diffusion Transformers) om de hoek kijken. De onderzoekers hebben een slimme oplossing bedacht die we kunnen vergelijken met het bouwen van een legpuzzel of het bakken van een taart.

1. Het Probleem: De Willekeurige Wolk

Stel je voor dat je een wolk van blokken hebt. Je kunt ze op elke manier stapelen en het resultaat is nog steeds een wolk. In AI noemen we dit "permutatie-symmetrie". Het betekent dat je dezelfde slimme AI kunt bouwen, maar de getallen erin staan in een heel andere volgorde. Voor een computer is het alsof je probeert een foto te maken van een wolk, maar elke keer als je een nieuwe foto maakt, zijn de wolkenblokken verplaatst. De computer ziet geen patroon en kan dus niet leren hoe je een nieuwe wolk moet "tekenen".

2. De Oplossing: De "Chef" die Alles Netjes Legt (GHN)

De onderzoekers gebruiken een slimme truc genaamd GHN (Graph HyperNetworks). Stel je dit voor als een super-chef die niet zelf de soep kookt, maar wel het recept schrijft.

Normale AI's (zoals SGD) koken de soep en gooien de ingrediënten erin zonder te kijken. Het resultaat is een rommelige soep.
De GHN-chef daarentegen heeft een vaste manier van werken. Hij legt de wortels altijd links en de aardappels rechts. Hij zorgt ervoor dat de blokken in de AI altijd netjes en op een logische manier gerangschikt zijn.

Dit is cruciaal. Omdat de GHN-chef alles netjes ordent, ontstaat er een patroon. De getallen in de AI zijn niet meer willekeurig; ze hebben een lokale structuur, alsof ze in een strakke rij staan.

3. De Magie: Het Legpuzzel (NNiT)

Nu komt het echte genie van NNiT. Omdat de GHN-chef alles netjes heeft gerangschikt, kunnen de onderzoekers de AI-blokken niet meer zien als één grote, rommelige hoop. Ze kunnen ze zien als kleine stukjes van een legpuzzel (of "patches", zoals ze in het paper zeggen).

Vroeger: Je probeerde een hele foto van 1000x1000 pixels te genereren. Als je de foto groter maakte (bijvoorbeeld 2000x2000), wist de computer niet wat hij moest doen.
Nu met NNiT: De computer leert hoe je een klein stukje van 10x10 pixels maakt.
- Wil je een grotere AI? Dan laat de computer gewoon meer van diezelfde puzzelstukjes maken.
- Het maakt niet uit of de AI nu klein of groot is; het principe blijft hetzelfde: "Maak een stukje, en nog een stukje, en nog een stukje."

Dit noemen ze "width-agnostic". Het betekent: "Onafhankelijk van de breedte". Je kunt de AI zo smal of zo breed maken als je wilt, en NNiT weet precies hoe hij de puzzelstukjes moet leggen.

4. Wat levert dit op? (De Robot)

De onderzoekers hebben dit getest op robots die taken moeten uitvoeren (zoals blokken stapelen of duwen).

Ze trainden NNiT op robots met een bepaalde grootte.
Vervolgens vroegen ze NNiT om een robot te maken die dubbel zo breed was, of een heel ander formaat had.
Resultaat: De oude methoden faalden volledig. De robots vielen om of konden niets doen.
NNiT: De robot deed de taak met succes! Omdat NNiT het patroon had geleerd in plaats van het exacte formaat, kon hij direct een nieuw, groter model "dromen" dat perfect werkte.

Samenvatting in één zin

NNiT is als een meester-architect die niet leert hoe je één specifiek huis bouwt, maar leert hoe je muren en ramen bouwt; hierdoor kan hij direct een klein huisje of een gigantisch kasteel ontwerpen zonder ooit een blauwdruk voor die specifieke grootte te hebben gezien.

Dit is een enorme stap voorwaarts omdat het betekent dat we in de toekomst AI's kunnen maken die zich direct aanpassen aan de hardware die we hebben, zonder dat we ze opnieuw hoeven te trainen.

Each language version is independently generated for its own context, not a direct translation.

Probleemstelling

Het genereren van volledige neurale netwerken via generatieve modellen (zoals diffusion-modellen) stuit op twee fundamentele obstakels:

Afhankelijkheid van vaste matrixdimensies: Traditionele methoden flattenen gewichten naar vaste vectoren. Dit koppelt het generatieve prior strikt aan de specifieke breedte van de lagen tijdens het trainen. Zodra de breedte van een laag verandert (bijv. een bredere of smallere architectuur), wordt de tokenisatie ongeldig en faalt het model om te generaliseren naar ongeziene topologieën.
Permutatie-symmetrie: In neurale netwerken (vooral MLP's) kunnen neurale eenheden binnen een laag in willekeurige volgorde worden gerangschikt zonder de input-output functie te veranderen. Dit zorgt ervoor dat aangrenzende gewichten in een matrix geen ruimtelijke correlatie hebben; ze lijken op ruis. Dit maakt het moeilijk om gewichten te modelleren als een gestructureerd veld (zoals bij afbeeldingen), wat essentieel is voor patch-based generatie.

Bestaande oplossingen proberen dit op te lossen door gewichten te canonicaliseren of in een latente ruimte te comprimeren, maar deze methoden blijven vaak kwetsbaar voor veranderingen in de breedte van de architectuur.

Methodologie

De auteurs introduceren NNiT (Neural Network Diffusion Transformers), een framework dat neurale netwerkgeneratie behandelt als een multimodaal sequentiemodelleertaak. De aanpak bestaat uit drie kerncomponenten:

1. Structurele Alignering via Graph HyperNetworks (GHN)

Om het probleem van permutatie-symmetrie op te lossen, gebruiken de auteurs Graph HyperNetworks (GHNs) met een Convolutional Neural Network (CNN) decoder als data-generator.

In plaats van gewichten te trainen via Stochastic Gradient Descent (SGD), die willekeurige permutaties produceert, gebruikt de GHN een grafische representatie van de architectuur.
De CNN-decoder introduceert een expliciete lokaliteitsbias. Hierdoor worden gewichten gegenereerd die consistente lokale ruimtelijke correlaties vertonen (bijv. zichtbare banden in de gewichtsverdeling), ongeacht de seed.
Dit creëert een "structuraal uitgelijnde" gewichtsruimte waarin gewichten kunnen worden behandeld als een continu ruimtelijk veld in plaats van onafhankelijke vectoren.

2. Patch-gebaseerde Tokenisatie (Width-Agnostic)

NNiT tokeniseert de uitgelijnde gewichtstensors in $p \times p$ patches (vergelijkbaar met Vision Transformers voor afbeeldingen).

Architectuur: Wordt gemodelleerd als een reeks discrete tokens (laagbreedtes).
Gewichten: Worden gemodelleerd als continue patches.
Width-Agnostic: Omdat gewichten als patches worden behandeld, betekent het vergroten van een laag simpelweg het genereren van extra patches, zonder de tokenisatieschema te veranderen. Dit maakt het model breedte-onafhankelijk.

3. Multimodaal Diffusion Transformer (DiT)

Het model gebruikt een Diffusion Transformer backbone die discrete architectuur-tokens en continue gewicht-patches in één enkele sequentie verwerkt.

Mixture of Noise Levels (MoNL): Het model wordt getraind om twee modi te hanteren:
- Gecombineerde generatie: Het leert de gezamenlijke verdeling $p(a, w)$ om zowel architectuur als gewichten tegelijk te genereren.
- Conditionele synthese: Het leert $p(w|a)$ om gewichten te genereren voor een specifieke, door de gebruiker opgegeven architectuur.
Adaptive Layer Norm (AdaLN-Zero): Hiermee wordt het model dynamisch gestuurd op basis van het ruisniveau van zowel de architectuur- als de gewicht-modi.

Belangrijkste Bijdragen

Structurale Alignering: Het aantonen dat GHNs met een CNN-decoder de permutatie-symmetrie effectief opheffen en een gewichtsruimte creëren met stabiele lokale correlaties, wat noodzakelijk is voor patch-based generatie.
Width-Agnostic Tokenisatie: De introductie van een patch-gebaseerde representatie voor gewichten, waardoor generatie mogelijk is voor architecturen met willekeurige breedtes zonder hertraining.
NNiT Framework: Een uniek multimodaal diffusion-model dat architectuur en parameters gezamenlijk modelleert, waardoor zowel co-design als conditionele synthese mogelijk is.

Resultaten

De methoden zijn geëvalueerd op ManiSkill3, een robuust robotica-benchmark voor manipulatietaken (zoals het pakken en stapelen van kubussen). De evaluatie focust op zero-shot generalisatie naar architecturen die niet tijdens het trainen zijn gezien.

Zero-Shot Generalisatie:
- NNiT: Bereikte een succespercentage van >85% op volledig nieuwe architectuurtopologieën (verschillende breedtes en dieptes).
- Baselines (SANE, D2NWG): Deze methoden faalden dramatisch op ongeziene breedtes. SANE bereikte slechts 0-2% succes, en D2NWG daalde naar 42-59% succes op complexe taken. Dit komt doordat hun vectorisatie-methoden niet schalen met veranderende breedtes.
Multimodale Synthese: NNiT kon succesvol volledige functionele beleidsregels genereren zonder een vooraf gedefinieerde architectuur (alleen op basis van de gezamenlijke verdeling), met successpercentages van 99-100% op bekende taken.
Kwaliteit: De gegenereerde netwerken presteerden op niveau met expert-netwerken die via traditionele training zijn verkregen, zonder dat er mode collapse (verlies van diversiteit) optrad.

Betekenis en Impact

Dit werk opent nieuwe wegen voor Meta-Learning en Embodied AI:

Efficiëntie: Het elimineert de noodzaak om voor elke nieuwe hardware-beperking of taak een nieuw netwerk te trainen. Een enkele generator kan gewichten synthetiseren die voldoen aan specifieke breedte- of rekenbudgetten.
Schaalbaarheid: Door het modelleren van netwerken als sequenties (waarbij diepte als tijdsdimensie en breedte als ruimtelijke dimensie wordt gezien), kan NNiT profiteren van efficiëntie-optimalisaties uit video-diffusiemodellen. Dit maakt het potentieel mogelijk om foundation modellen met miljarden parameters te synthetiseren.
Flexibiliteit: Het bewijst dat neurale netwerken niet hoeven te worden getraind, maar kunnen worden "gegenereerd" als een gestructureerd veld, wat de weg vrijmaakt voor snelle adaptatie aan veranderende omgevingen en hardware.

Kortom, NNiT doorbreekt de koppeling tussen generatieve modellen en vaste matrixdimensies, waardoor robuuste, breedte-onafhankelijke generatie van neurale netwerken mogelijk wordt.