Sintesi Tecnica: Sulla relazione tra gli outlier di attivazione e la morte delle feature negli Autoencoder Sparsi

1. Definizione del problema

Gli Autoencoder Sparsi (SAE) sono progettati per decomporre le attivazioni delle reti neurali in feature interpretabili e sparse, al fine di invertire il fenomeno della sovrapposizione (dove le reti codificano più concetti di quanti siano le dimensioni, sovrapponendo le direzioni). Tuttavia, un modo di fallimento critico noto come morte delle feature (feature death) si verifica quando le feature del dizionario apprese non si attivano mai su alcun input. Ciò spreca la capacità del dizionario e costringe le feature sopravvissute a compensare, potenzialmente reintroducendo la sovrapposizione.

Il documento identifica un'incoerenza significativa nelle prestazioni degli SAE: configurazioni identiche di SAE producono un numero di feature morte quasi nullo in modelli come GPT-2, ma oltre il 70% in modelli come AlphaFold3. Le tecniche di riattivazione esistenti (ad esempio, AuxK, ghost gradients, resampling) trattano la morte delle feature come un problema di dinamica di addestramento, tentando di iniettare segnali di gradiente nelle feature bloccate. Tuttavia, questi metodi non riescono a risolvere i tassi di morte nei casi più gravi, suggerendo che la causa principale non risieda nella dinamica di addestramento, ma nello stato di inizializzazione delle attivazioni.

2. Metodologia e quadro teorico

2.1 Outlier di attivazione a livello di dimensione

Gli autori propongono che gli outlier di attivazione a livello di dimensione siano la causa primaria della morte delle feature. A differenza degli outlier a livello di token (picchi in input specifici), questi sono dimensioni in cui l'entità della media dell'attivazione ($\mu$) è grande rispetto alla deviazione standard per token ($\sigma$). Ciò crea un offset quasi costante attraverso tutti gli input.

Gli autori formalizzano la gravità dell'outlier (outlier severity) come il rapporto:
$$\gamma = \frac{\|\mu\|}{\|\sigma\|}$$
dove $\mu$ è il vettore della media e $\sigma$ è il vettore della deviazione standard per token.

2.2 Meccanismo geometrico della morte

