Resumo Técnico: ASTOR – Aprendizado por Reforço Multi-Tarefa Guiado por Utilidade para LLMs de Código

Declaração do Problema

O aprendizado por reforço (RL) com recompensas verificáveis demonstrou ser eficaz para o pós-treinamento de Modelos de Linguagem de Grande Escala (LLMs) em tarefas de codificação. No entanto, o paradigma atual depende do treinamento de especialistas separados e específicos para cada tarefa (por exemplo, um modelo para geração de código, outro para geração de testes unitários). Essa abordagem incorre em custos de memória e computação que escalam linearmente com o número de tarefas. Por outro lado, confiar em um único especialista leva a uma degradação acentuada do desempenho quando o modelo é aplicado fora de seu domínio de treinamento específico.

Embora o Aprendizado por Reforço Multi-Tarefa (MTRL) ofereça uma solução unificada ao treinar um único modelo em todas as tarefas, os métodos existentes sofrem de duas limitações críticas devido ao tratamento uniforme de tarefas de codificação heterogêneas:

1. Agendamento de Dados ("O que aprender"): Os métodos atuais amostram de um pool de tarefas mistas ou seguem um currículo fixo, ignorando os efeitos inter-tarefas. Eles falham em considerar como o treinamento em uma tarefa (por exemplo, geração de código) pode impactar positiva ou negativamente outra (por exemplo, geração de testes unitários).
2. Otimização de Política ("Como aprender"): As abordagens existentes aplicam uma estratégia de atualização uniforme e um único coeficiente de regularização KL em todas as tarefas. Isso é subótimo porque as tarefas de codificação possuem requisitos de otimização divergentes; tarefas sintaticamente estritas (como a geração de testes unitários) exigem atualizações conservadoras para manter a executabilidade, enquanto tarefas semanticamente flexíveis (como a geração de mensagens de commit) beneficiam-se de uma exploração mais ampla.

Metodologia: ASTOR

Os autores propõem o ASTOR (framework de aprendizado por reforço multi-tarefa para código via coORdinação acionada por uTilidade), um framework que resolve esses desafios através de dois módulos acoplados centrados em um sinal unificado: Utilidade da Tarefa. A utilidade da tarefa captura tanto o potencial de aprendizado de uma tarefa quanto sua sinergia entre tarefas.

1. Agendamento de Dados Roteado por Utilidade Hierárquica

Este módulo aborda "o que aprender" alocando dinamicamente o orçamento de treinamento com base nas utilidades de tarefa estimadas.

• Estimativa de Utilidade em Dois Níveis:
  • Nível de Tarefa: A utilidade é definida por dois componentes:
    • Potencial de Aprendizado: Medido pela variância das recompensas de simulação (rollout). Alta variância indica um conjunto heterogêneo de resultados (sucessos e falhas), sugerindo sinais de otimização acionáveis.
    • Sinergia Inter-Tarefa: Medida pela similaridade do cosseno do gradiente entre a tarefa atual e outras tarefas. Similaridade positiva indica alinhamento construtivo, enquanto similaridade negativa sugere conflito.
  • Nível de Prompt: Dentro de uma tarefa, a utilidade é determinada pelo potencial de aprendizado atual (variância da recompensa) e pelo progresso histórico (melhoria na recompensa média).
• Alocação Hierárquica:
  • Alocação de Cota de Tarefa: O tamanho global do lote de treinamento é partitionado entre as tarefas usando um softmax escalado por temperatura sobre as utilidades de tarefa agregadas. Isso garante que tarefas com alto potencial e sinergia recebam mais dados de treinamento.
  • Priorização de Prompts: Dentro da cota alocada de cada tarefa, os prompts são selecionados via amostragem ponderada com base em sua utilidade no nível do prompt, priorizando amostras onde o modelo está melhorando ativamente.

2. Otimização de Política Adaptativa Calibrada por Utilidade

Este módulo aborda "como aprender" calibrando dinamicamente o coeficiente de regularização KL ($\beta_k$) para cada tarefa com base em sua utilidade atual.

