Mamba Différentielle

Differential Mamba

July 8, 2025

papers.authors: Nadav Schneider, Itamar Zimerman, Eliya Nachmani

cs.AI

papers.abstract

Les modèles séquentiels comme les Transformers et les RNNs allouent souvent une attention excessive à un contexte non pertinent, ce qui entraîne des représentations intermédiaires bruyantes. Cela dégrade les capacités des LLM en favorisant les hallucinations, en affaiblissant les capacités de traitement à long terme et de récupération, et en réduisant la robustesse. Des travaux récents ont montré qu'une conception différentielle peut atténuer ce problème dans les Transformers, améliorant ainsi leur efficacité dans diverses applications. Dans cet article, nous explorons si ces techniques, initialement développées pour les Transformers, peuvent être appliquées à Mamba, une architecture récente basée sur des couches d'espace d'état sélectives qui atteint des performances comparables aux Transformers avec une plus grande efficacité. Nous montrons qu'une adaptation naïve de la conception différentielle à Mamba est insuffisante et nécessite des modifications architecturales minutieuses. Pour résoudre ce problème, nous introduisons un nouveau mécanisme différentiel pour Mamba, validé empiriquement sur des benchmarks de modélisation du langage, démontrant des capacités de récupération améliorées et une performance supérieure à celle de Mamba standard. Enfin, nous menons des études d'ablation approfondies et des analyses empiriques pour justifier nos choix de conception et fournir des preuves que notre approche atténue efficacement le problème de surallocation dans les modèles basés sur Mamba. Notre code est disponible publiquement.

English

Sequence models like Transformers and RNNs often overallocate attention to irrelevant context, leading to noisy intermediate representations. This degrades LLM capabilities by promoting hallucinations, weakening long-range and retrieval abilities, and reducing robustness. Recent work has shown that differential design can mitigate this issue in Transformers, improving their effectiveness across various applications. In this paper, we explore whether these techniques, originally developed for Transformers, can be applied to Mamba, a recent architecture based on selective state-space layers that achieves Transformer-level performance with greater efficiency. We show that a naive adaptation of differential design to Mamba is insufficient and requires careful architectural modifications. To address this, we introduce a novel differential mechanism for Mamba, empirically validated on language modeling benchmarks, demonstrating improved retrieval capabilities and superior performance over vanilla Mamba. Finally, we conduct extensive ablation studies and empirical analyses to justify our design choices and provide evidence that our approach effectively mitigates the overallocation problem in Mamba-based models. Our code is publicly available.