Deep Hierarchical Learning with Nested Subspace Networks for Large Language Models

Dit artikel introduceert Nested Subspace Networks (NSNs), een nieuw architecturaal paradigma dat het mogelijk maakt om grote taalmodellen dynamisch en continu aan te passen aan verschillende rekentijd-begrotingen tijdens de inferentie, waardoor een flexibele afweging tussen prestaties en efficiëntie wordt bereikt zonder de noodzaak van het trainen van meerdere aparte modellen.

De kern

🚗 De "Slimme Auto" voor AI: Eén Model voor Elke Situatie

Stel je voor dat je een auto hebt die je elke dag gebruikt.

• Soms rijd je naar de supermarkt: je hebt dan een kleine, zuinige auto nodig (goedkoop, snel, weinig brandstof).
• Soms moet je een zware last vervoeren of een gevaarlijke bergweg nemen: dan heb je een grote, krachtige vrachtwagen nodig (veel vermogen, maar veel brandstof).

In de wereld van kunstmatige intelligentie (AI) is dit momenteel een groot probleem. Meestal heb je twee aparte modellen:

1. Een klein, snel model voor simpele vragen.
2. Een enorm, traag model voor moeilijke taken.

Het probleem is dat je deze twee aparte modellen moet trainen (leren). Dat kost enorm veel tijd, geld en energie. En als je midden in de dag merkt dat je batterij bijna leeg is (of je internettraject traag), kun je niet zomaar van de "vrachtwagen" op de "fiets" stappen. Je zit vast aan wat je hebt.

💡 De Oplossing: De "Chameleonsauto" (NSN)

De auteurs van dit paper (Paulius Rauba en Mihaela van der Schaar) hebben een nieuwe manier bedacht om dit op te lossen. Ze noemen het Nested Subspace Networks (NSN).

Stel je voor dat ze in plaats van twee aparte voertuigen, één super-chameleonsauto bouwen.

• Deze auto heeft een magische motor die je terwijl je rijdt (tijdens het gebruik) kunt aanpassen.
• Je kunt hem veranderen in een kleine, zuinige auto voor de supermarkt.
• Je kunt hem direct veranderen in een krachtige vrachtwagen voor de zware last.
• En alles daar tussenin: een sportwagen, een busje, etc.

Het geheim? Het is precies dezelfde auto, met dezelfde onderdelen, maar je gebruikt er op dat moment minder of meer van.

🧩 Hoe werkt het? (De Legpuzzel)

Om dit te begrijpen, moeten we kijken naar hoe een AI "denkt". Een AI bestaat uit lagen van wiskundige berekeningen. Stel je deze lagen voor als een grote legpuzzel.

• De oude manier: Om een snellere AI te maken, nam je de puzzel en snijd je er stukken uit (pruning) of bouw je een hele nieuwe, kleinere puzzel. Dat is lastig en je verliest vaak stukjes kennis.
• De nieuwe manier (NSN): Ze bouwen de puzzel anders. Ze maken één grote, perfecte puzzel, maar ze zorgen ervoor dat de stukjes in elkaar zitten als Russische poppetjes (vandaar de naam "Nested" of "Genest").
  • De kleinste poppetje (rank 1) bevat de allerbelangrijkste stukjes van de puzzel.
  • Het tweede poppetje (rank 2) zit eromheen en voegt iets extra's toe.
  • Het grootste poppetje (rank 100) bevat alles.

Als je een snelle AI nodig hebt, gebruik je alleen het kleinste poppetje. Als je een slimme AI nodig hebt, gebruik je het grootste. Omdat ze allemaal uit dezelfde basisstukken bestaan, werkt het naadloos.

🎓 De "Leraar" die slim lesgeeft

Nu is er een probleem: Hoe leer je deze auto om zowel als fiets als als vrachtwagen te rijden? Als je de auto alleen maar traint om een vrachtwagen te zijn, werkt hij als fiets heel slecht.

De auteurs hebben een slimme lesmethode bedacht:

1. Ze laten de auto tegelijkertijd oefenen als fiets (moeilijk, want weinig kracht) en als vrachtwagen (makkelijker, want veel kracht).
2. Ze gebruiken een slimme leraar (een wiskundig trucje genaamd "uncertainty-aware objective"). Deze leraar kijkt naar hoe goed de auto het doet in beide rollen.
   • Als de auto het als fiets erg slecht doet (hij is onzeker), geeft de leraar meer aandacht aan dat onderdeel.
   • Als hij het als vrachtwagen goed doet, geeft hij minder aandacht.
