Riepilogo Tecnico: Rational Sparse Autoencoder (RSAE)

Definizione del Problema

Gli Sparse Autoencoder (SAE) sono lo strumento standard per l'interpretabilità meccanicistica nei grandi modelli linguistici (LLM), decomponendo le attivazioni interne in combinazioni lineari sparse di dizionari di feature overcomplete. Tuttavia, le famiglie di SAE esistenti si affidano a non-linearità di codifica fisse e non smussate — specificamente ReLU (con regolarizzazione $\ell_1$), TopK e JumpReLU. Queste primitive di attivazione fisse introducono specifiche patologie:

• SAE ReLU: Soffrono di una contrazione dell'ampiezza delle feature attive e di una popolazione persistente di latenti "morti".
• SAE TopK: Sostituiscono i penalità morbide con vincoli di cardinalità rigidi, interrompendo il flusso del gradiente attraverso le feature inattive e richiedendo perdite di revival ausiliarie.
• SAE JumpReLU: Introducono soglie apprendibili ma richiedono surrogati di rilassamento continuo per la back-propagation attraverso i loro gate indicatori.

Questi meccanismi pre-codificati limitano la capacità del modello di adattarsi alla geometria delle pre-attivazioni osservate, potenzialmente distorcendo il compromesso tra fedeltà di ricostruzione e sparsità.

Metodologia

Gli autori propongono il Rational Sparse Autoencoder (RSAE), che sostituisce la funzione di attivazione del codificatore con una funzione razionale addestrabile. L'RSAE mantiene lo scheletro standard di un SAE (uno strato di pre-attivazione affine seguito da un blocco di attivazione sparsa e un decoder lineare) ma modifica la funzione di attivazione $\phi(\cdot)$ in una funzione razionale apprendibile applicata elemento per elemento.

Architettura

L'attivazione RSAE è definita come:
$$\phi(h) = C_{out} \cdot r_{(a,b)}\left(\frac{h}{C_{in}}\right)$$
dove $h$ è il vettore di pre-attivazione, $C_{in}$ e $C_{out}$ sono parametri di scala apprendibili, e $r_{(a,b)}(t) = P(t)/Q(t)$ è una funzione razionale con coefficienti $(a, b)$ addestrabili. I parametri di scala mappano la distribuzione della pre-attivazione nell'intervallo di progettazione della funzione razionale $[-1, 1]$ e mappano l'output nuovamente alla magnitudo attesa della feature.

Fondamento Teorico

Il documento stabilisce un'asimmetria espressiva tra le funzioni razionali e i gate piecewise-affine utilizzati negli attuali SAE:

1. Limite Superiore (Approssimazione Razionale): Utilizzando le funzioni segno di Zolotarev, gli autori dimostrano che le attivazioni razionali addestrabili possono approssimare le gate ReLU, JumpReLU e i gate TopK con soglia fornita con una dimensione polilogaritmica nell'accuratezza target $\epsilon$ (specificamente $O(\log(1/\epsilon)\log(1/\delta))$ per gate discontinui con margine $\delta$).
2. Limite Inferiore (Limite Piecewise-Affine): Al contrario, dimostrano che approssimare determinate mappe target razionali con $O(1)$-parametro utilizzando encoder a singolo strato piecewise-affine a output scalare (ReLU/JumpReLU/TopK) richiede $\Omega(\epsilon^{-1/2})$ coordinate attivate. Ciò suggerisce che le attivazioni razionali offrono una classe di rappresentazione più efficiente per l'encoder SAE.

Pipeline di Addestramento

L'RSAE è costruito tramite una pipeline a due stadi progettata per potenziare gli SAE "teacher" pre-addestrati:

1. Inizializzazione:
   • Scambio di Remez: Una funzione razionale viene adattata alla primitiva di attivazione del teacher (ReLU, JumpReLU o TopK) su una griglia densa sintetica utilizzando l'algoritmo di scambio di Remez rilassato per minimizzare l'errore $L_\infty$.
   • Calibrazione: I coefficienti razionali e i parametri di scala ($C_{in}, C_{out}$) vengono adattati alla distribuzione reale delle pre-attivazioni del teacher minimizzando l'errore quadratico tra l'output razionale e l'output del teacher.
   • Ereditarietà dei Pesi: I pesi del codificatore e del decoder vengono copiati parola per parola dal teacher pre-addestrato.
