PaperlayerView exact revision on arXiv

Source-linked AI summary

HMT: Hierarchical Memory Transformer for Efficient Long Context Language Processing

Zifan He, Yingqi Cao, Zongyue Qin, Neha Prakriya, Yizhou Sun, Jason Cong

arXiv:2405.06067v3cs.CLcs.LG

TL;DR

Device memory constraints limit transformer context windows, while long documents and ongoing question-answering can require enormous or even infinite inputs. HMT uses hierarchical memory with segment-level recurrence and relevant-history retrieval; the paper reports improved long-context processing and comparable or better quality than long-context LLMs with substantially fewer parameters and less inference memory.

  • Problem

    Device memory constraints lead most transformer models to restrict context windows, despite long-document and lifelong-question-answering inputs that can be enormous or infinite.

    Source: §1 · p. 1

  • Method

    HMT combines memory-augmented segment recurrence with interacting sensory, short-term, and long-term memory tiers.

    Source: §1 · p. 1

  • Results

    Across long-context evaluations, HMT improves processing; compared with long-context LLMs, it achieves comparable or higher metric results with 2–57× fewer parameters and 2.5–116× lower inference memory.

    Source: §1 · p. 2

  • Takeaways & Limitations

    HMT provides LLM accessibility for resource-constrained applications and represents a step forward for lifelong language tasks.

    Source: §6 · p. 9

    Training with BPTT creates a large computational graph, making HMT’s extra parameters memory-consuming to tune and impeding experiments on larger-scale models.

    Source: §7 · p. 9

Abstract

from arXiv · show

Transformer-based large language models (LLM) have been widely used in language processing applications. However, due to the memory constraints of the devices, most of them restrict the context window. Even though recurrent models in previous works can memorize past tokens to enable unlimited context and maintain effectiveness, they have ``flat'' memory architectures. Such architectures have limitations in selecting and filtering information. Since humans are good at learning and self-adjustment, we believe that imitating brain memory hierarchy is beneficial for model memorization. Thus, we propose the Hierarchical Memory Transformer (HMT), a novel framework that facilitates a model's long-context processing ability by imitating human memorization behavior. Leveraging memory-augmented segment-level recurrence, we organize the memory hierarchy by preserving tokens from early input segments, passing memory embeddings along the sequence, and recalling relevant information from history. Evaluating general language modeling, question-answering tasks, and the summarization task, we show that HMT consistently improves the long-context processing ability of existing models. Furthermore, HMT achieves a comparable or superior generation quality to long-context LLMs with $2 \sim 57\times$ fewer parameters and $2.5 \sim 116\times$ less inference memory, significantly outperforming previous memory-augmented models. Code on Github: https://github.com/OswaldHe/HMT-pytorch.

1 Introduction

HMT adds hierarchical, relevance-based memory to recurrently process long contexts, addressing limits in context length and information selection. The paper reports better long-context quality across model architectures and comparisons with long-context and memory-augmented models.

1 Introduction

  • HMT organizes learned memory tokens and current input tokens into interacting sensory, short-term, and long-term memory tiers.This hierarchical design is paired with retrieval that stores embeddings from prior iterations and searches for relevance to current segments.Source: §1 · pp. 1–2
  • HMT is a model-independent plug-and-play framework that can augment decoder-only backbones, including transformer-based and state-space models.The authors say joint training and fine-tuning applies to both newly introduced and original backbone parameters.Source: §1 · p. 2
  • HMT improves effectiveness by up to 25.5% in perplexity and 1.0% higher prediction accuracy over baseline models on long-context evaluations.Evaluations use Wikitext-103, PG-19, and PubMedQA with multiple contexts concatenated.Source: §1 · p. 2
  • With 2–57× fewer parameters and 2.5–116× lower inference memory, HMT achieves comparable or higher metric results than long-context large language models.Evaluated with SmolLM, OPT, and OpenLlamaV2 on LongBench.Source: §1 · p. 2
  • At the same or similar backbone size, HMT has better generation quality and lower memory complexity than specialized memory-augmented and hierarchical methods.Comparisons include RMT, LongMem, Memorizing Transformer, CCM, and HOMER.Source: §1 · p. 2

2 Related Works and Problem Formulation

