Sparse Crosscoders for diffing MoEs and Dense models

Dit artikel vergelijkt de interne representaties van Mixture-of-Experts (MoE) en dichte modellen met behulp van crosscoders en toont aan dat MoE's minder unieke, maar meer gespecialiseerde en actieve features ontwikkelen, terwijl dichte modellen informatie verspreiden over bredere, algemene features.

De kern

De Grote Wedstrijd: De "Alles-in-één" Chef vs. Het "Specialisten-Team"

Stel je voor dat je twee restaurants hebt die precies even groot zijn en evenveel geld kosten om te runnen, maar ze werken op een heel verschillende manier.

1. Restaurant A (Het Dichte Model): Dit is een superchef die alles zelf doet. Als er een bestelling binnenkomt, gebruikt deze chef elke vaardigheid die hij heeft: snijden, bakken, kruiden, serveren. Hij is een alleskunner, maar hij moet voor elke klant zijn hele brein gebruiken.
2. Restaurant B (Het MoE-model - Mixture of Experts): Dit is een team van gespecialiseerde chefs. Er is één chef die alleen pasta maakt, één die alleen desserts doet, en één die alleen vis bereidt. Als er een bestelling binnenkomt, kijkt een manager (de 'router') alleen naar de chef die het beste is voor die specifieke taak. De andere chefs doen niets. Dit is veel efficiënter: je gebruikt minder energie per klant, maar het team als geheel is enorm groot.

Het probleem: We weten precies hoe Restaurant A (de alleskunner) in zijn hoofd werkt. Maar bij Restaurant B (het specialistenteam) weten we niet echt wat er gebeurt in de hoofden van die individuele chefs. Werken ze samen? Heeft elke chef zijn eigen unieke trucjes? Of delen ze allemaal dezelfde basis?

De Oplossing: De "Twee-in-één" Vertaler

De onderzoekers (Marmik, Nishkal en Idhant) wilden dit mysterie oplossen. Ze bouwden een slimme vertaler, een Crosscoder.

Stel je voor dat je twee mensen hebt die in verschillende talen spreken. Je wilt weten: Welke woorden betekenen hetzelfde in beide talen, en welke woorden zijn uniek voor de ene taal?

Deze "Crosscoder" is een apparaat dat tegelijkertijd luistert naar wat er gebeurt in het hoofd van de alleskunner (Restaurant A) en wat er gebeurt in het hoofd van het specialistenteam (Restaurant B). Het probeert een lijst te maken van "concepten" (features) die beide gebruiken, en een lijst van concepten die alleen de een of de ander kent.

Wat hebben ze ontdekt?

Na het trainen van deze vertaler op een enorme hoeveelheid tekst (code, wetenschap en verhalen), kwamen ze tot drie verrassende conclusies:

1. De Alleskunner is een verzamelaar, de Specialist is een purist

• Restaurant A (Dicht): Deze chef heeft een enorme, diverse toolbox. Hij heeft duizenden unieke, specifieke gerechten in zijn hoofd. Hij is erg breed.
• Restaurant B (MoE): Het specialistenteam heeft veel minder unieke gerechten in totaal. Ze zijn veel selectiever. In plaats van duizenden kleine trucjes, hebben ze zich gefocust op een paar heel sterke, specifieke vaardigheden.
• De les: Door alleen de beste specialisten te kiezen, leert het model minder "dingen" in totaal, maar zijn die dingen wel extreem goed en specifiek.

2. Hoe vaak worden de trucjes gebruikt?

• De MoE-specialisten gebruiken hun unieke vaardigheden heel vaak. Als ze een taak hebben die bij hen past, gaan ze er voluit voor. Ze zijn "dicht" bij elkaar in hun activiteit.
• De Alleskunner gebruikt zijn unieke vaardigheden juist minder vaak. Hij verspreidt zijn kennis over een heel breed spectrum.
• De les: Specialisten zijn intensiever in hun werk; de alleskunner is meer verspreid.

3. De "Gedeelde" taal is lastig te vinden

