Representational Alignment Across Model Layers and Brain Regions with Multi-Level Optimal Transport

Probleemstelling

Bestaande methoden voor het vergelijken van neurale representaties (zowel tussen kunstmatige netwerken als tussen netwerken en de hersenen) maken gebruik van laag-voor-laag matching (layer-wise matching). Bij deze aanpak wordt elke laag van een bronnetwerk onafhankelijk gekoppeld aan de "beste" overeenkomstige laag in een doelnetwerk, gebaseerd op een gelijkaardigheidsmaatstaf (zoals RSA, CKA of lineaire voorspelbaarheid).

Deze aanpak heeft fundamentele beperkingen:

1. Rigiditeit: Het dwingt een één-op-één correspondentie af, wat faalt wanneer netwerken verschillende dieptes hebben of wanneer de representatie van één laag verspreid is over meerdere lagen in een ander netwerk.
2. Asymmetrie: De resultaten hangen af van de richting van de vergelijking (bron vs. doel).
3. Gebrek aan globale coherentie: Door elke match onafhankelijk te optimaliseren, wordt de globale activatiestructuur genegeerd, wat leidt tot overfitting op ruis en het missen van bredere hiërarchische patronen.
4. Geen uniek score: Er ontbreekt één enkele score die de algehele uitlijning van de volledige netwerken kwantificeert.

Methodologie: Multi-Level Optimal Transport (MOT)

De auteurs stellen Multi-Level Optimal Transport (MOT) voor, een unificerend kader dat globale, zachte (soft) uitlijningen infereert in plaats van harde één-op-één koppelingen.

Het Kader:
MOT werkt op twee niveaus binnen een Optimal Transport (OT) formulering:

1. Binnenste Niveau (Neuron-tot-Neuron):
   Voor elke mogelijke koppel van lagen $(\ell, m)$ wordt een OT-probleem opgelost om de dissimilariteit tussen de neuronen in die lagen te berekenen. Dit resulteert in een zachte koppelingsmatrix $Q_{\ell m}$ die aangeeft hoe sterk een neuron in laag $\ell$ correspondeert met neuronen in laag $m$. Dit maakt het mogelijk om neuronen te koppelen aan mengsels van neuronen in andere lagen.

2. Buitenste Niveau (Laag-tot-Laag):
   De kosten uit het binnenste niveau vormen een kostenmatrix $C$ tussen alle lagen van het bronnetwerk en het doelnetwerk. MOT lost een tweede OT-probleem op om een globale koppelingsmatrix $P$ te vinden.

   • Randvoorwaarden: Elke bronlaag moet 100% van zijn "massa" verdelen over de doellagen, en de totale massa die elke doellaag ontvangt, moet gebalanceerd zijn.
   • Resultaat: In plaats van één doel-laag te kiezen, mag een bronlaag zijn representatieve massa verdelen over meerdere doellagen. Dit lost automatisch diepte-mismatches op (bijv. een enkele laag in een ondiep netwerk kan corresponderen met een reeks lagen in een dieper netwerk).

Rotatie-invariantie (MOT + R):
Omdat veel neurale representaties rotatie-invariant zijn (de geometrie is hetzelfde, maar de coördinatenbasis verschilt), introduceren de auteurs een extensie die orthogonale rotatiematrices $R_{\ell m}$ optimaliseert. Dit minimaliseert de reconstructiefout terwijl het de geometrische equivalentie behoudt, wat cruciaal is voor Vision Transformers.

Evaluatie:
De kwaliteit van de uitlijning wordt gemeten door de oorspronkelijke activaties te reconstrueren op basis van de geleerde transportplannen en de correlatie met de werkelijke (gehouden) activaties te berekenen.

Belangrijkste Bijdragen

1. Principieel Globaal Uitlijningskader: MOT optimaliseert alle laag-correspondenties gezamenlijk, wat leidt tot symmetrische, globaal consistente toewijzingen met één enkele netwerk-niveau score.
2. Natuurlijke Afhandeling van Diepte-mismatches: Door zachte, veel-op-veel mapping toe te staan, kan MOT netwerken met verschillende dieptes uitlijnen zonder ongeschikte één-op-één koppelingen te forceren.
3. Rotatie-invariante Extensie: De MOT+R variant kan gemeenschappelijke representatieve kenmerken herkennen die in geroteerde subruimtes zijn ingebed.
4. Drie-niveau MOT: Een bewijs van concept waarbij MOT wordt uitgebreid naar drie niveaus om volledige trainings-trajecten van modellen over verschillende checkpoints uit te lijnen.

Resultaten

De methode werd geëvalueerd op drie domeinen:

1. Grote Taalmodellen (LLMs):

   • Vergelijkingen tussen modellen van verschillende families (LLaMA vs. Qwen) en schalen.
   • Resultaat: MOT behaalde hogere reconstructiescores dan pairwise baselines. De transportplannen toonden een duidelijke diagonale structuur (vroege lagen corresponderen met vroege lagen, diepe met diepe).
   • Insight: MOT onthulde dat lagen in ondiepere modellen hun massa verspreiden over meerdere opeenvolgende lagen in diepere modellen, wat suggereert dat diepere netwerken berekeningen verfijnen en verspreiden.

2. Visuele Cortex (fMRI):

   • Uitlijning van hersengebieden (V1-V4) tussen verschillende proefpersonen.
   • Resultaat: MOT herstelde de verwachte biologische correspondenties (bijv. V1 van persoon A naar V1 van persoon B) met een hoge structuur, terwijl pairwise OT willekeurige of ongeordende mapping produceerde. De scores waren vergelijkbaar met baselines, maar de interpretatie was superieur.

3. Vision Modellen (Vision Transformers):

   • Vergelijking van DINOv2 en ViT-MAE modellen.
   • Resultaat: De standaard MOT presteerde niet consistent beter dan pairwise OT. Echter, de MOT+R variant (met rotatie) behaalde aanzienlijk hogere scores en produceerde zeer schone, hiërarchische uitlijningen. Dit benadrukt de noodzaak van rotatie-invariantie voor Vision Transformers.

4. Trainings Trajectoires:

   • Drie-niveau MOT toonde een coherente uitlijning tussen checkpoints van twee Pythia-modellen, terwijl pairwise matching geen duidelijke traject-correspondentie kon vinden.

Betekenis en Conclusie

Dit paper introduceert een fundamentele verschuiving in hoe we neurale representaties vergelijken. In plaats van lokale, greedy matches te zoeken, biedt MOT een globaal geoptimaliseerde oplossing die de onderliggende hiërarchische structuur van zowel kunstmatige als biologische netwerken respecteert.

De belangrijkste implicaties zijn:

• Interpreteerbaarheid: MOT levert transportplannen die natuurlijke hiërarchische patronen blootleggen die door andere methoden worden gemist.
• Robuustheid: Het is minder gevoelig voor ruis en overfitting door het gebruik van globale randvoorwaarden.
• Toepasbaarheid: Het werkt effectief over verschillende schalen, architecturen en zelfs tussen biologische en kunstmatige systemen.
• Toekomst: De methode biedt een meetinstrument om theorieën over "universele representaties" en convergentie in intelligentie systemen empirisch te testen.

De auteurs erkennen wel dat de methode computationeel intensief is ($O(L^2 n^3 \log n)$), wat schaalbaarheid naar zeer grote modellen uitdagingen biedt, maar het biedt een nieuw paradigma voor representatieve analyse.