Understanding and Enhancing Mamba-Transformer Hybrids for Memory Recall and Language Modeling

TL;DR

This paper analyzes hybrid models combining state space models and attention mechanisms, revealing how architectural choices affect performance and proposing data augmentation techniques to improve memory recall in language modeling.

Contribution

It provides a detailed analysis of hybrid SSM-attention architectures, compares sequential and parallel designs, and introduces a data-centric training method to enhance recall capabilities.

Findings

Sequential hybrids excel on short contexts.

Parallel hybrids perform better on long contexts.

Data augmentation with paraphrases improves recall.

Abstract

Hybrid models that combine state space models (SSMs) with attention mechanisms have shown strong performance by leveraging the efficiency of SSMs and the high recall ability of attention. However, the architectural design choices behind these hybrid models remain insufficiently understood. In this work, we analyze hybrid architectures through the lens of memory utilization and overall performance, and propose a complementary method to further enhance their effectiveness. We first examine the distinction between sequential and parallel integration of SSM and attention layers. Our analysis reveals several interesting findings, including that sequential hybrids perform better on shorter contexts, whereas parallel hybrids are more effective for longer contexts. We also introduce a data-centric approach of continually training on datasets augmented with paraphrases, which further enhances recall while preserving other capabilities. It generalizes well across different base models and outperforms architectural modifications aimed at enhancing recall. Our findings provide a deeper understanding of hybrid SSM-attention models and offer practical guidance for designing architectures tailored to various use cases. Our findings provide a deeper understanding of hybrid SSM-attention models and offer practical guidance for designing architectures tailored to various use cases.