Utilisation de MemAlign pour améliorer l'évaluation de l'apprentissage automatique traditionnel dans le code Genie

Le Genie Code, annoncé récemment, est le partenaire d'IA autonome de Databricks, spécialement conçu pour le travail sur les données. Il a remplacé Databricks Assistant, tout en intégrant plusieurs agents et en offrant de nouveaux points d'intégration et de nouvelles capacités. Genie Code est profondément intégré à Unity Catalog, ce qui signifie qu'il comprend vos tables, colonnes, lignage, vues métriques et définitions métier (sémantique). Cette conscience contextuelle rend Genie Code beaucoup plus utile pour les praticiens de données que les chatbots génériques.

Lorsque Genie Code génère un notebook pour des tâches de ML traditionnelles, telles que "construire un modèle de prédiction de churn", nous nous attendons à ce qu'il produise un flux de travail prêt pour la production, incluant l'installation des bibliothèques Python appropriées, l'exploration et le prétraitement des données, l'entraînement, le réglage, l'enregistrement et le déploiement du modèle, ainsi que l'évaluation de ses performances. Nous nous attendons également à ce que chaque étape soit réellement informée par les données : par exemple, Genie Code doit comprendre que des classes déséquilibrées dans un problème de classification binaire entraînent des flux de travail et des métriques de succès radicalement différents.

Pour garantir que Genie Code suive systématiquement les meilleures pratiques natives de Databricks et évite, par exemple, de sauter la validation croisée, de ne pas remarquer les fuites de données ou une imputation de données incorrecte, nous avions besoin d'un moyen rigoureux de répondre à une question : Comment savoir si le code généré est réellement bon ? Le notebook généré dépendra grandement du problème que le client essaie de résoudre, et cela peut varier considérablement d'un client à l'autre, c'est donc une question très non triviale.

Dans cet article, nous allons expliquer comment nous avons construit un pipeline d'évaluation pour les capacités de ML traditionnelles de Genie Code et comment nous avons utilisé MemAlign (un nouveau framework d'alignement open-source dans MLflow) pour combler l'énorme fossé que nous avons trouvé entre les juges LLM et les experts humains. Les juges améliorés nous ont aidés à identifier et à corriger des lacunes dans les conseils ML de Genie Code que nous aurions autrement manquées.

Construction du cadre d'évaluation

Un cadre d'évaluation robuste est requis pour :

• Hillclimbing : quantifier comment les prompts, les outils, les compétences et les changements d'architecture affectent la sortie.
• Protection contre les régressions : S'assurer que l'amélioration de "l'entraînement du modèle" ne dégrade pas accidentellement "l'exploration des données".
• Benchmarking : Mesurer comment différents modèles de fondation (backends LLM) impactent la qualité du notebook.
• CI : Surveiller comment les changements dans la boucle agentique sous-jacente se répercutent sur les tâches ML finales.

L'évaluation des notebooks ML traditionnels est l'une des tâches d'évaluation les plus complexes car elle couvre l'évaluation de la qualité du code, des meilleures pratiques ML et des adaptations/personnalisations basées sur les données. Pour gérer une tâche aussi large et désordonnée que l'évaluation des notebooks ML, nous utilisons un LLM-as-a-judge - un "expert" LLM formé par des humains sur ce à quoi ressemble un bon notebook. Nous avons créé neuf juges qui sont invités à évaluer les notebooks ML selon neuf dimensions qui apparaissent dans la plupart des flux de travail ML :

Pour chaque dimension, nous avons rédigé des grilles de notation (réutilisées entre les évaluateurs humains et les juges LLM) qui attribuent un score de 1 à 3, et 0 pour "non applicable" :

• 3 (Bon) : Le notebook atteint un niveau élevé pour une dimension. Il démontre les meilleures pratiques, couvre la portée attendue et gère correctement les cas limites.
• 2 (Moyen) : Acceptable mais avec des lacunes. Les bases sont présentes, mais le notebook manque de raffinements qu'un praticien expérimenté attendrait.
• 1 (Mauvais) : Problèmes fondamentaux. Des étapes clés sont manquantes, incorrectes ou appliquées d'une manière qui conduirait à de fausses conclusions.
• N/A (Non applicable) : Cette dimension n'est pas applicable à ce prompt (par exemple, la dimension d'imputation des données ne peut pas être appliquée si l'ensemble de données n'a pas de valeurs manquantes).

