HMT : Hierarchical Memory Transformer for Efficient Long Context Language Processing (2025)

Abstract: 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.

## introduction transformer @xcite has demonstrated its strength in contextual learning and is utilized in various applications in language processing @xcite and computer vision @xcite . for a decoder-only transformer model, each transformer block contains a 1 https://github.com/oswaldhe/hmt-pytorch self-attention and a feedforward network module. an optimized self-attention layer has a quadratic computational and linear space complexity @xcite regarding the sequence length since it computes interactions between each token and all previous tokens in the sequence. to maintain the inference speed and satisfy memory requirements, most transformer models enforce maximum sequence length. for example, the llama 3 model is designed to process 8192 tokens @xcite and the llama 2 can process up to 4096 tokens @xcite . however, real-world applications involving long documents, such as book summarization @xcite and lifelong question-answering tasks @xcite , can have an enormous or even infinite stream of inputs. existing research attempts to build long context transformers using sparse attention @xcite @xcite , retrieval-augmented models @xcite , and recurrent sequence models @xcite @xcite . still, these models face at least one of two issues: (1) difficulty in adapting to future models due to a change in the core model architecture and (2) low effectiveness for long-range inputs under frequent context switching. in this work, we propose the hierarchical memory transformer (hmt), a novel framework to enable and augment models' long-context processing ability. hmt transforms models into a memory-augmented recurrent model that imitates the brain's memory hierarchy and human memorization behavior. it has the following unique features: hierarchical memorization: hmt mimics the memory hierarchy of the brain @xcite employing both learned memory tokens and current input tokens. hmt stratifies memory into sensory, short-term, and long-term, with interactions between each other. ## related works and problem formulation we will first discuss the existing efforts on longrange transformers and recurrent sequence models for infinitely long context language processing. then, we highlight a problem that is crucial in realworld applications. ## hierarchical memory transformer the main idea of hmt is to store information hierarchically and search for relevant information throughout the memory hierarchy. ## experiment we benchmark hmt with a variety of backbone models including smollm 135m @xcite , opt 350m, opt 2.7b @xcite , openllamav2 3b @xcite , rwkv 3b @xcite , and llama 2 7b @xcite , under the same memory constraint (i.e. same maximum context window). moreover, we test several models targeting long contexts (mamba 370m @xcite , yi-6b-200k @xcite , and mistral 7b @xcite ) to demonstrate the benefit hmt has on generation quality and memory consumption. we evaluate hmt with state-space models (rwkv and mamba) as backbones since we believe that models which can already process infinitely long inputs would benefit even further from hmt. all models mentioned are trained and assessed on 4 amd mi210 gpus, which can handle models up to 7b parameters. we further test hmt on 4 nvidia a100-80gb gpus for the qwen 2.5 14b model @xcite to justify its scalability to larger models and gain a consistent effectiveness boost. to tune the extra parameters introduced by hmt, we use the redpajamav2 (computer, 2023) dataset to pre-train each model. notice that hmt introduced new model hyperparameters on top of the backbone model (l, j, n , and k). a common configuration is @xmath0, @xmath1, @xmath2, and @xmath3, and we adjust these values for each model to achieve the best performance. to compare with previous works (rmt, longmem, memorizing transformer, ccm), we apply the same backbone models if the method is applicable to any model, or find a backbone model with a similar size if the method requires special architecture. for the long-context benchmark, we select subsets (narrativeqa, qasper, and multifieldqa-en for single document qa; hotpotqa, 2wikimqa, and musique for multi-document qa; govreport, qmsum, and multi-news for summarization; triviaqa for few-shot learning) from a widely acknowledged benchmark, longbench @xcite , and measure them against models reported in the longbench leaderboard. however, the maximum average document length of test sets in long-bench is shorter than 20k words, which is not very long for modern long-context models. to better understand hmt's long-context processing ability under various context scenarios, we further study hmt on crafted and controllable dataset samples. for crafted datasets, we derive from existing datasets to form long inputs. for general language tasks, models are tested for next token generation tasks with wikitext-103 @xcite (2-3k words per sample) and pg-19 (rae et al., 2019) datasets (69k words per sample on average). samples will be concatenated or split into chunks to form longer samples and investigate the relationships between input length and the effectiveness of the model. for question-answering tasks, we chose pubmedqa @xcite , which is a biomedical question-answering dataset with corresponding contexts. we artifact the dataset to assess hmt with multi-context inputs, described in appendix i. ## results and key observations in this section, we illustrate the main result of hmt. more ablation studies are in appendix e and g. ## conclusion we present hmt, a framework to augment models' long-range language processing ability with context switching. inspired by the brain's memory hierarchy, hmt imitates human memorization behavior by deploying hierarchical memory and the memory retrieval mechanism. hmt consistently improves the generation quality of the backbone models. compared with other long-context llms and memory-augmented models, hmt achieves higher generation quality with lower memory requirements. our model provides llm accessibility to resource-constrained applications and represents a step forward to lifelong language tasks.