Mamba Diferencial

Differential Mamba

July 8, 2025

Autores: Nadav Schneider, Itamar Zimerman, Eliya Nachmani

cs.AI

Resumo

Modelos de sequência como Transformers e RNNs frequentemente alocam atenção excessiva a contextos irrelevantes, resultando em representações intermediárias ruidosas. Isso prejudica as capacidades dos LLMs ao promover alucinações, enfraquecer habilidades de longo alcance e recuperação, e reduzir a robustez. Trabalhos recentes mostraram que o design diferencial pode mitigar esse problema em Transformers, melhorando sua eficácia em diversas aplicações. Neste artigo, exploramos se essas técnicas, originalmente desenvolvidas para Transformers, podem ser aplicadas ao Mamba, uma arquitetura recente baseada em camadas de espaço de estados seletivas que alcança desempenho equivalente ao dos Transformers com maior eficiência. Mostramos que uma adaptação ingênua do design diferencial ao Mamba é insuficiente e requer modificações arquiteturais cuidadosas. Para resolver isso, introduzimos um novo mecanismo diferencial para o Mamba, validado empiricamente em benchmarks de modelagem de linguagem, demonstrando capacidades aprimoradas de recuperação e desempenho superior ao Mamba padrão. Por fim, realizamos extensos estudos de ablação e análises empíricas para justificar nossas escolhas de design e fornecer evidências de que nossa abordagem mitiga efetivamente o problema de alocação excessiva em modelos baseados no Mamba. Nosso código está disponível publicamente.

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.