🤖 machine learning

Agent Memory Distillation: Empowering Small LLM Agents with Hierarchical Teacher Memory

This paper introduces Agent Memory Distillation (AMD), a training-free framework that enhances small language model agents by transferring hierarchical knowledge—comprising workflow strategies, subtask examples, and function conventions—from a large teacher agent, resulting in significant performance improvements across multiple tool-use benchmarks.

Original authors: Taeil Kim, Kangsan Kim, Sung Ju Hwang

Published 2026-08-10

📖 8 min read🧠 Deep dive

CC BY 4.0

Original authors: Taeil Kim, Kangsan Kim, Sung Ju Hwang

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are trying to teach a brand-new, eager apprentice how to fix a complex machine. You have a master mechanic who can solve any problem in seconds, but your apprentice is small, inexperienced, and prone to making mistakes. In the world of artificial intelligence, these "apprentices" are small language models—smart, but limited. They try to solve tasks by using digital tools, like APIs or code, but they often get stuck, forget the steps, or call the wrong tools. To help them, researchers have tried giving them a "memory," a notebook of past successes and failures to look back on. However, there's a catch: because the apprentice is so new, their notebook is mostly filled with pages of "I failed here" and "I got lost there." They don't have enough successful stories to learn from. This is the puzzle scientists are trying to solve: How do you give a small, struggling student the benefit of a master's experience without just dumping a giant, confusing textbook on their desk?

This is where the paper "Agent Memory Distillation" comes in. The researchers, led by Taeil Kim and colleagues from KAIST, propose a clever new way to transfer knowledge from a powerful "teacher" AI to a smaller "student" AI. They call their method Agent Memory Distillation (AMD). Instead of just handing the student a giant, messy pile of the teacher's notes, AMD acts like a master editor. It takes the teacher's successful journeys and breaks them down into three specific, easy-to-digest types of memories: a high-level Workflow (the big plan), a Subtask guide (the specific steps to get things done), and a Function reference sheet (how to use the tools correctly). The paper suggests that by organizing the teacher's knowledge this way, the small student can understand and use it much better than if they just tried to read the raw data.

The Problem: A Small Student with a Big Gap

Think of a small AI agent like a new intern at a tech company. When they try to do a job, like "log into this app and send a message," they often fail. If you ask them to keep a diary of what happened, their diary is mostly full of "I tried to log in, but I got an error" or "I clicked the wrong button." Because they fail so often, they don't have a good library of "how to succeed" stories to learn from.

Researchers tried a simple fix: "Let's just give the intern the master mechanic's diary!" They took the successful notes from a giant, super-smart AI (the teacher) and gave them to the small AI (the student). But this didn't work well. It was like giving a child a textbook written in advanced physics; the child just couldn't understand the complex language or figure out how to apply it. The "capability gap" was too wide. The student was too small to make sense of the teacher's big, abstract ideas.

The Solution: The Three-Part Memory Lunchbox

The authors realized that to bridge this gap, they couldn't just dump the teacher's memory on the student. They needed to "distill" it—like turning a complex soup into a concentrated, easy-to-swallow broth. They created a system that organizes the teacher's success stories into three distinct layers, each serving a different purpose:

Workflow Memory (The Map): This is the big picture. Before the student even starts the task, the system gives them a high-level map. It says, "First, you need to log in. Then, find the data. Finally, send the message." It doesn't give the code; it gives the strategy. It's like telling the apprentice, "To fix the engine, you first check the oil, then the spark plugs, then the battery." This helps the student know the order of operations before they get lost.
Subtask Memory (The Recipe): Once the student knows the plan, they need to know how to do each step. This memory provides concrete examples of how the teacher handled specific parts of the job. If the task is "fetch a list of items," the Subtask memory shows the student exactly how the teacher wrote the code to get that list, including how to handle pagination (getting page 1, then page 2, etc.). It's like handing the apprentice a recipe card that says, "Mix 2 cups of flour and 1 cup of sugar," rather than just saying "bake a cake."
Function Memory (The Tool Manual): This is the safety net. When the student tries to use a tool (like calling a specific function) and gets an error, the system instantly pulls up a "Function Memory" entry. This entry shows the teacher's correct way to use that specific tool, including the exact arguments needed. It's like a mechanic looking at a manual the moment they realize they are holding the wrench upside down.

How It Works in Practice

The process is a bit like a smart tutoring session.

