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

Taeil Kim, Kangsan Kim, Sung Ju Hwang

Agent Memory Distillation (AMD) transfers hierarchical, teacher-generated memory to small LLMs to boost tool-use performance.

How can we effectively transfer complex task-solving strategies from a large teacher agent to a small student agent without retraining?

Small language models struggle to build effective memory because their own task-solving attempts are often unsuccessful, leaving them with repositories dominated by failure. AMD is a training-free framework that distills successful trajectories from a large teacher agent into a hierarchical memory store, providing the student with task-level plans, subtask-level behavioral examples, and function-level error recovery. This approach consistently outperforms zero-shot baselines, with average accuracy gains of 27.2 percentage points on the AppWorld benchmark.

Paper Primer

AMD organizes teacher knowledge into three granularities: Workflow memory for high-level strategy, Subtask memory for concrete execution segments, and Function memory for reactive error correction. The system acts like a mentor providing a playbook: it gives the student a high-level strategy before starting, concrete examples for each step, and specific "fix-it" notes only when a tool call fails.

AMD significantly improves task success rates for small (4B–8B) student models across diverse tool-use benchmarks.

Average accuracy gains of 27.2 pp on AppWorld, 11.2 pp on BFCL V3, and 3.4 pp on ToolSandbox compared to zero-shot performance. Substantial performance boost, with some student models matching or exceeding the teacher's own accuracy.

The hierarchical structure is critical: Subtask memory provides the largest incremental gains, while Function memory offers targeted guidance that prevents the context inflation associated with naive, flat memory transfer.

Why does AMD use a hierarchical structure instead of just providing the student with the teacher's full successful trajectories?

Small models have limited in-context learning and instruction-following capacity; injecting full trajectories introduces noise that exceeds their comprehension. Hierarchical distillation breaks this knowledge into manageable, context-appropriate pieces.

Does the student model simply memorize the teacher's actions, or does it learn to generalize?

The student re-instantiates distilled decision-making patterns under its own inductive biases, allowing it to achieve performance that occasionally surpasses the teacher's own accuracy on specific benchmarks.

For developers of small-scale agents, AMD demonstrates that high-quality, structured teacher experience is a more effective lever for performance than model scaling or parameter-heavy fine-tuning.

Introduction

We expose why small LLM agents struggle and introduce Agent Memory Distillation to bridge the gap.

Memory has proven valuable for large agents, yet small LLM agents rarely generate enough successful trajectories to populate a useful memory store. Consequently, their memories are filled with failed attempts, limiting any performance boost from memory reuse.

**Figure 1.** **Motivation, Concept, and Results of AMD.** (A) Student-generated memory is limited by low task success rates. (B) Naive teacher memory transfer yields marginal gains due to the capability gap. (C) AMD transfers hierarchical memories spanning task, subtask, and function levels, making teacher knowledge accessible to small students. (D) AMD achieves significant accuracy gains across three benchmarks.

To close the capability gap we introduce Agent Memory Distillation (AMD), a training‑free framework that distills teacher experiences into three complementary memory types. Workflow memory encodes task‑level strategies, Subtask memory supplies concrete behavioral examples, and Function memory captures tool‑calling conventions, enabling small students to leverage teacher knowledge effectively.

Related Work

We situate our approach among prior LLM agents with memory and knowledge‑distillation work.

Prior work on LLM agents can be grouped into two strands: (1) augmenting agents with memory to reuse past experience, and (2) applying knowledge‑distillation techniques to compress teacher expertise into smaller students.

A zero‑shot baseline runs the student agent on a task without any injected memory or distilled knowledge, relying solely on its pretrained parameters.

The AMD Framework

3 Agent Memory Distillation

We present AMD, a novel framework that transfers teacher memory to a small student agent. We first formalize the problem setting in Section 3.1, then describe multi‑level memory generation by a teacher in Section 3.2 and memory retrieval and utilization by a student in Section 3.3.

3.1 Problem Formulation

