Résumé technique : Acquisition de données adaptative efficace via des représentations de croyances préentraînées (POLAR)

Énoncé du problème

L'acquisition adaptative de données englobe des problèmes d'apprentissage séquentiel tels que la conception expérimentale bayésienne (BED), l'optimisation bayésienne (BO) et l'apprentissage actif (AL). Dans ces contextes, un apprenant doit sélectionner séquentiellement des requêtes (designs) pour maximiser une utilité en aval (par exemple, le gain d'information, l'identification de l'optimum global ou la précision prédictive) sous un budget fini.

Les approches existantes pour apprendre des politiques d'acquisition tombent généralement dans deux catégories problématiques :

1. Méthodes basées sur la distribution postérieure : Elles reposent sur l'ajustement d'un modèle probabiliste aux données observées pour calculer une distribution postérieure, qui est ensuite utilisée pour évaluer une fonction d'acquisition (par exemple, l'amélioration attendue, le gain d'information attendu). Ces méthodes sont limitées par la qualité de l'approximation de la distribution postérieure. La spécification erronée du modèle ou les erreurs d'approximation peuvent fausser la règle d'acquisition, et ces erreurs se cumulent au fil des cycles séquentiels.
2. Méthodes d'apprentissage de politique directe : Elles cartographient directement les historiques d'observations vers des décisions (par exemple, via l'apprentissage par renforcement). Bien qu'elles évitent l'inférence postérieure explicite, elles doivent apprendre de zéro une cartographie complexe allant des données brutes aux décisions. Cela nécessite d'apprendre simultanément comment représenter l'ensemble de données et comment agir sur cette représentation, ce qui conduit souvent à une inefficacité de l'échantillonnage et à des difficultés d'entraînement.

Méthodologie : POLAR

Les auteurs introduisent POLAR (Policy Learning with Belief Representations), un cadre qui découple l'apprentissage de la représentation de l'apprentissage de la politique. L'idée centrale est que la décision séquentielle optimale dépend de l'historique des observations uniquement à travers un état de croyance (belief state). Plut lieu de l'apprendre à partir de zéro ou de s'appuyer sur des approximations de la distribution postérieure explicites, POLAR exploite des modèles de fondation prédictifs préentraînés comme encodeurs d'état de croyance.

Architecture

1. Encodeur de croyance (Backbone) : POLAR utilise un modèle de fondation tabulaire (TFM), tel que TabPFN ou TabICL, préentraîné sur diverses tâches synthétiques. Étant donné un historique d'observations $D_{t-1}$ et un candidat $x$, le TFM produit une représentation cachée conditionnée au candidat $r_\psi(x, D_{t-1})$.

   • Contrairement à la distribution prédictive finale produite par le TFM, cette représentation cachée est traitée comme la représentation de croyance. Elle encode la manière dont les prédictions changent avec les données (incertitude prédictive, structure locale, comportement dépendant du contexte) sans s'engager dans une approximation spécifique de la distribution postérieure.
   • Le cadre opère dans un contexte basé sur un pool (pool-based setting), où la politique sélectionne parmi un ensemble fini de candidats $C_t$. Le backbone traite l'historique une seule fois pour générer les représentations de tous les candidats du pool.

2. Tête de politique : Une tête de politique légère et apprenable $g_\eta$ cartographie la représentation de croyance $r_\psi$ vers un logit scalaire pour chaque candidat. Un softmax sur ces logits définit la politique stochastique $\pi_\eta$. Lors du déploiement, le candidat ayant le score le plus élevé est sélectionné.

3. Objectif d'entraînement :

   • Perte de politique : La politique est entraînée via des gradients de politique sur des simulations de parcours (rollouts). L'objectif maximise une utilité spécifique à la tâche (par exemple, le gain d'information attendu pour BED, la meilleure valeur terminale pour BO). Les auteurs utilisent une stratégie d'attribution de crédit locale (incréments à un pas) pour réduire la variance, plutôt qu'une fonction de valeur complète sur un long horizon.
   • Adaptation du Backbone (Optionnel) : Pour aligner le backbone préentraîné avec la distribution spécifique de la tâche en aval, une perte de prédiction supervisée auxiliaire ($L_{pred}$) est appliquée aux paramètres du backbone.
   • Découplage : Un choix de conception critique est que les gradients de politique ne sont pas propagés dans le backbone. Le backbone est mis à jour uniquement via la perte de prédiction supervisée, tandis que la tête de politique est mise à jour via le gradient de politique. Cela empêche les signaux d'apprentissage par renforcement à haute variance de déstabiliser les représentations préentraînées.

Fondement théorique

L'article fournit une vue décisionnelle unifiée, montrant que la politique bayésienne optimale $\pi^*$ est une fonction de la distribution postérieure pertinente pour la tâche $p(w|D_{t-1})$. Puisque la distribution postérieure est une fonction de l'historique, toute représentation $e_t$ permettant la récupération de la distribution postérieure (via un décodeur) est une représentation de croyance valide. POLAR postule que les TFM préentraînés fournissent de telles représentations qui sont suffisamment riches pour déterminer les actions optimales mais suffisamment structurées pour rendre la cartographie vers la tête de politique facile à apprendre, évitant ainsi les extrêmes de l'historique brut (trop complexe) ou de la distribution postérieure explicite (sujette aux erreurs d'approximation).

