Resumo Técnico: TraceNAS

Definição do Problema

A implantação de Grandes Modelos de Linguagem (LLMs) é limitada por contagens massivas de parâmetros e requisitos computacionais. Embora a poda estruturada ofereça um caminho para a implantação eficiente através da remoção de unidades arquitetônicas (ex: cabeças de atenção, blocos transformer), os métodos existentes enfrentam limitações significativas:

1. Dependências Locais vs. Globais: A maioria dos métodos de poda avalia a importância de camadas, cabeças ou canais isoladamente usando heurísticas locais (ex: magnitudes ponderadas por ativação ou reconstrução baseada em Hessian). Essas abordagens ignoram dependências estruturais globais complexas e a propagação de erros induzidos pela poda através do modelo.
2. Custo Computacional de Busca Consciente de Treinamento: Métodos que tentam capturar dependências globais geralmente dependem de Busca de Arquitetura de Redes Neurais (NAS) consciente de treinamento. Essas abordagens exigem atualizações de gradiente iterativas e fluxos contínuos de tokens não vistos para evitar o overfitting, tornando o processo de busca tão custoso computacionalmente quanto o próprio treinamento de recuperação pós-poda.
3. Inadequação dos Proxies Zero-Shot Existentes: Os atuais proxies livres de treinamento são projetados para identificar arquiteturas treináveis e inicializadas aleatoriamente, focando na treinabilidade (ex: estabilidade do gradiente, suavidade do landscape de perda). Eles falham em avaliar a "herança funcional" — a capacidade de um modelo podado de recuperar o desempenho ao preservar o landscape de perda pré-treinado.

Metodologia

O TraceNAS é proposto como um framework de NAS livre de treinamento projetado para otimizar conjuntamente a profundidade e a largura de LLMs sem treinamento durante a busca.

1. Espaço de Busca e Codificação

O framework trata o LLM pré-treinado como uma super-rede. O espaço de busca $G$ é definido como uma distribuição conjunta sobre:

• Profundidade ($d$): Uma máscara discreta sobre $L$ blocos transformer, permitindo redução de profundidade não uniforme enquanto mantém a continuidade do fluxo residual.
• Largura ($\kappa$): Razões de retenção de parâmetros para sub-blocos de atenção e MLP, permitindo esparsidade heterogênea através do modelo.

2. Realização Arquitetônica In-Place

Para evitar o overhead de memória de instanciar milhares de candidatos, o TraceNAS emprega uma estratégia de mascaramento in-place:

• Mascaramento de Largura: Utiliza um produto ponderado por ativação (semelhante ao Wanda) para gerar máscaras binárias para canais de peso. Esta heurística $O(d^2)$ identifica canais funcionalmente críticos enquanto preserva características outliers.
• Mascaramento de Profundidade: Ignora blocos inteiros se desativados, roteando as ativações através de conexões residuais.
• Execução: O framework modifica temporariamente as matrizes de peso do modelo base para computar gradientes em um conjunto de calibração mínimo, e então restaura os pesos originais, preservando o modelo base para iterações subsequentes.

3. O Proxy Zero-Shot: Correlação de Traço de Gradiente

A inovação central é um proxy invariante de escala $\Phi$ que mede o alinhamento funcional entre um candidato podado ($M_{sub}$) e a base pré-treinada ($M_{base}$).

• Traços de Gradiente: Gradientes são computados em um conjunto de calibração. Para garantir escalabilidade, estes são projetados em um subespaço de Baixo Rank (LoRA), aproveitando a baixa dimensionalidade intrínseca dos LLMs para reduzir o overhead de memória em ordens de magnitude.
• Métrica de Alinhamento: O proxy calcula o Coeficiente de Correlação de Pearson entre os traços de gradiente padronizados dos modelos podado e base para cada camada. Isso garante invariância de escala, focando no alinhamento direcional em vez de mudanças de magnitude.
• Agregação Ponderada por Esparsidade: A pontuação final $\Phi$ é uma soma ponderada de correlações por camada, onde os pesos correspondem às razões de retenção de parâmetros. Isso ancora a pontuação global em regiões de alta capacidade, evitando que o ruído de sub-blocos agressivamente podados distorça o ranking.

