MoxE: Mixture of xLSTM Experts with Entropy-Aware Routing for Efficient Language Modeling

TL;DR

MoxE combines xLSTM with a mixture of experts and entropy-aware routing to improve efficiency and scalability in large language models, balancing resource use and token handling.

Contribution

It introduces a novel entropy-based routing mechanism within a mixture of xLSTM experts, enhancing efficiency and balancing in large language models.

Findings

Significant efficiency improvements over existing models

Effective handling of rare and common tokens

Robust training with auxiliary entropy and group-wise losses

Abstract

This paper introduces MoxE, a novel architecture that synergistically combines the Extended Long Short-Term Memory (xLSTM) with the Mixture of Experts (MoE) framework to address critical scalability and efficiency challenges in large language models (LLMs). The proposed method effectively leverages xLSTM's innovative memory structures while strategically introducing sparsity through MoE to substantially reduce computational overhead. At the heart of our approach is a novel entropy-based routing mechanism, designed to dynamically route tokens to specialized experts, thereby ensuring efficient and balanced resource utilization. This entropy awareness enables the architecture to effectively manage both rare and common tokens, with mLSTM blocks being favored to handle rare tokens. To further enhance generalization, we introduce a suite of auxiliary losses, including entropy-based and group-wise balancing losses, ensuring robust performance and efficient training. Theoretical analysis and empirical evaluations rigorously demonstrate that MoxE achieves significant efficiency gains and enhanced effectiveness compared to existing approaches, marking a notable advancement in scalable LLM architectures.