Technische Samenvatting: Tensor-ähnlichkeit voor Mechanistische Interpretatie

Probleemstelling

Mechanistische interpretatie streeft ernaar neurale netwerken te ontleden in betekenisvolle, samenstellende onderdelen om hun interne berekeningen te begrijpen. Een kritieke voorwaarde voor dit vakgebied is het vermogen om rigoureus te verifiëren of twee ontbonden componenten (of twee verschillende modellen) dezelfde onderliggende functie implementeren. Huidige evaluatiemethoden lijden aan twee primaire beperkingen:

1. Empirische/gedragsmatige similariteit: Methoden die modeluitkomsten of activaties op specifieke datasets vergelijken, zijn blind voor mechanismen buiten de verdeling (out-of-distribution, OOD). Ze kunnen problemen zoals backdoors niet detecteren die alleen op specifieke inputs worden geactiveerd, of ze falen in generalisatie als de ontleding gebaseerd is op benaderende heuristieken in plaats van de ware functie.
2. Parameter-gebaseerde similariteit: Methoden die weegmatrices vergelijken (bijvoorbeeld matrix-cosinus-similariteit) zijn zeer gevoelig voor symmetrieën in de weegruimte (zoals permutatie en schaling). Ze falen in het herkennen van functionele equivalentie wanneer modellen op verschillende manieren zijn getraind of wanneer mechanismen over lagen zijn verdeeld in plaats van geconcentreerd in enkele neuronen.

Er ontbreekt een rigoureuze, dataset-onafhankelijke metriek die invariant is voor deze symmetrieën en in staat is globale functionele equivalentie vast te leggen.

Methodologie

De auteurs stellen Tensor Similarity voor, een op gewichten gebaseerde metriek ontworpen voor tensor-gebaseerde neurale netwerken (specifiek multilineaire modellen zoals bilineaire transformers). Deze benadering wisselt het gebruik van standaard niet-lineaire activatiefuncties in voor polynomiale (multilineaire) architecturen om algebraïsche analyse van gewichten mogelijk te maken.

Kernarchitecturale Verschuiving

De methode is gebaseerd op multilineaire modellen waarbij activatiefuncties worden vervangen door multiplicatieve gating (bijvoorbeeld bilineaire lagen en bilineaire attention).

• Bilineaire lagen: Mappen een opgeheven input $\tilde{x} = (1, x)$ af op een output via een polynoom van graad 2: $A(x) = D((L\tilde{x}) \odot (R\tilde{x}))$.
• Multilineaire representatie: Deze lagen kunnen worden uitgedrukt als een contractie van een gewichtstensor $A$ met $n$ kopieën van de input: $A(x) = A \cdot \text{in}^{\otimes n} \tilde{x}$.
• Diepe modellen: Diepe multilineaire netwerken worden opgebouwd als boomgestructureerde decomposities van globale tensoren, wat efficiënte berekening mogelijk maakt zonder de exponentieel grote globale tensor te materialiseren.

De Tensor Similarity Metriek

De metriek wordt gedefinieerd als een veralgemeend inproduct tussen twee gewichtstensoren $A$ en $B$, genormaliseerd door hun geïnduceerde normen:
$$\text{sim}_M(A, B) = \frac{\langle A | M | B \rangle}{\|A\|_M \|B\|_M}$$
waarbij $M$ een metriektensor is.

Om ervoor te zorgen dat de metriek voldoet aan de criteria van symmetrie-invariantie en dataset-onafhankelijkheid, leiden de auteurs een specifieke vorm voor $M$ af:

1. Symmetrisatie: Gebaseerd op de polarisatie-isomorfie stellen de auteurs vast dat twee multilineaire modellen functioneel equivalent zijn dan en slechts dan als hun gewichtstensoren hetzelfde representant hebben in de symmetrische deelruimte. De metriek projecteert tensoren op deze symmetrische deelruimte om vrijheidsgraden gerelateerd aan input-permutatiesymmetrieën te verwijderen.
2. Gaussische metriek: Om de metriek verder te verfijnen en te relateren aan gedragsmatige similariteit, definiëren de auteurs $M$ met behulp van de verwachting van de inputtensor onder een standaard Gaussische verdeling ($x \sim \mathcal{N}(0, I)$). Met behulp van Isserlis' stelling resulteert dit in een gesloten-vorm metriek die het inproduct van activaties over de volledige Gaussische inputruimte berekent, in plaats van een bemonsterde dataset.
3. Efficiënte berekening: Een op Gram gebaseerd recursief algoritme maakt de berekening van deze similariteit voor diepe netwerken mogelijk door lokale tensoren laag voor laag te verwerken, waardoor de noodzaak om de volledige globale tensor te instantiëren wordt vermeden.

Belangrijkste Bijdragen

