Logo
ScienceStack
Table of Contents
Fetching ...
Sign In

MossNet: Mixture of State-Space Experts is a Multi-Head Attention

Shikhar Tuli, James Seale Smith, Haris Jeelani, Chi-Heng Lin, Abhishek Patel, Vasili Ramanishka, Yen-Chang Hsu, Hongxia Jin

TL;DR

MossNet introduces a mixture-of-state-space-experts architecture to emulate linear multi-head attention within state-space models, addressing efficiency and expressiveness gaps in transformers and prior SSM/GRM approaches. The method extends MoE to both channel-mixing MLPs and time-mixing SSM kernels, with a theoretical mapping to linear $MHA$ and extensive experiments showing superior performance and scalability. Large-scale MossNet variants achieve strong results within comparable active parameter budgets and demonstrate favorable on-device and GPU profiling. The work provides a practical path toward efficient, high-performing recurrent LLMs.

Abstract

Large language models (LLMs) have significantly advanced generative applications in natural language processing (NLP). Recent trends in model architectures revolve around efficient variants of transformers or state-space/gated-recurrent models (SSMs, GRMs). However, prevailing SSM/GRM-based methods often emulate only a single attention head, potentially limiting their expressiveness. In this work, we propose MossNet, a novel mixture-of-state-space-experts architecture that emulates a linear multi-head attention (MHA). MossNet leverages a mixture-of-experts (MoE) implementation not only in channel-mixing multi-layered perceptron (MLP) blocks but also in the time-mixing SSM kernels to realize multiple "attention heads." Extensive experiments on language modeling and downstream evaluations show that MossNet outperforms both transformer- and SSM-based architectures of similar model size and data budgets. Larger variants of MossNet, trained on trillions of tokens, further confirm its scalability and superior performance. In addition, real-device profiling on a Samsung Galaxy S24 Ultra and an Nvidia A100 GPU demonstrate favorable runtime speed and resource usage compared to similarly sized baselines. Our results suggest that MossNet is a compelling new direction for efficient, high-performing recurrent LLM architectures.

MossNet: Mixture of State-Space Experts is a Multi-Head Attention

TL;DR

MossNet introduces a mixture-of-state-space-experts architecture to emulate linear multi-head attention within state-space models, addressing efficiency and expressiveness gaps in transformers and prior SSM/GRM approaches. The method extends MoE to both channel-mixing MLPs and time-mixing SSM kernels, with a theoretical mapping to linear and extensive experiments showing superior performance and scalability. Large-scale MossNet variants achieve strong results within comparable active parameter budgets and demonstrate favorable on-device and GPU profiling. The work provides a practical path toward efficient, high-performing recurrent LLMs.

Abstract

Large language models (LLMs) have significantly advanced generative applications in natural language processing (NLP). Recent trends in model architectures revolve around efficient variants of transformers or state-space/gated-recurrent models (SSMs, GRMs). However, prevailing SSM/GRM-based methods often emulate only a single attention head, potentially limiting their expressiveness. In this work, we propose MossNet, a novel mixture-of-state-space-experts architecture that emulates a linear multi-head attention (MHA). MossNet leverages a mixture-of-experts (MoE) implementation not only in channel-mixing multi-layered perceptron (MLP) blocks but also in the time-mixing SSM kernels to realize multiple "attention heads." Extensive experiments on language modeling and downstream evaluations show that MossNet outperforms both transformer- and SSM-based architectures of similar model size and data budgets. Larger variants of MossNet, trained on trillions of tokens, further confirm its scalability and superior performance. In addition, real-device profiling on a Samsung Galaxy S24 Ultra and an Nvidia A100 GPU demonstrate favorable runtime speed and resource usage compared to similarly sized baselines. Our results suggest that MossNet is a compelling new direction for efficient, high-performing recurrent LLM architectures.

Paper Structure

This paper contains 28 sections, 1 theorem, 12 equations, 4 figures, 13 tables.

Table of Contents

  1. Introduction
  2. Challenges and Motivation
  3. Our Contribution
  4. Method
  5. Preliminaries
  6. Mamba
  7. Mixture of Experts
  8. MossNet Architecture
  9. Remark 1
  10. Remark 2
  11. Remark 3
  12. Training and Evaluation Setup
  13. Experiments
  14. Downstream Language Modeling Performance
  15. Speed and Memory Profiling
  16. ...and 13 more sections

Key Result

Theorem 1

A mixture-of-expert implementation of $\bm{\bar{A}}$, $\bm{\bar{B}}$, and $\bm{C}$ is equivalent to a mixture-of-expert implementation of a linear multi-head attention.

Figures (4)

  • Figure 1: Simplified working schematic of the MossNet block. We implement MoE in channel mixing input, gate, and output projections and time mixing input-dependent SSM parameters $\bm{B}$, $\bm{C}$, and $\bm{\Delta}$.
  • Figure 2: (a) Perplexity and (b) commonsense average accuracy scaling for fairly-trained models.
  • Figure 3: (a) Memory consumption, (b) prefill speed, and (c) generation speed with context length for MossNet-8x200M+ and baselines on A100-80GB (FP16 precision, FlashAttention 2). Batch size set to 4.
  • Figure 4: (a) Memory consumption, (b) prefill speed, and (c) generation speed with context length for MossNet-8x200M+ and baselines on Samsung Galaxy S24 Ultra (Q8 precision). Batch size set to 1. Gray line plots depict performance without SWA implemented. Llama3-1.5B results not plotted for 32K context due to out-of-memory error.

Theorems & Definitions (2)

  • Theorem 1
  • proof : Proof