Resumo Técnico: Rumo a Embeddings de LM Orientados a Código para Busca de Arquitetura Neural Assistida por Surrogato

Declaração do Problema

A Busca de Arquitetura Neural (NAS) automatiza a descoberta de arquiteturas neurais de alto desempenho, mas enfrenta um gargalo computacional significativo: avaliar a aptidão de arquiteturas candidatas é dispendioso. Para mitigar isso, pesquisadores empregam modelos surrogatos (previsores de desempenho) para aproximar a aptidão sem treinamento completo. No entanto, as abordagens de surrogato existentes sofrem de duas limitações principais:

1. Rigidez de Representação: Codificações estruturais (por exemplo, matrizes de adjacência, codificações de caminho) são frequentemente adaptadas a espaços de busca específicos (como DAGs baseados em células) e lutam para generalizar para arquiteturas novas ou capturar hiperparâmetros complexos (por exemplo, tamanhos de kernel) sem engenharia personalizada.
2. Sobrecarga de Ajuste Fino: Tentativas recentes de usar Modelos de Linguagem (LMs) como regressores universais exigem converter arquiteturas em formatos de texto não padrão (por exemplo, strings de árvores de derivação ou ONNX-para-texto) e subsequentemente ajustar finamente o LM nessas representações estrangeiras. Esta fase de "aquecimento" nega os ganhos de eficiência de usar um surrogato, pois requer pré-treinamento dispendioso antes que a busca possa sequer começar.

Os autores postulam que essas ineficiências surgem porque os métodos atuais ignoram o fato de que as arquiteturas neurais existem naturalmente como código executável. Como os LMs são pré-treinados em corpora massivos de código, eles possuem um "viés indutivo" inerente para entender a lógica arquitetural que fica desalinhado quando alimentados com codificações de texto artificiais.

Metodologia: COLE

Os autores propõem Embeddings de LM Orientados a Código (COLE), uma estratégia que aproveita LMs pré-treinados congelados como extratores de recursos zero-shot para NAS.

1. Stringificação Nativa de Código

Em vez de converter arquiteturas em strings abstratas, o COLE mapeia o espaço de busca diretamente para definições de classes PyTorch (torch.nn.Module).

• NAS-Bench-201: Os autores analisam strings de topologia de grafo padrão e as convertem em classes Python válidas. Eles investigaram três níveis de verbosidade:
  • Método Auxiliar: Definindo operações complexas como classes separadas.
  • Em Linha: Definindo lógica dentro da classe usando nn.Sequential.
  • Auxiliar Excluído: Referenciando classes auxiliares sem definir sua implementação (mais eficiente em tokens).
• einspace: Um transpilador determinístico converte árvores de derivação de gramática livre de contexto diretamente em código PyTorch executável, mapeando nós não terminais (por exemplo, Branching, Routing) para classes de contêiner personalizadas e nós terminais para primitivas atômicas.

2. Extração de Embeddings Congelados

O sistema utiliza LMs pré-treinados congelados (especificamente CodeLlama Python 7B) para processar o texto bruto do código.

• Sem Ajuste Fino: Os parâmetros do LM são totalmente congelados, eliminando o custo computacional do pré-treinamento específico de domínio.
• Agrupamento (Pooling): As ativações do estado oculto da última camada são agregadas usando Mean Pooling através da sequência de tokens para gerar um vetor de embedding de dimensão fixa.

3. Cabeça de Surrogato

Os embeddings de alta dimensão (por exemplo, 4096 dimensões) são processados por um pipeline de regressão leve:

• Redução de Dimensionalidade: A Análise de Componentes Principais (PCA) reduz o embedding para 128 dimensões para evitar overfitting em regimes de poucos dados.
• Regressão: Um Perceptron Multicamadas (MLP) de 3 camadas prevê o desempenho da arquitetura.
• Função de Perda: O modelo é otimizado para a correlação de postos Kendall's Tau usando uma Perda de Hinge Pares. Isso prioriza a correta ordenação dos candidatos (crucial para algoritmos de busca evolutiva) sobre prever valores exatos de precisão absoluta.

