ModernBERT, mais pas que : tendance RAG chez LightOn

LightOn, à fond sur le RAG ?

Quelques semaines avant son IPO, l'entreprise française avait communiqué une avancée en la matière. Son nom : MonoQwen-2-VL. Sa nature : un modèle encodeur destiné au reranking - donc, entre autres, à assister la récupération de documents. Il s'agit d'une version LoRA de Qwen2-VL-2B-Instruct, entraînée sur MonoT5 (le modèle répond "vrai" si un document est pertinent vis-à-vis d'une requête ; "faux" sinon).

MonoQwen-2-VL est censé accompagner le développement des pipelines RAG "visuels". En d'autres termes, qui se passent de représentation textuelle... et qui évitent ainsi potentiellement la perte d'informations. Les méthodes ColPali et DSE (Document Screenshot Embeddings) s'inscrivent dans cette tendance. L'une et l'autre utilisent, pour le reranking, une version LoRA du LLM destiné aux tâches de génération. Une option qu'a également choisie LightOn.

ModernBERT, une autre contribution pour le RAG

Une deuxième annonce couleur RAG est intervenue vers la mi-décembre. LightOn a présenté ModernBERT. Il s'agit d'une version affinée de BERT (modèle encodeur de référence, utilisé entre autres pour la récupération de documents). La fenêtre de contexte a été nettement étendue (8k) et l'entraînement a fait la part belle au code, ouvrant la voie à des usages au sein des IDE.

ModernBERT s'inspire de l'approche Transformer++, que Meta avait introduite avec les modèles LLaMA 2. LightOn a ainsi remplacé un certain nombre de briques de BERT par des équivalents "modernes". En particulier, sur l'encodage positionnel (intégration du RoPE) et la fonction d'activation (GeGLU à la place de MLP). Il a aussi ajouté une couche de normalisation après embedding, afin de stabiliser l'entraînement.

ModernBERT tire également parti de la deuxième génération de l'attention flash. Il mêle ainsi attention "globale" et attention "locale" : l'input complet n'est traité que toutes les trois couches. Les autres utilisent une fenêtre glissante où chaque token ne "surveille" que les 128 plus proches.
Une autre amélioration réduisant l'usage de ressources de calcul concerne la technique dite de padding. Celle-ci consiste à repérer la phrase la plus longue dans une séquence, puis à compléter les autres avec des tokens insignifiants afin qu'elle aient toutes la même longueur. Une condition indispensable pour les traiter en lot. Avec ModernBERT, on effectue une suppression de ces tokens suivie d'une concaténation en mini-batchs... et - élément spécifique à l'implémentation de LightOn - on n'effectue ce processus qu'une fois plutôt qu'à la volée. Une approche rendue possible, en particulier, par la prise en charge de RoPE dans le mécanisme d'attention flash.

Avant MonoQwen-2-VL et ModernBERT, une implémentation maison de Mamba

En amont, LightOn avait notamment expérimenté l'architecture Mamba, alternative à Transformers. Il en a résulté un LLM nommé Mambaoutai.

L'implémentation d'origine a été modifiée sur plusieurs aspects. Il a par exemple été décidé d'utiliser un plus grand état caché (permettant au modèle de compresser davantage de contexte) et de ne pas lier la matrice d'embedding à la dernière couche de progression (une question de simplicité).

LightOn a par ailleurs opté pour le planificateur WSD (Warmup-Stable-Decay). Avec lui, le taux d'apprentissage ne suit pas une fonction cosinus. En conséquence, la performance n'est pas, dans les grandes lignes, tributaire de la définition préalable du nombre d'étapes. Une contrainte en moins, donc, si on souhaite surentraîner le modèle.

Autre choix : pondérer le poids de la perte sur les premiers tokens des séquences. Une façon, estime LightOn, de diminuer le bruit dans les signaux d'apprentissage, du fait qu'il est plus difficile pour un modèle décodeur de prédire ces premiers tokens.

Illustration générée par IA