Youjin Wang Yangjingyi Chen Jiahao Yan Jiaxuan Lu Xiao Sun

Résumé

Face à la croissance exponentielle des données, la modélisation de séquences longues est devenue de plus en plus essentielle dans des domaines tels que le traitement du langage naturel et la bioinformatique. Toutefois, les méthodes existantes sont confrontées à un compromis inhérent entre efficacité et consommation mémoire. Les réseaux de neurones récurrents souffrent du problème de disparition et d’explosion du gradient, ce qui les rend difficiles à mettre à l’échelle. Les Transformers peuvent modéliser des dépendances globales, mais leur complexité est quadratique. Récemment, les modèles à espace d’état sélectifs, tels que Mamba, ont démontré une haute efficacité avec une complexité temporelle linéaire O(n) et une inférence récurrente à coût constant O(1), tout en présentant toutefois une dégradation exponentielle de la mémoire à long terme. Dans ce travail, nous menons des dérivations mathématiques et une analyse informationnelle systématique afin d’élucider de manière approfondie le mécanisme de dégradation de la mémoire dans Mamba, répondant ainsi à une question fondamentale : quelle est la nature de la mémoire à long terme de Mamba, et comment ce modèle parvient-il à conserver l’information ? Pour quantifier la perte d’information clé, nous introduisons par ailleurs des métriques de fidélité de mémoire horizontale et verticale, capables de capturer la dégradation à la fois à l’intérieur des couches et entre elles. Inspirés par la manière dont les humains distillent et conservent les informations essentielles lorsqu’ils lisent des documents longs, nous proposons MemMamba, un nouveau cadre architectural intégrant un mécanisme de résumé d’état, ainsi qu’une attention inter-couche et inter-token, qui atténue le phénomène d’oubli à long terme tout en préservant une complexité linéaire. MemMamba obtient des améliorations significatives par rapport aux variantes existantes de Mamba et aux Transformers sur des benchmarks de séquences longues tels que PG19 et Passkey Retrieval, tout en offrant un gain de vitesse de 48 % en efficacité d’inférence. L’analyse théorique et les résultats expérimentaux démontrent que MemMamba réalise une percée dans le compromis complexité-mémoire, ouvrant ainsi une nouvelle voie pour la modélisation de séquences ultra-longues.