A lente logit é uma técnica em aprendizado de máquina interpretável usada para inspecionar as representações internas de transformadores baseados em grandes modelos de linguagem. Ela funciona ao pegar o estado oculto em uma determinada camada, aplicar a matriz de desembedding final do modelo (que mapeia estados ocultos para logits sobre o vocabulário) e, em seguida, aplicar um softmax para obter uma distribuição de probabilidade. Essa distribuição revela o que o modelo preveria se parasse naquela camada, oferecendo uma janela para como as previsões evoluem camada por camada. O método foi introduzido por nostalgebraist em 2020 e, desde então, tornou-se uma ferramenta padrão para estudar os cálculos internos de modelos como GPT-2 e GPT-3.
Diferentemente de classificadores de sondagem que exigem treinar modelos auxiliares, a lente logit usa a própria cabeça de saída do modelo, tornando-a uma abordagem de zero parâmetros e sem treinamento. Ela é particularmente eficaz para modelos com fluxos residuais, pois os estados ocultos em cada camada são aditivos e podem ser projetados diretamente. A técnica tem sido usada para identificar quando os modelos se comprometem com previsões específicas, detectar decodificação "tardia" versus "precoce" e analisar fenômenos como recuperação factual e raciocínio aritmético. No entanto, sua aplicabilidade é limitada a modelos com embeddings amarrados ou compartilhados, e pode não funcionar bem para modelos com normalização de camada complexa ou transformações não lineares.
Mecanismo e Implementação
A lente logit opera no fluxo residual de um transformador. Em um transformador típico, cada camada adiciona sua saída ao fluxo residual, então o estado oculto na camada \( l \) é a soma do embedding inicial e todas as saídas das camadas até \( l \). Os logits finais são calculados aplicando uma normalização de camada (se presente) e, em seguida, a matriz de desembedding \( W_U \). A lente logit aplica o mesmo desembedding aos estados ocultos intermediários, frequentemente após aplicar a normalização de camada final (ou uma transformação afim aprendida) para levar em conta a mudança de distribuição.
Para modelos com embeddings amarrados (onde o embedding de entrada e o desembedding de saída compartilham pesos), a projeção é direta. Para modelos com matrizes de desembedding separadas, a técnica ainda funciona desde que o desembedding esteja disponível. O resultado é uma sequência de distribuições de probabilidade sobre o vocabulário para cada camada, que pode ser visualizada como um mapa de calor ou usada para calcular métricas como a "diferença de logit" entre tokens candidatos.
Aplicações em Interpretabilidade
A lente logit tem sido aplicada a uma variedade de tarefas de interpretabilidade. Um uso comum é rastrear o desenvolvimento das previsões ao longo das camadas. Por exemplo, em um modelo solicitado a completar "A capital da França é __", a lente logit pode mostrar que as camadas iniciais atribuem alta probabilidade a "Paris" somente após várias camadas, revelando quando o modelo recupera o fato relevante. Pesquisadores usaram isso para identificar "cabeças de indução" e outros padrões de atenção que contribuem para comportamentos específicos.
Outra aplicação é na detecção de decodificação "tardia", onde um modelo pode inicialmente prever um token, mas depois mudar de ideia. Ao examinar a lente logit em cada camada, pode-se ver quando a previsão final é "travada" e quando ainda é flexível. Isso tem implicações para entender a confiança do modelo e para projetar intervenções.
A técnica também tem sido usada para estudar aritmética e raciocínio. Em modelos como GPT-3, a lente logit pode mostrar que camadas intermediárias representam somas ou transportes intermediários, fornecendo evidências para computação em múltiplas etapas. Isso levou a mais trabalhos em interpretabilidade mecanicista, como o framework de "cabeças de indução" e "circuitos" por Anthropic e outros grupos de pesquisa.
Limitações e Críticas
Apesar de sua utilidade, a lente logit tem várias limitações. Primeiro, ela assume que a matriz de desembedding é uma projeção significativa para todas as camadas, o que pode não se sustentar se os estados ocultos do modelo passarem por transformações significativas (por exemplo, normalização de camada) que não são levadas em conta. Alguns modelos, especialmente aqueles com arquiteturas de pré-normalização, exigem aplicar a normalização de camada final aos estados intermediários, mas isso nem sempre é suficiente.
Segundo, a lente logit é menos eficaz para modelos com grandes vocabulários ou para tarefas que exigem previsões de múltiplos tokens, pois ela mostra apenas a distribuição do próximo token. Ela também falha em capturar o estado completo do modelo, pois ignora a contribuição das subcamadas de atenção e feed-forward além de seu efeito aditivo.
Terceiro, a técnica pode produzir resultados enganosos se o modelo usar uma cabeça de saída separada que não está alinhada com o fluxo residual. Nesses casos, os logits projetados podem ser ruidosos ou pouco informativos. Pesquisadores propuseram variantes como a "lente ajustada" que aprendem uma transformação afim por camada para melhorar o alinhamento, mas essas exigem dados de treinamento adicionais.
Relação com Outros Métodos de Interpretabilidade
A lente logit faz parte de uma família mais ampla de técnicas de interpretabilidade que incluem classificadores de sondagem, correção de ativação e rastreamento causal. Classificadores de sondagem treinam um modelo linear ou não linear em estados ocultos para prever características específicas, mas exigem supervisão e podem não refletir a computação real do modelo. Em contraste, a lente logit usa a própria cabeça de saída do modelo, tornando-a mais direta.
A correção de ativação envolve intervir nos estados ocultos para ver como eles afetam a saída, o que pode ser usado para validar descobertas da lente logit. O rastreamento causal, popularizado por David Kaplan e outros, usa execuções corrompidas para identificar quais camadas são responsáveis por fatos específicos, e a lente logit pode complementar isso mostrando o token previsto em cada camada.
Outro método relacionado é a visão de "fluxo residual", que trata o transformador como uma série de atualizações aditivas. A lente logit é um ajuste natural para essa visão, pois projeta o fluxo residual em cada ponto. Isso inspirou ferramentas como o "Transformer Debugger" e outras bibliotecas de visualização.
Extensões e Variantes
Várias extensões da lente logit foram propostas. A "lente ajustada" aprende uma transformação afim por camada para melhor alinhar os estados ocultos com o desembedding, melhorando o desempenho em modelos onde a lente logit padrão falha. Outra variante é a "lente softmax", que aplica uma temperatura aos logits para afinar ou achatar a distribuição, facilitando a visualização de previsões de baixa probabilidade.
Para modelos com múltiplas cabeças de saída ou decodificadores, a lente logit pode ser aplicada a cada cabeça separadamente. Em modelos codificador-decodificador, a técnica pode ser aplicada aos estados ocultos do decodificador, mas é menos claro como aplicá-la ao codificador.
Trabalhos recentes também exploraram o uso da lente logit para analisar modelos multimodais, onde o vocabulário inclui tokens de imagem ou outras modalidades. No entanto, isso ainda é uma área ativa de pesquisa.
Considerações Computacionais
A lente logit é computacionalmente eficiente, pois requer apenas uma passagem direta e uma multiplicação de matrizes em cada camada. Para um modelo com \( L \) camadas e um tamanho de vocabulário \( V \), o custo adicional é \( O(L \cdot V \cdot d) \), onde \( d \) é a dimensão oculta. Isso é desprezível comparado ao custo da própria passagem direta.
No entanto, armazenar os logits para todas as camadas pode ser intensivo em memória, especialmente para grandes vocabulários. Na prática, pode-se calcular os tokens top-k em cada camada para reduzir o uso de memória. A técnica é implementada em várias bibliotecas de código aberto, incluindo a biblioteca TransformerLens e as ferramentas de interpretabilidade da Hugging Face.
Direções Futuras
À medida que os modelos de aprendizado profundo continuam a crescer em escala, ferramentas de interpretabilidade como a lente logit tornam-se cada vez mais importantes. Trabalhos futuros podem focar em adaptar a técnica a novas arquiteturas, como misturas de especialistas ou modelos de espaço de estado, e em integrá-la com pipelines de interpretabilidade automatizados. A lente logit também está sendo usada em pesquisa de segurança para detectar quando modelos estão "pensando" em saídas prejudiciais, o que pode informar estratégias de alinhamento.
No geral, a lente logit permanece uma ferramenta simples, porém poderosa, para espiar dentro da caixa preta das redes neurais, e suas aplicações provavelmente se expandirão à medida que o campo da interpretabilidade mecanicista amadurece.