Prior long-context approaches trade off attention cost, memory use, contextual modeling, and the ability to select relevant history. The problem formulation seeks adaptive processing of potentially unbounded streams while avoiding distraction from irrelevant past information.

2.1 Long Context Transformers

  • Sparse-attention methods extend receptive fields or retrieve selected tokens, but memory consumption can still rise with input length.Sliding windows can miss long-range interactions; global tokens, pooling, or top-K retrieval are used by other approaches.Source: §2.1 · p. 2

2.2 Recurrent Sequence Models

  • RNNs can memorize past information and support hardware-friendly accelerators, but have limited contextual-relationship learning compared with self-attention.Coarse-grain recurrence instead attends within segments and propagates compressed information between them.Source: §2.2 · p. 2
  • Segment recurrence propagates compressed states across inputs; Compressive Transformer and RMT use this strategy to retain past information.Repeated summarization can dilute earlier information, while concatenating summaries can sacrifice inference-memory efficiency.Source: §2.2 · p. 3
  • RWKV and Mamba offer efficient recurrent processing, but are described as limited in contextual relationships and filtering irrelevant information.Combining transformers with Mamba may mitigate this issue but reintroduces transformer scaling concerns.Source: §2.2 · p. 3

2.3 Problem Formulation: Adaptive Long-context Processing

  • The formulation targets infinitely long inputs and adaptive selection of relevant past information based on the input stream’s context or topic.Irrelevant context can distract the model.Source: §2.3 · p. 3
  • Memory bandwidth, capacity, and data-generation speed prevent hardware from reading long documents as a whole; encoding all history can also retain irrelevant information and degrade quality.Users may return to earlier topics or switch to topics relevant to past information.Source: §2.3 · p. 3

3 Hierarchical Memory Transformer

HMT organizes segment-level memory hierarchically: it retrieves relevant past information to augment each segment, then stores a memory embedding for later retrieval. Its workflow combines segment summarization, memory search, and memory updates.

3.1 Overall Workflow

  • For each segment, HMT encodes a summary, searches past memory for relevant information, augments the segment, and decodes it to produce logits and a memory embedding.The memory retrieval steps are representation extraction, memory search, and memory augmentation.Source: §3.1 · pp. 3–4, Figure 1 · p. 4

3.2 Memory Retrieval Mechanism

  • HMT uses the segment summary as a query over a sliding window of cached embeddings, producing a memorization prompt that carries information relevant to the current segment.The search applies normalized similarity scores directly to the memory embeddings; the accumulated memory covers prior segments, with greater information loss for older ones.Source: §3.2 · pp. 4–5
  • Representation encoding can run alongside inference, while memory search has O(N) time complexity and can also run in parallel when N is small.The paper describes the resulting runtime overhead as negligible.Source: §3.2 · p. 5

3.3 Hierarchical Memorization

  • HMT’s hierarchy combines the previous segment’s last k token embeddings as sensory memory with a segment summary and retrieved context, then caches a memory embedding for future retrieval.The memory embedding summarizes the augmented segment; the hidden embeddings used for logits are also produced by the backbone model.Source: §3.3 · p. 5

4 Experiment

The experiments compare HMT across backbone models and long-context settings, using both benchmark subsets and constructed longer inputs. The study spans language modeling, question answering, summarization, and few-shot learning.

4 Experiment

  • HMT is evaluated with varied transformer and state-space backbones, including models up to 14B parameters, and compared under the same maximum context window where applicable.The experiments also include long-context-targeted models and comparisons with prior memory-augmented methods using applicable or similarly sized backbones.Source: §4 · pp. 5–6
  • LongBench subsets cover single- and multi-document QA, summarization, and few-shot learning, but their test sets average fewer than 20k words per document.The paper adds crafted and controllable samples to examine performance across different input lengths and context scenarios.Source: §4 · p. 6
  • Constructed evaluations use Wikitext-103 and PG-19 for next-token generation, including PG-19 samples averaging 69k words, and adapted PubMedQA for multi-context question answering.Samples are concatenated or split into chunks to form longer inputs.Source: §4 · p. 6

5 Results and Key Observations

HMT improves long-context generation across backbone architectures and QA settings with a small parameter increase. It also compares favorably with long-context models and prior memory-augmented methods, including on inference memory and scalability.

