Visualizing LLM Latent Space Geometry Through Dimensionality Reduction

Resumo Técnico: Visualizando a Geometria do Espaço Latente de LLMs através da Redução de Dimensionalidade

Definição do Problema
Embora os Grandes Modelos de Linguagem (LLMs) alcancem um desempenho de estado da arte, seus processos internos de tomada de decisão permanecem opacos. A interpretabilidade mecanística busca realizar engenharia reversa nesses componentes internos para alcançar uma compreensão causal e granular. Um caminho promissor para isso é a geometria de características (feature geometry): analisar a estrutura e a organização das representações dentro de espaços latentes de alta dimensionalidade. No entanto, esses espaços são difíceis de interpretar diretamente devido à sua extrema dimensionalidade. O desafio reside em extrair, processar e visualizar esses estados latentes para descobrir como características abstratas são codificadas, transformadas e organizadas através das camadas do modelo sem perder informações estruturais salientes.

Metodologia
Os autores propõem um pipeline sistemático para capturar e visualizar geometrias de estados latentes em modelos Transformer de decodificador apenas (especificamente GPT-2 e LLaMa). A metodologia consiste em três estágios primários:

1. Pipeline de Extração de Latentes:

   • O pipeline processa dados de texto (do conjunto de dados PG-19) ou tokens individuais através da arquitetura Transformer.
   • Ele captura ativações em seis pontos distintos dentro de cada bloco Transformer: as saídas das camadas de normalização pré-atenção e pré-MLP, as saídas brutas dos módulos de atenção e MLP, e os estados do fluxo residual após a adição desses componentes (pré-adição e pós-adição).
   • Os autores adotam uma visão "centrada no fluxo residual", tratando o modelo como um canal de comunicação compartilhado onde os componentes leem de e escrevem em um espaço comum.

2. Processamento e Pré-processamento de Dados:

   • Para gerenciar a alta dimensionalidade e as restrições de memória, os estados latentes podem ser opcionalmente reduzidos (por exemplo, de 4096 para 512 dimensões) via PCA.
   • Etapas específicas de pré-processamento incluem a conversão de vetores para comprimento unitário para enfatizar a estrutura direcional sobre a magnitude, e a média de vetores ao longo de dimensões de amostra ou sequência para isolar fontes específicas de variabilidade (por exemplo, evolução camada a camada vs. ruído de nível de token).
   • Os autores excluem explicitamente a 0-ésima posição da sequência (token inicial) em muitas visualizações devido à sua norma anormalmente alta, que pode obscurecer outros padrões geométricos.

3. Redução de Dimensionalidade e Visualização:

   • Análise de Componentes Principais (PCA): Utilizada como uma linha de base linear para capturar direções de variância máxima. É aplicada sem centralização de média para preservar a posição da origem.
   • Uniform Manifold Approximation and Projection (UMAP): Utilizada como uma alternativa não linear para preservar vizinhanças locais e a estrutura do manifold global. Os autores utilizam a distância de cosseno para o UMAP para alinhar-se com a Hipótese da Representação Linear, que postula características como direções no espaço latente.
   • As visualizações são codificadas por cores por camada (azul para atenção, vermelho para MLP) e posição na sequência (gradiente de escuro para claro).

Principais Contribuições e Resultados

• Separação Geométrica Atenção vs. MLP:
  A descoberta mais significativa é uma separação geométrica clara e persistente entre as saídas dos componentes de atenção e os componentes Multi-Layer Perceptron (MLP) em camadas intermediárias. Em visualizações de dimensionalidade reduzida (tanto PCA quanto UMAP), as saídas de atenção ocupam consistentemente regiões distintas das saídas de MLP. Os autores observam que esse padrão não foi documentado em trabalhos anteriores.

• Alta Norma da Posição Inicial da Sequência:
  O estudo confirma e expande observações de que a norma do estado latente na 0-ésima posição da sequência é significativamente maior do que as posições subsequentes.

  • No GPT-2 (usando embeddings posicionais aprendidos), isso é esperado.
  • No LLaMa (usando Rotary Positional Encodings/RoPE), os autores descobrem que a alta norma não se limita ao token <BOS>, mas aplica-se à vasta maioria dos tokens do vocabulário quando aparecem como o token inicial. Isso sugere que o LLaMa desenvolveu um mecanismo complexo para detectar o início da sequência, distinto dos embeddings aprendidos triviais do GPT-2.

• Evolução por Camada e Dinâmica de Norma:
  Visualizações das normas dos estados latentes através das camadas revelam que as normas geralmente aumentam através das camadas médias e caem após a camada de normalização final. O token inicial exibe um "pico" único de norma que emerge após algumas camadas e se dissipa nas camadas finais.

• Geometrias de Embedding Posicional:

  • GPT-2: As visualizações confirmam a estrutura "helicoidal" dos embeddings posicionais aprendidos do GPT-2 no espaço de alta dimensão, visível através de vários pares de dimensões PCA.
  • LLaMa: Os padrões geométricos para o RoPE são menos distintos. Em dimensões iniciais de PCA, uma "cauda" de posições iniciais da sequência conduz a uma nuvem densa de posições posteriores. Em dimensões posteriores, esse padrão se dissolve em uma nuvem uniforme, sugerindo que a geometria posicional em modelos RoPE pode ser de baixa dimensionalidade relativa.

Significância e Alegações
O artigo afirma fornecer um framework robusto para visualizar a geometria interna de LLMs, tornando acessível o conceito abstrato de "geometria de características". Ao demonstrar a separação entre as saídas de atenção e MLP, o trabalho oferece uma nova perspectiva sobre como diferentes componentes Transformer contribuem para o fluxo residual compartilhado. Os autores posicionam este trabalho como um passo em direção à interpretabilidade principista, sugerindo que compreender essas estruturas geométricas é essencial para analisar como os modelos generalizam, raciocinam e abstraem. Os autores declaram explicitamente que suas descobertas sobre a separação atenção/MLP são inéditas de seu conhecimento e que os padrões observados nas normas do token inicial do LLaMa justificam investigações adicionais sobre os mecanismos de codificações posicionais relativas.

O trabalho não pretende resolver o problema completo da interpretabilidade, mas sim fornecer um "forte embasamento e exposição" e um conjunto de ferramentas (o pipeline de visualização) para facilitar pesquisas futuras sobre a dinâmica dos estados latentes através de camadas e modelos.