Data Mixing Agent: Learning to Re-weight Domains for Continual Pre-training

Título: Data Mixing Agent: Aprendendo a Re-pesquisar Domínios para Pré-treinamento Contínuo

1. O Problema

O Pré-treinamento Contínuo (Continual Pre-training) é um método eficaz para adaptar Grandes Modelos de Linguagem (LLMs) a novos campos específicos (como raciocínio matemático ou geração de código). No entanto, uma adaptação direta apenas aos dados do campo alvo frequentemente leva ao esquecimento catastrófico das capacidades originais do modelo (campo fonte) devido a grandes deslocamentos de distribuição.

A solução comum é criar misturas de dados de fontes (originais) e alvos para equilibrar o desempenho. Contudo, as estratégias existentes para re-pesquisar (re-weighting) esses domínios dependem de:

• Design manual baseado em heurísticas humanas ou intuição.
• Resultados empíricos específicos de tarefas.
• Métodos estáticos ou dinâmicos que não aprendem heurísticas generalizáveis.

Não existe um método automatizado que aprenda e generalize as melhores estratégias de mistura de dados para diferentes cenários sem retreinamento.

2. Metodologia: Data Mixing Agent

Os autores propõem o Data Mixing Agent, o primeiro framework baseado em modelo e end-to-end que aprende a re-pesquisar domínios. O problema é formalizado como um Processo de Decisão de Markov (MDP):

• Espaço de Estados ($S$): Representa a distribuição de dados de passos anteriores de re-pesagem.
• Espaço de Ações ($A$): A distribuição de probabilidade atual sobre os domínios (ex: 52 dimensões baseadas no classificador de domínio da Nvidia).
• Política ($f$): O próprio agente (um modelo leve) que decide a próxima distribuição de dados.
• Recompensa ($r$): Feedback de um ambiente de avaliação leve que mede o desempenho do modelo em campos alvo e fonte.

O Pipeline de Treinamento:

1. Amostragem de Trajetórias: O sistema gera milhares de trajetórias de mistura de dados aleatórias (com diferentes estratégias de re-pesagem).
2. Treinamento de Modelos Proxy: Pequenos modelos proxy são treinados do zero em cada trajetória amostrada.
3. Coleta de Feedback: Em cada passo de re-pesagem, os checkpoints dos modelos são avaliados em um ambiente leve (ex: subconjuntos de MMLU para capacidades gerais e MATH para raciocínio matemático).
4. Otimização com RL (Aprendizado por Reforço):
   • O agente é treinado para aprender as heurísticas contidas nessas trajetórias e feedbacks.
   • Utiliza-se o algoritmo Conservative Q-Learning (CQL). O CQL é escolhido porque é um algoritmo off-policy que evita a superestimação de valores Q para ações fora da distribuição, permitindo aprender tanto de trajetórias "boas" quanto "ruins".
   • O agente é inicialmente aquecido com Supervised Fine-Tuning (SFT) em trajetórias de alta qualidade, mas o CQL é crucial para refinar a política explorando todo o espaço de heurísticas.

Arquitetura do Agente:

• Baseado em Transformers (Decoder), extremamente leve (apenas 2,1 milhões de parâmetros).
• Recebe como entrada a trajetória histórica de distribuições e o feedback do ambiente para prever a distribuição ótima para o próximo passo.

3. Contribuições Principais

1. Primeiro Método Baseado em Agente: Propõem o primeiro método end-to-end e leve para re-pesagem de domínios em pré-treinamento contínuo, substituindo heurísticas manuais por um modelo aprendido.
2. Generalização Sem Retreinamento: O agente aprende heurísticas agnósticas a dados e modelos. Ele foi treinado em um cenário (raciocínio matemático) e aplicado com sucesso em:
   • Diferentes modelos alvo (LLaMA-3B, Pythia).
   • Diferentes espaços de domínio (2D baseado em fonte vs. 52D baseado em tópicos).
   • Novos campos alvo: Funcionou diretamente em Geração de Código sem nenhum ajuste de pesos, demonstrando transferibilidade de heurísticas.
3. Eficiência e Desempenho: O método supera baselines de ponta (como RegMix e DBL) em benchmarks de raciocínio matemático e gerais, alcançando um desempenho equilibrado com menos tokens de dados de origem.

4. Resultados Experimentais

Os experimentos foram realizados em tarefas de Raciocínio Matemático e Geração de Código.

• Desempenho Equilibrado: O DataAgentRL (agente com RL) superou significativamente as baselines.
  • No campo matemático, alcançou uma melhoria média de 3,02% sobre o RegMix e 3,53% sobre o DBL em benchmarks gerais e matemáticos combinados.
  • Reduziu drasticamente a variância de desempenho entre os diferentes benchmarks, indicando estabilidade.
• Eficiência de Dados: O agente aprendeu a usar menos dados do campo fonte para manter as capacidades originais, economizando até 2,14 bilhões de tokens de treinamento em comparação com métodos estáticos, enquanto atingia desempenho superior.
• Transferência de Campo: Ao aplicar o agente treinado em matemática diretamente na geração de código, ele superou métodos otimizados especificamente para código (como RegMix e DBL), provando que as heurísticas aprendidas são generalizáveis.
• Alinhamento com Intuição Humana: A análise das trajetórias mostrou que o agente aprendeu a aumentar dados de domínios benéficos (ex: "Ciência", "Saúde") e diminuir domínios prejudiciais (ex: "Moda e Beleza"), alinhando-se com descobertas humanas recentes (Wettig et al., 2025).

5. Significado e Impacto

O Data Mixing Agent representa um avanço significativo na automação do pré-treinamento de LLMs:

• Automação de Hiperparâmetros: Elimina a necessidade de design manual complexo de misturas de dados, que é custoso e propenso a viés.
• Escalabilidade: Por ser um modelo leve e generalizável, ele pode ser aplicado a novos cenários de pré-treinamento contínuo sem a necessidade de custosos processos de amostragem e treinamento de novos agentes.
• Eficiência Computacional: Reduz o custo total de treinamento ao otimizar o uso de dados, permitindo atingir o desempenho desejado com menos tokens de dados de origem.

Em resumo, o trabalho demonstra que é possível aprender heurísticas universais de mistura de dados através de Aprendizado por Reforço, criando agentes que guiam o pré-treinamento contínuo de forma mais inteligente, eficiente e equilibrada do que os métodos atuais.