Contribuições Principais

1. Superioridade do Código Nativo: Os autores demonstram que representar arquiteturas como código PyTorch bruto produz desempenho preditivo superior (Kendall's Tau) do que outras codificações baseadas em texto (ONNX-para-texto, árvores de derivação) ao usar LMs congelados.
2. Eliminação do Ajuste Fino: O COLE prova que o ajuste fino computacionalmente dispendioso não é estritamente necessário para embeddings eficazes de NAS; LMs prontos para uso atuam como extratores de recursos competitivos "prontos para uso".
3. Integração com SA-NAS: Integrar o COLE no algoritmo de busca assistida por surrogato BANANAS melhora significativamente a eficiência da busca em comparação com o uso de codificações de caminho estruturais padrão.
4. Diretrizes Práticas de Ajuste: O estudo documenta configurações eficazes, como o uso do modo "Auxiliar Excluído" para eficiência de tokens e PCA para redução de dimensionalidade, mostrando que o LM pode inferir semântica de operação mesmo sem detalhes completos de implementação na string de entrada.

Resultados Experimentais

O framework foi avaliado no NAS-Bench-201 e no espaço de busca einspace.

• Qualidade do Embedding:

  • No NAS-Bench-201, o COLE (usando CodeLlama 7B) alcançou um Kendall's Tau de 0.703 com 879 amostras de treinamento, superando codificações ONNX-para-texto (0.678) e strings de árvores de derivação (0.763 no einspace vs. 0.778 para código PyTorch).
  • Seleção de Modelo: Modelos centrados em código (CodeLlama, Yi Coder) consistentemente superaram modelos gerais apenas codificadores (ModernBERT, CodeBERT). Curiosamente, modelos menores (1.5B) performaram comparavelmente a modelos maiores de 7B/34B, sugerindo que escala massiva não é estritamente necessária.
  • Função de Perda: A Perda de Hinge Pares produziu melhor correlação de postos do que o Erro Quadrático Médio (MSE) padrão, apesar do MSE ter erro absoluto menor.

• Eficiência de Busca (BANANAS):

  • Substituir codificações de caminho por COLE no algoritmo BANANAS reduziu o orçamento de avaliação necessário para atingir desempenho quase ótimo.
  • CIFAR-100: O COLE atingiu dentro de 1% da arquitetura mais apta em 132 avaliações, uma redução de 34% em comparação com as 200 avaliações exigidas por codificações de caminho.
  • CIFAR-10: Alcançou uma redução de 13% (52 vs. 60 avaliações).
  • ImageNet16-120: Alcançou uma redução de 18% (164 vs. 200 avaliações).
  • Essas melhorias foram estatisticamente significativas ($p < 0.05$) na maior parte da trajetória de busca.

Significado e Alegações

O artigo afirma que o COLE estabelece uma base versátil e eficiente para avançar a NAS ao aproveitar o conhecimento existente de código em LMs pré-treinados.

• Agnóstico à Busca: Como qualquer arquitetura neural pode ser representada como código, o COLE não exige reengenharia de esquemas de codificação para novos espaços de busca, ao contrário de codificações estruturais.
• Capacidade de Início Frio: Ao remover a necessidade de ajuste fino, o COLE permite buscas de "início frio" onde o surrogato é eficaz imediatamente, abordando o paradoxo de eficiência de métodos anteriores de NAS baseados em LM.
• Generalizabilidade: A abordagem unifica a representação de topologia e hiperparâmetros em um único formato (código), permitindo que o LM interprete tanto o fluxo estrutural quanto os valores de parâmetros (por exemplo, kernel_size=3) naturalmente.

Os autores concluem que, embora sua avaliação atual seja limitada a benchmarks específicos (NAS-Bench-201, einspace), a metodologia oferece uma alternativa prática e de baixo custo a GNNs complexos ou transformers ajustados finamente, particularmente para cenários de NAS com recursos restritos.