3. Hierdoor leert de auto in één keer om in elke situatie goed te presteren.

🏆 Wat levert dit op?

De resultaten zijn indrukwekkend:

• Je kunt een enorm groot model (zoals een slimme taalcomputer) nemen dat al is getraind.
• Je past het chirurgisch aan (vervangt alleen de "motor" van de auto, niet de hele auto).
• Je krijgt een model dat 50% minder energie verbruikt, maar nog steeds 95% van de slimheid behoudt.
• Je kunt dit direct doen op je telefoon als de batterij laag is, of op een server als er veel vraag is.

🌍 Waarom is dit belangrijk?

Vroeger moest je kiezen: Of je hebt een slimme, dure AI, of je hebt een snelle, goedkope AI.
Met deze nieuwe methode (NSN) kun je alles in één.

Het is alsof je eindelijk een telefoon hebt die, afhankelijk van wat je doet, automatisch schakelt tussen een batterijbesparende modus en een gamemodus, zonder dat je de telefoon hoeft te vervangen of opnieuw te installeren. Dit maakt AI veel flexibeler, goedkoper en makkelijker te gebruiken in de echte wereld.

Kortom: Ze hebben een manier gevonden om één AI-model te maken dat zich aanpast aan elke situatie, net als een chameleonsauto, zonder dat je duizenden aparte modellen hoeft te bouwen.

Technische samenvatting

Probleemstelling

Grote neurale netwerken, zoals Large Language Models (LLMs), worden doorgaans getraind met een vast computatiebudget. Dit creëert een rigide afweging tussen prestaties en efficiëntie die niet geschikt is voor implementatie in omgevingen met beperkte middelen of dynamische eisen (bijv. batterijverbruik op mobiele apparaten, veiligheidskritieke taken die meer rekenkracht vereisen).

Bestaande oplossingen hebben twee grote nadelen:

1. Statische methoden: Technieken zoals netwerkknippen (pruning) of kennisdistillatie vereisen het trainen van aparte, gespecialiseerde modellen voor elk specifiek budget. Het creëren van een nieuw model voor een ander budget is computatief prohibitief en biedt geen flexibiliteit voor on-the-fly aanpassing.
2. Dynamische methoden: Bestaande dynamische netwerken (zoals "slimmable networks") vereisen vaak complexe trainingsschema's vanaf nul en zijn moeilijk toe te passen op reeds getrainde foundation models. Bovendien bieden ze vaak slechts een discrete set van operationele punten in plaats van een gladde, continue afweging.

Het doel is een enkele architectuur te ontwikkelen die onmiddellijke aanpasbaarheid, toepasbaarheid op bestaande modellen en granulaire controle over het rekenbudget biedt zonder extra fine-tuning-kosten bij het testen.

Methodologie: Nested Subspace Networks (NSNs)

De auteurs introduceren Nested Subspace Networks (NSNs), een nieuw architecturaal paradigma dat een continue hiërarchie van modellen binnen één enkele set gewichten realiseert.

1. Architectuur: Geneste Subruimtes
In plaats van traditionele lineaire lagen $f(x) = Wx + b$, repareert NSN de gewichtsmatrix $W$ via een lage-rang factorisatie $W \approx BA$.

• Er wordt een enkel paar factor-matrices $(A, B)$ gebruikt, waarbij $A \in \mathbb{R}^{R \times d_{in}}$ en $B \in \mathbb{R}^{d_{out} \times R}$ ($R$ is de maximale rang).
• Voor een specifieke rang $r$ wordt de effectieve gewichtsmatrix $W_r$ gevormd door alleen de eerste $r$ rijen van $A$ en de eerste $r$ kolommen van $B$ te gebruiken: $W_r = B_r A_r$.
• Geneste Eigenschap: Door deze constructie is het beeld (image) van de transformatie bij rang $r$ een strikte subruimte van het beeld bij rang $r+1$. Dit betekent dat het model bij een lagere rang een strikt subset is van het model bij een hogere rang.

2. Training: Onzekerheidsbewuste Doelfunctie
Het trainen van één set gewichten die optimaal presteert over de hele rang-hiërarchie is uitdagend. Een simpele truncatie van een hoog-rang model werkt niet goed. De auteurs lossen dit op met een multi-rang onzekerheidsbewuste trainingsstrategie:

