Résumé Technique : LaneRoPE

Énoncé du Problème

Les techniques de mise à l'échelle au moment de l'inférence en parallèle, telles que l'échantillonnage best-of-N, améliorent la précision des modèles de langage de grande taille (LLM) en générant plusieurs séquences conditionnées par le même prompt d'entrée. Bien que ces méthodes tirent parti de l'efficacité computationnelle de l'inférence par lots, elles fonctionnent traditionnellement avec un parallélisme non coordonné. Dans les approches standard, chaque séquence d'un lot est générée indépendamment, ce qui signifie qu'elles ne réutilisent pas les générations intermédiaires, les calculs ou les observations provenant d'autres séquences. Cette indépendance interdit l'exploitation des structures de problèmes décomposables, limite le partage des sorties intermédiaires et restreint la génération de réponses diversifiées et collaboratives.

Les approches « collaboratives » existantes tentent de remédier à cela en permettant une visibilité partagée sur les générations intermédiaires. Cependant, de nombreuses méthodes antérieures reposent sur la reparamétrisation des positions pour ordonner virtuellement plusieurs séquences en une seule longue séquence. Les auteurs soutiennent que cette approche est intrinsèquement un problème de génération multi-séquence où l'inférence doit refléter avec précision les positions le long de la dimension des tokens et de la dimension de la séquence (ou « voie »).

Méthodologie : LaneRoPE

L'article propose LaneRoPE, une méthode permettant la coordination et la collaboration entre $N > 1$ séquences au moment de la génération avec des modifications minimales de l'architecture sous-jacente du LLM. L'approche repose sur deux innovations techniques clés :

1. Masquage de l'Attention Inter-Séquence :
   LaneRoPE modifie le mécanisme d'attention causale standard pour permettre aux séquences d'échantillonnage d'être dépendantes les unes des autres. Au lieu de conditionner le token suivant $x^{(n)}_{i+1}$ uniquement sur son propre historique $x^{(n)}_{1:i}$, le modèle le conditionne sur les historiques de toutes les séquences $\{x^{(m)}_{1:i}\}_{m=1}^N$. Cela est réalisé via des masques d'attention causale inter-séquence, rendant les traces de raisonnement intermédiaires visibles à travers toutes les voies parallèles.

2. Encodage Positionnel Conscient de la Séquence (LaneRoPE) :
   L'Encodage Positionnel Rotatif (RoPE) standard ne peut pas distinguer les tokens au même pas de temps à travers différentes séquences car ils partagent des indices positionnels identiques. LaneRoPE étend le RoPE en introduisant une seconde matrice de rotation bloc-diagonale, $R_{\Omega, m}$, qui encode l'index de séquence $m$ (la « voie »).

   • Mécanisme : Les vecteurs de requête et de clé sont rotatés à la fois par l'index de token $i$ et l'index de voie $m$ : $f_{\{q,k\}}(x_i, i, m) = R_{\Omega, m} R_{\Theta, i} (W x_i + b)$.
   • Propriétés : Cette formulation traite la grille conjointe (token, séquence) comme une base de Fourier 2D. Crucialement, car la matrice de rotation de voie est orthogonale, les poids d'attention entre les tokens au sein de la même séquence restent inchangés, assurant la compatibilité ascendante. Pour $N=1$, LaneRoPE se simplifie en RoPE standard.
   • Stratégies d'Initialisation : Les auteurs fournissent des stratégies pour initialiser LaneRoPE à partir de modèles pré-entraînés afin d'imiter les comportements existants :
     • GroupThink : En fixant les fréquences de voie $\Omega = K\Theta$ (où $K$ est un écart de séquence), LaneRoPE récupère l'« ordre virtuel » de GroupThink.
     • Échantillonnage Indépendant : En augmentant les couches de projection avec un biais de Fourier, le modèle peut être initialisé pour supprimer l'attention inter-voies, simulant efficacement un échantillonnage indépendant.
     • Correction Sensible NTK : Pour atténuer les problèmes liés aux positions relatives virtuelles négatives dans l'initialisation GroupThink, les auteurs introduisent une fonction d'échelle qui atténue les rotations de voie de basse fréquence, préservant la capacité du modèle à encoder les positions absolues.

