Expansion Span: Combining Fading Memory and Retrieval in Hybrid State Space Models

• URL: http://arxiv.org/abs/2412.13328v1
• Date: Tue, 17 Dec 2024 20:55:42 GMT
• Title: Expansion Span: Combining Fading Memory and Retrieval in Hybrid State Space Models
• Authors: Elvis Nunez, Luca Zancato, Benjamin Bowman, Aditya Golatkar, Wei Xia, Stefano Soatto,
• Abstract summary: 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.
• Score: 59.607021334350385
• License:
• Abstract: The "state" of State Space Models (SSMs) represents their memory, which fades exponentially over an unbounded span. By contrast, Attention-based models have "eidetic" (i.e., verbatim, or photographic) memory over a finite span (context size). Hybrid architectures combine State Space layers with Attention, but still cannot recall the distant past and can access only the most recent tokens eidetically. Unlike current methods of combining SSM and Attention layers, we allow the state to be allocated based on relevancy rather than recency. In this way, for every new set of query tokens, our models can "eidetically" access tokens from beyond the Attention span of current Hybrid SSMs without requiring extra hardware resources. 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, thus expanding the eidetic memory span of the overall state. We call this reserved fraction of tokens the "expansion span," and the mechanism to retrieve and aggregate it "Span-Expanded Attention" (SE-Attn). To adapt Hybrid models to using SE-Attn, we propose a novel fine-tuning method that extends LoRA to Hybrid models (HyLoRA) and allows efficient adaptation on long spans of tokens. 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. We show that HyLoRA with SE-Attn is cheaper and more performant than alternatives like LongLoRA when applied to Hybrid models on natural language benchmarks with long-range dependencies, such as PG-19, RULER, and other common natural language downstream tasks.

Related papers

• Marconi: Prefix Caching for the Era of Hybrid LLMs [26.260418040965327]
  We present Marconi, the first system that supports efficient prefix caching with Hybrid LLMs. Across diverse workloads and Hybrid models, Marconi achieves up to 34.4$times$ higher token hit rates.
  arXiv  Detail & Related papers  (2024-11-28T21:10:20Z)
• SHERL: Synthesizing High Accuracy and Efficient Memory for Resource-Limited Transfer Learning [63.93193829913252]
  We propose an innovative METL strategy called SHERL for resource-limited scenarios. In the early route, intermediate outputs are consolidated via an anti-redundancy operation. In the late route, utilizing minimal late pre-trained layers could alleviate the peak demand on memory overhead.
  arXiv  Detail & Related papers  (2024-07-10T10:22:35Z)
• 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)
• Hierarchical Context Merging: Better Long Context Understanding for Pre-trained LLMs [61.40047491337793]
  We present Hierarchical cOntext MERging (HOMER), a new training-free scheme designed to overcome the limitations of large language models. HomeR uses a divide-and-conquer algorithm, dividing long inputs into manageable chunks. A token reduction technique precedes each merging, ensuring memory usage efficiency.
  arXiv  Detail & Related papers  (2024-04-16T06:34:08Z)
• 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)
• Effectively Modeling Time Series with Simple Discrete State Spaces [25.046268609133982]
  We introduce SpaceTime, a new state-space time series architecture for time series modeling. For expressivity, we propose a new SSM parameterization based on the companion matrix. For long horizon forecasting, we introduce a "closed-loop" variation of the companion SSM. For efficient training and inference, we introduce an algorithm that reduces the memory and compute of a forward pass with the companion matrix.
  arXiv  Detail & Related papers  (2023-03-16T17:08:21Z)
• Long Range Language Modeling via Gated State Spaces [67.64091993846269]
  We focus on autoregressive sequence modeling over English books, Github source code and ArXiv mathematics articles. We propose a new layer named Gated State Space (GSS) and show that it trains significantly faster than the diagonal version of S4.
  arXiv  Detail & Related papers  (2022-06-27T01:50:18Z)

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.