Table of Contents
Fetching ...

From Isolated Conversations to Hierarchical Schemas: Dynamic Tree Memory Representation for LLMs

Alireza Rezazadeh, Zichao Li, Wei Wei, Yujia Bao

TL;DR

MemTree tackles the long-standing memory bottleneck in large language models by introducing a dynamic, tree-structured memory that mirrors human schemas. Information is incrementally inserted via root-to-leaf traversals, with adaptive depth-aware thresholds guiding structure growth and parent-node content aggregation to maintain coherent hierarchies. Retrieval uses a collapse-based embedding similarity approach, enabling efficient online querying across the memory without full reconstructions. Across four benchmarks—multi-turn dialogues and multi-document question answering—MemTree outperforms online baselines and closely approaches offline methods, while enabling real-time updates and scalable memory management with interpretable hierarchical representations.

Abstract

Recent advancements in large language models have significantly improved their context windows, yet challenges in effective long-term memory management remain. We introduce MemTree, an algorithm that leverages a dynamic, tree-structured memory representation to optimize the organization, retrieval, and integration of information, akin to human cognitive schemas. MemTree organizes memory hierarchically, with each node encapsulating aggregated textual content, corresponding semantic embeddings, and varying abstraction levels across the tree's depths. Our algorithm dynamically adapts this memory structure by computing and comparing semantic embeddings of new and existing information to enrich the model's context-awareness. This approach allows MemTree to handle complex reasoning and extended interactions more effectively than traditional memory augmentation methods, which often rely on flat lookup tables. Evaluations on benchmarks for multi-turn dialogue understanding and document question answering show that MemTree significantly enhances performance in scenarios that demand structured memory management.

From Isolated Conversations to Hierarchical Schemas: Dynamic Tree Memory Representation for LLMs

TL;DR

MemTree tackles the long-standing memory bottleneck in large language models by introducing a dynamic, tree-structured memory that mirrors human schemas. Information is incrementally inserted via root-to-leaf traversals, with adaptive depth-aware thresholds guiding structure growth and parent-node content aggregation to maintain coherent hierarchies. Retrieval uses a collapse-based embedding similarity approach, enabling efficient online querying across the memory without full reconstructions. Across four benchmarks—multi-turn dialogues and multi-document question answering—MemTree outperforms online baselines and closely approaches offline methods, while enabling real-time updates and scalable memory management with interpretable hierarchical representations.

Abstract

Recent advancements in large language models have significantly improved their context windows, yet challenges in effective long-term memory management remain. We introduce MemTree, an algorithm that leverages a dynamic, tree-structured memory representation to optimize the organization, retrieval, and integration of information, akin to human cognitive schemas. MemTree organizes memory hierarchically, with each node encapsulating aggregated textual content, corresponding semantic embeddings, and varying abstraction levels across the tree's depths. Our algorithm dynamically adapts this memory structure by computing and comparing semantic embeddings of new and existing information to enrich the model's context-awareness. This approach allows MemTree to handle complex reasoning and extended interactions more effectively than traditional memory augmentation methods, which often rely on flat lookup tables. Evaluations on benchmarks for multi-turn dialogue understanding and document question answering show that MemTree significantly enhances performance in scenarios that demand structured memory management.

Paper Structure

This paper contains 66 sections, 2 theorems, 12 equations, 14 figures, 14 tables, 1 algorithm.

Key Result

Theorem 1

Assuming the data processed by MemTree satisfies the $\beta$-well-separated condition (see Appendix append:data-separation), the hierarchy maintained by MemTree achieves a revenue where $T^*$ is the optimal hierarchy maximizing the Moseley-Wang revenue.

Figures (14)

  • Figure 1: MemTree (subset) developed on the MultiHop dataset tang2024multihop. MemTree updates its structured knowledge when new information arrives, enhancing inference-time reasoning capabilities of LLMs.
  • Figure 2: Illustration of MemTree. MemTree represents knowledge schema via a dynamic tree. Both parent and leaf nodes archive textual content, summarizing information relevant to their respective levels. Upon receiving new information, the system begins traversal from the root node. If the new information is semantically akin to an existing leaf node under the current node, it is routed to that node. Conversely, if it diverges from all existing leaf nodes under the current node, a new leaf node is created under the current node, concluding the traversal. During this process, all ancestor nodes will integrate the new information into the higher-level summaries they maintain.
  • Figure 3: Visualization of the Learned MemTree Structure on the MultiHop RAG Dataset. Due to space limitations, we display only a small subtree from the entire tree (a larger subtree is depicted in Figure \ref{['fig:tree-banner']}). As we traverse deeper into the tree, the content stored in the nodes becomes increasingly specific. For instance, the three blue nodes shown in the bottom right corner begin with a general summary of the USMNT's 3-1 defeat to Germany, then branch into specific insights on individual performances and team dynamics, and ultimately delve into a detailed analysis of Gio Reyna's impact during the match. Note that all intermediate contents in the parent nodes are generated by MemTree during the node update step. This hierarchical organization demonstrates how MemTree efficiently stores and retrieves information, progressing from overarching concepts to specific details.
  • Figure 4: Depth-based Stats of MemTree learned on MultiHop RAG
  • Figure 5: Efficiency of MemTree vs. RAPTOR and GraphRAG: MemTree's top-down insertion strategy allows content aggregation and embedding updates to be parallelized on the CPU, significantly accelerating memory updates as memory grows. Despite its cumulative cost being approximately 1.4x higher than the offline algorithms (RAPTOR and GraphRAG), it remains manageable. Results are reported on the MultiHop dataset.
  • ...and 9 more figures

Theorems & Definitions (3)

  • Theorem 1: Approximation Guarantee of MemTree (Informal)
  • Theorem 1: Approximation Guarantee of MemTree (Informal)
  • Definition 1: Moseley-Wang Revenue moseley2017approximation