Proactive Injection: At the very start of a task, the student gets the Workflow and Subtask memories injected into their "brain" (their system prompt). They see the map and the recipes before they take a single step.
Reactive Injection: If the student makes a mistake and an error pops up, the system immediately grabs the Function memory related to that specific error and shows it to the student as a hint. This happens only when needed, so the student isn't overwhelmed with information when things are going well.

The Results: Small Models Getting Big Boosts

The researchers tested this method on three different benchmarks (AppWorld, BFCL V3, and ToolSandbox) using four different small student models (ranging from 4 billion to 8 billion parameters) and a powerful GPT-5-mini as the teacher.

The results were impressive. The small students, when equipped with this hierarchical memory, saw massive improvements:

On AppWorld, their accuracy jumped by an average of 27.2 percentage points.
On BFCL V3, they improved by 11.2 percentage points.
On ToolSandbox, they gained 3.4 percentage points.

In some cases, the small students didn't just get better; they actually performed as well as, or even better than, the giant teacher model itself! For example, on the AppWorld benchmark, a 4-billion-parameter student using AMD scored 49.40%, which was nearly identical to the teacher's 50.00%.

What the Data Tells Us

The paper digs deeper to understand why this worked so well.

The "Subtask" is King: They found that the Subtask memory (the concrete examples) was the most important part. It provided the biggest boost. This suggests that small models really need to see how to do things, not just what to do.
Size Matters (But Not Too Much): The 4-billion-parameter models benefited the most. They were smart enough to understand the teacher's memory but still had enough room to improve. The larger 8-billion models were already doing pretty well on their own, so the boost was smaller, though still positive.
The Teacher Must Be Compatible: Interestingly, having the strongest possible teacher didn't always mean the best results for the student. Sometimes, a teacher that was slightly less powerful but more "compatible" with the student's style worked better. It's not just about having a genius teacher; it's about having a teacher whose style the student can actually learn from.
Less is More: When they tried to give the student more memories (retrieving 5 instead of 1), performance actually dropped. The small models got confused by too much information. They needed just the right, high-quality hint, not a flood of data.

