Résumé technique : Échantillonnage guidé par l'influence pour l'adaptation de domaine des récupérateurs de texte

Énoncé du problème

Les systèmes de récupération dense en domaine ouvert à usage général sont généralement entraînés sur des mélanges vastes et éclectiques de corpus et de tâches de recherche. Un défi critique dans le développement de modèles de plongement (embeddings) robustes et généralisables consiste à déterminer comment échantillonner ces ensembles de données diversifiés pour l'entraînement. Les stratégies conventionnelles — telles que l'échantillonnage uniforme, l'échantillonnage proportionnel basé sur la population d'instances, ou les poids statiques définis par des experts — sont souvent sous-optimales. Elles ne parviennent pas à s'adapter dynamiquement à l'informativité spécifique des sous-ensembles de données nécessaires pour maximiser les performances sur un domaine cible. Bien que des approches dynamiques basées sur le gradient (par exemple, DDS, DoGE, DoReMi) aient été proposées pour apprendre les poids d'échantillonnage en ligne, elles font face à des limitations importantes : (1) elles souffrent d'instabilité et d'une variance élevée en raison des estimations de gradient stochastiques utilisées pour le calcul des récompenses, et (2) elles engendrent un surcoût computationnel substantiel lorsqu'elles nécessitent des modèles proxys ou des estimateurs coûteux.

Méthodologie : Inf-DDS

Les auteurs proposent Inf-DDS (Influence-guided Dynamic Data Sampling), un nouveau cadre piloté par l'apprentissage par renforcement, conçu pour ré-pondérer de manière adaptative les ensembles de données d'entraînement. L'innovation centrale réside dans le remplacement des récompenses basées sur le gradient, souvent bruitées, par des signaux de récompense basés sur l'influence dérivés de modèles proxys en ligne.

1. Formulation du problème

La tâche est formulée comme un problème d'optimisation bilatérale. L'objectif est d'apprendre une politique d'échantillonnage paramétrée $\psi$ qui ajuste la distribution de probabilité $P_D(i; \psi)$ sur un ensemble de jeux de données d'entraînement $\{D_i^{train}\}$ afin de maximiser la performance du modèle sur un ensemble de jeux de données de développement cibles $\{D_j^{dev}\}$.

• Boucle interne : Met à jour les paramètres du modèle $\theta$ en utilisant la distribution d'échantillonnage actuelle.
• Boucle externe : Met à jour les paramètres du scoreur $\psi$ pour optimiser la distribution d'échantillonnage en fonction de l'impact sur l'ensemble de développement.

2. Récompenses basées sur l'influence

Contra contrairement aux méthodes antérieures qui reposent sur des approximations de premier ordre ou de petits modèles proxys pour estimer l'influence des données, Inf-DDS utilise des modèles proxys en ligne pour calculer des scores d'influence exacts.

• Processus : À chaque étape $t$, les paramètres du modèle $\theta_t$ sont mis à jour en utilisant un jeu de données d'entraînement spécifique $D_i^{train}$ pendant $l$ étapes de gradient pour produire $\theta_{t+1}^i$.
• Métrique : L'influence est quantifiée par le changement de performance sur l'ensemble de développement : $\Delta M_j^i = M(\theta_{t+1}^i; d_j^{val}) - M(\theta_t; d_j^{val})$.
• Récompense : Le score d'influence final $I_i$ est la moyenne de ces changements de performance à travers tous les jeux de données de développement. Ce score sert de signal de récompense pour augmenter le poids des jeux de données qui améliorent de manière démontrable les performances en aval.

3. Optimisation et efficacité

Pour répondre aux enjeux de scalabilité computationnelle, Inf-DDS utilise une mise à jour méta de premier ordre de type Reptile :

• Réutilisation du gradient : Au lieu de maintenir des états d'optimiseur séparés pour chaque jeu de données, la méthode calcule une moyenne pondérée des paramètres mis à jour $\bar{\theta}_{t+1} = \sum p_i \theta_{t+1}^i$ (où $p_i$ sont les probabilités dérivées des scores d'influence) et effectue une seule étape de mise à jour. Cela permet de réutiliser un ensemble unique de gradients intermédiaires et d'états d'optimiseur, réduisant ainsi considérablement la surcharge mémoire.
• Sous-échantillonnage : Lorsque le nombre de domaines $M$ est élevé, le scoreur $\psi$ est mis à jour sur un sous-échantillon aléatoire $S$ de jeux de données pour réduire le calcul par itération, tout en maintenant un estimateur non biaisé pour l'objectif conditionnel.

