Résumé Technique : SPOKES – Optimiser la sélection de données de pré-entraînement diversifiées

1. Énoncé du Problème

La diversité des données est un facteur critique dans la construction des corpus de pré-entraînement, car elle réduit la redondance et améliore les performances en aval sous un budget de données fixe. Cependant, l'optimisation de la diversité est intrinsèquement difficile car il s'agit d'une propriété au niveau de l'ensemble (set-level property) dépendant des interactions entre les points de données plutôt que des exemples individuels. L'optimisation directe de la diversité est informatiquement intraitable en raison de la nature combinatoire de la sélection de sous-ensembles, particulièrement à l'échelle du pré-entraînement.

Les approches existantes reposent généralement sur des proxys ou des approximations, telles que le partitionnement grossier (par thèmes, compétences ou clusters non supervisés) ou la déduplication sémantique. Bien que ces méthodes assurent la représentation des clusters, elles reposent sur des partitions à basse résolution et échouent souvent à capturer l'information présente dans la longue traîne ou les connaissances sous-représentées. De plus, les méthodes actuelles traitent souvent les métriques de diversité (comme le score G-Vendi) comme des outils d'évaluation post-hoc plutôt que comme des objectifs d'optimisation. Le défi central adressé par ce travail est de savoir comment extraire de manière fiable des sous-ensembles diversifiés à partir de corpus de pré-entraînement à grande échelle pour améliorer les résultats en aval sans dépendre d'heuristiques qui ne parviennent pas à passer à l'échelle ou à capturer une diversité fine.

2. Méthodologie : SPOKES

Les auteurs introduisent SPOKES (Spokes: Optimizing for Diverse Pretraining Data Selection), un cadre évolutif qui optimise directement la diversité en utilisant le score G-Vendi dans un cadre de diversification probabiliste.

2.1 Objectif Central

La méthode formule la sélection de données comme un problème d'optimisation globale qui maximise conjointement la qualité et la diversité. Étant donné un ensemble de $n$ exemples, l'objectif est de sélectionner un sous-ensemble $S$ de taille $k$ qui maximise :
$$\max_{S \subseteq [n], |S|=k} \alpha \ln \left( \frac{1}{k} \sum_{i \in S} q_i \right) + (1 - \alpha) \ln \text{Vendi}(K_S)$$
Où :

• $q_i$ représente le score de qualité de l'exemple $i$ (par exemple, issu d'un classificateur FineWeb-Edu).
• $K_S$ est le noyau de similarité restreint au sous-ensemble sélectionné.
• $\alpha \in [0, 1]$ est un paramètre ajustable contrôlant le compromis entre qualité et diversité.

2.2 Diversité basée sur le Gradient (G-Vendi)

Contraal aux méthodes traditionnelles qui utilisent des plongements (embeddings) statiques, SPOKES opère dans l'espace des gradients. Il calcule les gradients $\nabla_\theta \ell(x_i)$ par rapport aux paramètres du modèle $\theta$ pour chaque point de donnée en utilisant un modèle proxy. Le score G-Vendi mesure la diversité en examinant l'entropie du spectre de la matrice de similarité construite à partir de ces gradients. Cette approche encourage l'orthogonalité entre les vecteurs de représentation, réduisant la redondance dans les directions de mise à jour et favorisant des étapes d'optimisation plus indépendantes et informatives.

2.3 Optimisation via la Descente de Gradient Exponentielle

Pour résoudre le problème de sélection de sous-ensembles combinatoires, SPOKES emploie une stratégie d'optimisation relaxée :

1. Relaxation Continue : Au lieu d'une sélection discrète, un vecteur de poids non négatif continu $w$ est assigné à chaque point de donnée, contraint au simplexe ($w \in \Delta_n$).
2. Construction du Noyau : Un noyau dépendant du poids $K(w)$ est défini où les entrées sont mises à l'échelle par $\sqrt{w_i w_j}$.
3. Optimisation : L'objectif est optimisé via une descente de gradient exponentielle. À chaque itération, les poids sont mis à jour via $w_i^{(t+1)} \propto w_i^{(t)} \exp(-\eta g_i^{(t)})$ et normalisés pour éviter l'effondrement sur quelques exemples à poids élevé.
4. Récupération Discrète : Après convergence, le sous-ensemble final est récupéré en sélectionnant les $k$ entrées les plus élevées du vecteur de poids appris.

2.4 Améliorations de la Scalabilité

Pour rendre cette approche réalisable à l'échelle du pré-entraînement, les auteurs implémentent plusieurs stratégies d'efficacité :

