TreeAgent: Structure-Aware Memory Agent for Processing Long Hierarchical Data
Wei Gu and colleagues at Zhejiang University and Alibaba Group present TreeAgent, a wrapper for recurrent memory agents that uses the hierarchy of JSON, HTML, YAML or JavaScript inputs to decide how to segment them and how to recall earlier content.
Ask this paper
Problem. Memory agents split long inputs into fixed token chunks, which can separate fields of one record, drop the structural position of a chunk, and overwrite details of earlier sibling records before they are needed.
Subtree segmentation. Segments are built from complete subtrees and groups of adjacent siblings, with proven upper bounds on segment count used to pick a token budget that matches fixed-length segmentation.
Paths and residual links. Each segment carries its structural path, and residual memory links let parent-level segments read evidence from child-level segments directly instead of through every intermediate memory update.
Results. Across eight retrieval and aggregation tasks with inputs from 8K to 1M tokens, TreeAgent raises the average score for all four memory agents it wraps, by 3.61% overall, with no extra training.
Abstract
Practical Web applications increasingly require agents to interpret long conversations or documents. For inputs exceeding the context windows of large language models (LLMs), memory agents process segments sequentially and retain past information in bounded memory. However, they typically divide inputs into flat token segments without preserving structure or hierarchy, causing three limitations for structured data such as API responses or HTML documents. First, fixed-length segmentation can split fields of the same record, breaking their relationships. Second, segments lack hierarchical information, obscuring their structural context and positions. Third, successive memory updates may discard details of earlier sibling records needed for reasoning on later segments. To address these limitations, we propose TreeAgent, a structure-aware agent that uses the input hierarchy to guide both segment-wise context processing and access to past information. Specifically, TreeAgent constructs segments from complete subtrees and groups of adjacent sibling subtrees to preserve local structural relationships. We provide theoretical guarantees on segment counts to guide token budget selection toward matching the segment count of fixed-length segmentation. We further attach structural paths to the segments for the agent to identify their ancestors and hierarchical positions. We also introduce residual memory links from the segments of low-level children to the segments of high-level parents, enabling direct information flows beyond successive memory updates. We evaluate TreeAgent with four representative memory agents on eight retrieval and aggregation tasks. The inputs are derived from real-world application data in JSON, HTML, YAML, and JavaScript and range from 8K to 1M tokens. TreeAgent improves the average score across the four datasets for all four agents, with 3.61% performance gain on average.