Contributions clés

1. Nouvel algorithme : La proposition d'Inf-DDS, une approche d'apprentissage par renforcement guidée par l'influence, qui apprend à ajuster les probabilités d'échantillonnage à travers divers jeux de données d'entraînement sans dépendre d'estimations de gradient bruitées pour les récompenses.
2. Efficacité computationnelle : La méthode est conçue pour être légère en termes de consommation de ressources GPU. En réutilisant les calculs de la distribution d'échantillonnage paramétrée pour mettre également à jour les paramètres du modèle, elle évite la lourde surcharge associée aux modèles proxys dans d'autres méthodes d'échantillonnage dynamique.
3. Validation empirique : Les auteurs valident Inf-DDS sur plusieurs benchmarks (BEIR, MLDR et le jeu de données d'embeddings Sentence-Transformers), démontrant qu'elle produit des trajectoires d'échantillonnage plus stables et atteint des performances supérieures ou comparables aux méthodes basées sur le gradient (MultiDDS, DoGE) et aux approches par modèle proxy (DoReMi).

Résultats expérimentaux

L'efficacité d'Inf-DDS a été évaluée dans trois scénarios distincts :

• BEIR (Adaptation de domaine) : Lors de l'optimisation pour l'ensemble de développement FEVER, Inf-DDS a systématiquement surpassé l'échantillonnage statique et MultiDDS à travers diverses températures d'initialisation. Elle a démontré des trajectoires d'échantillonnage stables, privilégiant correctement les jeux de données informatifs comme MS MARCO et FEVER, là où les méthodes basées sur le gradient présentaient des dynamiques instables.
• MLDR (Récupération de documents longs multilingues) : En utilisant le corpus BGE-M3, Inf-DDS a amélioré le modèle de base bge-m3-dense de 5,03 points NDCG@10 absolus. Notamment, elle a réussi à augmenter le poids des langues sous-représentées tout en évitant le surapprentissage sur les langues dominantes (anglais et chinois), atteignant la meilleure performance individuelle par langue dans 8 cas sur 13.
• Jeu de données d'embeddings Sentence-Transformers : Dans un contexte exigeant de 1 milliard de phrases parallèles réparties sur 32 domaines :
  • À partir d'une initialisation uniforme, Inf-DDS a atteint une performance proche du modèle Expert "prêt à l'emploi" (all-MiniLM-L6-v2), surpassant la référence uniforme de 1,83 point.
  • À partir des poids assignés par l'Expert, Inf-DDS a apporté une amélioration supplémentaire de 0,94 point NDCG@10 absolu.
  • En revanche, les méthodes basées sur le gradient (MultiDDS, DoGE) n'ont pas réussi à améliorer les références dans ces contextes.
  • Inf-DDS a réalisé ces gains avec 1,5× à 4× moins de calcul GPU que des méthodes comme DoReMi.

Signification et affirmations

L'article affirme qu'Inf-DDS offre une solution plus robuste et plus traitable sur le plan computationnel pour l'adaptation de domaine dans la récupération de texte. En passant de récompenses dérivées du gradient à des récompenses basées sur l'influence calculées via des modèles proxys en ligne, la méthode stabilise l'apprentissage des poids d'échantillonnage et fournit une estimation plus directe des sous-ensembles de données menant aux meilleurs minima atteignables. Les auteurs soulignent que, bien que l'algorithme ne converge pas toujours vers un optimum global, il délivre de manière fiable des améliorations substantielles à partir d'initialisations raisonnables. Ce travail souligne qu'une amélioration efficace des récupérateurs nécessite d'aller au-delà de la simple similitude de domaine ou des heuristiques statiques, nécessitant des stratégies adaptatives qui optimisent directement les métriques de performance en aval.