Vector Quantized Latent Concepts: A Scalable Alternative to Clustering-Based Concept Discovery

Resumo Técnico: Conceitos Latentes de Quantização Vetorial (VQLC)

Declaração do Problema

Os Grandes Modelos de Linguagem (LLMs) codificam informações semânticas ricas em seus estados ocultos, mas interpretar essas representações internas continua sendo um desafio. Embora os métodos de atribuição ao nível de token (ex: Gradientes Integrados, SHAP) identifiquem tokens de entrada salientes, eles oferecem uma visão limitada das estruturas semânticas de alto nível codificadas no modelo. A descoberta de conceitos latentes visa abordar isso, agrupando representações contextuais em conceitos discretos e interpretáveis. No entanto, as abordagens existentes enfrentam um compromisso fundamental entre escalabilidade e qualidade do conceito:

• Agrupamento Hierárquico (ex: LACOAT): Produz conceitos coerentes, mas sofre com custos de memória quadráticos ($O(N^2)$), limitando sua aplicação a conjuntos de dados pequenos.
• Agrupamento K-Means: Escala de forma eficiente, mas frequentemente gera conceitos menos semanticamente coerentes e fornece um compromisso mais fraco entre escalabilidade e fidelidade.

O artigo postula que é necessário um método que mantenha a escalabilidade dos métodos baseados em centroides enquanto alcança a coerência semântica do agrupamento hierárquico, especificamente para LLMs de grande escala.

Metodologia: Conceitos Latentes de Quantização Vetorial (VQLC)

Os autores propõem o VQLC, um framework de aprendizado de conceitos discretos que adapta a Quantização Vetorial (VQ) para mapear estados ocultos contínuos para um conjunto finito de códigos aprendidos. O framework opera sobre estados ocultos congelados extraídos de uma camada específica de um LLM.

Componentes da Arquitetura:

1. Codificador (Encoder): Um MLP residual leve mapeia representações de tokens contextuais ($h^{(\ell)}_{w_i}$) para um espaço de código. Ele combina uma projeção linear direta com um ramo de correção não linear para preservar o espaço original do modelo enquanto permite uma atribuição discreta estável.
2. Quantizador Vetorial (Vector Quantizer): Este módulo mantém um código de livros (codebook) aprendível $E = \{e_1, \dots, e_K\}$. Cada representação codificada é atribuída ao vetor do código mais próximo com base na distância euclidiana ao quadrado.
   • Inicialização: O codebook é inicializado usando uma travessia de "farthest-first" em representações de treinamento desduplicadas para garantir uma cobertura diversificada do espaço de representação.
   • Aprendizado: O codebook é atualizado usando uma Média Móvel Exponencial (EMA) das saídas do codificador atribuídas, em vez de passos diretos de gradiente, para garantir estabilidade. Um mecanismo de "recuperação de código morto" reinicializa códigos não utilizados usando saídas do codificador com os maiores erros de atribuição.
3. Decodificador (Decoder): Um decodificador residual reconstrói os estados ocultos originais a partir dos vetores quantizados. O objetivo de treinamento minimiza a perda de reconstrução ($L_{rec}$) e uma perda de compromisso ($L_{commit}$) para manter as saídas do codificador próximas aos seus vetores de código atribuídos.

Construção de Conceitos:
Após o treinamento, o modelo é congelado. Uma passagem de construção de conceitos agrega todas as representações de tokens codificadas atribuídas ao mesmo índice de código $k$ para formar um vetor de conceito $v_k$. Um conceito latente $c_k$ é definido pelo par $(v_k, \text{tokens atribuídos})$. Na inferência, novas representações de tokens são atribuídas ao codebook aprendido, e os conceitos correspondentes explicam a informação semântica codificada nessas representações.

Principais Contribuições