What This Means (And What It Doesn't)

The paper suggests that we don't need to train these small models from scratch or use expensive computing power to make them smarter. Instead, we can "distill" the experiences of larger models into a structured memory system that small models can actually use. This is a "training-free" approach, meaning the student model's internal weights don't change; it just learns to use the external memory better.

However, the authors are careful to note the limits. This method was tested on tasks involving code and structured tool use (like calling APIs). They aren't sure yet if it would work for tasks that require looking at pictures or writing creative stories from scratch. Also, the memory is "frozen"—it's built from the teacher's past successes and doesn't update in real-time as the student learns new things.

In the end, Agent Memory Distillation offers a promising new way to level the playing field. It suggests that even small, limited AI agents can become powerful problem-solvers if we just take the time to organize their teacher's wisdom into a format they can actually understand. It's a reminder that sometimes, the best way to learn isn't to read the whole encyclopedia, but to have a smart guide hand you the right page at the right time.

Technical Summary: Agent Memory Distillation (AMD)

Problem Statement

While memory systems have demonstrated significant potential for improving the performance of Large Language Model (LLM) agents by enabling the reuse of successful patterns and avoidance of past failures, their application to small language models (SLMs) remains largely unexplored and ineffective. Small agents (e.g., 4B–8B parameters) typically suffer from low task success rates, resulting in memory repositories dominated by unsuccessful trajectories. Consequently, self-generated memory offers limited utility.

Furthermore, a direct transfer of memory from a high-capacity "teacher" agent to a small "student" agent yields only marginal improvements. This is attributed to a capability gap: small models often lack the in-context learning ability and instruction-following capacity to interpret and execute high-level strategies provided by a teacher. For instance, a teacher's memory might suggest a strategy like "Log in first," but the student may lack the specific procedural knowledge to execute the login sequence itself. Existing knowledge distillation methods often rely on costly parameter updates (fine-tuning) or fail to address the specific granularity mismatch required for agent memory transfer.

Methodology: Agent Memory Distillation (AMD)

The authors propose Agent Memory Distillation (AMD), a training-free framework that transfers structured knowledge from a teacher agent to a student agent through a hierarchical memory system. Instead of a monolithic memory representation, AMD constructs three complementary memory types from successful teacher trajectories, organized by task granularity:

Workflow Memory (Task-Level):
- Content: Encodes high-level task-completion strategies, including key preconditions, decision rules, and common failure patterns.
- Format: Verbalized insights in natural language where concrete runtime values (IDs, credentials) are replaced with typed placeholders (e.g., <ID>, <EMAIL>).
- Injection: Proactively injected into the system prompt at the start of a task to provide a coherent high-level plan.
Subtask Memory (Intermediate Granularity):
- Content: Provides concrete behavioral examples bridging high-level plans and low-level tool calls. It captures semantically meaningful units of behavior (e.g., authenticating with a service, paginating through results).
- Format: Segments of the teacher's trajectory containing executable code or tool calls paired with observations, accompanied by a short natural language label and description.
- Injection: Proactively injected after the student decomposes the task into subtasks. The student retrieves the best-matching segment for each subtask label.
Function Memory (Function-Level):
- Content: Captures fine-grained tool invocation knowledge, including specific function schemas, argument constraints, and usage examples.
- Format: Concrete teacher examples of function calls, optionally augmented with API documentation.
- Injection: Reactively retrieved only when the student encounters a tool-calling error. The failing function name is used to index the memory, providing targeted corrective guidance without inflating the context during successful execution.

Inference Pipeline:
At inference time, the student agent ( $\pi_S$ ) first retrieves and injects Workflow and Subtask memories to guide initial planning. During execution, if a tool call fails, the system reacts by retrieving the relevant Function memory entry to correct the error. This process is governed by embedding-based retrieval (using cosine similarity) and deduplication mechanisms.

Key Contributions

Hierarchical Memory Structure: The paper introduces a novel three-tiered memory architecture (Workflow, Subtask, Function) specifically designed to bridge the capability gap between teacher and student agents, making teacher knowledge comprehensible to small models.
Training-Free Distillation: Unlike prior agent distillation works that rely on parameter updates, AMD operates entirely through memory injection, making it a scalable and cost-effective approach for enhancing small agents.
Proactive vs. Reactive Injection: The framework distinguishes between proactive planning support (Workflow/Subtask) and reactive error correction (Function), optimizing context usage and preventing information overload.
Systematic Investigation: This work represents the first systematic study of teacher-to-student memory transfer for small agents, explicitly addressing the limitations of direct transfer and the importance of teacher-student compatibility.

Experimental Results

The authors evaluated AMD on three tool-use benchmarks: AppWorld, BFCL V3, and ToolSandbox, using four student models (Qwen3-4B, Qwen3-8B, Gemma4-E4B, Llama3.1-8B) and GPT-5-mini as the teacher.

Performance Gains: AMD consistently outperformed zero-shot baselines and existing memory-based methods (ReasoningBank, MemP, SASM).
- AppWorld: Average accuracy gain of 27.2%p (e.g., Qwen3-4B improved from 14.88% to 49.40%).
- BFCL V3: Average accuracy gain of 11.2%p.
- ToolSandbox: Average accuracy gain of 3.4%p.
Teacher-Level Performance: On AppWorld, AMD enabled 4B and 8B student models to match or even surpass the performance of the GPT-5-mini teacher (e.g., Gemma4-E4B achieved 54.17% vs. teacher's 50.00%).
Interaction Efficiency: AMD significantly reduced the number of interaction turns required by students, bringing their behavior closer to the efficient trajectories of the teacher.
Ablation Findings:
- Subtask Memory contributed the largest performance gains, highlighting the critical need for concrete behavioral examples.
- Model Size: 4B models benefited most relatively, while larger models (8B+) saw diminishing returns as they approached teacher-level performance.
- Representation: Natural language was optimal for Workflow memory, while code-centric formats were essential for Subtask and Function memories.

Significance and Claims

The paper claims to provide a solid foundation for future research on agent memory distillation and the development of more capable small LLM agents. By identifying that direct memory transfer fails due to the capability gap, the authors demonstrate that structured, hierarchical knowledge transfer is necessary to make teacher experiences accessible to small students.

The authors modestly note that their work is currently limited to text-based tool-use benchmarks involving structured APIs and does not yet generalize to multimodal environments or open-ended code generation. Additionally, the framework relies on offline memory construction from a fixed set of teacher trajectories, meaning it cannot adapt to distribution shifts or incorporate the student's own test-time experiences. Finally, the effectiveness of the method depends on the compatibility between the teacher and student, suggesting that adaptive teacher selection remains an open problem.