From Weights to Features: SAE-Guided Activation Regularization for LLM Continual Learning

Sintesi Tecnica: Dai Pesi alle Caratteristiche: Regolarizzazione delle Attivazioni Guidata da SAE per l'Apprendimento Continuo di LLM

Problematica
L'apprendimento continuo (CL) nei Large Language Models (LLM) affronta la sfida dell'oblio catastrofico, in cui il fine-tuning sequenziale sovrascrive le rappresentazioni necessarie per i task precedenti. Sebbene i metodi basati sulla regolarizzazione come l'Elastic Weight Consolidation (EWC) siano uno standard per i modelli più piccoli, essi performano significativamente peggio negli LLM. Gli autori sostengono che questo fallimento sia strutturale, derivante dalla natura "polisemantica" degli LLM. A causa dell'ipotesi della sovrapposizione (superposition hypothesis), i singoli neuroni e i loro pesi associati codificano simultaneamente molteplici concetti non correlati. Di conseguenza, le stime di importanza nello spazio dei pesi (es. informazione di Fisher) sono troppo grossolane; proteggere un peso per preservare il Concetto A vincola involontariamente il Concetto B, creando una penalità non selettiva che soffoca la plasticità. Inoltre, gli esistenti metodi spesso non riescono a imporre esplicitamente la plasticità necessaria per apprendere nuovi task, portando all' "intransigenza", in cui il modello conserva la vecchia conoscenza ma non è in grado di adattarsi.

Metodologia
Il documento propone la regolarizzazione delle attivazioni guidata da SAE, un metodo che sposta il target della regolarizzazione dallo spazio dei parametri intrecciato a uno spazio di attivazione più selettivo utilizzando Sparse Autoencoder (SAE) pre-addestrati.

1. Spostamento del Sistema di Coordinate: Invece di penalizzare la deriva dei pesi, il metodo utilizza un SAE pre-addestato e congelato come un dizionario di feature monosemantiche fisse. Gli SAE decompongono le attivazioni dense del modello in direzioni di feature sparse che sono spesso più interpretabili e meno intrecciate rispetto ai singoli neuroni.
2. Formulazione dell'Ottimizzazione Vincolata: L'obiettivo di addestramento è derivato da un problema di ottimizzazione vincolata piuttosto che da un design di penalità ad hoc. Il modello deve minimizzare la perdita del task corrente soggetto a due vincoli:
   • Stabilità: Le feature irrilevanti per il task (bassa rilevanza) non devono deviare oltre un budget specifico ($\varepsilon_s$).
   • Plasticità: Le feature rilevanti per il task (alta rilevanza) devono deviare sufficientemente ($\ge \varepsilon_p$) per evitare soluzioni degenerate in cui il modello non impara nulla.
3. Funzione di Perdita: Utilizzando la rilassazione Lagrangiana e una penalità squared-hinge, gli autori derivano una perdita di addestramento composta da una perdita di protezione (che penalizza l'eccessiva deriva nelle feature protette) e una perdita di guida (che penalizza l'insufficiente deriva nelle feature rilevanti). L'EWC è mostrato essere un caso speciale di questo framework dove lo spazio di rappresentazione è quello dei pesi e il vincolo di plasticità è assente.
4. Costruzione della Maschera: Per ogni nuovo task, viene costruita una maschera di rilevanza continua ($M_t$) utilizzando dati del task corrente passati attraverso il modello base congelato e l'SAE. Questa maschera identifica le feature rilevanti tramite la magnitudo dell'attivazione e propaga la rilevanza alle feature vicine usando il k-NN sui vettori del decoder. Fondamentalmente, una volta costruita la maschera, tutti i dati del task vengono scartati; viene conservata solo la maschera compatta.
5. Addestramento: Durante l'addestramento, il modello corrente (con LoRA) e il modello base congelato sono entrambi codificati dall'SAE. La deriva delle feature è pesata dalla maschera per calcolare le perdite di protezione e di guida.

Contributi Chiave

• Regolarizzazione Guidata da SAE: Un metodo pratico di CL che regolarizza nello spazio delle attivazioni utilizzando le feature SAE. Non richiede dati di task precedenti, non memorizza attivazioni, non richiede parametri specifici per il task e non richiede routing al tempo di inferenza dopo la costruzione della maschera.
• Derivazione Principata: Il tradeoff stabilità-plasticità è formulato come un problema di ottimazione vincolata nello spazio delle feature SAE, producendo una perdita di addestramento squared-hinge tramite rilassazione Lagrangiana.
• Evidenza Empirica della Polisemanticità: Il documento fornisce prove meccanicistiche che la protezione nello spazio dei pesi è non selettiva. Gli esperimenti mostrano che le rappresentazioni rilevanti per il task sono linearmente separabili nello spazio delle feature SAE (AUC ~0.88) ma indistinguibili dal caso nello spazio dei pesi/neuroni (AUC ~0.50). Inoltre, proteggere i pesi per un task vincola le feature per un task successivo con un tasso del 91–96% rispetto al task originale, mentre la protezione nello spazio delle feature riduce questo vincolo collaterale al 43–61%.
• Scalabilità: Grazie alla minore dimensionalità dello spazio delle feature rispetto allo spazio dei parametri, lo storage per task è ridotto da gigabyte (ancore di peso) a sub-megabyte (maschere di feature).

Risultati
Il metodo è stato valutato su TRACE-5000 (8 task cross-domain) e MedCL (10 task biomedici) utilizzando Gemma-2 9B-it con LoRA.

• Performance: Su TRACE, la regolarizzazione guidata da SAE ha ottenuto la performance più forte tra gli approcci non architettonici (OP = 0.545), superando i regolarizzatori nello spazio dei pesi (EWC: 0.447, SI: 0.513), i metodi di proiezione del gradiente e i baseline di replay (anche con il 10% di dati). Su MedCL, ha superato tutti i metodi che non conservano esempi di task precedenti (OP = 0.510).
• Tradeoff Stabilità-Plasticità: A differenza di EWC, che ha ottenuto un alto mantenimento (ARA = 0.987) sacrificando severamente la plasticità (0.453), il metodo proposto ha mantenuto un'alta plasticità (0.567, comparabile al fine-tuning non protetto) pur mantenendo il 95% della performance appresa.
• Efficienza: Il metodo aggiunge un overhead computazionale minimo (~1.88s/step rispetto a 1.70s del baseline) e riduce drasticamente i requisiti di storage (412 KB per task rispetto a 6.5 GB per le ancore EWC).

Significatività e Rivendicazioni
Il documento sostiene che la scarsa performance della regolarizzazione nello spazio dei pesi negli LLM non è un problema di tuning degli iperparametri, ma una conseguenza strutturale della polisemanticità. Spostando il vincolo di regolarizzazione su una base di feature monosemantiche, il metodo ottiene una protezione più selettiva della conoscenza. Gli autori affermano che l'efficacia della regolarizzazione dipende criticamente dallo spazio di rappresentazione in cui vengono imposti i vincoli. Il metodo dimostra che è possibile ottenere una forte performance di apprendimento continuo senza conservare dati passati o modificare l'architettura del modello al tempo di inferenza, a condizione che la regolarizzazione operi su un sistema di coordinate allineato con la struttura concettuale interna del modello. Il documento conclude che la regolarizzazione basata su SAE offre un'alternativa più selettiva e scalabile ai tradizionali metodi nello spazio dei pesi per gli LLM.