We focus on multi‑turn tool‑use reasoning tasks, where an agent $\pi$ operates in an interactive environment by issuing a sequence of tool calls from a predefined set and receiving observations in return. Given a task s, the agent produces a trajectory $\tau$ = (a₁, o₁), …, ($a_T$, $o_T$), where $a_t$ ∈ F is a tool call and $o_t$ is the resulting observation at step t. We distinguish between a teacher agent $\pi_{T}$, backed by a large and capable language model, and a student agent $\pi_{S}$, backed by a small model (4 B or 8 B parameters). Given a task set, we run $\pi_{T}$ to collect a set of teacher trajectories $S_T$ = {$\tau_{T}$¹, …, $\tau_{T}$ᴺ}. We then use $\pi_{T}$ to construct a memory store which is subsequently transferred to $\pi_{S}$. At inference time, $\pi_{S}$ retrieves relevant memories from M to guide its reasoning on s. The goal of AMD is to maximize the task performance of $\pi_{S}$ through effective utilization of teacher‑generated memories:

$$ \max_{M}\; \mathbb{E}_{s\sim S}\big[ R\big(\pi_S(s; M)\big) \big] \tag{1} $$

3.2 Hierarchical Memory Generation

AMD constructs three types of memory from the teacher’s past experiences, organized at different levels of task granularity. AMD employs $\pi_{T}$ to construct each memory type from a subset of successful teacher trajectories, producing a memory store $f_n$, where $f_n$ denotes Workflow, Subtask, and Function memories.

**Workflow Memory** captures the teacher’s high‑level task‑completion strategy. For each successful trajectory $\tau_{T}$⁺, a verbalized insight is produced that describes the overall approach taken by the teacher in natural language, covering the apps and tools involved, key preconditions, decision rules for task completion, validation cues, and common failure patterns to avoid. The insight abstracts over concrete runtime values—identifiers, credentials, file paths, and other dynamic inputs are replaced with typed placeholders (e.g., <ID>, <EMAIL>, <`FILE_PATH`>) so that the memory remains applicable to future tasks with different specific inputs. Together with a natural‑language query that characterizes the task, the insight forms a workflow memory entry

$$ m^{\text{wf}}_i = (q_i, \text{ins}_i), $$

which is encoded into a dense vector for retrieval. The resulting memory bank provides the student agent with a high‑level task plan before execution begins, enabling it to identify the relevant tools and establish a coherent action sequence.

**Subtask Memory** provides concrete behavioral examples at an intermediate level of granularity, bridging the gap between high‑level workflow plans and low‑level tool calls. For each successful trajectory $\tau_{T}$⁺, we decompose the trajectory into a sequence of coherent subtask segments, where each segment corresponds to a semantically meaningful unit of the teacher’s behavior (e.g., authenticating with a service or executing a sequence of related API calls to fulfill one subtask). Segmentation is performed by a teacher LLM prompted to identify semantically coherent units in the trajectory, optionally guided by rule‑based heuristics that provide candidate breakpoint hints. Each segment $e_{i,k}$, which contains the concrete execution examples (tool calls or executable code paired with the corresponding observations from the teacher), is stored together with a label ℓ_{i,k} and a short natural‑language description $d_{i,k}$, forming a subtask memory entry

$$ m^{\text{st}}_{i,k} = (\ell_{i,k}, d_{i,k}, e_{i,k}). $$

Each description $d_{i,k}$ is encoded into a dense vector for retrieval, and all entries across all trajectories constitute the subtask memory bank.

**Function Memory** captures fine‑grained tool‑invocation knowledge at the level of individual function calls, built directly from the successful teacher trajectories and optionally augmented with the corresponding function documentation. For each successful trajectory $\tau_{T}$⁺, we extract its constituent function invocations, indexed by j. Each record stores the function name $f_{i,j}$ and a concrete teacher example $E_{i,j}$, and is optionally augmented with the API documentation doc($f_{i,j}$) that specifies the argument and response schema. The example $E_{i,j}$ comprises the executable invocation together with its surrounding context, along with any returned observation when available. This context makes the rationale for the call apparent, namely why the function was invoked at that step and what constraints govern its arguments. Together these form a function memory entry

$$ m^{\text{fn}}_{i,j} = (f_{i,j}, E_{i,j}, \text{doc}(f_{i,j})). $$

Unlike the workflow and subtask banks queried through dense‑vector similarity, function entries are indexed by …

Memory Retrieval and Injection

Retrieves and injects relevant memories into the agent’s context using cosine similarity and task‑aware queries.

At inference the agent often lacks the concrete examples needed to solve a new problem. By pulling past experiences that are semantically close to the current query, the agent can be given a concrete plan without additional training.

Pick the stored experiences whose vector embeddings are most aligned with the current query, and drop any that fall below a similarity cut‑off.

