Technische Samenvatting: Tree SAE: Het Leren van Hiërarchische Featurestructuren in Sparse Autoencoders

Probleemstelling

Sparse Autoencoders (SAE's) worden veel gebruikt om menselijk interpreteerbare features te extraheren uit ondoorzichtige taalmodellen. Echter, standaard SAE's behandelen features als een onafhankelijke verzameling, waardoor de inherente hiërarchische structuur van concepten uit de echte wereld wordt verwaarloosd. Deze nalatigheid leidt tot ongewenste fenomenen zoals feature-absorptie (waarbij fijnmazige features grove features onderdrukken), feature-splitsing (waarbij een enkel concept versplinterd raakt over meerdere features) en feature-compositie (waarbij modellen versplinterde of overmatig complexe "polysemantische" amalgamen leren).

Recente architecturen, zoals Matryoshka SAE's en Matching Pursuit SAE's, proberen deze problemen aan te pakken door multi-niveau reconstructie of conditionele orthogonaliteit in te voeren. Ze vertrouwen op Masked Cosine Similarity (MCS) om hiërarchische relaties te identificeren, wat gebaseerd is op de activatie-dekkingsvoorwaarde: een child-feature wordt beschouwd als een afstammeling van een parent als het alleen activeert wanneer de parent activeert.

De auteurs tonen aan dat activatie-dekking alleen onvoldoende is voor het identificeren van echte hiërarchische relaties. Zij leveren bewijs dat parent- en child-features kunnen voldoen aan de activatie-dekkingsvoorwaarde terwijl ze semantisch niet gerelateerd zijn (bijvoorbeeld een "tijd"-concept als parent dat onjuist wordt gekoppeld aan een "politiek"-concept als child, of een dichte "PCA-feature" die optreedt als parent voor niet-gerelateerde concepten). Dit leidt tot vals-positieven waarbij de geïdentificeerde hiërarchische paren geen conceptuele coherentie vertonen.

Methodologie

1. De Reconstructievoorwaarde

Om de beperkingen van activatie-dekking aan te pakken, introduceren de auteurs een nieuwe reconstructievoorwaarde. Deze voorwaarde dwingt een diepere functionele link af tussen hiërarchische niveaus.

• Definitie: Voor een parent-feature $f_p$ en een child-feature $f_c$ wordt de reconstructiescore $S_{res}$ gedefinieerd op basis van de uitlijning van hun decoder-vectoren ($d_p, d_c$) met de ware conceptvector van het kind ($d^*_c$).
• Mechanisme: De voorwaarde vereist dat zowel de parent- als de child-decoder-vectoren constructief bijdragen aan de reconstructie van het concept van het kind. Specifiek zorgt het ervoor dat de parent een brede conceptuele richting vastlegt die het kind verfijnt.
• Validatie: De auteurs tonen aan dat hoewel activatie-dekking kan worden voldaan door spurious paren, de reconstructiescore deze effectief filtert. Echte hiërarchische paren vertonen hoge reconstructiescores, terwijl niet-gerelateerde paren (zelfs met perfecte activatie-dekking) lage scores vertonen.

2. Tree SAE Architectuur

De auteurs stellen Tree SAE voor, een model dat ontworpen is om hiërarchische structuren direct binnen de featureverzameling te leren, in plaats van ze achteraf af te leiden.

• Boomstructuur: De featureverzameling is verdeeld in meerdere "privilege-lagen". Een feature in een hogere laag kan worden toegewezen als afstammeling van een feature in een lagere laag. Een imaginaire root-node in laag 0 dient als basislijn voor features zonder expliciete ouders.
• Toewijzingsvector: Een toewijzingsvector $a_l$ wijst parent-features toe aan child-features in elke laag $l$.
• Geregulariseerde Activering: Om de boomstructuur af te dwingen, wordt de activering van een child-feature $f_i$ geregulariseerd zodat deze alleen niet-nul is als de toegewezen parent-feature actief is. Dit volgt strikt de activatie-dekkingsvoorwaarde tijdens het trainen.
• Multi-niveau Reconstructieverlies: Het trainingsdoel omvat een multi-niveau reconstructieverlies. Voor een laag $l$ beperkt het verlies de cumulatieve reconstructie van alle voorgaande lagen (inclusief $l$) zodat deze de input nauwkeurig herstellen. Dit zorgt ervoor dat parent-vectoren algemene concepten vastleggen en child-vectoren deze verfijnen tot gespecialiseerde concepten.
• Dynamische Toewijzing: Om het probleem van "dode features" (features die nooit activeren) te mitigeren, introduceren de auteurs een dynamisch toewijzingsmechanisme. Dit algoritme wijst dode child-features opnieuw toe aan parent-features op basis van een "payoff"-metriek (afgeleid van het opgebouwde trainingsverlies), zodat child-features worden toegewezen aan ouders die hun leren effectief kunnen ondersteunen.

Belangrijkste Bijdragen

1. Identificatie van Beperkingen van Activatie-dekking: Het artikel levert concreet bewijs dat activatie-dekking alleen semantisch incoherente hiërarchische paren identificeert en stelt een nieuwe definitie voor die activatie-dekking combineert met de reconstructievoorwaarde.
2. Tree SAE Model: De introductie van een nieuwe SAE-architectuur die een hiërarchische boomstructuur direct binnen de featureverzameling encodeert, die dynamisch wordt geleerd tijdens het trainen.
3. Superieur Hiërarchisch Leren: Aantonen dat Tree SAE's aanzienlijk coherenter parent-child paren leren dan bestaande SAE's (inclusief Matryoshka en Matching Pursuit SAE's), terwijl ze de pathologieën van feature-splitsing, absorptie en compositie vermijden.
4. Geometrische Analyse: Het tonen van het nut van de expliciete boomstructuur bij het in kaart brengen van de geometrie van child-feature subruimtes en het begrijpen hoe taalmodellen concepten ontleden in grove en fijnmazige niveaus.

Experimentele Resultaten

De auteurs evalueerden Tree SAE's op GPT2-small over verschillende sparsiteitsniveaus ($L_0$) en woordenboekgroottes, en vergeleken ze met Top-k SAE's, Matryoshka SAE's en Matching Pursuit SAE's.

• Hiërarchische Metrieken: Tree SAE's ontdekten aanzienlijk meer hiërarchische paren die voldeden aan zowel de activatie-dekkings- als reconstructievoorwaarden in vergelijking met andere modellen. De interne boomstructuur bleek iets effectiever in het identificeren van deze paren dan MCS-analyse achteraf.
• Feature Pathologieën: Tree SAE's presteerden concurrerend met of beter dan state-of-the-art (SOTA) Matryoshka SAE's in het verminderen van feature-splitsing, absorptie en compositie.
• Reconstructie en Interpretatieerbaarheid: Tree SAE's behielden een concurrerend reconstructieverlies en verklaarde variantie, en reconstrueerden de originele latente variabelen nauwkeurig. Wat betreft interpretatieerbaarheid (gemeten via AutoInterp), kwamen Tree SAE's overeen met de prestaties van Matryoshka SAE's.
• Schalen: Tree SAE's schaalden goed met woordenboekgroottes (6k tot 49k), en presteerden vaak beter dan Matryoshka SAE's op hiërarchische metrieken bij kleinere schalen, terwijl ze vergelijkbare prestaties behielden bij grotere schalen.
• Diversiteit van Child Features: De geleerde child features vertoonden lage co-activeringspercentages (<3,0%), wat aangeeft dat ze verschillende aspecten van het parent-concept vertegenwoordigen in plaats van redundante informatie.

Betekenis en Claims

Het artikel beweert dat Tree SAE een belangrijke stap voorwaarts is in het begrijpen van de interne representaties van grote taalmodellen. Door verder te gaan dan onafhankelijke featureverzamelingen en expliciet hiërarchische structuren te modelleren, betogen de auteurs dat Tree SAE's:

• Bekende hiërarchische feature-pathologieën (splitsing, absorptie, compositie) mitigeren.
• Een nauwkeuriger kader bieden voor het identificeren van echte parent-child conceptrelaties, waardoor de vals-positieven die gebruikelijk zijn bij eerdere methoden worden vermeden.
• De geometrische analyse van hiërarchische concept-subruimtes mogelijk maken, wat nieuwe inzichten biedt in hoe modellen algemene concepten ontleden in gespecialiseerde sub-concepten.

De auteurs merken beperkingen op, specifiek dat de rigide boomstructuur kan leiden tot een hoger aantal dode features in vergelijking met andere SAE's, hoewel dynamische toewijzing helpt dit te mitigeren. Zij suggereren dat toekomstig werk zachte-gating mechanismen zou kunnen verkennen om het gebruik van features verder te verbeteren.