Mamba Diferencial

Differential Mamba

July 8, 2025

Autores: Nadav Schneider, Itamar Zimerman, Eliya Nachmani

cs.AI

Resumen

Los modelos de secuencia como los Transformers y las RNNs a menudo asignan excesiva atención a contextos irrelevantes, lo que genera representaciones intermedias ruidosas. Esto degrada las capacidades de los LLM al fomentar alucinaciones, debilitar las habilidades de largo alcance y recuperación, y reducir la robustez. Trabajos recientes han demostrado que el diseño diferencial puede mitigar este problema en los Transformers, mejorando su eficacia en diversas aplicaciones. En este artículo, exploramos si estas técnicas, desarrolladas originalmente para Transformers, pueden aplicarse a Mamba, una arquitectura reciente basada en capas de espacio de estados selectivas que logra un rendimiento comparable al de los Transformers con mayor eficiencia. Mostramos que una adaptación ingenua del diseño diferencial a Mamba es insuficiente y requiere modificaciones arquitectónicas cuidadosas. Para abordar esto, introducimos un nuevo mecanismo diferencial para Mamba, validado empíricamente en benchmarks de modelado de lenguaje, que demuestra capacidades de recuperación mejoradas y un rendimiento superior al de Mamba estándar. Finalmente, realizamos extensos estudios de ablación y análisis empíricos para justificar nuestras decisiones de diseño y proporcionar evidencia de que nuestro enfoque mitiga efectivamente el problema de asignación excesiva en modelos basados en Mamba. Nuestro código está disponible públicamente.

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.