Vector Quantized Latent Concepts: A Scalable Alternative to Clustering-Based Concept Discovery

Sintesi Tecnica: Vector Quantized Latent Concepts (VQLC)

Definizione del Problema

I Large Language Models (LLM) codificano ricche informazioni semantiche all'interno dei loro stati nascosti, eppure l'interpretazione di queste rappresentazioni interne rimane una sfida. Mentre i metodi di attribuzione a livello di token (ad esempio, Integrated Gradients, SHAP) identificano i token di input salienti, essi offrono una visione limitata delle strutture semantiche di alto livello codificate nel modello. La scoperta di concetti latenti mira ad affrontare questo problema raggruppando le rappresentazioni contestuali in concetti discreti e interpretabili. Tuttavia, gli approcci esistenti affrontano un compromesso fondamentale tra scalabilità e qualità del concetto:

• Clustering Gerarchico (ad es., LACOAT): Produce concetti coerenti ma soffre di costi di memoria quadratici ($O(N^2)$), limitandone l'applicazione a dataset di piccole dimensioni.
• K-Means Clustering: Scala efficientemente ma spesso produce concetti meno coerenti dal punto di vista semantico e fornisce un trade-off più debole tra scalabilità e fedeltà.

Il documento sostiene che sia necessario un metodo che mantenga la scalabilità dei metodi basati su centroidi pur raggiungendo la coerenza semantica del clustering gerarchico, specificamente per gli LLM su larga scala.

Metodologia: Vector Quantized Latent Concepts (VQLC)

Gli autori propongono VQLC, un framework di apprendimento di concetti discreti che adatta la Quantizzazione Vettoriale (VQ) per mappare gli stati nascosti continui in un insieme finito di codici discreti appresi. Il framework opera su stati nascosti congelati estratti da uno specifico strato di un LLM.

Componenti dell'Architettura:

1. Encoder: Un leggero MLP (Multi-Layer Perceptron) residuo mappa le rappresentazioni dei token contestuali ($h^{(\ell)}_{w_i}$) in uno spazio di codifica. Combina una proiezione lineare diretta con un ramo di correzione non lineare per preservare lo spazio originale del modello consentendo al contempo un'assegnazione stabile.
2. Vector Quantizer: Questo modulo mantiene un codebook apprendibile $E = \{e_1, \dots, e_K\}$. Ogni rappresentazione codificata viene assegnata al vettore del codebook più vicino in base alla distanza euclidea al quadrato.
   • Inizializzazione: Il codebook è inizializzato tramite una traversata "farthest-first" su rappresentazioni di addestramento de-duplicate per garantire una copertura diversificata dello spazio di rappresentazione.
   • Apprendimento: Il codebook viene aggiornato utilizzando una Media Mobile Esponenziale (EMA) degli output dell'encoder assegnati, piuttosto che tramite diretti passi di gradiente, per garantire stabilità. Un meccanismo di "dead code recovery" reinizializza i codici non utilizzati utilizzando gli output dell'encoder con i maggiori errori di assegnazione.
3. Decoder: Un decoder residuo ricostruisce gli stati nascosti originali a partire dai vettori quantizzati. L'obiettivo di addestramento minimizza la perdita di ricostruzione ($L_{rec}$) e una perdita di impegno ($L_{commit}$) per mantenere gli output dell'encoder vicini ai rispettivi vettori del codebook.

Costruzione del Concetto:
Dopo l'addestramento, il modello viene congelato. Un passaggio di costruzione del concetto aggrega tutte le rappresentazioni dei token codificate assegnate allo stesso indice di codice $k$ per formare un vettore di concetto $v_k$. Un concetto latente $c_k$ è definito dalla coppia $(v_k, \text{token assegnati})$. In fase di inferenza, le nuove rappresentazioni dei token vengono assegnate al codebook appreso, e i concetti corrispondenti spiegano l'informazione semantica codificata in tali rappresentazioni.

Contributi Chiave

