Differential Mamba

TL;DR

This paper introduces a novel differential mechanism for the Mamba architecture, enhancing its ability to mitigate attention overallocation and improve language modeling performance, especially in retrieval tasks.

Contribution

We adapt differential design techniques to Mamba, developing a new mechanism that improves its efficiency and effectiveness in language modeling tasks.

Findings

Enhanced retrieval capabilities in Mamba-based models

Superior performance over vanilla Mamba on benchmarks

Effective mitigation of attention overallocation issues

Abstract

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: https://github.com/NadavSc/Diff-Mamba