Compreender redes neuronais através de circuitos esparsos
Treinámos modelos para pensar em passos mais simples e rastreáveis, para que possamos compreender melhor como funcionam.
As redes neuronais são a base dos sistemas de IA mais avançados atualmente, mas continuam a ser difíceis de compreender. Não escrevemos estes modelos com instruções explícitas, passo a passo. Em vez disso, aprendem ajustando milhares de milhões de ligações internas, ou "pesos", até dominarem uma tarefa. Concebemos as regras do treino, mas não os comportamentos específicos que surgem, e o resultado é uma teia densa de ligações que nenhum humano consegue decifrar facilmente.
À medida que os sistemas de IA se tornam mais capazes e têm impacto real nas decisões no campo da ciência, educação e saúde, compreender como funcionam é essencial. A interpretabilidade refere-se aos métodos que nos ajudam a compreender porque é que um modelo produziu um determinado resultado. Existem muitas formas de o conseguirmos.
Por exemplo, os modelos de raciocínio são incentivados a explicar o seu processo até chegarem a uma resposta final. A interpretabilidade da cadeia de pensamento utiliza estas explicações para monitorizar o comportamento do modelo. Isto é imediatamente útil: as cadeias de pensamento dos modelos de raciocínio atuais parecem fornecer informações relevantes relativamente a comportamentos preocupantes, como a mentira. No entanto, confiar totalmente nesta característica é uma estratégia frágil, que pode deixar de funcionar com o tempo.
Por outro lado, a interpretabilidade mecanicista, que é o foco deste trabalho, procura reverter completamente a engenharia dos cálculos de um modelo. Até agora, tem sido menos útil de imediato, mas, em princípio, poderia oferecer uma explicação mais completa do comportamento do modelo. Ao tentar explicar o comportamento do modelo no nível mais detalhado, a interpretabilidade mecanicista pode fazer menos suposições e aumentar a nossa confiança. Mas o caminho desde os detalhes mais básicos até às explicações de comportamentos complexos é muito mais longo e difícil.
A interpretabilidade apoia vários objetivos fundamentais como, por exemplo, permitir uma melhor supervisão e fornecer sinais de alerta precoces sobre comportamentos inseguros ou estrategicamente desalinhados. Complementa também os nossos outros esforços de segurança, como a supervisão escalável, o treino em situações adversas e as simulações de ataque (red-teaming).
Neste trabalho, demonstramos que podemos frequentemente treinar modelos de forma a torná-los mais fáceis de interpretar. Consideramos o nosso trabalho como um complemento promissor à análise pós-hoc de redes densas.
Esta é uma aposta muito ambiciosa; há um longo caminho a percorrer entre o nosso trabalho e a compreensão completa dos comportamentos complexos dos nossos modelos mais poderosos. Ainda assim, para comportamentos simples, descobrimos que os modelos esparsos treinados com o nosso método contêm pequenos circuitos independentes que são compreensíveis e suficientes para executar o comportamento. Isto sugere que pode existir um caminho viável para treinar sistemas maiores cujos mecanismos possamos compreender.
Trabalhos anteriores sobre interpretabilidade mecanicista começaram com redes densas e entrelaçadas, tentando desenredá-las. Nestas redes, cada neurónio individual está ligado a milhares de outros neurónios. A maioria dos neurónios parece desempenhar várias funções distintas, tornando a compreensão aparentemente impossível.
Mas e se treinássemos redes neuronais não emaranhadas, com muito mais neurónios, mas em que cada neurónio tivesse apenas algumas dezenas de ligações? Então, talvez a rede resultante seja mais simples e mais fácil de compreender. Esta é a principal hipótese de investigação do nosso trabalho.
Com este princípio em mente, treinámos modelos de linguagem com uma arquitetura muito semelhante aos modelos existentes, como o GPT‑2, com uma pequena modificação: obrigamos a que a grande maioria dos pesos do modelo seja zero. Isto restringiu o modelo a usar apenas um número muito reduzido das ligações possíveis entre os seus neurónios. Esta é uma alteração simples que, defendemos, clarifica substancialmente os cálculos internos do modelo.
Nas redes neuronais densas normais, cada neurónio está ligado a todos os neurónios da camada seguinte. Nos nossos modelos esparsos, cada neurónio liga-se apenas a alguns neurónios da camada seguinte. Esperamos que isto torne os neurónios, e a rede como um todo, mais fáceis de compreender.
Queremos medir até que ponto os cálculos dos nossos modelos esparsos estão desagregados. Considerámos vários comportamentos simples do modelo e verificámos se conseguíamos isolar as partes do modelo responsáveis por cada comportamento, que designamos por circuitos.
Selecionámos cuidadosamente um conjunto de tarefas algorítmicas simples. Para cada uma, reduzimos o modelo ao circuito mais pequeno que ainda consegue executar a tarefa e examinámos quão simples é esse circuito. (Para mais informações, veja o nosso artigo(abre numa nova janela).) Concluímos que, ao treinar modelos maiores e mais esparsos, conseguimos criar modelos cada vez mais capazes com circuitos cada vez mais simples.
Representamos graficamente a interpretabilidade em função da capacidade nos vários modelos (quanto mais perto do canto inferior esquerdo, melhor). Para um tamanho fixo de modelo esparso, aumentar a esparsidade — definindo mais pesos como zero — reduz a capacidade, mas aumenta a interpretabilidade. Aumentar o tamanho do modelo desloca esta fronteira para fora, sugerindo que podemos construir modelos maiores que sejam simultaneamente capazes e interpretáveis.
Para tornar isso concreto, considere-se uma tarefa em que um modelo treinado em código Python precisa de completar uma string com o tipo de aspas correto. Em Python, 'hello' deve terminar com uma aspa simples e "hello" deve terminar com uma aspa dupla. O modelo consegue resolver isto memorizando qual o tipo de aspa que abriu a string e reproduzindo-a no final.
Os nossos modelos mais interpretáveis parecem conter circuitos desagregados que implementam exatamente esse algoritmo.

