Résumé technique : Quand la distance distrait : Biais de distance de représentation dans la perte BT (BT-Loss) pour les modèles de récompense

1. Énoncé du problème

Les modèles de récompense (RM) sont centraux dans l'apprentissage par renforcement à partir de rétroaction humaine (RLHF), servant de substitut aux préférences humaines pour guider l'alignement des grands modèles de langage (LLM). L'objectif standard pour l'entraînement de ces modèles est la perte de Bradley-Terry (BT), qui apprend à partir de données par paires consistant en des réponses « choisies » ($y_w$) et « rejetées » ($y_l$).

Les auteurs identifient une faille critique dans la dynamique de mise à jour de la perte BT standard : le Biais de distance de représentation. Alors que le gradient de la BT est censé être piloté par l'erreur de prédiction (l'ampleur de l'erreur de classement), l'analyse révèle que la norme du gradient dépend de deux composantes distinctes :

1. Erreur de prédiction : La différence entre les récompenses prédites pour les réponses choisie et rejetée.
2. Distance de représentation : La distance entre les deux réponses dans l'espace de sortie de la dernière couche ($\|h_w - h_l\|$).

Ce couplage crée un signal d'apprentissage fallacieux où l'ampleur de la mise à jour est dictée par la distance entre les représentations des deux éléments, plutôt que uniquement par l'ampleur de l'erreur commise par le modèle.

• Paires à faible distance : Les paires qui sont sémantiquement ou structurellement similaires (par exemple, un extrait de code correct versus un avec une seule erreur logique) ont souvent de faibles distances de représentation. Sous la perte BT standard, ces paires reçoivent des mises à jour extrêmement faibles, même lorsque le modèle les classe mal. Cela est problématique car les distinctions fines sont souvent les plus critiques pour des tâches comme le raisonnement et le codage.
• Paires à grande distance : Les paires qui sont très différentes (par exemple, un refus versus une réponse à une requête malveillante) ont naturellement de grandes distances de représentation. Ces paires reçoivent des mises à jour disproportionnellement fortes, pouvant éclipser les signaux d'apprentissage des paires plus subtiles et fines.

Par conséquent, le modèle sous-entraîne les distinctions fines et difficiles où il rencontre le plus de difficultés, tout en surestimant les différences superficielles.

2. Méthodologie : NormBT

Pour remédier à ce biais, les auteurs proposent NormBT, un schéma de normalisation par paire adaptatif qui redimensionne la contribution du gradient de chaque paire de préférences.

Base théorique

Les auteurs décomposent la norme du gradient BT $\|\nabla_\theta L_{BT}\|$ et montrent qu'elle est proportionnelle à :
$$\|\nabla_\theta L_{BT}\| \approx |\sigma(d) - 1| \cdot \|h_w - h_l\|$$
où $d$ est la différence de récompense et $\|h_w - h_l\|$ est la distance de représentation. Pour que la taille de la mise à jour dépende principalement du terme d'erreur de prédiction $|\sigma(d) - 1|$, les auteurs introduisent un poids $w_i$ qui met à l'échelle inverse la distance de représentation.

Détails d'implémentation

Le calcul direct de la norme exacte du gradient pour la pondération est coûteux en termes de calcul. Au lieu de cela, NormBT utilise la distance de représentation de la dernière couche ($\|h_w - h_l\|$) comme un substitut, qui est montré comme étant un substitut exact pour la tête de score linéaire et une forte approximation pour le modèle complet.

La fonction de perte proposée est :
$$\mathcal{L}_{NormBT}(\theta) = -\mathbb{E}_{(x, y_w, y_l) \sim \mathcal{D}} \left[ \tilde{w}(y_w, y_l) \cdot \log \sigma(r_w - r_l) \right]$$

Où le poids $\tilde{w}$ est défini par :
$$\tilde{w}(y_w, y_l) = \frac{\mu_t}{\|h_w - h_l\| + \epsilon}$$

Choix de conception clés :

