MoM: Linear Sequence Modeling with Mixture-of-Memories

• URL: http://arxiv.org/abs/2502.13685v1
• Date: Wed, 19 Feb 2025 12:53:55 GMT
• Title: MoM: Linear Sequence Modeling with Mixture-of-Memories
• Authors: Jusen Du, Weigao Sun, Disen Lan, Jiaxi Hu, Yu Cheng,
• Abstract summary: We introduce a novel architecture called Mixture-of-Memories (MoM) MoM utilizes multiple independent memory states, with a router network directing input tokens to specific memory states. MoM performs exceptionally well on recall-intensive tasks, surpassing existing linear sequence modeling techniques.
• Score: 9.665802842933209
• License:
• Abstract: Linear sequence modeling methods, such as linear attention, state space modeling, and linear RNNs, offer significant efficiency improvements by reducing the complexity of training and inference. However, these methods typically compress the entire input sequence into a single fixed-size memory state, which leads to suboptimal performance on recall-intensive downstream tasks. Drawing inspiration from neuroscience, particularly the brain's ability to maintain robust long-term memory while mitigating "memory interference", we introduce a novel architecture called Mixture-of-Memories (MoM). MoM utilizes multiple independent memory states, with a router network directing input tokens to specific memory states. This approach greatly enhances the overall memory capacity while minimizing memory interference. As a result, MoM performs exceptionally well on recall-intensive tasks, surpassing existing linear sequence modeling techniques. Despite incorporating multiple memory states, the computation of each memory state remains linear in complexity, allowing MoM to retain the linear-complexity advantage during training, while constant-complexity during inference. Our experimental results show that MoM significantly outperforms current linear sequence models on downstream language tasks, particularly recall-intensive tasks, and even achieves performance comparable to Transformer models. The code is released at https://github.com/OpenSparseLLMs/MoM and is also released as a part of https://github.com/OpenSparseLLMs/Linear-MoE.

Related papers

• MEMO: Fine-grained Tensor Management For Ultra-long Context LLM Training [24.066283519769968]
  Large Language Models (LLMs) have been trained using extended context lengths to foster more creative applications. We propose MEMO, a novel framework for fine-grained activation memory management. MeMO achieves an average of 1.97x and 1.80x MFU compared to Megatron-LM and DeepSpeed.
  arXiv  Detail & Related papers  (2024-07-16T18:59:49Z)
• B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory [91.81390121042192]
  We develop a class of models called B'MOJO to seamlessly combine eidetic and fading memory within an composable module. B'MOJO's ability to modulate eidetic and fading memory results in better inference on longer sequences tested up to 32K tokens.
  arXiv  Detail & Related papers  (2024-07-08T18:41:01Z)
• LongVQ: Long Sequence Modeling with Vector Quantization on Structured Memory [63.41820940103348]
  Self-attention mechanism's computational cost limits its practicality for long sequences. We propose a new method called LongVQ to compress the global abstraction as a length-fixed codebook. LongVQ effectively maintains dynamic global and local patterns, which helps to complement the lack of long-range dependency issues.
  arXiv  Detail & Related papers  (2024-04-17T08:26:34Z)
• MSPipe: Efficient Temporal GNN Training via Staleness-Aware Pipeline [8.889825826072512]
  Memory-based Temporal Graph Neural Networks (MTGNNs) are a class of temporal graph neural networks that utilize a node memory module to capture and retain long-term temporal dependencies. Existing optimizations for static GNNs are not directly applicable to MTGNNs due to differences in training paradigm, model architecture, and the absence of a memory module. We propose MSPipe, a general and efficient framework for MTGNNs that maximizes training throughput while maintaining model accuracy.
  arXiv  Detail & Related papers  (2024-02-23T05:57:22Z)
• CAMELoT: Towards Large Language Models with Training-Free Consolidated Associative Memory [38.429707659685974]
  Large Language Models (LLMs) struggle to handle long input sequences due to high memory and runtime costs. We introduce an associative memory module which can be coupled to any pre-trained (frozen) attention-based LLM without re-training. This architecture, which we call CAMELoT, demonstrates superior performance even with a tiny context window of 128 tokens.
  arXiv  Detail & Related papers  (2024-02-21T01:00:17Z)
• Blockwise Parallel Transformer for Large Context Models [70.97386897478238]
  Blockwise Parallel Transformer (BPT) is a blockwise computation of self-attention and feedforward network fusion to minimize memory costs. By processing longer input sequences while maintaining memory efficiency, BPT enables training sequences 32 times longer than vanilla Transformers and up to 4 times longer than previous memory-efficient methods.
  arXiv  Detail & Related papers  (2023-05-30T19:25:51Z)
• A Model or 603 Exemplars: Towards Memory-Efficient Class-Incremental Learning [56.450090618578]
  Class-Incremental Learning (CIL) aims to train a model with limited memory size to meet this requirement. We show that when counting the model size into the total budget and comparing methods with aligned memory size, saving models do not consistently work. We propose a simple yet effective baseline, denoted as MEMO for Memory-efficient Expandable MOdel.
  arXiv  Detail & Related papers  (2022-05-26T08:24:01Z)
• LaMemo: Language Modeling with Look-Ahead Memory [50.6248714811912]
  We propose Look-Ahead Memory (LaMemo) that enhances the recurrence memory by incrementally attending to the right-side tokens. LaMemo embraces bi-directional attention and segment recurrence with an additional overhead only linearly proportional to the memory length. Experiments on widely used language modeling benchmarks demonstrate its superiority over the baselines equipped with different types of memory.
  arXiv  Detail & Related papers  (2022-04-15T06:11:25Z)
• Memformer: A Memory-Augmented Transformer for Sequence Modeling [55.780849185884996]
  We present Memformer, an efficient neural network for sequence modeling. Our model achieves linear time complexity and constant memory space complexity when processing long sequences.
  arXiv  Detail & Related papers  (2020-10-14T09:03:36Z)

This list is automatically generated from the titles and abstracts of the papers in this site.

This site does not guarantee the quality of this site (including all information) and is not responsible for any consequences.