Differential Mamba

papers.abstract

Sequenzmodelle wie Transformers und RNNs überbeanspruchen häufig die Aufmerksamkeit für irrelevante Kontexte, was zu verrauschten Zwischendarstellungen führt. Dies beeinträchtigt die Fähigkeiten von LLMs, indem es Halluzinationen fördert, die langreichweitigen und Abruffähigkeiten schwächt und die Robustheit verringert. Jüngste Arbeiten haben gezeigt, dass differenzielle Gestaltung dieses Problem in Transformern mildern kann, wodurch deren Effektivität in verschiedenen Anwendungen verbessert wird. In diesem Artikel untersuchen wir, ob diese Techniken, die ursprünglich für Transformer entwickelt wurden, auf Mamba angewendet werden können, eine neuere Architektur, die auf selektiven Zustandsraum-Schichten basiert und Transformator-ähnliche Leistung mit höherer Effizienz erreicht. Wir zeigen, dass eine naive Anpassung der differenziellen Gestaltung auf Mamba unzureichend ist und sorgfältige architektonische Modifikationen erfordert. Um dies zu adressieren, führen wir einen neuartigen differenziellen Mechanismus für Mamba ein, der empirisch anhand von Sprachmodellierungs-Benchmarks validiert wurde und verbesserte Abruffähigkeiten sowie überlegene Leistung gegenüber dem Standard-Mamba demonstriert. Schließlich führen wir umfangreiche Ablationsstudien und empirische Analysen durch, um unsere Designentscheidungen zu rechtfertigen und Belege dafür zu liefern, dass unser Ansatz das Problem der Überbeanspruchung in Mamba-basierten Modellen effektiv mildert. Unser Code ist öffentlich verfügbar.

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.