Visualizing LLM Latent Space Geometry Through Dimensionality Reduction

Technische Samenvatting: Visualiseren van de Geometrie van de Latente Ruimte van LLM's via Dimensiereductie

Probleemstelling
Hoewel Large Language Models (LLM's) de staat van de kunst bereiken qua prestaties, blijft hun interne besluitvormingsproces opaak. Mechanistische interpreteerbaarheid streeft ernaar om deze interne componenten te reverse-engineeren om een fijnmazig, causaal begrip te verkrijgen. Een veelbelovende weg hiervoor is feature-geometrie: het analyseren van de structuur en organisatie van representaties binnen hoogdimensionele latente ruimtes. Echter, deze ruimtes zijn direct moeilijk te interpreteren vanwege hun extreme dimensionaliteit. De uitdaging ligt in het extraheren, verwerken en visualiseren van deze latente toestanden om te ontdekken hoe abstracte kenmerken worden gecodeerd, getransformeerd en georganiseerd over de lagen van het model heen, zonder daarbij essentiële structurele informatie te verliezen.

Methodologie
De auteurs stellen een systematische pijplijn voor om de geometrie van latente toestanden in decoder-only Transformer-modellen (specifiek GPT-2 en LLaMa) te vangen en te visualiseren. De methodologie bestaat uit drie primaire fasen:

1. Latente Extractie-pijplijn:

   • De pijplijn verwerkt tekstgegevens (van de PG-19 dataset) of individuele tokens door de Transformer-architectuur.
   • Het legt activaties vast op zes afzonderlijke punten binnen elke Transformer-block: de outputs van de pre-attention en pre-MLP normalisatielagen, de ruwe outputs van de attention- en MLP-modules, en de residual stream-toestanden na de additie van deze componenten (pre-add en post-add).
   • De auteurs hanteren een "residual stream-centrische" visie, waarbij het model wordt beschouwd als een gedeeld communicatiekanaal waar componenten uit een gemeenschappelijke ruimte lezen en naar schrijven.

2. Dataverwerking en Preprocessing:

   • Om de hoge dimensionaliteit en geheugenbeperkingen te beheersen, worden latente toestanden optioneel gereduceerd (bijv. van 4096 naar 512 dimensies) via PCA.
   • Specifieke preprocessing-stappen omvatten het omzetten van vectoren naar een eenheidslengte om de nadruk te leggen op de directionele structuur boven de magnitude, en het middelen van vectoren over sample- of sequentiedimensies om specifieke bronnen van variabiliteit te isoleren (bijv. laag-gewijze evolutie versus token-niveau ruis).
   • De auteurs sluiten de 0-de sequentiepositie (het initiële token) expliciet uit bij veel visualisaties vanwege de anomalie van de extreem hoge norm, die andere geometrische patronen kan maskeren.

3. Dimensiereductie en Visualisatie:

   • Principal Component Analysis (PCA): Wordt gebruikt als een lineaire baseline om richtingen van maximale variantie te vangen. Het wordt toegepast zonder gemiddelde-centrering om de positie van de oorsprong te behouden.
   • Uniform Manifold Approximation and Projection (UMAP): Wordt gebruikt als een niet-lineair alternatief om lokale buurten en de globale manifold-structuur te behouden. De auteurs gebruiken cosinusafstand voor UMAP om aan te sluiten bij de Linear Representation Hypothesis, die stelt dat features zich manifesteren als richtingen in de latente ruimte.
   • Visualisaties zijn gecodeerd met kleur op basis van laag (blauw voor attention, rood voor MLP) en sequentiepositie (gradiënt van donker naar licht).

Belangrijkste Bijdragen en Resultaten

• Geometrische Separatie tussen Attention versus MLP:
  De meest significante bevinding is een duidelijke, persistente geometrische scheiding tussen de outputs van attention-componenten en Multi-Layer Perceptron (MLP)-componenten in intermediaire lagen. In gedimensiereduceerde visualisaties (zowel PCA als UMAP) bezetten attention-outputs consistent andere regio's dan MLP-outputs. De auteurs merken op dat dit patroon in eerdere werken niet is gedocumenteerd.

• Hoge Norm van de Initiële Sequentiepositie:
  De studie bevestigt en breidt observaties uit dat de norm van de latente toestand op de 0-de sequentiepositie aanzienlijk groter is dan bij opeenvolgende posities.

  • In GPT-2 (met geleerde positional embeddings) is dit te verwachten.
  • In LLaMa (met Rotary Positional Encodings/RoPE) vinden de auteurs dat de hoge norm niet beperkt is tot het <BOS> token, maar geldt voor het overgrote deel van de vocabulaire-tokens wanneer zij als het initiële token verschijnen. Dit suggereert dat LLaMa een complex mechanisme heeft ontwikkeld om het begin van de sequentie te detecteren, wat verschilt van de triviale geleerde embeddings van GPT-2.

• Laag-gewijze Evolutie en Norm-dynamiek:
  Visualisaties van de latente toestand-normen over de lagen heen laten zien dat de normen over het algemeen toenemen door de middelste lagen en dalen na de laatste normalisatielaag. Het initiële token vertoont een unieke "spike" in de norm die verschijnt na enkele lagen en verdwijnt in de laatste lagen.

• Positional Embedding Geometrieën:

  • GPT-2: De visualisaties bevestigen de "helische" structie van GPT-2's geleerde positional embeddings in de hoogdimensionele ruimte, zichtbaar over diverse PCA-dimensieparen.
  • LLaMa: De geometrische patronen voor RoPE zijn minder uitgesproken. In vroege PCA-dimensies leidt een "staart" van initiële sequentieposities naar een dichte wolk van latere posities. In latere dimensies lost dit patroon op in een uniforme wolk, wat suggereert dat de positionele geometrie in RoPE-modellen relatief laag-dimensionaal kan zijn.

Betekenis en Claims
Het paper beweert een robuust framework te bieden voor het visualiseren van de interne geometrie van LLM's, waardoor het abstracte concept van "feature-geometrie" toegankelijk wordt gemaakt. Door de scheiding tussen attention- en MLP-outputs aan te tonen, biedt het werk een nieuw perspectief op hoe verschillende Transformer-componenten bijdragen aan de gedeelde residual stream. De auteurs positioneren dit werk als een stap naar een principiële interpreteerbaarheid, waarbij zij suggereren dat het begrijpen van deze geometrische structuren essentieel is voor het analyseren van hoe modellen generaliseren, redeneren en abstraheren. De auteurs verklaren expliciet dat hun bevindingen over de attention/MLP-scheiding nieuw zijn naar hun weten en dat de geobserveerde patronen in de LLaMa-initiële token verdere investigatie behoeven met betrekking tot de mechanismen van relatieve positionele encodings.

Het werk claimt niet het volledige probleem van interpreteerbaarheid op te lossen, maar biedt eerder een "sterke achtergrond en expositie" en een set instrumenten (de visualisatie-pijplijn) om verder onderzoek naar de dynamiek van latente toestanden over lagen en modellen heen te faciliteren.