Mamba Differenziale

Differential Mamba

July 8, 2025

Autori: Nadav Schneider, Itamar Zimerman, Eliya Nachmani

cs.AI

Abstract

I modelli sequenziali come i Transformer e le RNN spesso allocano eccessivamente l'attenzione a contesti irrilevanti, portando a rappresentazioni intermedie rumorose. Ciò compromette le capacità dei LLM favorendo allucinazioni, indebolendo le capacità di gestione di contesti a lungo raggio e di recupero delle informazioni, e riducendo la robustezza. Recenti lavori hanno dimostrato che un design differenziale può mitigare questo problema nei Transformer, migliorandone l'efficacia in varie applicazioni. In questo articolo, esploriamo se queste tecniche, originariamente sviluppate per i Transformer, possano essere applicate a Mamba, una recente architettura basata su strati di spazio degli stati selettivi che raggiunge prestazioni paragonabili ai Transformer con maggiore efficienza. Mostriamo che un adattamento ingenuo del design differenziale a Mamba è insufficiente e richiede modifiche architetturali accurate. Per affrontare questo problema, introduciamo un nuovo meccanismo differenziale per Mamba, validato empiricamente su benchmark di modellazione del linguaggio, dimostrando capacità di recupero migliorate e prestazioni superiori rispetto alla versione base di Mamba. Infine, conduciamo ampi studi di ablazione e analisi empiriche per giustificare le nostre scelte progettuali e fornire evidenze che il nostro approccio mitiga efficacemente il problema di sovraallocazione nei modelli basati su Mamba. Il nostro codice è pubblicamente disponibile.

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.