Resumo Técnico: Modelos de Linguagem de Regressão para Código

Declaração do Problema

O artigo aborda o desafio da regressão código-métrica: prever resultados numéricos de execuções de código (como footprint de memória, latência ou precisão) diretamente a partir de código-fonte ou representações de grafos computacionais. Historicamente, essa tarefa tem sido dificultada pela natureza aberta das linguagens de programação e pela dependência de engenharia de características pesada e específica de domínio. Métodos anteriores frequentemente exigiam a conversão de código em representações de grafos rígidas (por exemplo, árvores de sintaxe, matrizes de adjacência) ou vetores tabulares por meio de contagem manual de comandos, o que limitava a transferibilidade para novas tarefas, linguagens ou tipos de kernel. A dificuldade central reside em aprender modelos de regressão precisos sobre dados altamente não tabulares e baseados em grafos, sem perder a flexibilidade de generalizar através de domínios diversos.

Metodologia

Os autores propõem um Modelo de Linguagem de Regressão (RLM), um framework unificado que trata a regressão código-métrica como um problema padrão de previsão do próximo token, alinhando-se ao paradigma moderno de Modelos de Linguagem de Grande Escala (LLM).

Arquitetura e Treinamento

• Codificador-Decodificador Unificado: O RLM utiliza um codificador congelado (especificamente T5Gemma-S) para processar o texto de entrada e um decodificador treinável para gerar saídas numéricas. O codificador é pré-treinado em dados de linguagem padrão, enquanto apenas o decodificador sofre retropropagação, reduzindo significativamente os custos de treinamento.
• Regressão Texto-a-Texto: Diferentemente das abordagens tradicionais que mapeiam entradas para vetores de comprimento fixo, o RLM aceita representações de texto bruto de código (por exemplo, Python, C++, grafos ONNX) e gera valores numéricos como sequências de tokens.
• Tokenização Numérica Especializada: Para lidar com a ampla faixa dinâmica de métricas (de $10^{-2}$ a $10^6$) sem normalização, o modelo usa uma tokenização explícita dígito a dígito. Números são representados usando tokens especiais para sinal, expoente e mantissa (por exemplo, <+><-><1><7><2><5> para $72.5$). Isso evita instabilidade numérica e a necessidade de escalonamento específico do conjunto de dados.
• Modelagem Multi-Objetivo: Aproveitando a natureza autoregressiva do decodificador, o modelo pode prever condicionalmente múltiplas métricas sequencialmente (por exemplo, $p(y_{latência} | y_{precisão}, x)$), capturando correlações inerentes entre objetivos como precisão e latência de hardware.

Fontes de Dados

O modelo é treinado em um conjunto de dados diversificado e unificado, compreendendo:

1. Programação de Alto Nível: Código de APPS (problemas do LeetCode), CodeNet (14M de amostras em 37 linguagens) e KernelBook (kernels GPU Triton).
2. Busca de Arquitetura Neural (NAS): 520K arquiteturas únicas de benchmarks como NASBench-101/201, FBNet e Once-for-All, representadas como representações intermediárias ONNX.

Contribuições Principais

1. Modelo Unificado de Propósito Geral: O artigo demonstra que um único RLM pode simultaneamente prever o uso de memória em múltiplas linguagens (Python, C++, etc.), latência de kernel GPU e precisão/latência de redes neurais, eliminando a necessidade de engenharia de características específica de tarefa.
2. Superioridade sobre Métodos Baseados em Grafos: Apesar de ler representações de texto ONNX densas, o RLM supera as linhas de base de Redes Neurais em Grafos (GNN) mais avançadas em benchmarks de NAS (por exemplo, alcançando um Kendall-Tau de 0,46 em espaços de design onde as GNNs anteriormente dominavam).
3. Eficiência e Escalabilidade: O uso de um codificador congelado permite um treinamento economicamente viável. Os autores mostram que um modelo relativamente pequeno de 300M de parâmetros alcança alto desempenho, e a escalabilidade para 600M de parâmetros gera ganhos adicionais.
4. Previsão sem Normalização: A estratégia de tokenização dígito a dígito permite que o modelo preveja valores através de ordens de grandeza sem normalização manual, uma vantagem significativa sobre cabeças de regressão baseadas em MSE.

Resultados Experimentais

• Desempenho de Código: No conjunto de dados APPS, o RLM alcança uma correlação de posto de Spearman ($\rho$) de 0,93 para previsão de memória. No CodeNet, alcança um $\rho$ médio > 0,5 em 17 linguagens, com C++ atingindo 0,748.
• Desempenho em NAS: O RLM alcança um Kendall-Tau médio de 0,46 em cinco espaços de design clássicos de NAS, superando métodos anteriores SOTA como FLANT e codificadores baseados em GNN. Ele prevê com sucesso tanto a precisão quanto latências específicas de hardware (por exemplo, em Pixel3, Eyeriss, CPU Intel, GPU Nvidia) simultaneamente.
• Estudos de Ablação:
  • Codificadores Congelados: Codificadores pré-treinados melhoram significativamente a convergência e o desempenho final em comparação com inicialização aleatória.
  • Decodificador vs. Cabeças de Regressão: A saída baseada em decodificador (0,800 Spearman $\rho$) supera substancialmente cabeças de regressão explícitas (0,478 sem normalização, 0,717 com normalização).
  • Tokenização: Tokenizadores personalizados e compactos para grafos ONNX melhoram o desempenho (0,723 vs. 0,533 com tokenizadores padrão).
  • Aprendizado em Contexto (ICL): Embora o ICL com modelos grandes (por exemplo, GPT-5) melhore com mais exemplos, ele permanece consistentemente abaixo do RLM treinado em pesos e incorre em custos de inferência significativamente maiores devido às limitações da janela de contexto.

Significado e Alegações

O artigo afirma abrir caminho para simplificar a regressão de grafos computacionais em um problema genérico de previsão do próximo token. Ao demonstrar que modelos de linguagem podem efetivamente realizar regressão sobre representações de texto puro de código e grafos, o trabalho sugere que o campo pode se afastar da engenharia de características rígida e manual em direção a uma abordagem mais flexível e unificada.

Os autores posicionam isso como uma ferramenta valiosa para:

• Busca de Programas: Acelerar a seleção de soluções de código eficientes.
• Co-design Hardware-Software: Prever latência em diversas plataformas de hardware.
• Otimização de Compiladores: Estimar métricas de desempenho para passes de compilador.

O artigo permanece modesto quanto a implicações mais amplas, observando que, embora o método mostre promessa para prever resultados numéricos a partir de código bruto, estender isso para prever resultados de experimentos científicos completos permanece uma questão em aberto para trabalhos futuros. A contribuição principal é estabelecer a capacidade geral dos modelos de linguagem de servirem como regressores universais para código e grafos computacionais sem pós-processamento ou extração de características especializada.