1. Een principieel op gewichten gebaseerde metriek: Het paper introduceert een metriek die invariant is voor symmetrieën in de weegruimte (permutaties, schaling) en onafhankelijk is van de inputdataset. Het verifieert functionele equivalentie formeel door te controleren of twee modellen dezelfde polynomiale logica implementeren.
2. Gesloten-vorm, dataloze analyse: In tegenstelling tot gedragsmatige metrieken vereist tensor-similariteit geen bemonsterde data. Het integreert over de volledige Gaussische inputruimte, waardoor het OOD-mechanismen (zoals backdoors) kan detecteren die onzichtbaar zijn voor standaard prestatiemetrieken.
3. Lokalisatiecapaciteiten: De metriek kan worden "gesneden" om zich te richten op specifieke outputdimensies of "diff"-tensoren om functionele veranderingen tussen checkpoints te isoleren, waardoor het toewijzen van specifieke mechanismen aan specifieke outputs mogelijk wordt.
4. Schaalbaarheid: Het recursieve algoritme zorgt ervoor dat de metriek computationeel hanteerbaar is voor diepe netwerken, en vermijdt exponentiële of kwadratische convergentiegrenzen die geassocieerd zijn met Monte Carlo-bemonstering.

Resultaten

De auteurs valideren tensor-similariteit over vier verschillende domeinen, waarbij ze de superioriteit ervan ten opzichte van bestaande metrieken (CKA, matrix-cosinus, gedragsmatige similariteit) aantonen:

1. Catastrofaal Vergeten (Visie): In een progressieve trainingsopstelling op SVHN lokaliseert tensor-similariteit precies het verlies aan kennis voor specifieke cijfers (bijvoorbeeld cijfer '9') wanneer de dataset wordt gereduceerd. Het toont een scherpe daling in similariteit voor de betrokken klasse terwijl het stabiel blijft voor anderen, terwijl gedragsmatige metrieken ambiguë signalen tonen.
2. Grokking (Modulaire Aritmetiek): Bij taken voor modulaire optelling volgt tensor-similariteit de continue reorganisatie van het model naar specifieke Fourier-frequenties. Het onthult dat structurele reorganisatie doorgaat zelfs nadat de nauwkeurigheid naar perfectie is gesprongen, en biedt een meer leesbaar beeld van trainingsdynamiek dan nauwkeurigheidscursussen alleen.
3. Backdoor-detectie (OOD): Wanneer een backdoor wordt geïnjecteerd in SVHN (geactiveerd door een zwarte ruit-patch), falen standaard gedragsmatige metrieken op schone data om de verandering te detecteren. Tensor-similariteit markeert echter de backdoor uitsluitend op basis van schone data en lokaliseert de verandering tot de specifieke aangevallen klasse ('9'), wat het vermogen aantoont om OOD-mechanismen te detecteren.
4. Taalmodellering: Toegepast op een twee-laags bilineaire attention-transformer getraind op The Pile, onthult tensor-similariteit duidelijke blokkenstructuren en scherpe overgangen in trainingsregimes die door bestaande metrieken worden vervaagd. Het volgt de vorming van complexe n-gramstatistieken effectiever dan gedragsmatige proxies.

Betekenis en Claims

Het paper positioneert tensor-similariteit als een fundamenteel hulpmiddel voor mechanistische interpretatie. De primaire betekenis ligt in het reduceren van het probleem van het meten van similariteit en het verifiëren van trouw van een empirisch benaderingsprobleem tot een opgelost algebraïsch probleem.

• Trouw: Door te vertrouwen op de algebraïsche eigenschappen van multilineaire modellen, zorgt de metriek ervoor dat ontledingen de ware onderliggende functie vastleggen in plaats van benaderende heuristieken.
• Robuustheid: Het biedt sterkere claims over robuustheid door toe te passen op de volledige verdeling, waardoor mechanismen die alleen op specifieke inputs worden geactiveerd (zoals backdoors) effectief worden uitgesloten.
• Beperkingen: De auteurs erkennen nederig dat de methode momenteel afhankelijk is van multilineaire modellen, wat de meeste state-of-the-art architecturen met niet-polynomiale niet-lineariteiten uitsluit. Ze betogen echter dat multilineaire modellen dienen als een cruciaal testbed voor interpretatie en dat er bewijs is dat ze de prestaties van standaardarchitecturen kunnen evenaren.
• Toekomstperspectief: Het werk suggereert dat de beperking tot multilineaire modellen een noodzakelijke trade-off is voor hanteerbare algebraïsche analyse, maar het baant de weg voor toekomstige toewijzingsmethoden (bijvoorbeeld het berekenen van similariteit tussen paden in diepe transformers) en ondersteunt het beeld dat multilineaire architecturen transparanter zijn dan gewone diepe netwerken.