TA-Mem: Tool-Augmented Autonomous Memory Retrieval for LLM in Long-Term Conversational QA

TL;DR

TA-Mem introduces an adaptive, tool-augmented memory retrieval framework for LLMs, enhancing long-term conversational QA by dynamically extracting and utilizing structured memories, leading to improved performance and flexibility.

Contribution

The paper presents a novel framework combining adaptive memory extraction and multi-tool retrieval for LLMs, addressing limitations of static retrieval methods.

Findings

Significant performance improvements on LoCoMo dataset.

Effective adaptive memory extraction and retrieval demonstrated.

Enhanced flexibility in handling diverse question types.

Abstract

Large Language Model (LLM) has exhibited strong reasoning ability in text-based contexts across various domains, yet the limitation of context window poses challenges for the model on long-range inference tasks and necessitates a memory storage system. While many current storage approaches have been proposed with episodic notes and graph representations of memory, retrieval methods still primarily rely on predefined workflows or static similarity top-k over embeddings. To address this inflexibility, we introduced a novel tool-augmented autonomous memory retrieval framework (TA-Mem), which contains: (1) a memory extraction LLM agent which is prompted to adaptively chuck an input into sub-context based on semantic correlation, and extract information into structured notes, (2) a multi-indexed memory database designed for different types of query methods including both key-based lookup and similarity-based retrieval, (3) a tool-augmented memory retrieval agent which explores the memory autonomously by selecting appropriate tools provided by the database based on the user input, and decides whether to proceed to the next iteration or finalizing the response after reasoning on the fetched memories. The TA-Mem is evaluated on the LoCoMo dataset, achieving significant performance improvements over existing baseline approaches. In addition, an analysis of tool use across different question types also demonstrates the adaptivity of the proposed method.