Structured Agent Distillation for Large Language Model

TL;DR

This work tackles the deployment bottlenecks of large language model–based agents by introducing Structured Agent Distillation (SAD), a framework that compresses high-capacity ReAct-style teachers into smaller students while preserving both reasoning fidelity and action execution. SAD explicitly segments agent trajectories into [REASON] and [ACT] spans and trains with span-specific supervision and curriculum learning, using binary masks to apply distinct CoT- and action-focused losses. Across ALFWorld, WebShop, and HotPotQA-ReAct, SAD consistently outperforms token-level distillation in task success, reasoning efficiency, and CoT coherence, while achieving substantial compression. Ablation and scaling studies show that span-level alignment is crucial for robust, deployable agents, with promising prospects for extension to multimodal and memory-enabled settings.

Abstract

Large language models (LLMs) exhibit strong capabilities as decision-making agents by interleaving reasoning and actions, as seen in ReAct-style frameworks. Yet, their practical deployment is constrained by high inference costs and large model sizes. We propose Structured Agent Distillation, a framework that compresses large LLM-based agents into smaller student models while preserving both reasoning fidelity and action consistency. Unlike standard token-level distillation, our method segments trajectories into {[REASON]} and {[ACT]} spans, applying segment-specific losses to align each component with the teacher's behavior. This structure-aware supervision enables compact agents to better replicate the teacher's decision process. Experiments on ALFWorld, HotPotQA-ReAct, and WebShop show that our approach consistently outperforms token-level and imitation learning baselines, achieving significant compression with minimal performance drop. Scaling and ablation results further highlight the importance of span-level alignment for efficient and deployable agents.

Figures (9)

• Figure 1: Structured Agent Distillation framework.
• Figure 2: Scaling behavior of student agents across model sizes. Top-Left: Task Success Rate (%). Top-Right: Reasoning Length (tokens). Bottom-Left: Chain-of-Thought (CoT) Match Rate (%). Bottom-Right: Latency (steps per episode). Structured Agent Distillation consistently outperforms MiniLLM-OPT baselines and better approaches teacher performance as model capacity increases.
• Figure 3: Illustration of the teacher-student interaction in our Structured Agent Distillation framework. The teacher model generates both the chain-of-thought (CoT) ott2023thoughtsource reasoning steps and the corresponding action trajectory based on the task input. The student model learns to imitate these trajectories by minimizing both reasoning-level and action-level distillation losses. This joint supervision allows the student to capture both high-level problem-solving strategies and low-level task executions.
• Figure 4: Illustration of the segmented trajectory structure used in our Structured Agent Distillation framework. The teacher’s trajectory is divided into a reasoning segment (language-based, multi-step inference) and an action segment (tool invocation or final structured decision). Each segment is used to compute a dedicated imitation loss for the student model.
• Figure 5: An overview of the loss flow in our Structured Agent Distillation framework. The student model learns from the teacher's trajectory by decomposing the learning signals into reasoning and action objectives. These loss terms are aggregated to guide effective student optimization.