• Mecanismo: O coeficiente KL para a tarefa $k$ é decomposto em uma base específica da tarefa ($\beta_{base,k}$) e um multiplicador dinâmico condicionado à utilidade da tarefa:
  $$\beta_k(t) = \beta_{base,k} \cdot (1 + \lambda \cdot U(t)(T_k))$$
• Lógica de Adaptação:
  • Tarefas com alta utilidade (sinais de aprendizado fortes) recebem um $\beta_k$ maior, apertando a restrição KL para evitar ultrapassagem quando o sinal do gradiente está concentrado.
  • Tarefas com baixa utilidade recebem um coeficiente menor, permitindo atualizações mais agressivas onde a política tem espaço para melhorar.
• Sistema de Recompensa: O framework suporta quatro tarefas heterogêneas que abrangem o fluxo de trabalho de desenvolvimento de software: Predição de I/O de Código (correção de execução), Geração de Código (taxa de aprovação), Geração de Teste Unitário (cobertura) e Geração de Mensagem de Commit (similaridade textual). Cada uma utiliza um sinal de recompensa verificável e baseado em execução.

Principais Contribuições

1. Primeiro Estudo de MTRL para LLMs de Código: O artigo formula o RL multi-tarefa para código em torno de duas decisões acopladas ("o que aprender" e "como aprender") e introduz a utilidade da tarefa como um sinal unificado que governa ambas.
2. Framework ASTOR: Um framework de coordenação que apresenta:
   • Agendamento de Dados Roteado por Utilidade Hierárquica: Aloca orçamento entre tarefas e prioriza prompts com base no potencial de aprendizado e sinergia.
   • Otimização de Política Adaptativa Calibrada por Utilidade: Ajusta dinamicamente a regularização KL por tarefa para corresponder ao estado de treinamento de tarefas heterogêneas.
3. Validação Empírica: Extensos experimentos demonstrando que um único modelo ASTOR supera tanto especialistas específicos para tarefas quanto as bases de referência MTRL existentes.

Resultados Experimentais

Os autores avaliaram o ASTOR em dois modelos base (Qwen2.5-Coder-7B e Qwen3-8B) em quatro tarefas de codificação representativas.

• vs. Aprendizado de Tarefa Única (STL): O ASTOR alcança forte desempenho multi-tarefa, rivalizando ou superando especialistas específicos para tarefas. No Qwen2.5-Coder-7B, supera o melhor modelo de tarefa única em 9,0% no desempenho médio geral. No Qwen3-8B, a melhoria é de 9,5%.
• vs. Bases de Referência MTRL: O ASTOR supera significativamente três bases de referência: Aprendizado Conjunto (mistura uniforme), Aprendizado de Currículo (ordem sequencial fixa) e Mesclagem de Modelos (integração post-hoc de especialistas).
  • As melhorias sobre a base de referência MTRL mais forte (Aprendizado Conjunto) variam de 7,5% a 19,2% entre modelos e métricas.
  • Especificamente, no Qwen2.5-Coder-7B, o ASTOR supera o Aprendizado Conjunto em 12,8%.
• Estudos de Ablação: Remover o módulo de agendamento de dados (alocação de tarefa ou priorização de prompts) ou o módulo de otimização de política (regularização dinâmica ou específica da tarefa) leva a uma degradação significativa do desempenho, confirmando que ambos os componentes são essenciais para coordenar tarefas heterogêneas.

Significado e Alegações

O artigo afirma que o ASTOR valida a eficácia da coordenação acionada por utilidade para o RL multi-tarefa de código. Ao tratar tarefas de forma heterogênea — alocando recursos com base no potencial de aprendizado e sinergia, e restringindo atualizações com base em necessidades específicas da tarefa — um único modelo pode alcançar capacidades diversas sem os custos de escalonamento linear de implantar múltiplos especialistas.

Os autores posicionam este trabalho como um passo em direção a sistemas de codificação assistida por IA eficientes e unificados que podem lidar com todo o fluxo de trabalho de desenvolvimento de software (compreensão, implementação, teste e documentação) dentro de um único modelo, reduzindo barreiras de implantação enquanto mantêm ou melhoram o desempenho em todas as etapas. O trabalho reconhece limitações quanto à cobertura de tarefas/idiomas e ao ajuste manual dos coeficientes KL base, sugerindo esses pontos como áreas para trabalhos futuros.