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

1. Il Problema

L'interpretabilità meccanica mira a comprendere come i Large Language Models (LLM) rappresentino e processino le informazioni. Un approccio promettente utilizza i Transcoder (basati sull'apprendimento di dizionari) per approssimare i livelli MLP delle reti neurali tramite combinazioni lineari sparse di feature interpretabili. Questo permette di costruire Grafici di Attribuzione delle Feature, dove i nodi sono le feature e gli archi rappresentano influenze causali.

Tuttavia, esistono due limiti principali che ostacolano l'adozione pratica su larga scala:

1. Ridondanza e Dimensione: I grafici di attribuzione generati dai transcoder standard diventano rapidamente enormi (centinaia di nodi) e ridondanti, poiché feature simili emergono in più livelli, rendendo l'analisi difficile.
2. Mancanza di Infrastruttura Open-Source: Sebbene i Cross-Layer Transcoders (CLT) risolvano il problema della ridondanza condividendo le feature tra i livelli (producendo grafici più compatti), la loro addestramento è computazionalmente costoso a causa dell'elevato numero di parametri. Attualmente, non esistono framework open-source completi che supportino l'addestramento end-to-end dei CLT su larga scala, l'analisi automatizzata delle feature e la visualizzazione interattiva. Le soluzioni esistenti sono frammentate o proprietarie (es. Anthropic).

2. Metodologia: CLT-Forge

Gli autori introducono CLT-Forge, una libreria open-source unificata progettata per abilitare l'addestramento scalabile e l'analisi di interpretabilità dei CLT. L'architettura si basa su quattro pilastri tecnici:

A. Addestramento Scalabile e Condivisione delle Feature

• Cross-Layer Transcoders: A differenza dei transcoder standard che ricostruiscono l'output di un livello MLP usando solo le feature dello stesso livello, i CLT permettono alle feature estratte da livelli precedenti di contribuire alla ricostruzione dei livelli successivi. Questo riduce la ridondanza.
• GPU Sharding (Feature-wise): Poiché i CLT hanno un numero di parametri che scala quadraticamente con il numero di livelli e la dimensione del modello, non possono essere ospitati su una singola GPU. CLT-Forge implementa uno sharding delle feature: i diversi GPU gestiscono sottoinsiemi di feature, aggregando gli output alla fine del forward pass. Questo è più efficiente rispetto al Fully Sharded Data Parallel (FSDP) per questa architettura specifica.
• Caching delle Attivazioni Compresso: L'addestramento richiede l'accesso alle attivazioni di tutti i livelli su dataset massicci. Per evitare l'uso di terabyte di storage, il sistema pre-calcola le attivazioni, le quantizza (supportando modalità int8, int4, int2) e le comprime usando codifica entropica (zstd). Questo riduce l'uso di storage di un fattore 4-12x con una perdita minima di qualità (2-3%).

B. Pipeline di Interpretabilità Automatizzata (AutoInterp)

Il framework integra una pipeline parallela per analizzare le feature apprese:

1. Identifica le sequenze di input che attivano maggiormente ciascuna feature.
2. Calcola statistiche di riepilogo (token più attivi, ecc.).
3. Genera automaticamente spiegazioni in linguaggio naturale utilizzando un LLM (prompting) basato sui dati di attivazione.
4. Tutto ciò avviene in un'unica passata parallela, riducendo drasticamente l'I/O su disco rispetto alla scrittura di milioni di file separati.

C. Calcolo del Grafico di Attribuzione

Il sistema si integra nativamente con la libreria Circuit-Tracer per calcolare i punteggi di attribuzione tra le feature di diversi livelli e posizioni. Questo permette di:

• Costruire grafici di attribuzione per specifiche frasi.
• Potare (prunare) il grafico per isolare i circuiti rilevanti.
• Eseguire interventi (interventions) per testare la causalità.

D. Interfaccia Visuale

È stata sviluppata un'interfaccia interattiva basata su Python Dash che permette agli utenti di:

• Esplorare le feature, le sequenze attive e le spiegazioni generate.
• Visualizzare i grafici di attribuzione.
• Creare cluster di feature e definire strutture di ordine superiore.
• Eseguire interventi sui cluster e vedere gli effetti in tempo reale.

3. Contributi Chiave

1. Infrastruttura di Addestramento Scalabile: Supporto per l'addestramento di CLT su modelli grandi (es. LLaMA 1B) tramite sharding delle feature e caching delle attivazioni quantizzate, rendendo fattibile l'addestramento su hardware disponibile.
2. Integrazione End-to-End: Unificazione di addestramento, analisi delle feature (AutoInterp), calcolo dei grafici di attribuzione e visualizzazione in un unico flusso di lavoro coerente.
3. Efficienza e Accessibilità: Riduzione delle barriere all'ingresso per la ricerca sull'interpretabilità meccanica, fornendo strumenti che prima erano disponibili solo in laboratori industriali chiusi o frammentati.
4. Supporto al Fine-tuning: Implementazione di tecniche di low-rank fine-tuning per adattare CLT pre-addestrati a nuovi domini senza costi computazionali proibitivi.

4. Risultati

• Performance di Ricostruzione: Su GPT-2, l'implementazione raggiunge una varianza spiegata di ~0.8 con un L0 totale di 100, un punteggio di sostituzione (replacement score) di ~0.8 e una completezza del grafico di ~0.95, allineandosi con i lavori precedenti.
• Efficienza di Memoria: Lo sharding delle feature e la quantizzazione delle attivazioni hanno permesso di addestrare un CLT per LLaMA 3.2 1B (con un fattore di espansione di 48, ~1.5 milioni di feature) su 8 GPU, riducendo l'footprint di storage da ~20TB a ~4TB.
• Validazione Qualitativa: Il paper dimostra la capacità di identificare circuiti interpretabili (es. il meccanismo "greater-than" o relazioni antonimiche) e di visualizzare le interazioni tra feature in modo coerente con la letteratura esistente.

5. Significato e Impatto

CLT-Forge rappresenta un passo fondamentale per democratizzare l'interpretabilità meccanica su larga scala.

• Riduzione della Ridondanza: Dimostra che i CLT possono produrre rappresentazioni molto più compatte e interpretabili rispetto ai transcoder standard.
• Riproducibilità: Fornisce un framework standardizzato che permette alla comunità di ricerca di riprodurre, estendere e confrontare studi su CLT, un'area precedentemente dominata da framework proprietari.
• Scalabilità: Abilita l'analisi di modelli di dimensioni industriali (miliardi di parametri) con risorse di calcolo accessibili, aprendo la strada a una comprensione più profonda dei meccanismi interni degli LLM moderni.

In sintesi, CLT-Forge non è solo un tool di addestramento, ma un ecosistema completo che trasforma l'analisi dei CLT da un processo frammentato e costoso in un flusso di lavoro unificato, scalabile e interattivo.