Pipeline d'Entraînement

Pour susciter des interactions inter-voies, les auteurs proposent deux stratégies de fine-tuning :

1. Fine-Tuning Supervisé (SFT) : Un jeu de données est synthétisé où un LLM joue le rôle de $N$ assistants différents collaborant sur un problème. Les assistants sont incités à voir les pensées partielles des autres et à éviter le travail redondant. Le jeu de données final se compose de traces individuelles de ces conversations, forçant le modèle à apprendre les dépendances inter-séquences.
2. Optimisation Kahneman-Tversky (KTO) : Le modèle est entraîné sur des lots de traces de raisonnement (certaines correctes, d'autres incorrectes) générées indépendamment par le modèle de base. La perte KTO maximise la vraisemblance des complétions correctes et minimise celle des incorrectes, encourageant le modèle à apprendre une distribution conjointe non triviale où les voies parallèles peuvent se corriger ou se compléter mutuellement.

Résultats

Les auteurs évaluent LaneRoPE sur des benchmarks de raisonnement mathématique (MATH500, AMC 23, AIME 24/25) en utilisant des modèles DeepSeek-R1-Distill-Qwen (1,5B et 7B paramètres).

• Performance : LaneRoPE démontre des améliorations fiables de la précision par rapport aux baselines séquentielles et à l'échantillonnage indépendant. Les modèles 7B montrent des gains constants, en particulier lors de l'utilisation de l'entraînement KTO avec une initialisation sensible NTK et des fréquences apprenables.
• Comparaison : LaneRoPE surpasse généralement la baseline « Bridge » (une autre méthode de génération interdépendante) sur les modèles 7B, attribué à la plus grande capacité de LaneRoPE à partager l'information et à son architecture plus efficace.
• Petits Modèles : Le modèle 1,5B montre des résultats mitigés, bénéficiant moins de l'initialisation GroupThink non entraînée et montrant une mise à l'échelle plus faible avec l'augmentation du nombre de voies, suggérant une limitation de la taille du modèle de base plutôt que de l'architecture.
• Efficacité : LaneRoPE introduit une surcharge négligeable. Les mesures de temps d'inférence montrent une augmentation d'environ 6 % par rapport aux transformateurs standard, principalement due au pré/post-traitement plutôt qu'au mécanisme d'attention lui-même. En revanche, la baseline Bridge entraîne une augmentation allant jusqu'à 25 % du temps d'inférence en raison de couches supplémentaires.

Signification et Revendications

L'article affirme que LaneRoPE représente une étape prometteuse vers l'activation du raisonnement collaboratif parallèle en exploitant fondamentalement les positions relatives des tokens à l'intérieur et à l'extérieur d'une séquence.

• Efficacité Architecturale : Contrairement aux travaux antérieurs nécessitant des moteurs d'exécution spécialisés ou des mécanismes de branchement complexes, LaneRoPE nécessite des modifications architecturales minimales et introduit une surcharge d'inférence négligeable, le rendant adapté à une intégration rapide dans les pipelines LLM existants.
• Flexibilité : La formulation est suffisamment expressive pour englober les techniques existantes (comme l'échantillonnage indépendant et GroupThink) tout en permettant un fine-tuning ultérieur pour découvrir de nouvelles stratégies collaboratives.
• Efficacité des Données : La méthode peut être débloquée avec très peu de paramètres supplémentaires (< 0,5 % par rapport aux modèles de base) en utilisant un SFT simple ou KTO avec des données synthétiques.

Les auteurs reconnaissent des limites, notant qu'ils n'emploient pas encore de schéma d'entraînement dédié pour fusionner les sorties de plusieurs voies et que leurs expériences actuelles se concentrent sur de petits lots ($N \le 4$). Ils suggèrent que des travaux futurs pourraient explorer l'apprentissage par renforcement avec feedback vérifiable (RLVF) pour renforcer davantage les comportements collaboratifs.