Résumé technique : ASTOR – Apprentissage par renforcement multi-tâches guidé par l'utilité pour les LLM de code

Énoncé du problème

L'apprentissage par renforcement (RL) avec récompenses vérifiables s'est avéré efficace pour le post-entraînement des grands modèles de langage (LLM) sur des tâches de codage. Cependant, le paradigme actuel repose sur l'entraînement de spécialistes RL distincts et spécifiques à chaque tâche (par exemple, un modèle pour la génération de code, un autre pour la génération de tests unitaires). Cette approche engendre des coûts de mémoire et de calcul qui augmentent linéairement avec le nombre de tâches. À l'inverse, s'en remettre à un seul spécialiste entraîne une dégradation brutale des performances lorsque le modèle est appliqué en dehors de son domaine d'entraînement spécifique.

Bien que l'apprentissage par renforcement multi-tâches (MTRL) offre une solution unifiée en entraînant un modèle unique sur toutes les tâches, les méthodes existantes souffrent de deux limitations critiques dues à leur traitement uniforme des tâches de codage hétérogènes :

1. Planification des données (« Quoi apprendre ») : Les méthodes actuelles échantillonnent à partir d'un pool de tâches mixte ou suivent un curriculum fixe, ignorant les effets inter-tâches. Elles ne tiennent pas compte de la manière dont l'entraînement sur une tâche (par exemple, la génération de code) peut avoir un impact positif ou négatif sur une autre (par exemple, la génération de tests unitaires).
2. Optimisation de la politique (« Comment apprendre ») : Les approches existantes appliquent une stratégie de mise à jour uniforme et un seul coefficient de régularisation KL sur toutes les tâches. Cela est sous-optimal car les tâches de codage ont des exigences d'optimisation divergentes ; les tâches strictes sur le plan syntaxique (comme la génération de tests unitaires) nécessitent des mises à jour conservatrices pour maintenir l'exécutabilité, tandis que les tâches flexibles sur le plan sémantique (comme la génération de messages de commit) bénéficient d'une exploration plus large.

Méthodologie : ASTOR

Les auteurs proposent ASTOR (framework d'apprentissage par renforcement multi-tâches pour le code via une coORdination pilotée par l'utilité), un cadre qui résout ces défis grâce à deux modules couplés centrés sur un signal unifié : l'utilité de la tâche. L'utilité de la tâche capture à la fois le potentiel d'apprentissage d'une tâche et sa synergie inter-tâches.

1. Planification des données hiérarchique à routage par utilité

Ce module répond à la question « quoi apprendre » en allouant dynamiquement le budget d'entraînement basé sur les utilités de tâches estimées.

• Estimation de l'utilité à deux niveaux :
  • Niveau tâche : L'utilité est définie par deux composantes :
    • Potentiel d'apprentissage : Mesuré par la variance des récompenses de défilement (rollout). Une variance élevée indique un ensemble hétérogène de résultats (succès et échecs), suggérant des signaux d'optimisation actionnables.
    • Synergie inter-tâches : Mesurée par la similarité cosinus des gradients entre la tâche actuelle et les autres tâches. Une similarité positive indique un alignement constructif, tandis qu'une similarité négative suggère un conflit.
  • Niveau prompt : Au sein d'une tâche, l'utilité est déterminée par le potentiel d'apprentissage actuel (variance de la récompense) et les progrès historiques (amélioration de la récompense moyenne).
• Allocation hiérarchique :
  • Allocation de quota par tâche : La taille globale du lot d'entraînement est répartie entre les tâches en utilisant un softmax à température échelonnée sur les utilités de tâches agrégées. Cela garantit que les tâches à fort potentiel et synergie reçoivent plus de données d'entraînement.
  • Priorisation des prompts : Au sein du quota alloué à chaque tâche, les prompts sont sélectionnés par échantillonnage pondéré basé sur leur utilité au niveau du prompt, privilégiant les échantillons où le modèle s'améliore activement.

2. Optimisation de la politique adaptative calibrée par l'utilité

Ce module répond à la question « comment apprendre » en calibrant dynamiquement le coefficient de régularisation KL ($\beta_k$) pour chaque tâche en fonction de son utilité actuelle.