Compute cosine similarity: $\text{cos\_sim}(q,m_1)=0.99$, $\text{cos\_sim}(q,m_2)=0.71$, $\text{cos\_sim}(q,m_3)=0.71$.

All three exceed the $0.7$ threshold, so none are discarded.

Rank by similarity: $m_1$ (0.99) first, $m_2$ and $m_3$ tie (0.71). Keep the top‑2, i.e., $m_1$ and $m_2$.

The retrieved memories are $m_1$ and $m_2$, which will be injected into the prompt.

Even a modest threshold can filter out completely unrelated memories while still preserving multiple plausible candidates.

How does this retrieval differ from a standard nearest‑neighbor search over raw text?

Standard nearest‑neighbor search compares raw strings, which is brittle to wording variations. Here we first embed both query and memories into a semantic vector space, then compare using cosine similarity, so retrieval is robust to lexical differences and focuses on meaning.

Once the relevant memories are selected, they must be placed into the agent’s context before it begins acting. The paper distinguishes two proactive injection strategies: one for a high‑level workflow plan and another for detailed sub‑tasks.

Insert a concise, high‑level plan derived from the top‑k retrieved memories at the very start of the prompt, giving the agent a roadmap before any tool calls.

Encode the instruction and each workflow entry with the same embedding model.

Compute cosine similarity; $w_1$ scores 0.92, $w_2$ scores 0.85.

Select $w_1$ as the top‑1 entry.

Prepend “Propose times → create calendar event → email invites” to the system prompt.

The injected plan gives the agent a concrete sequence to follow, reducing the need for trial‑and‑error during execution.

Why not simply let the agent generate its own plan without injecting a retrieved workflow?

Without an injected plan the agent must discover a viable sequence from scratch, which often leads to inefficient or incorrect tool usage. The retrieved workflow provides a proven template that guides the agent toward successful execution.

Break the overall instruction into a short list of subtask labels, query the memory store for each label, and stitch the best matching examples together as part of the prompt.

Encode each subtask label and all subtask memories.

Compute cosine similarity; each label retrieves its corresponding entry with similarity >0.8.

Check for duplicates – none found.

Append the three retrieved snippets to the system prompt after the workflow plan.

Providing concrete code snippets for each subtask reduces the agent’s need to synthesize low‑level details on the fly.

How does this subtask‑wise retrieval avoid the redundancy that could blow up the prompt size?

After each label retrieves its best match, the system checks for identical memory entries across labels and removes duplicates, guaranteeing that each piece of information appears at most once.

When a tool call fails, look up past failure cases indexed by function name, rank them by relevance to the current task, and append the most helpful example as a hint to the error message.

Identify the failing function name “`fetch_url`”.

Retrieve both records $r_1$ and $r_2$.

Rank by similarity to the instruction; $r_1$ wins.

Format $r_1$ as a hint block and append it to the error message.

By presenting the most relevant past fix, the agent can immediately apply a proven remedy instead of guessing.

Why not always inject function memories pre‑emptively rather than reactively?

Pre‑emptive injection would bloat the prompt with many rarely used hints, wasting context space. Reactively adding only the needed hint keeps the prompt compact while still providing targeted guidance when a failure occurs.

Experimental Results

AMD injects teacher experience into small agents, dramatically boosting accuracy and efficiency.

Recall that small LLM agents lack successful experience; AMD distills teacher trajectories into a hierarchical memory that is injected at inference time.

Agent Memory Distillation (AMD) yields large accuracy gains over zero‑shot baselines across all benchmarks.

Table 1 shows average gains of 27.2 pp on AppWorld, 11.2 pp on BFCL V3, and 3.4 pp on ToolSandbox.

**Figure 3.** Interaction steps on two benchmarks. AMD brings the student's turn count closer to the teacher's, particularly where the zero-shot gap is large.

**Table 1.** Main results across three benchmarks. $\Delta$ shows absolute accuracy gains of AMD over zero-shot.

Ablation Studies

Ablation analysis reveals which memory components and teacher agents drive performance gains.

We isolate the impact of each memory type and each teacher model by measuring performance changes on the AppWorld and BFCL V3 benchmarks.

Adding Subtask memory to the workflow yields a +25.0 %p improvement for Qwen3‑4B on AppWorld over workflow alone.

Table 2 shows Qwen3‑4B accuracy rising from 30.36 % (WF) to 35.50 % (WF + ST).

Function memory provides a modest additional gain of +5.1 %p when combined with workflow and subtask memories.