• Stop-Gradient : Le terme de distance est calculé avec un opérateur de blocage de gradient (stop-gradient) pour empêcher le poids de normalisation lui-même d'être optimisé.
• Stabilisation par EMA : Pour éviter la dérive d'échelle pendant l'entraînement (car les magnitudes d'embedding peuvent varier), les auteurs utilisent une moyenne mobile exponentielle (EMA) de la distance de représentation moyenne par lot ($\mu_t$) pour normaliser les poids. Cela garantit que les poids restent stables et adaptables aux changements de distribution.
• Intégration légère : NormBT ne nécessite aucun changement architectural, préserve le fondement probabiliste de la BT et n'entraîne qu'un surcoût computationnel négligeable.

3. Principales contributions

1. Analyse du gradient : L'article fournit une analyse théorique et empirique détaillée montrant que les normes de gradient de la BT sont déterminées conjointement par l'erreur de prédiction et la distance de représentation, menant à des signaux d'apprentissage biaisés.
2. Identification du biais : Les auteurs démontrent que ce couplage fait que les paires à faible distance (souvent les plus critiques pour les tâches fines) reçoivent des mises à jour faibles, tandis que les paires à grande distance dominent le signal d'apprentissage.
3. Proposition de NormBT : Ils introduisent NormBT, une modification directe (drop-in) qui normalise les mises à jour par la distance de représentation, garantissant que les signaux d'apprentissage se concentrent sur les erreurs de prédiction.
4. Validation empirique : Des expériences approfondies sur différents backbones de LLM (Gemma-2b-it, Llama-3.2-3B) et jeux de données (Unified-Feedback, Skywork-Reward) montrent des améliorations de performance constantes.

4. Résultats expérimentaux

Les auteurs ont évalué NormBT par rapport à la ligne de base BT standard et à plusieurs variantes (BT + Margin, BT + Label Smoothing) sur le benchmark RewardBench et sur des tâches de sélection Best-of-N (BoN) en aval.

• Performance sur RewardBench : NormBT a systématiquement surpassé la ligne de base BT dans tous les réglages. Les gains les plus importants ont été observés dans la catégorie Reasoning (Raisonnement), où la précision s'est améliorée de plus de 5 % en moyenne. Cela corrobore l'hypothèse que les tâches de raisonnement reposent sur des distinctions fines (faibles distances de représentation) qui sont supprimées par la BT standard.
• RLHF en aval (BoN) : Dans les expériences de sélection Best-of-N, les modèles entraînés avec NormBT ont obtenu des « scores d'or » (qualité des réponses sélectionnées) plus élevés par rapport aux bases BT.
• Études d'ablation :
  • NormBT a surpassé d'autres techniques d'entraînement comme la LayerNorm et le Gradient Clipping, qui n'ont pas réussi à traiter le problème de pondération relative entre les paires.
  • La méthode s'est avérée complémentaire, améliorant les performances même lorsqu'elle est appliquée sur d'autres variantes de la BT (ex: NormBT + Margin).
  • Le retrait de la stabilisation par EMA ou l'utilisation d'autres métriques de similitude (similitude cosinus, embeddings avec pooling moyen) a entraîné une dégradation des performances, validant les choix de conception spécifiques.

5. Signification et affirmations

L'article affirme que NormBT offre une solution fondée sur des principes et légère à une limitation structurelle de la perte BT standard. En découplant la magnitude de la mise à jour de la distance de représentation, NormBT garantit que le modèle de récompense apprend de manière plus fidèle les distinctions fines, qui sont souvent les plus difficiles mais aussi les plus critiques pour un alignement de haute qualité.

Les auteurs soulignent que leur travail ne traite pas directement de la qualité des données ou du biais des annotateurs, mais améliore l'efficacité et la fidélité des signaux d'apprentissage extraits des données par paires existantes. Ils positionnent NormBT comme une amélioration fondamentale pour le paradigme de modélisation de récompense discriminative, pouvant mener à des signaux de récompense plus fiables dans les applications de RLHF en aval, particulièrement dans les domaines exigeant un raisonnement précis et une détection subtile d'erreurs.