Couches d'espace d'état simplifiées pour la modélisation de séquences

Les modèles utilisant des couches de séquence à espace d’état structuré (S4) ont atteint des performances de pointe sur les tâches de modélisation de séquences à longue portée. Une couche S4 combine des modèles d’espace d’état linéaires (SSM), le cadre HiPPO et l’apprentissage profond afin d’obtenir une haute performance. Nous nous appuyons sur la conception de la couche S4 pour introduire une nouvelle couche d’espace d’état, la couche S5. Alors qu’une couche S4 repose sur de nombreux SSM indépendants à entrée unique et sortie unique, la couche S5 utilise un seul SSM à entrées multiples et sorties multiples. Nous établissons un lien entre S5 et S4, et exploitons cette relation pour développer une initialisation et une paramétrisation adaptées au modèle S5. Le résultat est une couche d’espace d’état capable d’exploiter efficacement des balayages parallèles largement implémentés, permettant à S5 de rivaliser avec la performance computationnelle de S4, tout en atteignant des performances de pointe sur plusieurs tâches de modélisation de séquences à longue portée. S5 atteint en moyenne 87,4 % sur le benchmark Long Range Arena, et 98,5 % sur la tâche la plus difficile, Path-X.