Résumé technique : SimReg - Atteindre des performances supérieures lors du pré-entraînement grâce à la régularisation de similarité des embeddings

1. Énoncé du problème

Les grands modèles de langage (LLM) reposent fondamentalement sur la prédiction du token suivant, un processus qui peut être considéré comme un problème de classification sur les représentations combinées du contexte précédent. Bien que le pré-entraînement avec une fonction de perte d'entropie croisée ait permis des avancées remarquables, les auteurs identifient une limitation critique : la nature dépendante du contexte des embeddings de tokens entraîne une forte variance intra-classe et une similarité inter-classe.

Contrairement à la classification d'images, où un objet spécifique (par exemple, un chat) correspond systématiquement à une étiquette et un embedding stables, les tokens linguistiques sont prédits en fonction de caractéristiques contextuelles diverses. Par conséquent, les embeddings prédisant le même token (par exemple, « murs » dans différentes phrases) peuvent varier considérablement en raison de leurs contextes distincts. L'article soutient que la perte d'entropie croisée standard échoue à imposer une forte cohérence à ces embeddings. L'analyse empirique révèle que, bien que l'entropie croisée contraigne rapidement la similarité contrastive au début de l'entraînement, elle atteint un plateau une fois un seuil de séparabilité de base atteint. Elle ne continue pas à favoriser l'élargissement des marges de classification ni à imposer une hétérogénéité suffisante parmi les embeddings de tokens, même lorsque la taille du modèle augmente. Cette limitation entrave l'efficacité de l'apprentissage de représentations et ralentit la convergence.

2. Méthodologie : SIMREG

Pour résoudre ces problèmes, les auteurs proposent SIMREG (Similarity Regularization), une fonction de perte conçue pour régulariser explicitement la similarité des embeddings lors du pré-entraînement à grande échelle.

Mécanisme central

Pour chaque token d'une séquence, la méthode partitionne tous les tokens en groupes positifs et négatifs sur la base de leurs étiquettes de vérité terrain :

• Groupe positif ($P_i$) : Tokens partageant la même étiquette de vérité terrain que le token ancre.
• Groupe négatif ($N_i$) : Tokens ayant des étiquettes différentes.

La perte SIMREG vise à minimiser la distance entre les embeddings du groupe positif tout en maximisant la séparation par rapport au groupe négatif. La perte pour le $i$-ème token est formulée comme suit :
$$L_i = L_{ce}^i + \lambda \cdot \text{softplus}(L_{sr}^i)$$
Où $L_{sr}^i$ est le terme de régularisation de similarité :
$$L_{sr}^i \triangleq \log \sum_{j \in N_i} \phi_{i,j} - \log \sum_{j \in P_i} \phi_{i,j}$$
Ici, $\phi$ représente la similarité cosinus entre les embeddings, mise à l'échelle par un coefficient de température $\tau$. Pour garantir la stabilité et la validité, la méthode inclut :

• Auto-similarité : Ajouter le token ancre à son propre groupe positif pour s'assurer qu'au moins une paire positive existe, même pour des tokens uniques dans une séquence.
• Mise à l'échelle Softplus : Garantir que la perte de régularisation reste non négative.
• Moyenne au niveau du groupe : Calculer la perte au niveau du groupe plutôt qu'au niveau de l'échantillon pour équilibrer les contributions entre les tokens et maintenir la stabilité sur de longues périodes d'entraînement.

Détails de l'implémentation

• Calcul par blocs (Chunk-wise) : Pour gérer la complexité $O(n^2)$ du calcul des similarités au sein de longues séquences et pour prendre en charge le parallélisme de séquence (SP) sans frais de communication excessifs, les auteurs divisent les séquences en blocs. La perte est calculée internement au sein des blocs, réduisant la complexité à $O(n^2/b)$ où $b$ est le nombre de blocs.
• Placement : Les études empiriques suggèrent que l'application de SIMREG à la dernière couche du modèle est la plus efficace, car les couches intermédiaires encodent des informations contextuelles entrelacées, tandis que la dernière couche projette ces informations dans des espaces sémantiques distincts utilisés pour la prédiction.