• De onderzoekers dachten eerst dat ze makkelijk zouden kunnen zien welke vaardigheden beide restaurants delen. Maar het bleek lastig. De standaardmethode dacht dat ze veel deelden, terwijl ze eigenlijk totaal verschillende manieren gebruikten om hetzelfde te zeggen.
• Ze moesten hun vertaler (de Crosscoder) aanpassen. Ze moesten de "specialisten" (de unieke vaardigheden) harder straffen als ze te vaak werden gebruikt, zodat de "gemeenschappelijke" vaardigheden er echt uit sprongen. Uiteindelijk slaagden ze erin om 87% van wat er in de hoofden gebeurde, te verklaren.

Waarom is dit belangrijk?

Vroeger dachten we dat als je een model groter maakt (meer parameters), het gewoon "slimmer" wordt op dezelfde manier. Dit onderzoek laat zien dat hoe je het model bouwt (specialisten vs. alleskunner) de manier waarop het leert, fundamenteel verandert.

• Dichte modellen leren een brede, algemene kennis.
• MoE-modellen leren een diepe, gefocuste kennis door zich te specialiseren.

Samenvatting in één zin

De onderzoekers hebben ontdekt dat "specialisten" (MoE-modellen) minder unieke ideeën hebben maar dieper en frequenter op die ideeën focussen, terwijl "alleskunners" (dichte modellen) een breder, maar minder intensief palet aan ideeën hebben. Hun nieuwe "vertaler" (Crosscoder) helpt ons eindelijk te zien hoe die twee werelden van binnen verschillend zijn.

Technische samenvatting

Probleemstelling

Mixture of Experts (MoE) modellen zijn populair geworden voor het schalen van Large Language Models (LLMs) omdat ze parameters efficiënter gebruiken door slechts een subset van "experts" te activeren per token. Hoewel dense modellen (waarbij alle parameters actief zijn) uitgebreid zijn onderzocht op hun interne structuur en interpretatie, blijft de interne werking van MoE-modellen slecht begrepen.

De kernvraag is of de intuïties en bevindingen die gelden voor dense modellen ook opgaan voor MoE-modellen. Specifiek is het onduidelijk:

• Of experts unieke feature-representaties ontwikkelen.
• Hoe de routing-strategie de specialisatie van features beïnvloedt.
• Hoe de geleerde representaties van MoE-modellen verschillen van die van dense modellen, zelfs wanneer ze dezelfde hoeveelheid actieve parameters hebben.

Methodologie

De auteurs voeren een systematische vergelijking uit tussen een dense model en een MoE-model met een identieke hoeveelheid actieve parameters.

1. Modeltraining:

   • Er zijn twee 5-laagse modellen getraind op een dataset van ongeveer 1 miljard tokens (bestaande uit code, wetenschappelijke teksten en Engelse verhalen).
   • Beide modellen zijn getraind met dezelfde Cross Entropy loss; het MoE-model gebruikte additionally een Switch load balancing loss.
   • Beide modellen zijn 2 epochs getraind.