Pour donner une idée de la granularité, voici la grille spécifique que nous utilisons pour la dimension "imputation des données" :

Parallèlement aux juges, nous maintenons un ensemble de cas de test d'évaluation qui couvrent une gamme de tâches ML (classification, régression, prévision), différentes tailles de jeux de données, domaines et niveaux de complexité. Chaque cas de test comprend un prompt utilisateur qui indique à Genie Code la tâche ML qu'il est censé résoudre sur le jeu de données spécifié ("J'ai des données de passagers dans les tables titanic_train_table et titanic_test_table. Pouvez-vous déterminer qui a survécu ?"). La boucle d'évaluation consiste à utiliser Genie Code pour générer un ou plusieurs notebooks pour chaque cas de test, puis à évaluer chaque notebook selon toutes les dimensions applicables.

Évaluation du système d'évaluation

En utilisant des juges LLM, au lieu d'humains, pour évaluer les artefacts de Genie Code, nous avons essentiellement échangé un problème difficile contre un autre : le juge prêt à l'emploi est incompétent dans la tâche à accomplir et désaligné avec les évaluations humaines. Notre problème est de faire en sorte que les scores des juges LLM correspondent à ceux des évaluateurs humains.

L'ensemble d'évaluation pour l'appréciation des juges LLM contient 50 notebooks générés par Genie Code ("cas de test") où des experts humains ont noté chaque dimension applicable, fournissant à la fois un score et une courte justification servant de vérité terrain. Dans les zones grises entre deux scores, les évaluateurs étaient autorisés à exprimer leur propre jugement, mais les schémas étaient rédigés de manière à ce que cela soit rarement le cas.

La mesure de l'alignement humain-machine est l'erreur absolue moyenne (MAE) entre les scores dans chaque dimension. Les résultats ont été mitigés, certaines dimensions ont montré un fort alignement (4 dimensions avaient une MAE de <= 0,10), tandis que d'autres ont révélé un désaccord significatif :

• Entraînement du modèle : MAE de 0,680
• Utilisation du modèle : MAE de 0,562
• Imputation des données : MAE de 0,474
• Exploration des données : MAE de 0,407

Cet écart existe parce que les humains et les LLM n'interprètent pas la même grille de notation de la même manière. Alors qu'un évaluateur humain peut repérer une stratégie d'imputation subtilement erronée ou une boucle d'entraînement qui "fonctionne" mais est logiquement unsound, un juge LLM manque souvent cette nuance technique. Nous avons également constaté que le juge souffrait d'un biais de positivité classique - il était simplement trop "poli" et cela l'empêchait d'obtenir des résultats objectifs.

Il est devenu évident qu'étant donné la même grille de notation, les juges LLM et les humains ne produiraient pas les mêmes résultats - un désalignement. C'est exactement le scénario pour lequel MemAlign a été conçu.

Utilisation de MemAlign pour l'alignement

MemAlign est un framework au sein de MLflow qui peut, avec une très petite quantité de feedback en langage naturel humain, effectuer un alignement entre les évaluateurs humains et les juges LLM. Ceci est réalisé grâce à deux types de "mémoires" formées en lisant le feedback humain :

• Mémoire sémantique stocke des directives généralisées - des règles distillées à partir des commentaires qui s'appliquent largement
• Mémoire épisodique stocke des exemples spécifiques - des cas où le juge s'est trompé, conservés comme points d'ancrage pour les décisions futures

Au moment de l'inférence, MemAlign construit un contexte de travail en extrayant toutes les directives sémantiques et en récupérant les exemples épisodiques les plus pertinents pour l'entrée actuelle. Le juge charge tout cela dans son contexte, avec la grille d'évaluation d'origine, et utilise les connaissances accumulées pour donner un score plus précis à tous les futurs notebooks.