3. Analyse théorique

L'article fournit une justification théorique de l'efficacité de SIMREG en analysant son impact sur les marges de classification.

• Élargissement des marges : Les auteurs démontrent que la perte d'entropie croisée est bornée supérieurement par une fonction exponentielle de la marge de classification ($m_i$). En minimisant la perte SIMREG, le centre positif de l'embedding d'un token se déplace vers les échantillons de même classe, et le centre négatif s'éloigne.
• Borne de convergence : Les preuves théoriques montrent que ce processus augmente la marge de classification ($m_k$) d'au moins $\delta \|\epsilon^+\| + \sqrt{2}L_P \gamma$, où $\epsilon^+$ est la perturbation provenant de la perte de similarité. Cet élargissement de la marge entraîne une réduction multiplicative de la perte d'entropie croisée, accélérant théoriquement le processus de pré-entraînement.

4. Contributions clés

1. Proposition de SIMREG : Une nouvelle fonction de perte de régularisation de similarité d'embeddings qui encourage explicitement la cohérence intra-classe et la séparation inter-classe lors du pré-entraînement à grande échelle, répondant aux limites de l'entropie croisée seule.
2. Insights théoriques : Une analyse détaillée montrant comment SIMREG élargit les marges de multi-classification, réduisant ainsi la borne supérieure de la perte d'entropie croisée et améliorant la dynamique d'optimisation.
3. Implémentation évolutive : Une stratégie de calcul par blocs permettant une application efficace de la régularisation de similarité dans des contextes de parallélisme de séquence, la rendant réalisable pour des modèles à grande échelle.
4. Évaluation complète : Des expériences extensives sur des architectures denses (LLaMA-350M à 7B) et des mélanges d'experts (Mixtral-8×1B).

5. Résultats expérimentaux

Les auteurs ont évalué SIMREG sur des benchmarks standards en utilisant des modèles allant de 350M à 7B paramètres et le modèle MoE Mixtral-8×1B.

• Accélération de l'entraînement : SIMREG accélère systématiquement la convergence de l'entraînement. Sur les modèles LLaMA, il réalise une accélération de plus de 30 % (jusqu'à 38 % sur les modèles 7B) pour atteindre le même niveau de perte d'entropie croisée par rapport à la référence.
• Performance en aval : La méthode améliore la performance moyenne en zéro-shot sur les tâches en aval de plus de 1 % sur les benchmarks standards (par exemple, +1,14 % sur LLaMA-1,3B, +1,22 % sur LLaMA-7B).
• Séparabilité des embeddings : L'analyse montre que, tandis que l'entropie croisée seule conduit à un angle inter-classe moyen d'environ 61,3 degrés, l'ajout de SIMREG augmente cet angle à environ 74 degrés, indiquant une séparabilité des caractéristiques nettement meilleure.
• Efficacité : La perte SIMREG complète introduit une augmentation de moins de 2 % du temps d'exécution et moins de 1 % d'augmentation de la mémoire GPU. La variante par blocs est presque neutre en termes de ressources.
• Sensibilité aux hyperparamètres : La méthode est robuste sur une large gamme d'hyperparamètres. La température optimale $\tau$ est stable à 0,01, tandis que $\lambda$ s'adapte efficacement à la dimension d'embedding du modèle (environ proportionnelle à $\sqrt{d}$).

6. Signification et affirmations

L'article affirme que SIMREG offre une stratégie simple mais largement applicable pour améliorer l'efficacité et la généralisation du pré-entraînement des LLM. En imposant explicitement la cohérence des embeddings, il surmonte la stagnation de l'apprentissage de représentations observée dans l'entraînement standard par entropie croisée. Les auteurs positionnent ce travail comme une avancée pratique nécessitant un coût computationnel supplémentaire minimal tout en offrant des gains substantiels tant en vitesse d'entraînement qu'en performance des tâches en aval. Ils soulignent que la régularisation de cohérence, précédemment explorée principalement dans le fine-tuning ou des contextes à plus petite échelle, est hautement efficace lorsqu'elle est étendue à la phase de pré-entraînement des modèles à grande échelle.