Contributions clés

1. Vue unifiée : Les auteurs formalisent l'acquisition de données adaptative comme l'apprentissage de politiques sur des représentations de croyance, comblant ainsi le fossé entre les approches basées sur l'historique brut et celles basées sur la distribution postérieure.
2. Architecture découplée : Ils proposent une architecture simple qui utilise des modèles de fondation préentraînés comme encodeurs de croyance fixes (ou affinés) et entraîne une tête de politique distincte, séparant l'apprentissage de la représentation de la prise de décision.
3. Efficacité de l'échantillonnage : Les résultats empiriques démontrent que POLAR atteint des performances de pointe à travers diverses tâches (BED, BO, AL) tout en nécessitant nettement moins d'échantillons d'entraînement (jusqu'à $100\times$ de moins) par rapport aux méthodes amorties existantes.

Résultats expérimentaux

Les auteurs évaluent POLAR sur trois bancs d'essai principaux :

1. Conception expérimentale bayésienne (BED) :

   • Tâches : Localisation (2D et 5D) et Élasticité de substitution constante (CES).
   • Résultats : POLAR surpasse de manière significative les méthodes amorties de pointe (DAD, RL-BOED, ALINE). Dans la tâche de localisation 5D, POLAR dépasse la performance de pointe de RL-BOED et DAD en utilisant 10 000$\times$ moins d'échantillons d'entraînement et dépasse ALINE avec 100$\times$ moins d'échantillons.
   • Ablations :
     • Backbone figé vs affiné : L'affinage du backbone avec une perte de prédiction auxiliaire produit de meilleures performances que de le laisser figé, bien que la version figée soit déjà compétitive.
     • Gradients découplés vs couplés : Autoriser les gradients de politique à mettre à jour le backbone (couplé) dégrade gravement les performances, confirmant la nécessité du découplage.
     • Préentraîné vs à partir de zéro : Entraîner l'architecture à partir de zéro converge beaucoup plus lentement et donne de moins bons résultats, soulignant la valeur du modèle de fondation préentraîné.

2. Optimisation bayésienne (Optimisation d'hyperparamètres) :

   • Benchmark : HPO-B (Hyperparameter Optimisation Benchmark).
   • Résultats : POLAR obtient le regret moyen le plus bas et le meilleur rang moyen à travers six espaces de recherche, surpassant les processus gaussiens classiques, les Meta-GPs et les méthodes amorties de bout en bout comme NAP. Il démontre que l'utilisation d'un backbone TFM préentraîné produit des performances plus robustes que l'apprentissage de représentations à partir de zéro.

3. Optimisation moléculaire :

   • Benchmark : DOCKSTRING (amarrage moléculaire de haute dimension).
   • Résultats : Dans un espace discret de 512 dimensions, POLAR obtient le regret le plus bas par rapport aux GP avec noyaux Tanimoto, aux forêts aléatoires et à une ligne de base TabICL avec une fonction d'acquisition conçue à la main. Cela confirme que l'apprentissage direct de la règle d'acquisition à partir de la représentation de croyance surpasse les stratégies d'acquisition fixes.

Signification et affirmations

L'article affirme que POLAR représente une étape importante dans la scalabilité et l'efficacité de l'acquisition de données amortie. En passant du paradigme consistant à apprendre conjointement les représentations et les décisions à celui d'agir sur des représentations de croyance préentraînées, POLAR atténue l'inefficacité de l'échantillonnage et l'instabilité de l'entraînement associés à l'apprentissage de politique directe et aux erreurs d'approximation des méthodes basées sur la distribution postérieure.

Les auteurs soulignent que leur approche ne consiste pas à concevoir un nouvel estimateur ou un nouvel algorithme d'entraînement, mais plutôt à réutiliser un modèle prédictif préentraîné comme encodeur de croyance. Cela permet à une petite tête de politique d'être entraînée avec des gradients de politique standards, atteignant des performances supérieures avec des coûts de calcul d'entraînement considérablement réduits. Le cadre est présenté comme une solution unifiée applicable à BED, BO et AL, ne différant que par l'utilité spécifique à la tâche utilisée pour l'entraînement.