Polar probe linearly decodes semantic structures from LLMs

Resumo Técnico: A Sonda Polar Decodifica Linearmente Estruturas Semânticas de LLMs

Formulação do Problema
O artigo aborda a questão fundamental de como os Modelos de Linguagem de Grande Escala (LLMs) vinculam conceitos para formar estruturas semânticas complexas. Embora trabalhos anteriores tenham demonstrado que estruturas sintáticas (por exemplo, dependências sujeito-verbo) podem ser decodificadas linearmente das ativações de LLMs usando princípios geométricos (Hewitt & Manning, 2019; Diego-Simon et al., 2024), permanece incerto se princípios semelhantes se aplicam à vinculação semântica. Os autores hipotetizam que os LLMs representam a existência e o tipo de relações entre entidades não por meio de mecanismos complexos e não lineares, mas via um código geométrico simples: a distância euclidiana entre os embeddings de entidades codifica a existência de uma relação, enquanto a direção relativa codifica o tipo de relação.

Metodologia
Para testar essa hipótese, os autores introduzem uma Sonda Polar, uma transformação linear treinada para mapear as ativações de camadas de LLMs em um espaço de sonda onde o grafo relacional pode ser recuperado.

1. Conjuntos de Dados: O estudo utiliza um conjunto de dados sintético abrangendo cinco domínios semânticos distintos:

   • Domínios Euclidianos: Ordinalidade variável (variáveis matemáticas), layouts espaciais (arranjos de objetos 2D) e papéis temáticos (interações sociais).
   • Domínios Não Euclidianos: Árvores genealógicas (relações não comutativas, muitos-para-muitos) e mapas de metrô (caminhos mais curtos baseados em rede).
   • Cada amostra consiste em uma descrição em linguagem natural de um grafo relacional. Para impedir que o modelo aprenda correlações de nível superficial, a ordem das relações e entidades é aleatorizada, e múltiplas descrições textuais são geradas para o mesmo grafo.

2. Arquitetura da Sonda:

   • Perda Estrutural ($L_s$): Otimiza a correlação de posto de Spearman entre a matriz de distância par a par prevista e a matriz de distância de caminho mais curto de verdade fundamental do grafo semântico.
   • Perda Angular ($L_a$): Otimiza o alinhamento entre o vetor de diferença de dois embeddings de entidade ($\delta_{ij}$) e um vetor protótipo aprendido ($p_r$) específico para o tipo de relação $r$.
   • A sonda é treinada independentemente em cada camada de vários LLMs (Llama3-8B, OLMo-7B e a família Pythia).

3. Avaliação:

   • Desempenho de Decodificação: Medido via $\rho$ de Spearman para existência de relação (ordenação de distância) e tipo de relação (alinhamento de direção).
   • Generalização: Testado em nomes de entidades e formas superficiais de relações fora da distribuição (OOD).
   • Intervenção Causal: Os autores intervêm nos estados ocultos direcionando as ativações ao longo das direções de protótipo aprendidas para observar mudanças no desempenho subsequente de Perguntas e Respostas (QA).
   • Correlação: Os erros da sonda são correlacionados com a probabilidade de logit de respostas corretas em tarefas de QA.

Principais Resultados

• Decodificabilidade Linear: As verdadeiras estruturas semânticas podem ser recuperadas linearmente das ativações de LLMs. A Sonda Polar decodifica com sucesso tanto a existência quanto o tipo de relação em todos os cinco domínios.
• Emergência de Camada: O código geométrico é mais proeminente nas camadas intermediárias de LLMs pré-treinados (atingindo o pico por volta das camadas 12–15 para o Llama3-8B). Em modelos inicializados aleatoriamente, o desempenho de decodificação permanece em níveis de acaso, indicando que o código emerge durante o pré-treinamento.
• Tamanho do Modelo e Treinamento: O desempenho escala com o tamanho do modelo (de 14M para 6,9B de parâmetros) e aumenta com as etapas de pré-treinamento, sugerindo que o código polar é uma capacidade aprendida dependente da capacidade do modelo e da duração do treinamento.
• Generalização: As sondas generalizam para novas entidades e formas superficiais de relações, embora o desempenho degrade modestamente com entidades OOD e mais significativamente com tipos de relação OOD. O desempenho também declina à medida que o número de entidades no grafo aumenta, provavelmente devido à expansão do espaço combinatório.
• Influência Causal: Direcionar as ativações do modelo ao longo das direções de protótipo aprendidas desloca causalmente a probabilidade de respostas corretas em tarefas de QA, confirmando que essas representações geométricas são funcionalmente relevantes para o raciocínio do modelo.
• Disjunção de Subespaços: Os subespaços aprendidos para diferentes domínios semânticos são amplamente disjuntos no espaço de ativação, com a notável exceção de layouts espaciais e ordinalidade variável, que mostram sobreposição significativa (apoiando a hipótese de superposição).

Significado e Alegações
O artigo alega demonstrar que os LLMs aprendem a construir estruturas semânticas complexas vinculando representações através de um princípio geométrico simples. Especificamente, estende o escopo da sondagem linear além da sintaxe para uma classe mais ampla de grafos direcionados e rotulados.

• Além da Sintaxe: Diferentemente de trabalhos anteriores limitados a árvores sintáticas, esta abordagem decodifica com sucesso estruturas semânticas que não são necessariamente estruturadas em árvore e incluem topologias não euclidianas (por exemplo, árvores genealógicas, mapas de metrô), embora com algumas limitações de precisão para relações não comutativas.
• Além da Recuperação de Conhecimento: O estudo visa estruturas composicionais que não poderiam ser memorizadas durante o treinamento (devido a entidades e relações aleatorizadas), fornecendo evidências de como os modelos realizam aprendizado em contexto e vinculação.
• Tensão Simbólico-Connectionista: As descobertas sugerem uma resolução para a tensão entre abordagens simbólicas e connectionistas, propondo que estruturas simbólicas podem ser representadas dentro de sistemas vetoriais via subespaços lineares e relações geométricas.

Os autores permanecem modestos quanto às limitações do método, notando que a Sonda Polar depende de distâncias lineares e é adaptada para grafos euclidianos. Consequentemente, não pode representar perfeitamente relações inteiras não comutativas, muitos-para-muitos ou distâncias complexas de caminho mais curto encontradas em domínios não euclidianos, embora capture coerência local dentro deles.