Target-Oriented Pretraining Data Selection via Neuron-Activated Graph

1. Problema e Motivação

O pré-treinamento de Grandes Modelos de Linguagem (LLMs) é intensivo em recursos computacionais. Uma estratégia eficaz para melhorar o desempenho em tarefas específicas é a seleção de dados de pré-treinamento orientada a objetivos (target-oriented). No entanto, existem desafios significativos nas abordagens atuais:

• Definição Ambígua de "Qualidade": Métodos tradicionais focam em métricas gerais de qualidade (como perplexidade ou classificadores de "educação"), que muitas vezes não se alinham com capacidades específicas necessárias para tarefas downstream (ex: raciocínio matemático vs. compreensão narrativa).
• Abordagens "Caixa-Preta": Métodos anteriores orientados a objetivos (como BETR) frequentemente utilizam representações de embeddings ou perdas correlacionadas com desempenho, que são difíceis de interpretar e diagnosticar. Eles dependem de sinais internos opacos do modelo.
• Desalinhamento: Há uma lacuna entre os dados selecionados e as capacidades reais que o modelo precisa adquirir, levando a ganhos subótimos ou ineficiências computacionais.

O objetivo deste trabalho é criar um framework sem treinamento adicional (training-free), interpretável e eficiente para selecionar dados que ativem especificamente as capacidades desejadas em um LLM.

2. Metodologia: NAG-based Ranking

Os autores propõem o NAG-based Ranking (Classificação baseada em Grafo Ativado por Neurônios), que seleciona dados com base na similaridade de padrões de ativação neuronal em relação a exemplos-alvo.

A. Impacto do Neurônio (Neuron Impact)

Em vez de usar embeddings finais, o método quantifica a contribuição de neurônios individuais durante a inferência.

• Para uma projeção $W$ em uma camada, cada coluna é tratada como um neurônio.
• O impacto de um neurônio $N_k$ é definido pela magnitude da mudança na saída da camada se esse neurônio for desativado (zerado).
• Matematicamente: $Imp(N_k | h_{in}) = ||h_{in}^T W_{:,k}||_2$. Isso identifica quais neurônios são mais críticos para processar uma entrada específica.

B. Grafo Ativado por Neurônios (NAG)

Para cada entrada de texto, constrói-se um NAG, que é uma representação compacta e estruturada:

1. Identifica-se os neurônios de alto impacto em cada camada do modelo.
2. Seleciona-se os top-$K$ neurônios por camada.
3. O NAG é definido como a coleção desses índices de neurônios através de todas as camadas: $NAG(c) = \{(l, k) | l \in \text{camadas}, k \in \text{top-K}\}$.

C. Seleção de Dados via Similaridade

O processo de seleção compara candidatos com um conjunto pequeno de exemplos-alvo ($D_{target}$):

1. Perfil Alvo: Agrega-se os NAGs dos exemplos-alvo para criar um perfil de ativação (frequência de ativação de neurônios).
2. Ranking: Para cada candidato no pool de dados, calcula-se a similaridade entre seu NAG e o perfil alvo. A similaridade é baseada na interseção de neurônios ativados (semelhante ao coeficiente de Jaccard).
3. Seleção: Os dados são classificados pela similaridade e os top-$r_f$ (ex: 20%) são selecionados para o pré-treinamento.

Vantagem Chave: O método não requer treinamento de classificadores auxiliares e utiliza sinais interpretáveis diretamente de um LLM "off-the-shelf" (pronto para uso).

3. Contribuições Principais

1. Framework Interpretável: Diferente de métodos baseados em embeddings finais, o NAG mapeia a "trajetória computacional" do modelo através de todas as camadas, capturando o "esqueleto funcional" necessário para a tarefa.
2. Sem Treinamento Adicional: Elimina a necessidade de treinar classificadores auxiliares ou distilar modelos, reduzindo o custo e a complexidade do pipeline de seleção.
3. Análise Causal: Os autores provam que os neurônios selecionados pelo NAG são críticos. Desativar apenas 0,12% dos neurônios do modelo (aqueles selecionados pelo NAG) causa um colapso de desempenho de 23,5%, enquanto a desativação aleatória tem efeito insignificante.
4. Robustez Multi-Objetivo: O método mantém sua eficácia em cenários com múltiplos objetivos simultâneos, onde outras abordagens (como BETR) tendem a degradar.

4. Resultados Experimentais

Os experimentos foram realizados em 6 benchmarks (ARC-Challenge, HellaSwag, TriviaQA, MMLU, XStoryCloze, XWinograd) com modelos pré-treinados a partir do zero (1.2B parâmetros) usando um pool de 150B tokens.

• Desempenho Geral: O NAG-based Ranking melhorou o pré-treinamento orientado a objetivos em 4,9% em média em comparação com a amostragem aleatória.
• Comparação com SOTA: Superou o método state-of-the-art (BETR) em 5,3% de precisão no benchmark HellaSwag e em média 1,0% em todos os benchmarks.
• Cenário Multi-Objetivo: Em configurações com múltiplos alvos, o NAG superou os baselines em 1,1% e 4,1%, demonstrando melhor generalização.
• Independência do Modelo: O método funcionou consistentemente com diferentes modelos base (Qwen, Llama, SmolLM) para extração de NAG, indicando que é agnóstico à arquitetura.
• Combinação com Qualidade Geral: Ao combinar NAG com classificadores de qualidade geral (FineWeb-Edu), houve ganhos adicionais, sugerindo que o NAG captura sinais complementares.

5. Significado e Conclusão

O trabalho estabelece que a seleção de dados para pré-treinamento pode ser guiada diretamente pela mecânica interna do modelo (neurônios específicos) em vez de representações superficiais.

• Interpretabilidade: O NAG revela que tarefas específicas dependem de um subconjunto esparso e altamente crítico de neurônios distribuídos por todas as camadas do modelo, não apenas na camada final.
• Eficiência Computacional: O método atua como um "multiplicador de computação", permitindo que modelos alcancem o mesmo desempenho com menos dados de treinamento (ganhos de eficiência de 1,27x a 2,42x relatados).
• Futuro: Abre caminho para o uso de sinais de nível de neurônio para entender, direcionar e otimizar as capacidades aprendidas durante o pré-treinamento em larga escala, oferecendo uma alternativa mais transparente e eficiente aos métodos baseados em "caixa-preta".

Em resumo, o NAG-based Ranking oferece uma nova perspectiva para a seleção de dados, alinhando o pré-treinamento com as necessidades específicas da tarefa através da análise direta da arquitetura neuronal do modelo.