CLT-Forge: A Scalable Library for Cross-Layer Transcoders and Attribution Graphs

1. O Problema

A Interpretabilidade Mecanística visa compreender como os Grandes Modelos de Linguagem (LLMs) representam e processam informações. Abordagens recentes, baseadas em dictionary learning (aprendizado de dicionário) e Transcoders, permitem representar a computação do modelo através de características (features) esparsas e interpretáveis, gerando Gráficos de Atribuição (Feature Attribution Graphs).

No entanto, existem limitações críticas que impedem a escalabilidade prática dessas abordagens:

• Redundância e Complexidade: Os gráficos de atribuição gerados por transcoders padrão tornam-se excessivamente grandes e redundantes (com centenas de nós), pois características semelhantes emergem em múltiplas camadas, dificultando a análise.
• Custo Computacional dos CLTs: Os Transcoders de Camada Cruzada (Cross-Layer Transcoders - CLTs) foram propostos para resolver a redundância, compartilhando características entre camadas enquanto mantêm a decodificação específica de cada camada. Isso resulta em gráficos mais compactos. Contudo, os CLTs possuem um número massivo de parâmetros (crescimento quadrático em relação ao número de camadas), tornando seu treinamento e análise em larga escala computacionalmente proibitivos.
• Falta de Infraestrutura Unificada: Não existiam frameworks open-source que suportassem o treinamento end-to-end de CLTs em escala industrial. As soluções existentes careciam de componentes essenciais como treinamento distribuído eficiente, cache de ativações comprimido e pipelines integrados de interpretação automatizada e visualização.

2. Metodologia e Solução (CLT-Forge)

O artigo apresenta o CLT-Forge, uma biblioteca open-source unificada projetada para o treinamento escalável e a análise de interpretabilidade de CLTs. A metodologia é dividida em quatro componentes principais:

A. Infraestrutura de Treinamento Escalável

• Sharding de GPU (Feature-wise Sharding): Para contornar as limitações de memória de uma única GPU, o CLT-Forge implementa um sharding baseado em características (features). As características são divididas entre múltiplas GPUs, onde cada worker processa um subconjunto de características, agregando os resultados no final do forward pass. Isso supera as limitações do Fully Sharded Data Parallel (FSDP) para arquiteturas de CLT.
• Cache de Ativações Comprimido: O treinamento requer acesso a ativações de todas as camadas em grandes conjuntos de dados. Para evitar o armazenamento de terabytes de dados em float16, o sistema utiliza:
  • Quantização Simétrica por Camada: Reduz a precisão para int8, int4 ou int2.
  • Codificação Entropica (zstd): Comprime os dados quantizados.
  • Resultado: Redução de armazenamento de ~20TB para ~4TB (para um modelo LLaMA 1B com 300M tokens) com apenas uma queda de 2-3% na qualidade de reconstrução.

B. Treinamento do Modelo

O objetivo de treinamento visa reconstruir as saídas das camadas MLP (Multi-Layer Perceptron) a partir das ativações de entrada, impondo representações esparsas. A função de perda inclui:

1. Erro de Reconstrução: Diferença entre a saída real e a reconstruída.
2. Regularização de Esparsidade: Uso de JumpReLU e agendamento de esparsidade para garantir que poucas características sejam ativas.
3. Regularização de "Features Mortas": Penaliza características que nunca são ativadas.
   O framework suporta Low-Rank Fine-Tuning (LoRA) para adaptar CLTs pré-treinados a novos dados sem retreinamento completo e custoso.

C. Pipeline de Interpretabilidade Automatizada (AutoInterp)

O sistema integra um pipeline paralelo para análise post-hoc:

• Identifica as sequências de entrada que mais ativam cada característica.
• Calcula estatísticas de resumo (tokens mais ativos, etc.).
• Gera explicações em linguagem natural usando LLMs (via prompting) baseadas nessas estatísticas.
• Todo o processo é paralelizado na dimensão das características, permitindo a análise de milhões de features em uma única passagem.

D. Gráficos de Atribuição e Visualização

• Circuit-Tracer: Integração nativa com a biblioteca Circuit-Tracer para calcular e podar os gráficos de atribuição de características, capturando influências causais entre camadas.
• Interface Visual Interativa: Uma interface baseada em Python (Dash) que permite:
  • Explorar características individualmente.
  • Visualizar gráficos de atribuição.
  • Realizar interventions (intervenções) no nível de clusters de características.
  • Agrupar características para formar estruturas de nível superior.

3. Contribuições Principais

1. Primeira Biblioteca Unificada para CLTs: Oferece o primeiro framework open-source completo para treinamento escalável e análise de interpretabilidade de Transcoders de Camada Cruzada.
2. Eficiência de Treinamento: Suporte nativo para sharding de características em GPU e cache de ativações quantizado/comprimido, permitindo treinar CLTs em modelos como LLaMA 1B com ~1,5 milhão de características em hardware acessível (8x GPUs de 80GB).
3. Pipeline Integrado: Une treinamento, análise de atribuição, geração de explicações automatizadas e visualização interativa em um único fluxo de trabalho, eliminando a necessidade de combinar múltiplas bibliotecas fragmentadas.
4. Ferramentas de Análise Avançada: Integração com Circuit-Tracer para intervenções e visualização de circuitos causais complexos.

4. Resultados e Avaliação

• Desempenho em GPT-2: A implementação alcançou resultados comparáveis ao estado da arte, com ~0,8 de variância explicada (em L0=100), pontuação de substituição (replacement score) de ~0,8 e completude do gráfico de ~0,95 em prompts padrão.
• Eficiência de Memória: O feature-sharding demonstrou ser significativamente mais eficiente em memória do que o DDP tradicional, permitindo o treinamento de modelos maiores.
• Validação Qualitativa: O pipeline foi validado qualitativamente na geração de gráficos de atribuição interpretáveis (ex: análise da frase "O oposto de 'grande' é...") e na capacidade de realizar intervenções no nível de clusters.
• Aplicações: O código já foi utilizado em estudos de replicação de mecanismos (ex: "greater-than") e análises multilíngues.

5. Significado e Impacto

O CLT-Forge representa um avanço significativo para a comunidade de Interpretabilidade Mecanística ao:

• Democratizar o Acesso: Remove barreiras técnicas e computacionais para o uso de CLTs, que anteriormente eram restritos a laboratórios industriais com frameworks proprietários (como o da Anthropic).
• Facilitar a Pesquisa Reprodutível: Fornece uma base padronizada para treinar, analisar e visualizar modelos interpretáveis em escala.
• Acelerar a Descoberta: Ao permitir a visualização interativa e a intervenção em circuitos complexos de forma eficiente, acelera a compreensão de como os LLMs processam informações e tomam decisões.

O trabalho aponta para futuros desenvolvimentos, como a inclusão de mapas de atenção na atribuição e a otimização de arquiteturas para reduzir ainda mais o custo computacional. O código está disponível publicamente sob licença MIT.