Table 2 reports Qwen3‑4B accuracy increasing from 30.36 % (WF + ST) to 35.50 % (WF + ST + FN).

GPT‑5.5 as teacher produces the highest student accuracy for Qwen3‑8B (58.93 %).

Table 3 lists 58.93 % for the Qwen3‑8B student when paired with the GPT‑5.5 teacher.

GPT‑5‑mini as teacher yields the best Qwen3‑4B student performance (49.40 %).

Table 3 shows 49.40 % accuracy for Qwen3‑4B with the GPT‑5‑mini teacher.

**Table 2.** Ablation study across models and benchmarks. WF, ST, and FN denote workflow, subtask, and function memory.

**Table 3.** Effect of teacher agents on AppWorld. GPT-5-mini shows the strongest transfer effectiveness for Qwen3-4B student.

Scaling and Model Effects

Student size and memory design shape AMD performance on AppWorld.

AMD yields a peak accuracy gain of +34.52 %p at a 4 B student model.

Figure 4 shows the gain curve across model sizes.

**Figure 4:** Effect of student model size on AMD performance. Accuracy increases with model size, while accuracy gain peaks at 4B.

Increasing the retrieval count harms Subtask memory, dropping accuracy from 49.40 % to 33.34 % at k = 5.

Results in §5.4 show a monotonic decline as more memories are retrieved.

Pure‑text Subtask memory attains only 23.21 % accuracy, far below the 49.40 % achieved with code‑centric format.

Table 4 contrasts the two representations.

When all three memory types use natural‑language representations, overall accuracy collapses to 26.19 %.

Table 4 reports the combined effect.

**Table 4.** Effect of memory representation on AppWorld. Our design choice achieves the best accuracy.

The qualitative case study in §5.6 demonstrates how the three memory tiers jointly resolve successive failure modes, culminating in a fully successful transaction.

Algorithmic Details

Algorithm 1 details the AMD inference pipeline, showing proactive and reactive memory injection steps.

This section spells out the full AMD inference pipeline, distinguishing proactive workflow injection from reactive function‑memory injection.

Algorithm 1 – AMD student agent inference with hierarchical memory

Additional Experiments

Ablation studies reveal each memory type’s contribution and robustness under disjoint evaluation.

We first quantify how each memory type—Workflow, Subtask, and Function—affects performance across benchmarks, then illustrate concrete case studies, and finally test robustness when memory is distilled from disjoint tasks.

Adding Function memory to the WF+ST configuration harms LLaMA3.1‑8B on AppWorld, dropping accuracy from 30.36 % to 27.38 %.

Table 2 shows the decrease when FN memory is added to the WF+ST setup for LLaMA3.1‑8B.

For Gemma4‑E4B on AppWorld, injecting Function memory raises accuracy from 30.36 % to 40.48 %.

Table 2 reports the improvement when FN memory is added to the WF baseline for Gemma4‑E4B.

Overall, Subtask memory delivers the biggest incremental boost when layered on top of Workflow memory, while Function memory adds modest gains except for the LLaMA3.1‑8B anomaly. The Student‑only memory variant remains close to zero‑shot performance, confirming that teacher‑generated trajectories are essential.

**Figure 5.** Effect of retrieval count $k$ on AMD performance. Accuracy at $k=1$ is near-optimal for all memory types, and increasing $k$ generally degrades performance.

Disjoint‑evaluation protocol for Qwen3‑4B student agents.

Conclusion and Limitations

We discuss AMD’s remaining challenges and future research directions.

We briefly restate the core idea: AMD transfers teacher experiences to small agents via a hierarchical memory that is injected at inference time.

AMD builds three complementary memory stores—Workflow memory for high‑level planning, Subtask memory for concrete behavioral references, and Function memory for fine‑grained tool‑invocation guidance.

Across three benchmarks and four student models, AMD consistently beats zero‑shot and other memory‑based baselines, with some students reaching or exceeding teacher‑level performance.

Ablation studies confirm that each memory type adds a distinct benefit and that effective distillation requires matching memory complexity to the student’s capacity.

Limitations: first, AMD has only been tested on text‑based tool‑use tasks involving fixed Python APIs, so its applicability to multimodal settings or open‑ended coding tasks is unknown.

Second, the memory is pre‑computed from a static set of teacher trajectories and cannot adapt to a student’s own successes, failures, or distribution shifts at test time.

Third, performance hinges on the quality and compatibility of teacher trajectories; stronger teachers do not always translate into larger gains, leaving adaptive teacher selection as an open problem.