• Het probleem wordt benaderd als multi-task learning, waarbij elke rang een taak is met een verschillende intrinsieke moeilijkheidsgraad.
• Er worden leerbare variance-parameters $\sigma^2_k$ (of log-variances $s_k = \log \sigma^2_k$) geïntroduceerd voor elke rang $k$.
• De totale loss-functie combineert de cross-entropy loss van een "anker-rang" (meestal de maximale rang) en een "variant-rang" (een willekeurige lagere rang), gewogen door de onzekerheid:
  $$\mathcal{L}_{total} = (e^{-s_{\tilde{R}}} \mathcal{L}_{CE}(\tilde{R}) + s_{\tilde{R}}) + (e^{-s_{r}} \mathcal{L}_{CE}(r) + s_{r})$$
• Deze methode zorgt ervoor dat rangen met hogere onzekerheid (moeilijker te leren, vaak lagere rangen) automatisch minder gewicht krijgen in de gradiënt-update, terwijl rangen met lagere onzekerheid meer invloed hebben. Dit stabiliseert het trainen van de hiërarchie.

3. Post-hoc Toepassing op LLMs
NSNs kunnen chirurgisch worden toegepast op bestaande, voorgeprogrammeerde LLMs:

• De lineaire lagen in de MLP-blokken van een transformer worden vervangen door NSN-lagen.
• De initiële factor-matrices $(A, B)$ worden geïnitieerd via Singular Value Decomposition (SVD) van de oorspronkelijke gewichten. Hierdoor wordt de informatie van het getrainde model behouden en is slechts een korte fine-tuning fase nodig.

Kernbijdragen

1. NSN Architectuur: Een nieuwe architectuur die een continue hiërarchie van modellen mogelijk maakt binnen één set gewichten, gebaseerd op de geneste subruimte-eigenschap.
2. Trainingsstrategie: Een theoretisch onderbouwde, onzekerheidsbewuste trainingsdoelstelling die het gezamenlijk optimaliseren van meerdere rangen mogelijk maakt zonder instabiliteit.
3. Theoretische Garanties: Bewijzen dat de prestatie-computatie frontier glad en voorspelbaar is, zelfs voor rangen die niet expliciet tijdens het trainen zijn geoptimaliseerd (interpolatie), mits een "energy decay" aanname geldt voor de geleerde basisvectoren.
4. Empirische Validatie: Uitgebreide experimenten die aantonen dat NSNs succesvol kunnen worden toegepast op grote pre-getrainde modellen (Pythia, GPT-Neo, Gemma, Qwen) met minimale prestatieverlies.

Resultaten

• Gladde Frontier: NSNs bieden een gladde, voorspelbare afweging tussen rekenkosten (FLOPs) en nauwkeurigheid. Bijvoorbeeld, een enkel NSN-aangepast model kan een reductie van 50% in inference FLOPs bereiken met slechts een daling van 5 procentpunten in nauwkeurigheid.
• Interpolatie: Modellen presteren stabiel en monotoon op interpolatie-rangen (rang-waarden die niet expliciet in de training zijn gebruikt), wat bevestigt dat de geneste structuur werkt zoals bedoeld.
• Energieverval: Experimenten tonen aan dat de geleerde basisvectoren een duidelijke afname in "energie" (norm) vertonen naarmate de index toeneemt. Dit bevestigt de theoretische aanname dat de belangrijkste informatie in de lagere rangen zit.
• Vergelijking: In tegenstelling tot LoRA (dat statisch is voor een vaste rang) of slimmable networks (die vaak van nul moeten worden getraind), kunnen NSNs direct worden toegepast op bestaande modellen en bieden ze een continue spectrum van opties.

Betekenis en Impact

Dit werk stelt een fundamentele verschuiving voor in hoe we grote modellen deployen. In plaats van een verzameling van statische, gespecialiseerde modellen te onderhouden, stelt NSNs organisaties in staat om één enkel adaptief model te gebruiken dat dynamisch kan schalen tussen snelheid en nauwkeurigheid afhankelijk van de beschikbare middelen (bijv. batterijstatus, netwerkvertraging, of veiligheidsvereisten).

De methode is bijzonder waardevol omdat:

• Het post-hoc toepasbaar is op de huidige ecosysteem van foundation models zonder dat deze opnieuw getraind hoeven te worden.
• Het granulaire controle biedt, waardoor het mogelijk is om precies het juiste evenwicht te vinden voor specifieke gebruiksscenario's.
• Het de weg vrijmaakt voor de volgende generatie van adaptieve foundation modellen die efficiënter en robuuster zijn in real-world omgevingen.

Samenvattend introduceren NSNs een krachtig raamwerk dat de rigiditeit van huidige deep learning-deployments doorbreekt en een nieuwe standaard zet voor flexibele, resource-bewuste AI-systemen.