All'inizializzazione (con bias zero), la pre-attivazione per la feature $i$ è $z_i = w_i \cdot x$. Decomponendo l'input $x$ nella sua media e nella sua variazione ($x = \mu + (x - \mu)$), la pre-attivazione diventa:
$$z_i = \underbrace{w_i \cdot \mu}_{\text{spostamento costante}} + \underbrace{w_i \cdot (x - \mu)}_{\text{segnale dipendente dall'input}}$$

Quando $\gamma$ è elevato, lo spostamento costante domina il segnale. Ciò porta a due distinti modi di fallimento all'inizializzazione:

1. Morte per ReLU (Dead-by-ReLU): Le feature con pesi anti-allineati con $\mu$ ($w_i \cdot \mu \ll 0$) ricevono pre-attivazioni permanentemente negative. Poiché la funzione ReLU restituisce zero per valori negativi, queste feature non si attivano mai e non ricevono alcun gradiente.
2. Morte per TopK (Dead-by-TopK): Negli SAE TopK, le feature con spostamenti positivi moderati possono superare la ReLU, ma perdono la competizione contro le feature con grandi spostamenti positivi. Solo le $k$ feature più allineate con $\mu$ vincono la selezione TopK, lasciando le altre morte.

2.3 Derivazione analitica

Il documento deriva espressioni in forma chiusa per la probabilità di morte delle feature basandosi esclusivamente su $\gamma$:

• Morte per ReLU: $P(\text{dead}) = \Phi\left(\frac{-C}{\gamma}\right)$, dove $\Phi$ è la CDF della normale standard e $C$ dipende dal numero di campioni di valutazione.
• Morte per TopK: $P(\text{dead}) \approx \Phi\left(\frac{t_k - C}{\gamma}\right)$, dove $t_k$ è il quantile corrispondente alla soglia TopK.

Queste formule prevedono che all'aumentare di $\gamma$, il tasso di morte aumenti monotonicamente, avvicinandosi al 50% per ReLU e a $1 - k/n$ per TopK in condizioni di estrema gravità.

2.4 Setup sperimentale

Gli autori validano queste affermazioni attraverso 454 combinazioni modello-layer che spaziano tra modelli linguistici, di visione, proteici e genomici. Essi:

• Calcolano $\gamma$ su attivazioni reali (utilizzando la LayerNorm per token).
• Addestrano SAE con e senza perdite ausiliarie (AuxK).
• Conducono esperimenti sintetici in cui $\gamma$ è controllato per isolare la causalità.
• Testano il centramento della media (sottrazione di $\mu$ dagli input o inizializzazione del bias a $\mu$) come strategia di mitigazione.

3. Contributi chiave

1. Diagnostica della gravità dell'outlier ($\gamma$): Gli autori derivano $\gamma$ dai primi principi, dimostrando che esso determina analiticamente la frazione di feature morte all'inizializzazione. Questa diagnostica correla fortemente con i tassi di morte osservati in diversi modelli (Spearman $\rho = 0.89$ per TopK, $\rho = 0.82$ per ReLU) senza alcun fitting.
2. Causa geometrica della morte: Il documento stabilisce che la morte delle feature nei modelli ad alto $\gamma$ è un problema di inizializzazione geometrica, non di dinamica di addestramento. La media di attivazione crea spostamenti di pre-attivazione indipendenti dall'input che predeterminano il destino delle feature prima dell'inizio dell'addestramento.
3. Analisi della dinamica di recupero: Gli autori caratterizzano due distinti percorsi di riattivazione:
   • Rapido (Fast): Le feature morte per TopK si riattivano man mano che le feature concorrenti si rimpiccioliscono durante l'addestramento.
   • Lento (Slow): Le feature morte per ReLU si riattivano solo quando il bias dell'SAE impara lentamente ad assorbire l'offset della media $\mu$. Questo processo è proibitivamente lento ad alti valori di $\gamma$, spiegando perché l'addestramento da solo spesso non riesce a recuperare le feature.
4. Validazione tramite centramento della media: Il documento dimostra che inizializzare il bias dell'SAE con la media dell'attivazione (effettivamente il centramento della media) elimina la morte indotta dagli outlier all'inizializzazione. Ciò conferma il meccanismo geometrico e fornisce una razionale fondata su quando questo pre-processing è necessario.

4. Risultati

• Correlazione: $\gamma$ predice accuratamente i tassi di morte iniziali attraverso 454 combinazioni modello-layer. Esperimenti sintetici confermano una correlazione di rango perfetta ($\rho=1.0$) tra l'aumento di $\gamma$ e l'aumento dei tassi di morte.
• Limitazioni dell'addestramento: Negli scenari ad alto $\gamma$ (es. $\gamma \ge 30$), l'addestramento standard non riesce a recuperare le feature anche dopo 2 milioni di step. Il processo di apprendimento del bias è il collo di bottiglia; le feature con pre-attivazioni permanentemente negative rimangono morte perché il bias non può apprendere l'ampio offset abbastanza velocemente.
• Limitazioni di AuxK: AuxK accelera il recupero delle feature morte per TopK fornendo segnali di gradiente, ma non accelera l'apprendimento del bias necessario per il recupero per ReLU. Di conseguenza, AuxK aiuta in scenari con $\gamma$ moderato ma fallisce nel risolvere la morte ad alto $\gamma$.
• Efficacia del centramento della media:
  • Sui dati sintetici ($\gamma=40$), il centramento della media riduce le feature morte quasi a zero e abbassa l'errore di ricostruzione di un ordine di grandezza.
  • Sui modelli reali (es. ESM3, AlphaFold3), il centramento della media riduce la morte iniziale da >80% a quasi zero.
  • Qualità delle feature: Gli SAE con centramento della media producono feature più monosemantiche. Su ESM3, un SAE centrato con 2.048 feature recupera più concetti biologici (100) rispetto a un SAE baseline con 8.192 feature (73).
• Morte residua: Il centramento della media non elimina tutta la morte. In layer con rango estremamente basso (es. Evo1), dove la varianza è concentrata in pochissime direzioni, le feature muoiono ancora a causa della competizione TopK. Gli autori suggeriscono il whitening tramite PCA o l'Active Subspace Initialization come rimedi per questi casi specifici.

5. Significato e affermazioni

Il documento sostiene di fornire una base fondata per comprendere e mitigare la morte delle feature negli SAE, spostando il paradigma dalla "correzione delle dinamiche di addestramento" alla "correzione della geometria di inizializzazione".

• Guida pratica: Gli autori propongono un flusso di lavoro diagnostico semplice: calcolare $\gamma$ per un layer. Se $\gamma$ è alto, il centramento della media (inizializzare il bias alla media dell'attivazione) è necessario per prevenire la morte delle feature. Se $\gamma$ è basso, l'addestramento standard è sufficiente.
• Efficienza: Il centramento della media elimina la necessità di tecniche di riattivazione computazionalmente costose (come AuxK) nei regimi ad alto $\gamma$ e consente a dizionari più piccoli di raggiungere la stessa o migliore interpretabilità rispetto a quelli più grandi e non centrati.
• Approfondimento teorico: Il lavoro chiarisce che la variabilità nelle prestazioni degli SAE tra diversi domini (linguaggio vs biologia) è guidata dalla geometria delle attivazioni dei modelli sorgente, non dall'architettura dell'SAE stesso.

Il documento conclude che, per la maggior parte dei modelli, inizializzare il bias dell'SAE alla mediana geometrica delle attivazioni è un passaggio a costo zero e ad alto impatto che garantisce che la dimensione effettiva del dizionario corrisponda alla dimensione desiderata, consentendo così un'interpretabilità meccanicistica più affidabile attraverso diverse famiglie di modelli.