Implementation Details

Implementation details covering evaluation protocols, memory prompts, and retrieval mechanisms.

We evaluate two disjoint protocols. In cross‑split evaluation we reserve 30 % of each benchmark for evaluation and distill memory only from the remaining 70 % training split. In self‑excluded retrieval we distill over the full benchmark but forbid each task from retrieving its own distilled memory, preserving a larger pool while still preventing self‑reuse.

To test stability we repeat the zero‑shot and AMD configurations five times with the same student model. Table 5 shows that AMD consistently yields higher mean success rates and that the standard deviations stay around one percentage point, indicating stable gains across runs.

**Figure 14.** Sub-task Memory Segmentation Prompt for AppWorld Tasks

Workflow memory is generated from a prompt (Figure 13) that asks the teacher LLM to produce a short high‑level insight describing the overall task strategy, key preconditions, decision rules, validation cues, and common failure patterns, with concrete API calls written as app.`function_name` placeholders.

Function memory stores the function name together with a concrete example from a successful teacher trajectory; for richly documented benchmarks we also attach the argument and response schema, while for minimal tool suites we keep only the example call and surrounding context.

All memory entries are embedded with OpenAI’s text-embedding-3-small model and retrieved via cosine similarity. A minimum similarity threshold $\delta$ discards low‑scoring candidates, and we set $k = 1$ to fetch the top‑1 entry for each of workflow, sub‑task, and function memories. Deduplication across sub‑task segments prevents multiple injections of the same segment.

AMD consistently outperforms the zero‑shot baseline across benchmarks, achieving up to a 24.19‑point gain in success rate.

See Table 5 for mean success rates and Table 6 for protocol‑wise comparisons; AMD’s gains are stable across five repeated runs.

Memory Examples

Concrete examples illustrate how each memory type is built and used.

We now walk through concrete, head‑simulable instances of the three memory types introduced earlier, together with the prompts that generate them.

Step 1 – Retrieve passwords and log into Venmo (supervisor.`show_account_passwords` → venmo.login).

Step 2 – Paginate through pending requests: while True, fetch page i via `show_received_payment_requests`(status='pending', `page_index`=i); stop when the page is empty.

Step 3 – Filter each request’s

Step 4 – For each kept request, call `approve_payment_request`(`request_id`, `access_token`) and verify the response contains the word “message”.

Step 5 – After approvals, call `show_venmo_balance`, iterate over possible keys (balance, amount, `venmo_balance`, `current_balance`) until a numeric value is found, then withdraw that amount to the card whose last four digits are 8907.

Step 6 – Call supervisor.`complete_task`() to finish.

Injecting this workflow memory reduces the agent’s step count from >40 (exhausting the budget) to 6, and eliminates common bugs such as missing pagination or incorrect date parsing.

Step 1 – Agent receives

Step 2 – Subtask memory supplies a pattern: replace “Z” with “+00:00” before calling

Step 3 – Agent applies the fix, successfully obtains a

With the subtask memory, the agent resolves the timestamp mismatch in a single step, freeing the remaining budget for the withdrawal subtask.

Step 1 – Agent checks for

Step 2 – Function memory provides a loop over candidate keys (balance, amount, `venmo_balance`, `current_balance`) and selects the first present.

Step 3 – The loop finds

Step 4 – Agent proceeds to withdraw the amount, completing the task.

By iterating over possible keys, the function memory makes the balance extraction robust to API schema variations.

The prompts that create these memories are themselves concise templates that the teacher agent uses to distill successful trajectories.

Given a task instruction and a successful Python trajectory, produce a short paragraph that captures the high‑level strategy, the key APIs used, a validation cue, and a common failure pattern.

Analyze a successful trajectory and split it into coherent sub‑tasks, each labeled with a short verb phrase and a one‑sentence description.

Given a high‑level task, the student first logs into all required apps, then enumerates up to six sub‑tasks that will be used as retrieval queries.

Case study: approving all pending Venmo payment requests for the current month and withdrawing the remaining balance demonstrates the cascading benefit of each memory type.

Analysis shows that Workflow Memory narrows the request set, Subtask Memory fixes date‑parsing bugs, and Function Memory resolves response‑schema mismatches, together turning a failure that exhausts the 40‑step budget into a successful execution in under 25 steps.

Read the original paper

Open the simplified reader on Paperglide

Browse all simplified papers