4. Busca Evolutiva

O TraceNAS emprega uma estratégia evolutiva para navegar no espaço de busca discreto:

• Inicialização: Uma população é gerada com base em priors de importância global.
• Evolução: A população passa por crossover e mutação. A evolução de profundidade usa mutações de bit-flip e crossover de intervalo contíguo, enquanto a evolução de largura usa interpolação e jitter Gaussiano.
• Seleção: Os candidatos são classificados unicamente pelo proxy $\Phi$, selecionando modelos com maior potencial de recuperação de desempenho durante o pré-treinamento contínuo (CPT) pós-poda.

Principais Contribuições

1. NAS Escalável Livre de Treinamento: O TraceNAS fornece um framework unificado para poda conjunta de profundidade e largura, permitindo a descoberta de arquiteturas não uniformes sem treinamento durante a busca.
2. Proxy de Herança Funcional: Introduz um novo proxy baseado em gradiente e invariante de escala que avalia o alinhamento dos traços de gradiente entre os modelos podado e pré-treinado. Ele alcança alta correlação com a acurácia downstream (Spearman $\rho = 0.94$, Kendall $\tau = 0.82$).
3. Eficiência: O framework reduz o overhead de busca utilizando subespaços de gradiente de baixo rank e mascaramento in-place, permitindo a descoberta de modelos de alta fidelidade em uma única GPU.

Resultados Experimentais

O TraceNAS foi avaliado no Llama-2-7B, Llama-3.1-8B e Qwen-2.5-14B, podando-os para 2.7B, 4.6B e 8.4B parâmetros, respectivamente.

• Eficiência de Busca: O TraceNAS identificou arquiteturas ótimas em 8.5 horas de GPU (em uma A100), representando uma redução de 10× no tempo de busca em comparação com baselines conscientes de treinamento como o DarwinLM. Requereu apenas 98M de tokens de busca, uma redução de 10× em relação ao DarwinLM.
• Desempenho:
  • Llama-2-7B: O TraceNAS alcançou a maior acurácia média (62.81%) em oito benchmarks de raciocínio, superando o DarwinLM (62.57%) e o ShearedLLaMA (62.63%). Notavelmente, o desempenho do ShearedLLaMA caiu significativamente (para 60.86%) quando restrito ao mesmo orçamento de tokens usado pelo TraceNAS.
  • Llama-3.1-8B & Qwen-2.5-14B: O TraceNAS superou baselines de poda uniforme e métodos conscientes de treinamento (DarwinLM, E3-Pruner) em tarefas de raciocínio complexas (ex: ARC-Challenge, BoolQ) enquanto requereu significativamente menos tokens de recuperação (20B vs. 94B para o Minitron).
• Validação do Proxy: O proxy demonstrou alta estabilidade através de vários hiperparâmetros (tamanho de amostra, comprimento de contexto e rank LoRA), confirmando que o alinhamento do traço de gradiente captura efetivamente a herança do modelo mesmo sob configurações restritas.

Significância e Alegações

O artigo afirma que o TraceNAS aborda o gargalo crítico do custo computacional na poda estruturada. Ao mudar o foco de "treinabilidade" (aprender do zero) para "herança" (preservar o landscape de perda pré-treinado), o TraceNAS permite a descoberta de modelos podados não uniformes de alto desempenho com gasto mínimo de recursos.

Os autores enfatizam que seu método:

• Elimina o custo proibitivo do treinamento durante a busca, tornando a compressão de modelos em larga escala acessível.
• Captura dependências estruturais globais que heurísticas locais perdem, garantindo a preservação de capacidades complexas de raciocínio.
• Reduz a pegada de carbono da compressão de modelos ao reduzir drasticamente as horas-GPU necessárias para a busca de arquitetura.

O trabalho posiciona o TraceNAS como uma fundação escalável para implantar LLMs eficientes em ambientes com recursos limitados, mantendo a integridade funcional dos modelos densos pré-treinados.