2. Crosscoders toepassing:

   • In plaats van traditionele Sparse Autoencoders (SAE's) die per model apart worden getraind, gebruiken de auteurs Crosscoders. Dit zijn een variant van SAE's die twee activatieruimtes (van het dense en het MoE-model) gezamenlijk modelleren.
   • Het doel is om een gezamenlijke set van sparse features $f_i(x)$ te leren die beide activatieruimtes reconstrueren via model-specifieke decoder-weights ($W^{MoE}$ en $W^{Dense}$).
   • BatchTopK met expliciete gedeelde features: Om de interpretatie te verbeteren, gebruiken ze de BatchTopK variant. Hierbij wordt een subset van features expliciet gedefinieerd als "gedeeld" (shared) met gekoppelde decoder-parameters, terwijl andere features "exclusief" zijn voor één model. Dit helpt om de druk op de exclusieve features te verminderen en monosemantische (betekenisvol en eenduidig) features te bevorderen.

3. Analyse van Feature Specialisatie:

   • Om te bepalen of een feature gedeeld of model-specifiek is, berekenen de auteurs de $\Delta_{norm}$: een maatstaf voor het verschil in de norm van de decoder-vectoren tussen de twee modellen.
   • $\Delta_{norm} \approx 0.5$ betekent een gedeelde feature.
   • $\Delta_{norm} \approx 0$ betekent een feature exclusief voor het MoE-model.
   • $\Delta_{norm} \approx 1$ betekent een feature exclusief voor het dense model.

Belangrijkste Bevindingen en Resultaten

• Uitdaging bij Standaard Crosscoders: Een standaard crosscoder-objective overschatte de hoeveelheid gedeelde structuur, zelfs wanneer de cosine-similariteit van de decoder-weights dicht bij nul lag. Dit suggereert dat de activatieruimtes van twee onafhankelijk getrainde modellen (dense vs. MoE) te divergent zijn voor de standaard methode.
• Aangepaste Regularisatie: De auteurs ontdekten dat de eerder gepubliceerde verhouding voor regularisatie ($\lambda_s/\lambda_f \approx 0.1-0.2$) niet werkte voor hun setting. Omdat de modellen vanaf nul onafhankelijk zijn getraind (in plaats van een base vs. fine-tuned scenario), was een sterkere regularisatie voor gedeelde features nodig. Een verhouding van ongeveer 0.7 bleek nodig om effectief te onderscheiden tussen model-specifieke features.
• Prestaties: Met de aangepaste BatchTopK crosscoder werd een fractionele variantie-uitgelegdheid (fractional variance explained) van ongeveer 87% bereikt over de activaties van beide modellen.
• Verschil in Feature-aantallen:
  • Het dense model leerde aanzienlijk meer unieke (exclusieve) features dan het MoE-model (3.226 dense-only features vs. 910 MoE-only features).
  • Het MoE-model ontwikkelde dus minder, maar meer gespecialiseerde features.
• Activatie-dichtheid:
  • MoE-specifieke features vertoonden een hogere activatie-dichtheid dan de gedeelde features.
  • Dense-specifieke features vertoonden een lagere activatie-dichtheid dan de gedeelde features.
  • Dit verschilt van eerdere studies (bijv. bij base vs. fine-tuned modellen), waar beide model-specifieke sets een hogere dichtheid hadden dan de gedeelde features.
• Geen Trimodale Distributie: In tegenstelling tot eerdere crosscoder-studies die een duidelijke trimodale verdeling zagen (MoE-only, Shared, Dense-only), zagen de auteurs hier geen scherpe trimodale structuur in de $\Delta_{norm}$ verdeling.

Bijdragen en Significantie

1. Systematische Vergelijking: Dit is een van de eerste werken dat de interne representaties van MoE- en dense-modellen direct vergelijkt met een gelijke actieve parameter-count, gebruikmakend van geavanceerde interpretatietechnieken.
2. Inzicht in Architecturale Verschillen: De studie toont aan dat MoE-modellen informatie organiseren via meer gespecialiseerde en gefocuste representaties, terwijl dense modellen informatie verspreiden over bredere, meer algemene features. Sparsity in MoE's lijkt lokale specialisatie te stimuleren.
3. Methodologische Vooruitgang: Het paper demonstreert dat crosscoders kunnen worden aangepast voor het vergelijken van fundamenteel verschillende architecturen (niet alleen fine-tuning), maar benadrukt dat de regularisatie-parameters (zoals de $\lambda$-verhouding) moeten worden aangepast aan de mate van divergentie tussen de modellen.
4. Toekomstige Richting: De resultaten vormen een basis voor toekomstig mechanistisch interpretatieonderzoek bij sparse architecturen. De auteurs wijzen erop dat er nog kwalitatieve analyse nodig is om de semantische betekenis van de ontdekte features te valideren.

Kortom, het paper levert bewijs dat MoE-modellen niet simpelweg "dense modellen met minder actieve parameters" zijn, maar een fundamenteel andere manier van informatieverwerking en feature-specialisatie hanteren.