• Apprendimento Discreto Scalabile: VQLC introduce un framework di quantizzazione vettoriale per la scoperta di concetti latenti che evita la dipendenza quadratica dalla memoria del clustering gerarchico.
• Trade-off Ottimizzato: Il metodo raggiunge un favorevole equilibrio tra fedeltà (qualità del concetto) e scalabilità, mantenendo un uso della memoria quasi costante all'aumentare delle dimensioni del dataset.
• Valutazione Esaustiva: Il documento fornisce una valutazione multifasica attraverso 12 configurazioni dataset-modello, tra cui:
  • Scalabilità: Analisi del tempo di esecuzione e della memoria.
  • Fedeltà: Misurata tramite la variazione della confidenza del modello e delle etichette di predizione dopo l'ablazione della direzione del concetto tramite proiezione ortogonale.
  • LLM-as-a-Judge: Valutazione tramite l'uso di molteplici LLM (Claude, Gemini, DeepSeek) per classificare l'allineamento del concetto con le predizioni del modello.
  • Analisi Qualitativa: Visualizzazione di nuvole di parole dei concetti e casi studio.
  • Confronto con SAE: Contrasto con gli Autoencoder Sparsi (SAE) per evidenziare le differenze tra spiegazioni a livello di feature e a livello di concetto.

Risultati

• Scalabilità: Sul dataset ERASER con il modello Qwen, VQLC scala linearmente con il numero di token. Mentre l'uso di memoria di K-Means cresceva da 1,8GB a 17,5GB e LACOAT falliva a 200k token (richiedendo circa 229GB), VQLC è rimasto quasi costante, aumentando solo da 1,94GB a 2,54GB a 500k token.
• Fedeltà: VQLC ha ottenuto la massima variazione media di confidenza in 6 su 12 configurazioni dataset-modello. Ha dimostrato i guadagni più chiari sui modelli decoder-only (Llama-2, Qwen), superando i baseline in 5 su 6 configurazioni decoder, pur rimanendo competitivo con K-Means e LACOAT sui modelli basati su encoder (RoBERTa).
• Valutazione LLM: Nelle classifiche LLM-as-a-judge, VQLC ha ottenuto il miglior rango medio in 8 su 12 configurazioni, suggerendo che i suoi concetti siano più allineati ai task rispetto a quelli di K-Means o LACOAT.
• Confronto con SAE: A differenza degli SAE, che producono un gran numero di feature sparse con bassi tassi di attivazione, VQLC genera un set più piccolo di concetti (ad es., ~400) che sono attivi su una frazione significativamente maggiore di esempi di test. È stato rilevato che i concetti di VQLC codificano più informazioni rilevanti per il task nella configurazione del modello Qwen.
• Risultati Qualitativi: Le visualizzazioni hanno mostrato che i concetti di VQLC erano semanticamente coerenti (ad es., cluster distinti per "Giocatori di Baseball" vs "Football" in Sports, o "Farmaci Clinici" vs "Infrastrutture di Rete" in Sci/Tech). Nei casi studio, VQLC ha fornito spiegazioni più precise per le predizioni del modello rispetto ai cluster più ampi di K-Means o LACOAT.

Significato e Rivendicazioni

Il documento afferma che VQLC offre un'alternativa pratica e scalabile alla scoperta di concetti basata su clustering per gli LLM. Affronta con successo il collo di bottiglia della memoria del clustering gerarchico senza sacrificare la coerenza semantica dei concetti scoperti. Gli autori sottolineano che VQLC è particolarmente efficace per i modelli decoder-only, dove identifica le direzioni dei concetti rilevanti per il task in modo più efficace rispetto ai baseline esistenti.

Il lavoro posiziona VQLC non come un sostituto degli strumenti di interpretabilità meccanicistica come gli SAE, ma come un approccio complementare che organizza le rappresentazioni in concetti latenti a livello di task piuttosto che in feature monosemantiche. Gli autori segnalano i limiti, inclusa la sensibilità del framework agli iperparametri (dimensione del codebook, peso dell'impegno) e l'attuale focalizzazione sui compiti di classificazione, suggerendo che l'estensione di questo approccio al ragionamento generativo token-per-token sia una direzione per il lavoro futuro.