2. Fine-Tuning: Tutti i parametri (pesi, bias e coefficienti razionali) vengono sbloccati e ottimizzati congiuntamente sotto l'obiettivo di ricostruzione $\ell_1$-regolarizzato standard.

Contributi Chiave

1. Architettura RSAE: Introduzione di un autoencoder sparso con un'attivazione razionale addestrabile, basata sulla teoria dell'approssimazione adattata al codificatore SAE superficiale.
2. Gap di Espressività Teorica: Dimostrazione che le attivazioni razionali forniscono approssimazioni compatte per le gate fisse degli SAE, mentre gli encoder piecewise-affine richiedono significativamente più coordinate attivate per approssimare semplici target razionali.
3. Algoritmo di Addestramento a Due Stadi: Una procedura pratica di inizializzazione e fine-tuning che permette agli RSAE di essere "inseriti" per sostituire gli SAE pre-addestrati con un overhead trascurabile.
4. Validazione Empirica: Valutazione completa su tre modelli linguistici open-weight (GPT-2 small, Pythia-160m, Gemma-2-2B) e tutte e tre le famiglie di attivazione di base.

Risultati

Gli autori valutano l'RSAE rispetto ai baseline pre-addestrati sulle attivazioni dello stream residuo. I risultati indicano miglioramenti rigorosi dopo il fine-tuning:

• Metriche di Ricostruzione: L'RSAE raggiunge costantemente un MSE inferiore, una $\ell_0$ inferiore (sparsità) e una frazione più alta di feature "vive" rispetto ai modelli teacher in tutte le configurazioni testate.
• Metriche Downstream: L'RSAE migliora il comportamento downstream, mostrando una riduzione della degradazione della cross-entropia ($\Delta CE$) e un recupero della perdita ($LR$) più elevato quando intercetta lo stream residuo.
• Interpretabilità: Gli esperimenti di probing sparso su Pythia-160m mostrano che l'RSAE preserva l'interpretabilità a livello di feature. L'accuratezza del probe a dizionario completo su otto compiti di classificazione binaria rimane comparabile a quella del teacher, senza degradazioni significative.
• Efficienza: L'aggiornamento aggiunge solo una manciata di parametri scalari per autoencoder. L'intera pipeline (inizializzazione + 2.000 passi di fine-tuning) gira in pochi minuti su una singola GPU consumer (NVIDIA RTX 5090).
• Robustezza: I guadagni sono costanti attraverso diversi modelli host, famiglie di attivazione baseline e un ampio intervallo di livelli di sparsità baseline ($\ell_0 \in [51, 683]$).

Significato e Rivendicazioni

Il documento afferma che l'RSAE offre una modifica "drop-in" per gli SAE pre-addestrati che migliora rigorosamente la fedeltà del modello a parità di sparsità senza sacrificare l'interpretabilità. Sostituendo le non-linearità fisse con funzioni razionali addestrabili, l'RSAE si adatta alla specifica geometria delle pre-attivazioni del modello host, superando i limiti delle primitive di attivazione codificate rigidamente.

Gli autori sottolineano che questi miglioramenti sono ottenuti con un costo computazionale trascurabile (minuti su una singola GPU) e un overhead di parametri minimo. Posizionano l'RSAE come una strategia di aggiornamento universale per gli SAE esistenti, capace di migliorare sia la qualità della ricostruzione che le prestazioni dei task downstream mantenendo le feature monosemantiche richieste per l'interpretabilità meccanicistica. Il documento nota che, sebbene la valutazione copra tre modelli e tre famiglie di attivazione, l'approccio è generale e potrebbe potenzialmente essere esteso ad altri modelli di frontiera o varianti architettoniche (ad esempio, Gated, BatchTopK, Matryoshka SAEs).