5.1 Impacts on Backbone Models

  • Across 2k–100k-token inputs, HMT consistently improves generation quality for OPT 2.7B, RWKV 3B, and OpenLlamaV2 3B.HMT adds 0.5%–1.3% parameters (1.77M–33.5M); on PG-19 it also improves Qwen 2.5 14B effectiveness by 10.0%.Source: §5.1 · p. 6, Figure 4 · p. 7
  • HMT improves long-answer contextual reasoning by 9.48% in perplexity and short-answer accuracy by 1.0% over the backbone on PubMedQA.The samples contain 2 to 10 contexts; short-answer advantages are significant when samples have more contexts.Source: §5.1 · p. 6
  • HMT improves OPT 2.7B perplexity over RMT by 13.0% on Wikitext-103 and 3.98% on PG-19.For RWKV, HMT improves effectiveness by 16.5% on Wikitext-103 and 9.96% on PG-19, while RMT worsens it on Wikitext-103.Source: Figure 2 · p. 7, Figure 3 · p. 7

5.2 Comparison to Long Context Models

  • HMT-augmented small models achieve comparable or better LongBench metrics than long-context models with lower inference memory, but can be weaker on long responses.The LongBench subsets cover single- and multi-document QA, summarization, and few-shot learning.Source: §5.2 · p. 7, Table 3 · p. 8
  • 33.9 GB VRAM lets HMT + Yi-6B-200K process 30k tokens, versus 44.8 GB for Yi-6B-SW-5.2K, with a 2% effectiveness improvement.The HMT configuration uses a 512-token segment length.Source: §5.2 · p. 7

5.3 Comparison to Memory-augmented and Hierarchical Methods

  • HMT outperforms the compared memory-augmented methods in generation quality and has the lowest memory complexity among prior methods.The comparisons include Memorizing Transformer, LongMem, and CCM-concat using same- or similar-sized backbones.Source: §5.3 · pp. 8–9, Table 5 · p. 9
  • 9.9% lower perplexity on PG-19 versus HOMER with YaRN; HMT maintains constant peak memory as input length grows.The benefit becomes more substantial as input length increases.Source: §5.3 · p. 9, Figure 7 · p. 9

6 Conclusion

HMT augments long-range language processing with context switching through hierarchical memory and memory retrieval. The authors report improved generation quality and higher quality with lower memory requirements than other long-context and memory-augmented models.

6 Conclusion

  • HMT combines hierarchical memory with a memory retrieval mechanism to augment models’ long-range language processing with context switching.The approach is inspired by the brain’s memory hierarchy and human memorization behavior.Source: §6 · p. 9
  • HMT consistently improves the generation quality of its backbone models.Source: §6 · p. 9
  • Compared with other long-context LLMs and memory-augmented models, HMT achieves higher generation quality with lower memory requirements.The authors describe this as supporting LLM accessibility for resource-constrained applications and as a step toward lifelong language tasks.Source: §6 · p. 9

7 Limitations and Ongoing Works

The authors identify retrieval overhead as a concern when the memory bank grows, BPTT training as memory-consuming, and HMT’s single long-term-memory level as a possible constraint on information access efficiency.

7 Limitations and Ongoing Works

  • Retrieval overhead can become significant as the number of saved memory embeddings grows, especially when they span different physical memory hierarchies.The authors report negligible overhead at N = 300, sufficient for 100k-token samples, and leave intelligent memory prefetching for future work.Source: §7 · p. 9
  • BPTT’s large computational graph makes tuning HMT’s extra parameters memory-consuming and impedes experiments on larger-scale models.The authors identify extending BPTT depth without memory overhead as future work.Source: §7 · p. 9
  • HMT uses only one level of long-term memory; the authors suggest multiple levels as a possible way to improve information access efficiency.Source: §7 · p. 9

8 Ethical Statements

The paper discusses privacy concerns, related memory approaches, comparisons with LongMem and Unlimiformer, and training configurations for HMT, RMT, and baselines.

8 Ethical Statements

  • The authors note that memorization may raise privacy-leakage concerns in conversations with language-model agents and suggest offline edge deployment as a possible resolution.Source: §8 · p. 9

A Other Related Works

B Comparison to LongMem

