MixFormer: Linear Transformer with Mixture of Memory Experts

MixFormer proposes a novel linear Transformer architecture that integrates a Mixture-of-Memory-Experts (MoE) mechanism to address limitations in existing State Space Models (SSMs) for ultra-long sequence modeling. The paper states: existing SSMs suffer from limited input adaptivity and constrained memory capacity, leading to information loss when modeling ultra-long sequences. To address these, MixFormer "maintains differentiated memory states through multiple collaborating memory experts and employs a novel Time-Aware Linear Attention (TALA) mechanism, which leverages learnable exponential decay functions and positional biases to dynamically update memory. This design enables the model to selectively reinforce important historical information while effectively mitigating memory dilution, substantially improving long-range dependency modeling."

The paper identifies two key limitations of current SSMs: (1) Existing approaches typically update the memory state using fixed coefficients or a limited number of learnable parameters, resulting in insufficient input adaptivity, and (2) "Most current SSMs employ a single memory state matrix to store historical information. Under ultra-long context inputs, this single-state architecture inherently limits storage capacity, making it difficult to prevent early critical information from being overwritten or diluted by subsequent inputs."

MixFormer introduces the MoE into state memory modeling for the first time, where multiple memory experts collaboratively construct the final memory state through weighted fusion. Within each expert, TALA dynamically adjusts the memory update strategy according to the current input, enabling the model to reinforce important historical information while progressively attenuating less relevant content. The multi-expert parallel architecture breaks the capacity limitation of a single memory state by forming a distributed memory storage, with each expert employing a differentiated memory strategy via the TALA mechanism, focusing on distinct temporal patterns such as long-range dependencies or short-term dynamics.

The main contributions are: (1) We propose a novel linear Transformer with mixture-of-memory-experts (MixFormer). To the best of our knowledge, this is the first work to incorporate multi-memory-experts into State Space Models (SSMs). (2) "We propose a Time-Aware Linear Attention (TALA) mechanism based on linear attention, which integrates an exponential decay function with learnable pair-wise position biases to significantly enhance the model’s capability for long-range dependency modeling. (3) Experimental results on both text and image tasks demonstrate that MixFormer achieves strong performance while supporting efficient parallel training and fast autoregressive inference."

The TALA mechanism supports three computation modes: parallel representation for training, recurrent representation for inference, and chunkwise recurrent representation for efficient long-sequence training. The MoME mechanism extends the memory state by maintaining multiple memory experts and dynamically activates a subset of experts based on the current input and constructs the current memory state by weighted aggregation of their outputs.

Experiments on the Long-Range Arena (LRA) benchmark show MixFormer achieves the best average accuracy (56.72%) compared to all other efficient Transformer models, with top performance on ListOps (37.35%), Text (66.56%), and Retrieval (60.50%). On image generation tasks (MNIST and CIFAR-10), MixFormer produces high-quality samples with sharp boundaries and minimal noise and successfully reconstructs occluded images with well-preserved structural integrity, clear boundaries, and minimal noise.

Interpretability analysis reveals that softmax attention produces explicit token-pair attention scores, while linear attention constructs an implicit global attention matrix. The causal mask visualization clearly reveals the temporal evolution of memory attention values—an effect that arises from our exponential decay mechanism and learnable pairwise interaction biases. The MoME analysis shows that when the current token is 'Paris' and the context is 'The Eiffel Tower is located in the city of', the MoME mechanism exhibits significantly stronger expert activations for the tokens 'Eiffel', 'Tower', and 'city', confirming that the MoME mechanism achieves dynamic selection and persistent retention of crucial historical information through functional specialization among experts.

The paper concludes that MixFormer provides a promising solution for efficient Transformer design with strong modeling capabilities for ultra-long contexts, with future work planned to scale up the model parameters to further enhance its performance and extend the application of MixFormer to a broader range of tasks and domains.

Improvements to AI Systems:

1. Adaptive Memory Management for Long-Context Models

• Integrate a Mixture-of-Memory-Experts (MoME) layer into existing Transformer or SSM architectures, replacing single-state memory with multiple specialized experts. Each expert uses a Time-Aware Linear Attention (TALA) mechanism with learnable exponential decay and pairwise positional biases.

• Resulting capability: The system can selectively reinforce critical historical tokens (e.g., key entities, events) while attenuating irrelevant noise, preventing memory dilution over sequences of millions of tokens. It dynamically allocates memory capacity based on input content, avoiding fixed-capacity bottlenecks.

2. Distributed Memory Storage for Ultra-Long Sequences

• Implement parallel memory experts that each track distinct temporal patterns (e.g., short-term dynamics, long-range dependencies, periodic structures). Use weighted fusion to combine expert outputs into a final memory state.

• Resulting capability: The system can retain and retrieve early critical information (e.g., a user’s initial instruction in a multi-hour conversation, or a rare fact in a 10M-token document) without overwriting, even when later inputs are highly dense. This enables coherent reasoning over entire books, codebases, or multi-session dialogues.

3. Input-Adaptive Memory Update via TALA

• Replace fixed or limited-parameter memory update rules with TALA’s exponential decay and learnable pairwise biases, which adjust memory update strength per token based on its relevance to the current context.

• Resulting capability: The system can automatically reinforce important tokens (e.g., a named entity like “Eiffel Tower” when predicting “Paris”) and progressively forget irrelevant ones, improving next-token prediction, question answering, and summarization on long inputs without manual tuning.

4. Hybrid Training/Inference Modes

• Adopt TALA’s three computation modes: parallel (for fast GPU training), recurrent (for autoregressive inference with O(1) memory per step), and chunkwise recurrent (for efficient training on very long sequences by processing in chunks).

• Resulting capability: The system can train on sequences of arbitrary length without quadratic attention cost, and perform real-time generation (e.g., streaming text, code completion, or dialogue) with constant memory footprint, even on edge devices.

5. Functional Specialization for Interpretable Memory

• Use MoME’s expert activation patterns as an interpretability tool: visualize which experts fire for specific tokens (e.g., “Paris” activates experts tuned to geographic or landmark-related patterns).

• Resulting capability: The system can explain its memory decisions—e.g., why it retained a certain fact—by showing expert activation heatmaps. This aids debugging, trust, and user-facing transparency in applications like legal document analysis or medical record summarization.

6. Improved Performance on Long-Range Benchmarks

• Apply MixFormer’s architecture directly to tasks requiring long-range reasoning, such as document classification, retrieval, and image generation with global context.

• Resulting capability: The system achieves higher accuracy on Long-Range Arena (e.g., 56.72% average, with 66.56% on Text and 60.50% on Retrieval) compared to prior efficient Transformers, while maintaining linear complexity. It can also generate high-fidelity images (MNIST, CIFAR-10) with sharp boundaries and reconstruct occluded images by leveraging long-range spatial dependencies.

7. Scalable Future Extensions

• Use the MoME+TALA design as a drop-in replacement for attention in existing large models (e.g., LLMs, vision transformers) to extend their context windows without retraining from scratch.

• Resulting capability: The system can be scaled to larger parameter counts and new domains (e.g., video understanding, multi-modal long-form reasoning) while preserving efficient training and inference, enabling deployment in resource-constrained environments.