Differentiële Mamba

Differential Mamba

July 8, 2025

Auteurs: Nadav Schneider, Itamar Zimerman, Eliya Nachmani

cs.AI

Samenvatting

Sequentiemodellen zoals Transformers en RNN's besteden vaak te veel aandacht aan irrelevante context, wat leidt tot ruis in de tussenliggende representaties. Dit vermindert de capaciteiten van LLM's door hallucinaties te bevorderen, het vermogen voor langeafstandsrelaties en retrieval te verzwakken, en de robuustheid te verminderen. Recent onderzoek heeft aangetoond dat differentieel ontwerp dit probleem in Transformers kan mitigeren, waardoor hun effectiviteit in diverse toepassingen verbetert. In dit artikel onderzoeken we of deze technieken, oorspronkelijk ontwikkeld voor Transformers, kunnen worden toegepast op Mamba, een recente architectuur gebaseerd op selectieve state-space lagen die Transformer-niveau prestaties behaalt met grotere efficiëntie. We laten zien dat een naïeve aanpassing van differentieel ontwerp aan Mamba onvoldoende is en zorgvuldige architectonische aanpassingen vereist. Om dit aan te pakken, introduceren we een nieuw differentieel mechanisme voor Mamba, empirisch gevalideerd op taalmodelleringsbenchmarks, dat verbeterde retrieval-capaciteiten en superieure prestaties ten opzichte van standaard Mamba aantoont. Tot slot voeren we uitgebreide ablatiestudies en empirische analyses uit om onze ontwerpkeuzes te rechtvaardigen en bewijs te leveren dat onze aanpak het probleem van overmatige aandachtstoewijzing in Mamba-gebaseerde modellen effectief mitigeert. Onze code is publiekelijk beschikbaar.

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.