Hungry Hungry Hippos : Vers une modélisation du langage basée sur des modèles à espace d’états

Les modèles à espace d’état (SSM) ont démontré des performances de pointe dans la modélisation de séquences dans certaines modalités, mais peinent à rivaliser avec les modèles à attention en modélisation du langage. En outre, malgré une croissance quasi linéaire de la complexité en fonction de la longueur de la séquence (au lieu d’une croissance quadratique comme pour les Transformers), les SSM restent plus lents en pratique en raison d’une utilisation médiocre des ressources matérielles. Dans cet article, nous faisons des progrès dans la compréhension de l’écart d’expressivité entre les SSM et les modèles à attention en modélisation du langage, ainsi que dans la réduction de la barrière matérielle entre ces deux approches. Premièrement, nous utilisons des tâches de modélisation du langage synthétique pour explorer cet écart. Nous constatons que les SSM existants éprouvent des difficultés à deux capacités fondamentales : rappeler des tokens antérieurs dans la séquence et comparer des tokens à différentes positions. Pour mieux comprendre l’impact de ces limites sur la modélisation du langage, nous proposons une nouvelle couche SSM, nommée H3, conçue explicitement pour ces deux fonctionnalités. H3 atteint les performances des modèles à attention sur les langages synthétiques et se rapproche de 0,4 point de perplexité (PPL) des Transformers sur OpenWebText. En outre, un modèle hybride de 125 millions de paramètres combinant H3 et attention, qui conserve seulement deux couches à attention, surpasse les Transformers sur OpenWebText de 1,0 point de PPL. Ensuite, afin d’améliorer l’efficacité d’entraînement des SSM sur les architectures matérielles modernes, nous introduisons FlashConv. FlashConv utilise un algorithme FFT fusionné par blocs pour améliorer l’efficacité sur des séquences allant jusqu’à 8K, et propose un nouvel algorithme de passage d’état exploitant les propriétés récurrentes des SSM afin de permettre une mise à l’échelle efficace sur des séquences plus longues. FlashConv permet une accélération de 2× sur le benchmark Long-Range Arena et permet aux modèles hybrides de générer du texte 2,4 fois plus vite que les Transformers. En utilisant FlashConv, nous avons échelonné des modèles hybrides H3-attention jusqu’à 2,7 milliards de paramètres sur le corpus Pile, obtenant des résultats prometteurs : une perplexité inférieure à celle des Transformers, et une performance supérieure sur la majorité des tâches du benchmark SuperGLUE en apprentissage zéro et à faibles exemples.