Juste lire deux fois : combler le fossé de rappel pour les modèles de langage récurrents

Des grands modèles de langage récurrents qui concurrencent les Transformers en perplexité de modélisation du langage émergent à un rythme rapide (par exemple, Mamba, RWKV). De manière enthousiasmante, ces architectures utilisent une quantité constante de mémoire lors de l'inférence. Cependant, en raison de la mémoire limitée, les modèles récurrents ne peuvent pas se souvenir et utiliser toutes les informations dans de longs contextes, ce qui conduit à une qualité fragile de l'apprentissage en contexte (ICL). Un défi clé pour les modèles efficaces est de sélectionner quelles informations stocker versus lesquelles ignorer. Dans ce travail, nous observons que l'ordre dans lequel l'information est montrée au modèle impacte la difficulté de sélection. Pour formaliser cela, nous montrons que la difficulté du rappel d'information se réduit à la difficulté d'un problème appelé disjonction d'ensembles (SD), un problème fondamental en complexité de communication qui nécessite qu'un algorithme en flux (par ex., modèle récurrent) décide si les ensembles d'entrée sont disjoints. Nous montrons empiriquement et théoriquement que la mémoire récurrente requise pour résoudre SD change avec l'ordre des ensembles, c'est-à-dire, si le plus petit ensemble apparaît en premier dans le contexte. Notre analyse suggère que, pour réduire la dépendance à l'ordre des données, nous pouvons mettre l'information dans le bon ordre ou traiter les invites de manière non causale. Dans cette optique, nous proposons : (1) JRT-Prompt, où le contexte est répété plusieurs fois dans l'invite, montrant ainsi au modèle tous les ordres de données. Cela apporte une amélioration moyenne de 11,0 ± 1,3 points, calculée sur 16 modèles récurrents et 6 tâches ICL, avec un débit 11,9 fois supérieur à FlashAttention-2 pour le pré-remplissage lors de la génération (longueur 32k, taille de lot 16, NVidia H100). Nous proposons ensuite (2) JRT-RNN, qui utilise une attention linéaire préfixe non causale pour traiter les invites et fournit 99 % de la qualité Transformer à 360M paramètres, 30B tokens, et 96 % à 1,3B paramètres, 50B tokens en moyenne sur les tâches, avec un débit 19,2 fois supérieur à FA2 pour le pré-remplissage.