Exemplo de circuito num transformador esparso que prevê se uma string deve terminar com aspas simples ou duplas. Este circuito utiliza apenas cinco canais residuais (linhas cinzentas verticais), dois neurónios MLP na camada 0, e um canal de atenção para a chave de consulta e um canal de valor na camada 10. O modelo (1) codifica as aspas simples num canal residual e as aspas duplas noutro; (2) utiliza uma camada MLP para converter isto num canal que deteta qualquer aspa e outro que classifica entre aspas simples e duplas; (3) utiliza uma operação de atenção para ignorar os tokens intermédios, encontrar a aspa anterior e copiar o seu tipo para o token final; e (4) prevê a aspa de fecho correspondente.
Na nossa definição, as ligações exatas mostradas acima são suficientes para realizar a tarefa — se removermos o resto do modelo, este pequeno circuito ainda funcionará. Elas são também necessárias — eliminar estas poucas ligações faz com que o modelo falhe.
Também analisámos alguns comportamentos mais complicados. Os nossos circuitos para estes comportamentos (por exemplo, a ligação variável mostrada abaixo) são mais difíceis de explicar completamente. Mesmo assim, podemos ainda obter explicações parciais relativamente simples que são preditivas do comportamento do modelo.
Outro exemplo de circuito, com menos detalhes. Para determinar o tipo de uma variável chamada current, uma operação de atenção copia o nome da variável para o token set() quando esta é definida, e outra operação posterior copia o tipo do token set() para uma utilização subsequente da variável, permitindo ao modelo inferir o próximo token correto.
Este trabalho é um passo inicial rumo a um objetivo maior: tornar os cálculos dos modelos mais fáceis de compreender. Porém, ainda há um longo caminho a percorrer. Os nossos modelos esparsos são muito menores do que os modelos de ponta e grande parte dos seus cálculos continua por interpretar.
De seguida, esperamos escalar as nossas técnicas para modelos maiores e explicar melhor o comportamento desses modelos. Ao enumerar os padrões dos circuitos subjacentes a raciocínios mais complexos em modelos esparsos capazes, poderíamos desenvolver uma compreensão que nos ajude a direcionar melhor as investigações sobre modelos de ponta.
Para superar a ineficiência do treino de modelos esparsos, vemos dois caminhos a seguir. Uma opção é extrair circuitos esparsos de modelos densos existentes, em vez de treinar modelos esparsos a partir do zero. Os modelos densos são fundamentalmente mais eficientes de implementar do que os modelos esparsos. O outro caminho é desenvolver técnicas mais eficientes para treinar modelos com interpretabilidade, que possam ser mais fáceis de implementar em produção.
Note-se que as nossas conclusões aqui apresentadas não garantem que esta abordagem se estenda a sistemas mais capazes, mas estes resultados iniciais são promissores. O nosso objetivo é expandir gradualmente a parte do modelo que conseguimos interpretar de forma fiável e criar ferramentas que facilitem a análise, a depuração e a avaliação de sistemas futuros.
Autores
Leo Gao, Achyuta Rajaram, Jacob Coxon, Soham V. Govande, Bowen Baker, Dan Mossing