• Gradients Tronqués : Au lieu de calculer les gradients complets du modèle, la méthode les approxime en utilisant uniquement les dernières couches $n$ du transformer (typiquement 2 couches). Cela réduit considérablement le coût computationnel tout en préservant la structure du noyau de similarité.
• Projection de Johnson-Lindenstrauss (JL) : Pour gérer les vecteurs de gradient de haute dimension (ex: 200 millions de paramètres), les auteurs appliquent des projections aléatoires (matrices de Rademacher) pour réduire la dimensionnalité. Les résultats empiriques montrent que même des projections de faible dimension (ex: $k=1024$) maintiennent une haute stabilité du sous-ensemble (95,5 % de chevauchement avec les références de haute dimension).
• Parallélisation : Le jeu de données peut être partitionné en sous-ensembles disjoints optimisés indépendamment, offrant une accélération linéaire.

3. Contributions Clés

1. Cadre SPOKES : Introduction d'une technique scalable et rigoureuse pour l'optimisation de la diversité qui opère efficacement à l'échelle des jeux de données de pré-entraînement modernes, dépassant les heuristiques de clustering.
2. Optimisation Directe de G-Vendi : Démonstration que le score G-Vendi peut être directement optimisé via la descente de gradient exponentielle pour extraire des sous-ensembles hautement diversifiés, plutôt que de servir simplement de métrique post-hoc.
3. Optimisation Conjointe Qualité-Diversité : Preuve que l'optimisation conjointe de la qualité et de la diversité produit des résultats supérieurs par rapport à l'optimisation d'un seul objectif ou à l'utilisation de baselines standards.
4. Améliorations d'Efficacité : Validation que le calcul des gradients tronqués et la réduction de dimensionnalité via la projection JL permettent de passer à l'échelle sur des corpus massifs sans coûts de calcul prohibitifs.

4. Résultats Expérimentaux

La méthode a été évaluée sur deux corpus de pré-entraînement à grande échelle : DCLM (sous-ensemble Dolmino, filtré de manière agressive) et FineWeb (filtré de manière moins agressive). Des modèles ont été pré-entraînés à partir de zéro sur une architecture LLaMA-1B avec 100 milliards de tokens.

4.1 Métriques de Diversité

• Augmentation du Score G-Vendi : Sur un sous-ensemble de 500k échantillons, SPOKES a obtenu une augmentation de +489 points du score G-Vendi par rapport à l'échantillonnage aléatoire. En revanche, la déduplication sémantique (SemDeDup) n'a obtenu qu'une augmentation de +7 points.
• Diversité des Batches : Tout au long de l'entraînement, les batches sélectionnés par SPOKES ont maintenu des scores de diversité moyenne systématiquement plus élevés (passant de 315 à 520 sur DCLM) par rapport à l'échantillonnage aléatoire.

4.2 Performance en Aval

L'évaluation a été menée sur 10 benchmarks anglais de la suite OLMES.

• Diversité Seule : SPOKES (diversité uniquement) a amélioré la performance moyenne en aval de +0,4 (DCLM) et +0,5 (FineWeb) par rapport à l'échantillonnage aléatoire.
• Optimisation Conjointe : Les meilleurs résultats ont été obtenus en optimisant conjointement la qualité et la diversité. SPOKES a atteint des gains de +1,5 (DCLM) et +1,4 (FineWeb) points par rapport à l'échantillonnage aléatoire, surpassant toutes les baselines, y compris la déduplication sémantique et le filtrage de qualité uniquement.
• Forces Complémentaires : Les résultats indiquent que la qualité et la diversité capturent des propriétés distinctes mais complémentaires. Alors que le filtrage de qualité excelle sur les benchmarks intensifs en connaissances (ex: MMLU), l'optimisation de la diversité améliore le raisonnement de bon sens (ex: HellaSwag). L'optimisation conjointe comble ce fossé, offrant des performances compétitives sur toutes les tâches.

5. Signification et Revendications

L'article affirme que l'optimisation directe de la diversité est à la fois traitable à grande échelle et pratiquement bénéfique. En s'élochant des approximations grossières et des heuristiques, SPOKES démontre que :

• La diversité est un axe critique : Même sans supervision explicite de la qualité, l'optimisation de la diversité seule produit des gains constants, suggérant qu'une représentation de gradient diversifiée de la distribution des données améliore la généralisation.
• Les métriques de qualité sont insuffisantes : Les métriques de qualité standard ne capturent pas pleinement l'utilité des données. L'optimisation conjointe de la diversité peut fournir des signaux d'entraînement complémentaires qui mènent à de meilleures performances en aval, même lorsque le sous-ensemble sélectionné possède un score de qualité moyenne inférieur à celui d'une baseline de qualité pure.
• Scalabilité : La méthode proposée extrait avec succès des sous-ensembles denses et hautement diversifiés à partir de corpus massifs, offrant une alternative viable aux pipelines de sélection existants qui reposent sur la déduplication sémantique ou l'échantillonnage aléatoire.

Les auteurs positionnent SPOKES comme un outil de construction de sous-ensembles de haute qualité et diversifiés, soulignant que les bénéfices de la diversité se traduisent de manière fiable à travers différents jeux de données avec des régimes de filtrage et des profils de bruit variés.