La propriété clé qui a fait ressortir MemAlign est ses hautes performances en utilisant seulement un petit nombre d'exemples. C'est parce que MemAlign distille efficacement l'apprentissage à partir de riches signaux d'apprentissage dans les commentaires en langage naturel, et les intègre dans le système de double mémoire.

Voici un exemple de certains extraits de mémoire sémantique générés pour la dimension « imputation de données », comblant les lacunes de la grille d'évaluation que nous avions définie en fournissant généralement des points d'ancrage, des exemples et des contre-exemples :

De plus, comme mentionné précédemment, la mémoire sémantique reflétée dans l'invite est complétée par des exemples pertinents de la mémoire épisodique du juge au moment de la notation, donnant ainsi au juge encore plus de contexte pour interpréter les instructions optimisées.

Conception de l'expérience

Validation croisée K-Fold

En suivant le paradigme d'entraînement-test ML, nous avons appliqué la validation croisée K-fold (K=4) sur 50 cas de test (notebooks), évitant ainsi la fuite de données et la nécessité d'étiqueter un ensemble de test distinct. Pour chaque pli, nous avons fait ce qui suit :

1. Phase d'entraînement : MemAlign a aligné le juge en utilisant les traces des autres plis pour obtenir le juge.
2. Phase d'évaluation : A évalué les notebooks du pli ⁱ avec le juge.

Intervalles de confiance par bootstrap

Pour calculer les intervalles de confiance sans données étiquetées supplémentaires, nous avons généré 100 échantillons bootstrap avec remplacement à partir des 50 d'origine. En répétant cela 10 000 fois et en suivant la MAE entre les scores humains et machine, nous avons calculé les intervalles de confiance pour l'alignement humain-machine avec un IC à 95 % définissant un changement statistiquement significatif.

Implémentation

Le pipeline d'évaluation est implémenté comme un seul extrait MLflow qui orchestre l'ensemble du processus :

L'optimiseur MemAlign est capable d'aligner les juges LLM en fonction des traces des cas de test en seulement quelques lignes de code. Nous avons utilisé ce nouveau juge « aligné » pour calculer la nouvelle MAE. L'alignement d'un juge sur une seule dimension prend environ 25 secondes par pli, donc l'alignement lui-même n'est pas un goulot d'étranglement.

Résultats

Trois dimensions sur 9 ont montré une amélioration statistiquement significative :

• Entraînement du modèle amélioré de 0,500 MAE (0,680 → 0,180), une réduction de 74 %
• Utilisation du modèle améliorée de 0,438 MAE (0,562 → 0,125), une réduction de 78 %
• Imputation de données améliorée de 0,421 MAE (0,474 → 0,053), une réduction de 89 %

Ces 3 dimensions font partie des 4 dimensions initiales qui étaient fortement désalignées. Un faible alignement initial est indicatif d'une compréhension fondamentalement différente des grilles d'évaluation partagées par les LLM et les humains, et la mémoire injectée par MemAlign semble fournir suffisamment de contexte pour les mettre « sur la même longueur d'onde ».

• Évaluation des métriques et journalisation MLflow étaient déjà bien alignées (MAE < 0,10 à l'origine), et leur dégradation n'est pas statistiquement significative (bruit expérimental)
• Exploration des données a montré une légère régression (-0,130), mais pas statistiquement significative compte tenu de son intervalle de confiance [-0,33, +0,09]. Cette dimension présentait la plus forte variance inter-évaluateurs, et ce bruit a empêché MemAlign de s'améliorer (et l'a peut-être même gêné).

Expérience avec la mémoire sémantique uniquement

