Дифференциальная Мамба

Differential Mamba

July 8, 2025

Авторы: Nadav Schneider, Itamar Zimerman, Eliya Nachmani

cs.AI

Аннотация

Последовательные модели, такие как Transformers и RNN, часто избыточно распределяют внимание на нерелевантный контекст, что приводит к зашумленным промежуточным представлениям. Это ухудшает возможности крупных языковых моделей (LLM), способствуя возникновению галлюцинаций, ослаблению способностей к обработке длинных зависимостей и извлечению информации, а также снижению устойчивости. Недавние исследования показали, что дифференциальный подход может смягчить эту проблему в Transformers, повышая их эффективность в различных приложениях. В данной работе мы исследуем, могут ли эти методы, изначально разработанные для Transformers, быть применены к Mamba — новой архитектуре, основанной на селективных слоях пространства состояний, которая достигает производительности уровня Transformers с большей эффективностью. Мы показываем, что простое адаптирование дифференциального подхода к Mamba недостаточно и требует тщательных архитектурных модификаций. Для решения этой проблемы мы предлагаем новый дифференциальный механизм для Mamba, эмпирически проверенный на бенчмарках языкового моделирования, демонстрирующий улучшенные способности к извлечению информации и превосходную производительность по сравнению с базовой версией Mamba. Наконец, мы проводим обширные исследования и эмпирические анализы, чтобы обосновать наши проектные решения и предоставить доказательства того, что наш подход эффективно смягчает проблему избыточного распределения внимания в моделях на основе Mamba. Наш код доступен публично.

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.