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

TL;DR

This work analyzes hybrid models that combine state-space models (SSMs) with attention, focusing on how architectural design choices affect memory recall and language modeling. It introduces a unified experimental setup and evaluates three axes—long-context language modeling, commonsense reasoning, and memory recall—across SSM-only, Transformer, and hybrid variants. The key findings show that parallel hybrids with merge-attention excel in long-context recall, while sequential hybrids train more stably and suit short contexts; a data-centric approach using paraphrased continual pretraining substantially improves recall with minimal harm to reasoning, generalizing across base models up to 2.8B parameters. Collectively, these insights offer practical guidance for tailoring hybrid architectures to task context and demonstrate that paraphrase-based data augmentation can outperform architectural modifications aimed at boosting recall, with broad applicability and scalability.

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.

Figures (21)

• Figure 1: Comparison of different architectural designs: SSM, Transformer, Sequential Hybrid, and Parallel Hybrid. Each architecture consists of stacked blocks that incorporate Mamba and Attention layers. The key difference lies in how these layers are arranged: SSM uses only Mamba layers, Transformer uses only Attention layers, while the hybrid models combine both. Sequential Hybrid stacks Mamba and Attention layers within each block, whereas Parallel Hybrid applies them in parallel and aggregates their outputs. Feedforward (FF) layers are omitted in the hybrid models for clarity, as it varies by design.
• Figure 2: Comparison of different model architectures on Commonsense Reasoning (y-axis) vs. Recall Ability (x-axis). Commonsense Reasoning and Recall Ability are measured using answer accuracy. The models compared included Mamba-only, SWA-only, Hybrid (Sequential), and Hybrid (Parallel). For details of each model, see Figure \ref{['app_fig:scatter_names']} in Appendix \ref{['app_subsec: correlation']}.
• Figure 3:
• Figure 4:
• Figure 6: Performance of best performing models from each architecture in commonsense reasoning and recall ability, where divided by length of context.