La structure de double mémoire de MemAlign nous a amené à nous demander si les deux contribuent réellement à l'alignement du juge. En particulier, la mémoire épisodique est censée aider le juge en lui donnant un ensemble des notebooks annotés les plus similaires comme point de référence (en utilisant la recherche du voisin le plus proche). Mais que se passe-t-il si les notebooks récupérés (voisins les plus proches) ne sont pas réellement similaires à celui actuel, mais seulement les moins dissemblables ? Les charger dans le contexte du juge pourrait embrouiller les choses plutôt que d'aider. L'espace de problèmes que nous évaluons (notebooks ML) est très large, et nous avions initialement émis l'hypothèse qu'un ensemble de 50 notebooks ne suffirait tout simplement pas pour obtenir un ensemble de mémoires suffisamment dense pour que le juge puisse s'en souvenir.

Sans mémoire épisodique, l'image se dégrade considérablement :

• Entraînement du modèle s'améliore toujours (+0,420), mais le gain est plus faible que le +0,500 avec MemAlign complet, et la MAE alignée est de 0,260 contre 0,180.
• Utilisation du modèle perd complètement la signification statistique - l'amélioration passe de +0,438 à +0,294, l'intervalle de confiance traversant maintenant zéro.
• Imputation de données passe d'une réduction d'erreur de 89 % à aucune amélioration - la MAE alignée est égale à la MAE d'origine (0,455).
• Journalisation MLflow et évaluation des métriques régressent en fait de manière significative. Sans exemples épisodiques pour ancrer le juge, les directives distillées seules introduisent du bruit dans des dimensions qui étaient déjà bien calibrées, faisant passer la journalisation MLflow de 0,062 à 0,396 MAE.

C'était le contraire de ce que nous attendions. Nous avions initialement émis l'hypothèse que notre ensemble annoté épars finirait par confondre le juge, mais presque toutes les dimensions se sont détériorées sans mémoire épisodique. La seule exception était l'exploration des données, où la suppression des exemples épisodiques a pu aider - sans les notebooks spécifiques sur lesquels nos annotateurs étaient en désaccord, le juge n'avait que les directives distillées, et un signal moins bruyant pour travailler.

La leçon à retenir : même lorsque vos entrées sont volumineuses et désordonnées, la mémoire épisodique améliore considérablement les performances du juge. Les mémoires sémantique et épisodique sont toutes deux intégrales au fonctionnement de MemAlign.

Conclusion : Combler le fossé expert

Juger si un agent de codage fait son travail est déjà assez difficile, tandis qu'évaluer un partenaire IA autonome dans la construction et l'exécution de flux de travail ML traditionnels est d'un autre niveau de complexité. En raison de l'itération rapide sur les produits d'IA, il n'y a tout simplement pas assez de temps pour que des experts surveillent l'« intégration continue » de l'agent. La seule solution évolutive viable sont les juges LLM - mais nous avons toujours besoin d'un jury humain pour garder le juge LLM sous contrôle.

En appliquant MemAlign, nous avons réduit l'erreur du juge de 74 à 89 % sur les dimensions où cela importait le plus. Mais, comme pour tout travail ML/LLM, le résultat n'est aussi bon que les informations que vous y mettez, alors assurez-vous que l'étiquetage est compétent.

Points à retenir :

• Mesurez votre système de mesure : Un système bruyant n'est pas idéal pour l'évaluation, et tant que nous n'avions pas investi le temps et les ressources nécessaires pour valider et améliorer les juges, nous ne pouvions pas faire confiance à notre système d'évaluation.
• Les rubriques ne suffisent pas à elles seules : Il existe des différences subtiles entre la façon dont un humain perçoit les instructions et la façon dont un LLM les perçoit. Ces différences doivent être prises en compte, et des outils d'alignement comme MemAlign sont un moyen efficace de combler le fossé.
• Qualité de l'étiquetage > quantité : Lorsque les annotateurs humains ne sont pas d'accord entre eux (comme nous l'avons vu dans notre régression d'exploration de données), l'alignement n'a pas de signal cohérent à apprendre.

MemAlign est livré avec MLflow et il a fonctionné pour nous avec seulement ~50 exemples étiquetés. Si vos juges LLM ne correspondent pas à vos experts, cela vaut la peine d'y consacrer un après-midi.

(Cet article de blog a été traduit à l'aide d'outils basés sur l'intelligence artificielle) Article original