C Comparison to Unlimiformer

  • HMT uses cross-attention over cached hidden embeddings, unlike Unlimiformer’s kNN over encoded segments, and encodes each segment with its previous memory embedding.HMT targets decoder-only models without injecting modules into the backbone; Unlimiformer targets encoder-decoder models and injects retrieval modules.Source: C Comparison to Unlimiformer · p. 14

D HMT, RMT, and Baseline Training Details and Hyperparameters

  • HMT uses 300 or 400 long-term memory embeddings and 32 sensory tokens, and summarizes half a segment for representation extraction.Training uses two stages, with S1 trained on two segments and S2 at the maximum BPTT unroll depth the GPUs can manage; experiments average three random seeds.Source: D HMT, RMT, and Baseline Training… · p. 15

E Ablation Study

The ablations examine memory retrieval, partial-segment representation extraction, cached-memory size, and multi-stage training. They report benefits from retrieval and caching more embeddings, negligible impact from partial summarization, and better performance with two-stage retrieval training.

Impact of memory retrieval mechanism

  • With memory retrieval, HMT outperforms the version without retrieval at every tested input length; effectiveness improves or tends to improve as input length grows.The comparison covers OPT 350M and OPT 2.7B on Wikitext-103 inputs up to 100k tokens.Source: E Ablation Study · p. 15, Figure 8 · p. 15

Impact of summarizing partial segment in memory retrieval

  • The impact of summarizing half a segment rather than extracting from the whole segment is negligible.Partial extraction is used to overlap or reduce the previous segment’s inference time with representation extraction for the next segment.Source: E Ablation Study · pp. 15–16, Figure 9 · p. 15

Impact of limited cached memory embeddings

  • For HMT + Llama 2 7B on 100k-token Wikitext-103 samples, caching more memory embeddings marginally improves effectiveness.Retrieval favors recent embeddings, but long-context retrieval also occurs; 6.5% of segments retrieve memory tokens within two segments.Source: E Ablation Study · p. 16, Figure 10 · p. 16, Figure 11 · p. 16

Multi-stage training

  • 902 s of multi-stage training outperforms 1680 s of single-stage training for HMT + OPT 2.7B with memory retrieval.Source: Figure 12 · p. 16

F Multi-stage Training

HMT training is split into two stages: first training without memory retrieval, then adding retrieval and training with the longest BPTT unroll the GPUs can handle.

F Multi-stage Training

  • Multi-stage training achieves better effectiveness and lower total training time than single-stage training within the same number of steps.Stage 1 runs faster per iteration (1.15 s) than Stage 2 (3.36 s).Source: F Multi-stage Training · p. 17
  • Stage 1 trains without memory retrieval using BPTT with 2 segments unrolled, then Stage 2 adds retrieval and unrolls up to 15 segments in this experiment.Stage 1 saves a checkpoint that the retrieval mechanism loads and extends in Stage 2.Source: F Multi-stage Training · p. 17

G HMT Memory Retrieval Behavior

HMT’s retrieval behavior follows context-switching patterns and avoids recalling inserted irrelevant segments; its effectiveness also depends on sensory-memory size.

G HMT Memory Retrieval Behavior

H Relationships Between Effectiveness and Size of Sensory Memory

I Dataset Construction for PubMedQA

The appendix describes how PubMedQA long-context samples are constructed, analyzes gradient stability in HMT and RMT, and details training-memory optimizations.

I Dataset Construction for PubMedQA

  • PubMedQA long-context samples concatenate contexts from selected question-context-answer tuples, then append each tuple’s question and answer to the shared context.Controlling the number of tuples determines the fraction of useful context for each question; the pqa_artificial subset is split 75%/15%/10% for training, validation, and test.Source: I Dataset Construction for PubMedQA · pp. 18–19

J Gradient Stability in HMT and RMT

K Distributed Training with Memory Consumption Optimization

  • Training HMT can unroll up to 15 segments using CPU/GPU offloading and distributed optimizer states and gradients; LoRA rank 8 fits 7B models on four MI210 GPUs.The authors identify GPU memory used by parameters, intermediate data, and optimization data as constraints on BPTT depth.Source: K Distributed Training with Memory… · p. 21

Related papers

Loading 2405.06067v3…