CLT-Forge: A Scalable Library for Cross-Layer Transcoders and Attribution Graphs

Probleemstelling

Mechanistische interpretabiliteit van Large Language Models (LLMs) richt zich op het begrijpen van hoe deze modellen informatie verwerken. Een veelgebruikte aanpak is het gebruik van transcoders (gebaseerd op dictionary learning) om de berekeningen van MLP-lagen (Multi-Layer Perceptrons) te benaderen als spaarse, lineaire combinaties van interpreteerbare features. Dit leidt tot feature attribution graphs, waarin knopen features voorstellen en randen causale invloeden weergeven.

Er zijn echter twee grote beperkingen aan deze bestaande methoden:

1. Redundantie en schaalbaarheid: In standaard transcoders ontstaan vergelijkbare features vaak in meerdere lagen, wat leidt tot enorme, redundante attributie-grafieken met honderden knopen die moeilijk te analyseren zijn.
2. Gebrek aan geïntegreerde tooling: Hoewel Cross-Layer Transcoders (CLTs) dit redundantieprobleem oplossen door features over lagen te delen (wat leidt tot compactere grafieken), ontbreekt er een open-source framework dat schaalbaar trainen en post-hoc interpretatie ondersteunt. Bestaande tools zijn gefragmenteerd, ontberen efficiënte distributie (zoals model sharding) en hebben geen geïntegreerde visualisatie of geautomatiseerde interpretatiepijplijnen.

Methodologie: CLT-Forge

De auteurs introduceren CLT-Forge, een open-source bibliotheek die een end-to-end workflow biedt voor het trainen en interpreteren van CLTs. De kern van de methode bestaat uit de volgende technische componenten:

1. Cross-Layer Transcoders (CLTs)

In tegenstelling tot standaard transcoders die features per laag decoderen, laten CLTs features toe die uit eerdere lagen komen bijdragen aan de reconstructie van latere lagen.

• Formulering: De output van een MLP in laag $\ell'$ wordt gereconstrueerd als een som van gedecodeerde features uit alle lagen $\ell \leq \ell'$.
• Voordeel: Dit reduceert redundantie aanzienlijk en resulteert in dichter, interpreteerbaarder attributie-grafieken.
• Uitdaging: De parametercount groeit kwadratisch met het aantal lagen en de modeldimensie, wat training op grote schaal zeer rekenintensief maakt.

2. Schaalbare Training Infrastructuur

Om CLTs op grote schaal (bijv. LLaMA 1B) te trainen, implementeert CLT-Forge:

• GPU Sharding: In plaats van alleen Data Parallelism (DDP), gebruikt de library feature-wise sharding. Features worden verdeeld over meerdere GPU's, waardoor elke worker een subset van de features verwerkt. Dit omzeilt geheugenbeperkingen die optreden bij Fully Sharded Data Parallel (FSDP) voor deze specifieke architectuur.
• Compressed Activation Caching: Omdat het trainen toegang vereist tot activa-ties over grote datasets (honderden miljoenen tokens), worden deze vooraf berekend en opgeslagen. Om de opslag (die anders terabytes zou beslaan) te reduceren, wordt er gebruikgemaakt van:
  • Symmetrische per-laag kwantisatie (int8, int4, int2).
  • Zstd entropy coding.
  • Dit reduceert de opslag met een factor 4–12× met slechts een verwaarloosbaar verlies in reconstructiekwaliteit (2–3%).

3. Geautomatiseerde Interpretatie (AutoInterp)

De library bevat een geïntegreerde pijplijn om de betekenis van features te achterhalen:

• Parallelle Verwerking: Berekeningen worden parallel over de feature-dimensie uitgevoerd.
• Top-K Tracking: De meest activerende sequenties en tokens worden direct op de GPU bijgehouden om het geheugengebruik laag te houden.
• LLM-gegenereerde Uitleg: Op basis van deze statistieken worden automatisch prompts gegenereerd voor een LLM om natuurlijke taal-uitingen van de features te produceren.

4. Attributie Grafieken en Visualisatie

• Circuit-Tracer Integratie: De library is naadloos geïntegreerd met de Circuit-Tracer bibliotheek om attributie-scores te berekenen en grafieken te prunen.
• Interactieve Interface: Een Python-based Dash-interface stelt onderzoekers in staat om attributie-grafieken te visualiseren, clusters van features te vormen, en interventies (manipulatie van features) uit te voeren om causale effecten te testen.

Belangrijkste Bijdragen

1. Eerste Unified Library: CLT-Forge is de eerste open-source bibliotheek die schaalbaar CLT-training combineert met post-hoc interpretatie, attributie-analyse en visualisatie in één workflow.
2. Efficiënte Schaalbaarheid: Door middel van feature-sharding en geoptimaliseerde activatie-caching (kwantisatie + compressie) wordt het mogelijk om CLTs te trainen op modellen van 1B+ parameters, wat eerder beperkt was tot gesloten industriële frameworks.
3. Geautomatiseerde Workflow: Het biedt een volledig geautomatiseerde pijplijn voor het genereren van feature-uitingen (AutoInterp) en het berekenen van attributie-grafieken, wat de barrière voor onderzoekers verlaagt.
4. Flexibele Visualisatie: Een uitbreidbare interface voor het verkennen van features en het uitvoeren van interventies, als een alternatief voor bestaande, minder flexibele tools zoals Neuronpedia.

Resultaten

• Kwaliteit: Op GPT-2 bereikt de implementatie vergelijkbare resultaten met eerdere werken: ~0.8 verklaarde variantie bij een totale L0 van 100, een vervangingsscore van ~0.8 en grafische volledigheid van ~0.95.
• Efficiëntie: De feature-sharding benadering verbetert de geheugenefficiëntie aanzienlijk ten opzichte van DDP, waardoor training op 8x80GB GPU's mogelijk is voor een LLaMA 1B CLT met een expansiefactor van 48 (ongeveer 1,5 miljoen features).
• Opslag: De compressietechnieken reduceren de opslagbehoefte voor activa-ties van ~20TB naar ~4TB voor een LLaMA 1B model op 300M tokens, met minimaal verlies aan prestaties.
• Validatie: De volledige pijplijn is kwalitatief gevalideerd, inclusief het genereren van interpreteerbare grafieken die overeenkomen met bekende circuits (bijv. antoniemen) en meertalige analyses.

Betekenis en Impact

CLT-Forge is een mijlpaal voor het veld van mechanistische interpretabiliteit. Het democratiseert de toegang tot geavanceerde CLT-analyses, die tot nu toe grotendeels beperkt waren tot gesloten labs zoals Anthropic. Door een gestroomlijnde, end-to-end oplossing te bieden, stelt het onderzoekers in staat om:

• De interne werking van grote taalmodellen op schaal te bestuderen.
• Reproduceerbaar onderzoek te doen met betrekking tot feature-interacties.
• Sneller te experimenteren met interventies en causale verklaringen.

De library vormt een cruciale brug tussen theoretische concepten (zoals CLTs) en praktische toepassing, en legt de basis voor toekomstig werk gericht op nog efficiëntere architecturen en uitgebreidere attributiemethoden (zoals het meenemen van attention-maps).