• Mécanisme : Le coefficient KL pour la tâche $k$ est décomposé en une base spécifique à la tâche ($\beta_{base,k}$) et un multiplicateur dynamique conditionné par l'utilité de la tâche :
  $$\beta_k(t) = \beta_{base,k} \cdot (1 + \lambda \cdot U(t)(T_k))$$
• Logique d'adaptation :
  • Les tâches à haute utilité (signaux d'apprentissage forts) reçoivent un $\beta_k$ plus élevé, resserrant la contrainte KL pour éviter les dépassements lorsque le signal de gradient est concentré.
  • Les tâches à faible utilité reçoivent un coefficient plus faible, permettant des mises à jour plus agressives là où la politique a de la marge d'amélioration.
• Système de récompense : Le cadre prend en charge quatre tâches hétérogènes couvrant le flux de travail de développement logiciel : Prédiction d'E/S de code (exactitude d'exécution), Génération de code (taux de réussite), Génération de tests unitaires (couverture) et Génération de messages de commit (similarité textuelle). Chacun utilise un signal de récompense vérifiable basé sur l'exécution.

Contributions clés

1. Première étude du MTRL pour les LLM de code : L'article formule le RL multi-tâches pour le code autour de deux décisions couplées (« quoi apprendre » et « comment apprendre ») et introduit l'utilité de la tâche comme signal unifié régissant les deux.
2. Framework ASTOR : Un cadre de coordination comprenant :
   • Planification des données hiérarchique à routage par utilité : Alloue le budget entre les tâches et priorise les prompts en fonction du potentiel d'apprentissage et de la synergie.
   • Optimisation de la politique adaptative calibrée par l'utilité : Ajuste dynamiquement la régularisation KL par tâche pour correspondre à l'état d'entraînement des tâches hétérogènes.
3. Validation empirique : Des expériences extensives démontrant qu'un modèle ASTOR unique surpasse à la fois les spécialistes spécifiques à une tâche et les bases de référence MTRL existantes.

Résultats expérimentaux

Les auteurs ont évalué ASTOR sur deux modèles de base (Qwen2.5-Coder-7B et Qwen3-8B) à travers quatre tâches de codage représentatives.

• Vs. Apprentissage mono-tâche (STL) : ASTOR atteint de solides performances multi-tâches, rivalisant avec ou dépassant les spécialistes spécifiques à une tâche. Sur Qwen2.5-Coder-7B, il surpasse le meilleur modèle mono-tâche de 9,0 % en performance moyenne globale. Sur Qwen3-8B, l'amélioration est de 9,5 %.
• Vs. Bases de référence MTRL : ASTOR surpasse significativement trois bases de référence : Apprentissage conjoint (mélange uniforme), Apprentissage par curriculum (ordre séquentiel fixe) et Fusion de modèles (intégration post-hoc de spécialistes).
  • Les améliorations par rapport à la base de référence MTRL la plus forte (Apprentissage conjoint) varient de 7,5 % à 19,2 % selon les modèles et les métriques.
  • Plus précisément, sur Qwen2.5-Coder-7B, ASTOR surpasse l'Apprentissage conjoint de 12,8 %.
• Études d'ablation : La suppression soit du module de planification des données (allocation de tâche ou priorisation des prompts), soit du module d'optimisation de la politique (régularisation dynamique ou spécifique à la tâche) entraîne une dégradation significative des performances, confirmant que les deux composants sont essentiels pour coordonner les tâches hétérogènes.

Importance et affirmations

L'article affirme qu'ASTOR valide l'efficacité de la coordination pilotée par l'utilité pour le RL multi-tâches de code. En traitant les tâches de manière hétérogène — en allouant des ressources en fonction du potentiel d'apprentissage et de la synergie, et en contraignant les mises à jour en fonction des besoins spécifiques à la tâche — un modèle unique peut atteindre des capacités diverses sans les coûts de mise à l'échelle linéaire liés au déploiement de multiples spécialistes.

Les auteurs positionnent ce travail comme une étape vers des systèmes d'aide au codage unifiés et efficaces capables de gérer l'intégralité du flux de travail de développement logiciel (compréhension, implémentation, test et documentation) au sein d'un seul modèle, réduisant les barrières de déploiement tout en maintenant ou en améliorant les performances à tous les stades. Le travail reconnaît les limitations concernant la couverture des tâches/langages et le réglage manuel des coefficients KL de base, suggérant ces aspects comme des pistes pour des travaux futurs.