Visualizing LLM Latent Space Geometry Through Dimensionality Reduction

Sintesi Tecnica: Visualizzare la Geometria dello Spazio Latente degli LLM attraverso la Riduzione della Dimensionalità

Enunciato del Problema
Sebbene i Large Language Models (LLM) raggiungano prestazioni allo stato dell'arte, i loro processi decisionali interni rimangono opachi. L'interpretabilità meccanicistica cerca di l'ingegnerizzare a ritroso questi componenti interni per ottenere una comprensione causale e granulare. Una via promettente per questo è la geometria delle feature: analizzare la struttura e l'organizzazione delle rappresentazioni all'interno di spazi latenti ad alta dimensionalità. Tuttavia, questi spazi sono difficili da interpretare direttamente a causa della loro estrema dimensionalità. La sfida consiste nell'estrarre, elaborare e visualizzare questi stati latenti per scoprire come le feature astratte vengono codificate, trasformate e organizzate attraverso i layer del modello senza perdere informazioni strutturali salienti.

Metodologia
Gli autori propongono una pipeline sistematica per catturare e visualizzare le geometrie degli stati latenti in modelli Transformer decoder-only (specificamente GPT-2 e LLaMa). La metodologia consiste in tre fasi principali:

1. Pipeline di Estrazione del Latente:

   • La pipeline elabora dati testuali (dal dataset PG-19) o singoli token attraverso l'architettura Transformer.
   • Cattura le attivazioni in sei punti distinti all'interno di ogni blocco Transformer: gli output dei layer di normalizzazione pre-attention e pre-MLP, gli output grezzi dei moduli di attention e MLP, e gli stati del residual stream dopo l'aggiunta di questi componenti (pre-add e post-add).
   • Gli autori adottano una visione "centrata sul residual stream", trattando il modello come un canale di comunicazione condiviso dove i componenti leggono da e scrivono in uno spazio comune.

2. Elaborazione e Pre-elaborazione dei Dati:

   • Per gestire l'alta dimensionalità e i vincoli di memoria, gli stati latenti sono opzionalmente ridotti (ad esempio, da 4096 a 512 dimensioni) tramite PCA.
   • I passaggi specifici di pre-elaborazione includono la conversione dei vettori in lunghezza unitaria per enfatizzare la struttura direzionale rispetto alla magnitudo, e la media dei vettori attraverso le dimensioni di campione o sequenza per isolare specifiche fonti di variabilità (ad esempio, l'evoluzione layer-wise rispetto al rumore a livello di token).
   • Gli autori escludono esplicitamente la 0-esima posizione della sequenza (il token iniziale) in molte visualizzazioni a causa della sua norma anomalamente elevata, che può oscurare altri pattern geometrici.

3. Riduzione della Dimensionalità e Visualizzazione:

   • Analisi delle Componenti Principali (PCA): Utilizzata come baseline lineare per catturare le direzioni di massima varianza. Viene applicata senza centratura della media per preservare la posizione dell'origine.
   • Uniform Manifold Approximation and Projection (UMAP): Utilizzata come alternativa non lineare per preservare i vicinati locali e la struttura globale del manifold. Gli autori utilizzano la distanza del coseno per UMAP per allinearsi con la Linear Representation Hypothesis, la quale postula le feature come direzioni nello spazio latente.
   • Le visualizzazioni sono codificate per colore in base al layer (blu per l'attention, rosso per l'MLP) e alla posizione nella sequenza (gradiente da scuro a chiaro).

Contributi Chiave e Risultati

• Separazione Geometrica Attention vs. MLP:
  Il risultato più significativo è una chiara e persistente separazione geometrica tra gli output dei componenti di attention e quelli dei componenti Multi-Layer Perceptron (MLP) nei layer intermedi. Nelle visualizzazioni a dimensionalità ridotta (sia PCA che UMAP), gli output dell'attention occupano costantemente regioni distinte dagli output dell'MLP. Gli autori osservano che questo pattern non è stato documentato in lavori precedenti.

• Alta Norma della Posizione Iniziale della Sequenza:
  Lo studio conferma ed espande le osservazioni secondo cui la norma dello stato latente alla 0-esima posizione della sequenza è significativamente maggiore rispetto alle posizioni successive.

  • In GPT-2 (utilizzando embedding posizionali appresi), ciò è previsto.
  • In LLaMa (utilizzando Rotary Positional Encodings/RoPE), gli autori scoprono che l'alta norma non è limitata al token <BOS>, ma si applica alla stragrande maggioranza dei token del vocabolario quando appaiono come token iniziale. Ciò suggerisce che LLaMa abbia sviluppato un meccanismo complesso per rilevare l'inizio della sequenza, distinto dai semplici embedding appresi di GPT-2.

• Evoluzione Layerwise e Dinamiche della Norma:
  Le visualizzazioni delle norme degli stati latenti attraverso i layer rivelano che le norme generalmente aumentano attraverso i layer centrali e scendono dopo l'ultimo layer di normalizzazione. Il token iniziale presenta uno "spike" unico nella norma che emerge dopo alcuni layer e si dissipa nei layer finali.

• Geometrie degli Embedding Posizionali:

  • GPT-2: Le visualizzazioni confermano la struttura "elicoidale" degli embedding posizionali appresi di GPT-2 nello spazio ad alta dimensione, visibile attraverso varie coppie di dimensioni PCA.
  • LLaMa: I pattern geometrici per RoPE sono meno distinti. Nelle prime dimensioni PCA, una "coda" di posizioni iniziali della sequenza conduce in una nuvola densa di posizioni successive. Nelle dimensioni successive, questo pattern si dissolve in una nuvola uniforme, suggerendo che la geometria posizionale nei modelli RoPE possa essere relativamente a bassa dimensionalità.

Significato e Rivendicazioni
Il documento sostiene di fornire un framework robusto per visualizzare la geometria interna degli LLM, rendendo accessibile il concetto astratto di "geometria delle feature". Dimostrando la separazione tra gli output di attention e MLP, il lavoro offre una nuova prospettiva su come diversi componenti Transformer contribuiscano al residual stream condiviso. Gli autori posizionano questo lavoro come un passo verso un'interpretabilità principiata, suggerendo che comprendere queste strutture geometriche sia essenziale per analizzare come i modelli generalizzano, ragionano e astraggono. Gli autori dichiarano esplicitamente che le loro scoperte riguardanti la separazione attention/MLP sono nuove a loro conoscenza e che i pattern osservati nelle norme del token iniziale di LLaMa meritano ulteriori indagini sui meccanismi di encoding posizionale relativo.

Il lavoro non pretende di risolvere il problema completo dell'interpretabilità, ma fornisce un "forte background ed esposizione" e un insieme di strumenti (la pipeline di visualizzazione) per facilitare la ricerca ulteriore sulla dinamica degli stati latenti attraverso i layer e i modelli.