Weights to Code: Extracting Interpretable Algorithms from the Discrete Transformer

1. O Problema

A extração de algoritmos visa sintetizar programas executáveis diretamente a partir de modelos treinados em tarefas algorítmicas, permitindo a descoberta de algoritmos de novo sem depender de código escrito por humanos. Embora métodos baseados em Redes Neurais Recorrentes (RNNs), como o MIPS (Mechanistic-Interpretability-based Program Synthesis), tenham tido sucesso, a aplicação dessa paradigmática à arquitetura Transformer enfrenta barreiras significativas.

O principal obstáculo é o emaranhamento de representações (conhecido como superposition ou superposição). Nos Transformers padrão, as características são codificadas em direções sobrepostas e não ortogonais, tornando as representações internas altamente entrelaçadas e polissêmicas. Isso impede a recuperação direta de expressões simbólicas discretas e legíveis por humanos, que exigem mapeamentos entrada-saída desentrelaçados para a regressão simbólica.

2. Metodologia: O Discrete Transformer

Os autores propõem o Discrete Transformer, uma arquitetura especializada projetada para preencher a lacuna entre representações contínuas e lógica simbólica discreta.

Arquitetura

O modelo é estruturado como um grafo computacional com especialização funcional rigorosa, alinhado ao modelo computacional RASP (Restricted Access Sequence Processing):

• Stream Residual Numérico: Em vez de vetores latentes densos, o modelo utiliza um conjunto de variáveis escalares explícitas, preservando o histórico computacional como coordenadas desentrelaçadas.
• Leitura Discretizada: Um mecanismo de projeção aprendida seleciona variáveis de entrada do stream residual usando amostragem controlada por temperatura.
• Numerical Attention (Roteador): Atua estritamente como um operador de roteamento de informações. Utiliza Piecewise Linear Encoding (PLE) para projetar escalares e uma atenção "hard" (determinística) para identificar padrões de roteamento (ex: "copiar o valor do token anterior").
• Numerical MLP (Operador): Responsável exclusivamente por transformações aritméticas e lógicas elementares. É decomposto em sub-módulos independentes e limitados em complexidade para facilitar a extração de expressões fechadas.
• Cabeça de Saída Linear: Agrega os passos computacionais discretos para gerar a previsão final.

Processo de Treinamento e Extração

1. Annealing de Temperatura: Durante o treinamento, uma temperatura ($\tau$) controla a função de amostragem. Começa alta (espaço contínuo para exploração) e é reduzida gradualmente para zero (espaço discreto para exploração), forçando o modelo a convergir para um grafo computacional totalmente discreto.
2. Pipeline de Extração Modular: Após a convergência, o modelo é tratado como um grafo onde:
   • Atenção: Analisada via Teste de Hipóteses para identificar padrões de roteamento (ex: "Offset Fixo" ou "Extremos em Janela").
   • MLP: Analisado via Regressão Simbólica (usando ferramentas como PySR) para inferir expressões aritméticas exatas a partir das entradas e saídas dos sub-módulos.
   • Montagem Global: Os primitivos extraídos são integrados através da cabeça de saída linear, gerando código Python executável e legível.

3. Contribuições Principais

• Arquitetura Discrete Transformer: Uma nova arquitetura alinhada ao RASP que impõe desentrelaçamento funcional para conectar otimização contínua e lógica discreta.
• Pipeline de Extração Híbrido: Combinação de teste de hipóteses (para roteamento) e regressão simbólica (para aritmética), estendendo a síntese baseada em interpretabilidade mecânica para variáveis contínuas.
• Controle via Vieses Indutivos: Demonstração de que restrições arquitetônicas podem ser usadas para guiar a descoberta de algoritmos específicos, permitindo intervenções no espaço de soluções.

4. Resultados

Os autores avaliaram o modelo em um conjunto diversificado de tarefas algorítmicas, incluindo aritmética linear, composição não-linear e dinâmicas contínuas.

• Desempenho Preditivo: O Discrete Transformer alcançou desempenho comparável ao método MIPS baseado em RNNs, com erro quadrático médio (MSE) próximo de zero na maioria das tarefas.
• Fidelidade da Extração: Os programas extraídos (código Python) resolveram as tarefas com precisão de 1.00 (100%) na maioria dos casos, demonstrando alta fidelidade na recuperação da lógica subjacente.
• Recuperação de Identidades Não-Lineares: O modelo conseguiu recuperar expressões algébricas exatas para operações complexas. Por exemplo, para a tarefa de paridade ($x_t \oplus x_{t-1}$), a regressão simbólica derivou a fórmula exata $y_t = x_t + x_{t-1} - 2x_t x_{t-1}$.
• Dinâmicas Contínuas: Diferente do MIPS (limitado a variáveis discretas/latentes), o Discrete Transformer lidou nativamente com variáveis contínuas, permitindo a extração de algoritmos para tarefas como "média móvel exponencial" e "altura de queda livre", que envolvem números de ponto flutuante.
• Análise de Dinâmica de Treinamento: Observou-se uma transição de fase clara: o modelo primeiro atinge a convergência funcional (aprendendo o mapeamento) e, posteriormente, ocorre a cristalização estrutural (discretização completa), garantindo uma transição robusta de exploração para exploração.

5. Significado e Impacto

Este trabalho estabelece o Discrete Transformer como um framework robusto para a descoberta de algoritmos sem demonstração e para a interpretabilidade de Transformers.

• Interpretabilidade Mecânica: Oferece uma via para entender "o que" os Transformers estão aprendendo, transformando pesos opacos em código executável e verificável.
• Descoberta Científica: A capacidade de extrair equações fechadas e lógica simbólica a partir de dados contínuos abre novas possibilidades para a descoberta científica automatizada, onde o modelo pode propor novas leis físicas ou lógicas independentes de dados humanos prévios.
• Controllabilidade: A arquitetura permite que pesquisadores intervenham no processo de descoberta (ex: forçando o uso de atenção em vez de MLP), oferecendo um ambiente controlável para testar hipóteses sobre como algoritmos podem ser aprendidos e representados.

Em resumo, o artigo supera a barreira da superposição nos Transformers, permitindo que modelos modernos não apenas prevejam resultados, mas revelem explicitamente a lógica algorítmica e as equações matemáticas que governam seus comportamentos.