B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory

• URL: http://arxiv.org/abs/2407.06324v1
• Date: Mon, 8 Jul 2024 18:41:01 GMT
• Title: B'MOJO: Hybrid State Space Realizations of Foundation Models with Eidetic and Fading Memory
• Authors: Luca Zancato, Arjun Seshadri, Yonatan Dukler, Aditya Golatkar, Yantao Shen, Benjamin Bowman, Matthew Trager, Alessandro Achille, Stefano Soatto,
• Abstract summary: 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.
• Score: 91.81390121042192
• License: http://creativecommons.org/licenses/by/4.0/
• Abstract: We describe a family of architectures to support transductive inference by allowing memory to grow to a finite but a-priori unknown bound while making efficient use of finite resources for inference. Current architectures use such resources to represent data either eidetically over a finite span ("context" in Transformers), or fading over an infinite span (in State Space Models, or SSMs). Recent hybrid architectures have combined eidetic and fading memory, but with limitations that do not allow the designer or the learning process to seamlessly modulate the two, nor to extend the eidetic memory span. We leverage ideas from Stochastic Realization Theory to develop a class of models called B'MOJO to seamlessly combine eidetic and fading memory within an elementary composable module. The overall architecture can be used to implement models that can access short-term eidetic memory "in-context," permanent structural memory "in-weights," fading memory "in-state," and long-term eidetic memory "in-storage" by natively incorporating retrieval from an asynchronously updated memory. We show that Transformers, existing SSMs such as Mamba, and hybrid architectures such as Jamba are special cases of B'MOJO and describe a basic implementation, to be open sourced, that can be stacked and scaled efficiently in hardware. We test B'MOJO on transductive inference tasks, such as associative recall, where it outperforms existing SSMs and Hybrid models; as a baseline, we test ordinary language modeling where B'MOJO achieves perplexity comparable to similarly-sized Transformers and SSMs up to 1.4B parameters, while being up to 10% faster to train. Finally, we show that B'MOJO's ability to modulate eidetic and fading memory results in better inference on longer sequences tested up to 32K tokens, four-fold the length of the longest sequences seen during training.

Related papers

• MoM: Linear Sequence Modeling with Mixture-of-Memories [9.665802842933209]
  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.
  arXiv  Detail & Related papers  (2025-02-19T12:53:55Z)
• Expansion Span: Combining Fading Memory and Retrieval in Hybrid State Space Models [59.607021334350385]
  Hybrid architectures combine State Space layers with Attention, but still cannot recall the distant past and can access only the most recent tokens eidetically. We describe a method to expand the memory span of the hybrid state by "reserving" a fraction of the Attention context for tokens retrieved from arbitrarily distant in the past. We show that SE-Attn enables us to efficiently adapt pre-trained Hybrid models on sequences of tokens up to 8 times longer than the ones used for pre-training.
  arXiv  Detail & Related papers  (2024-12-17T20:55:42Z)
• Memory Layers at Scale [67.00854080570979]
  This work takes memory layers beyond proof-of-concept, proving their utility at contemporary scale. On downstream tasks, language models augmented with our improved memory layer outperform dense models with more than twice the budget, as well as mixture-of-expert models when matched for both compute and parameters. We provide a fully parallelizable memory layer implementation, demonstrating scaling laws with up to 128B memory parameters, pretrained to 1 trillion tokens, comparing to base models with up to 8B parameters.
  arXiv  Detail & Related papers  (2024-12-12T23:56:57Z)
• HMT: Hierarchical Memory Transformer for Efficient Long Context Language Processing [33.720656946186885]
  Hierarchical Memory Transformer (HMT) is a novel framework that facilitates a model's long-context processing ability. HMT consistently improves the long-context processing ability of existing models.
  arXiv  Detail & Related papers  (2024-05-09T19:32:49Z)
• 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)
• 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)
• Kanerva++: extending The Kanerva Machine with differentiable, locally block allocated latent memory [75.65949969000596]
  Episodic and semantic memory are critical components of the human memory model. We develop a new principled Bayesian memory allocation scheme that bridges the gap between episodic and semantic memory. We demonstrate that this allocation scheme improves performance in memory conditional image generation.
  arXiv  Detail & Related papers  (2021-02-20T18:40:40Z)
• End-to-End Egospheric Spatial Memory [32.42361470456194]
  We propose a parameter-free module, Egospheric Spatial Memory (ESM), which encodes the memory in an ego-sphere around the agent. ESM can be trained end-to-end via either imitation or reinforcement learning. We show applications to semantic segmentation on the ScanNet dataset, where ESM naturally combines image-level and map-level inference modalities.
  arXiv  Detail & Related papers  (2021-02-15T18:59:07Z)
• 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.