• Aprendizado Discreto Escalável: O VQLC introduz um framework de quantização vetorial para a descoberta de conceitos latentes que evita a dependência de memória quadrática do agrupamento hierárquico.
• Compromisso Otimizado: O método alcança um equilíbrio favorável entre fidelidade (qualidade do conceito) e escalabilidade, mantendo o uso de memória quase constante à medida que o tamanho do conjunto de dados aumenta.
• Avaliação Abrangente: O artigo fornece uma avaliação multifacetada em 12 configurações de dataset-modelo, incluindo:
  • Escalabilidade: Análise de tempo de execução e memória.
  • Fidelidade: Medida pela mudança na confiança do modelo e nos rótulos de predição após a ablação da direção do conceito via projeção ortogonal.
  • LLM-as-a-Judge: Avaliação usando múltiplos LLMs (Claude, Gemini, DeepSeek) para classificar o alinhamento do conceito com as predições do modelo.
  • Análise Qualitativa: Visualização de nuvens de palavras de conceitos e estudos de caso.
  • Comparação com SAEs: Contraste com Autoencoders Esparsos (SAEs) para destacar as diferenças entre explicações de nível de característica (feature-level) e de nível de conceito.

Resultados

• Escalabilidade: No dataset ERASER com o modelo Qwen, o VQLC escala linearmente com a contagem de tokens. Enquanto o uso de memória do K-Means cresceu de 1,8GB para 17,5GB e o LACOAT falhou em 200k tokens (exigindo ~229GB), o VQLC permaneceu quase constante, aumentando apenas de 1,94GB para 2,54GB em 500k tokens.
• Fidelidade: O VQLC alcançou a maior mudança média de confiança em 6 de 12 configurações de dataset-modelo. Demonstrou ganhos mais claros em modelos decoder-only (Llama-2, Qwen), superando os baselines em 5 de 6 configurações de decoder, enquanto permaneceu competitivo com K-Means e LACOAT em modelos baseados em encoder (RoBERTa).
• Avaliação por LLM: Nas classificações de LLM-as-a-judge, o VQLC alcançou o melhor ranking médio em 8 de 12 configurações, sugerindo que seus conceitos são mais alinhados com tarefas do que os do K-Means ou LACOAT.
• Comparação com SAEs: Ao contrário dos SAEs, que produzem um grande número de características esparsas com baixas taxas de ativação, o VQLC gera um conjunto menor de conceitos (ex: ~400) que são ativos em uma fração significativamente maior de exemplos de teste. Os conceitos do VQLC foram encontrados como codificadores de mais informações relevantes para a tarefa na predição no cenário do modelo Qwen.
• Achados Qualitativos: As visualizações mostraram que os conceitos do VQLC eram semanticamente coerentes (ex: clusters distintos para "Jogadores de Beisebol" vs. "Futebol" em Esportes, ou "Drogas Clínicas" vs. "Infraestrutura de Rede" em Ciência/Tecnologia). Em estudos de caso, o VQLC forneceu explicações mais precisas para as predições do modelo comparado aos clusters mais amplos do K-Means ou LACOAT.

Significância e Alegações

O artigo afirma que o VQLC oferece uma alternativa prática e escalável para a descoberta de conceitos baseada em agrupamento para LLMs. Ele aborda com sucesso o gargalo de memória do agrupamento hierárquico sem sacrificar a coerência semântica dos conceitos descobertos. Os autores enfatizam que o VQLC é particularmente eficaz para modelos decoder-only, onde identifica direções de conceitos relevantes para a tarefa de forma mais eficaz que os baselines existentes.

O trabalho posiciona o VQLC não como um substituto para ferramentas de interpretabilidade mecanística como os SAEs, mas como uma abordagem complementar que organiza representações em conceitos latentes de nível de tarefa, em vez de características monossêmicas. Os autores observam limitações, incluindo a sensibilidade do framework a hiperparâmetros (tamanho do codebook, peso de compromisso) e o foco atual em tarefas de classificação, sugerindo que estender esta abordagem para o raciocínio generativo token a token é uma direção para trabalhos futuros.