AtomSD: Procedural Atom-Guided Self-Distillation for LLM Agents
Authors: Bohan Yu, Pengfei Cao, Kai Tian, Chen Han, Shi-Yang Li, Chenxi Zhou, Junlin Liu, Cao Liu, Ke Zeng, Jun Zhao, Kang Liu
Organizations: School of Advanced Interdisciplinary Sciences, University of Chinese Academy of Sciences · The Key Laboratory of Cognition and Decision Intelligence for Complex Systems, Institute of Automation, Chinese Academy of Sciences · School of Artificial Intelligence, University of Chinese Academy of Sciences · Meituan
Large language models (LLMs) possess broad factual knowledge and strong reasoning capabilities, but solving complex tasks also requires procedural knowledge about how to act. Existing approaches often represent procedural knowledge through trajectories or bundled skills, limiting flexible reuse across interaction histories. This paper introduces Procedural Atoms, fine-grained and independently applicable units that represent Intent, When, Do, and Avoid guidance. Building on this representation, we propose AtomSD, an Atom-guided Self-Distillation framework in which the learner continually extracts Procedural Atoms from its own successful trajectories and distills them into the evolving policy through Atom-conditioned supervision and task-reward optimization. Comprehensive experiments show that AtomSD consistently outperforms strong reward-based and distillation methods, exceeding the state-of-the-art baseline by up to 9.31 points in average performance. Further analyses demonstrate that the proposed Procedural Atoms provide effective guidance, highlighting the value of atom-level procedural knowledge for policy learning.
Figures & tables
Figure 1: Procedural Atoms enable fine-grained reuse and composition of procedural knowledge.
Figure 2: Overview of AtomSD. Procedural Atoms are continually abstracted from the student’s own successful trajectories and distilled back into the evolving policy through Atom-guided self-distillation.
Table 1: Main results across four benchmarks. denotes methods requiring no additional training; denotes methods that update model parameters. Higher is better in every column. Avg. is the average of the six reported scores; green arrows ( ↑ ) indicate our absolute gains over the highest baseline Avg. within each backbone. Best , Second-best , and Third-best scores are highlighted within each backbone.
Figure 4Figure 5
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
Benchmark
Task characteristics
Task types
Training
Evaluation
ALFWorld
Household interaction with object search, navigation, and prerequisite-aware manipulation
6 families
3,553
140 seen / 134 unseen
ScienceWorld
Multi-step scientific experimentation involving transformations, measurement, circuits, and biological reasoning
30 types
3,592
194 seen / 211 unseen
WebShop
Interactive product search, attribute matching, option selection, and purchase
Open-ended
10,587
200 test
StableToolBench
Multi-turn tool use involving API selection, parameterization, and execution
3 groups
10,129
169 solvable
Total
27,861
1,048
Appendix
Table 2: Dataset characteristics and task counts used for training and evaluation.
Parameter
Setting
Shared training settings
Backbones
Qwen3.5-4B / Qwen3.5-9B
Tasks per iteration
48 (12 per benchmark)
Rollouts per task
4
Optimization mini-batch
12 tasks / 48 trajectories
Optimization epochs
1
Appendix
Table 3: Shared training hyperparameters and AtomSD-specific settings.
Method
GRPO reward loss
Distillation weight
Teacher update
Distillation configuration
GRPO
Yes
—
—
Task-reward optimization only.
OPSD
No
1.0
Fixed
Sampled-token reverse-KL surrogate with successful same-task demonstrations.
GRPO+OPSD
Yes
0.001
Fixed
Same demonstration-conditioned distillation objective as OPSD.
SDPO
No
1.0
EMA ( α=0.95 )
Distribution-level self-distillation over the top- K tokens and the aggregated tail probability, with K=100 .
Table 4: Method-specific configurations for the training baselines.
Model
Stage
Hist.
N
Smax (%)
Dα (%)
Dβ (%)
Dαβ (%)
4B
Early
Same
175
36.9
68.4
64.9
67.4
Cross
144
16.9
70.8
72.9
75.0
Late
Same
175
21.2
52.7
50.0
56.3
Cross
144
5.8
47.1
46.1
44.2
9B
Early
Same
173
45.2
59.0
56.0
54.4
Cross
149
25.9
69.9
68.7
61.9
Appendix
Table 5: Procedural Atom internalization across model scales, training stages, and history settings. N denotes the number of analysis units; Smax measures maximum-token support change; Dα , Dβ , and Dαβ measure guidance-direction agreement for Atom ωα , Atom ωβ , and their composition, respectively.
Procedural Pattern
Core Dependency
Representative Case
Observation-Conditioned Branching
A downstream action depends on newly observed information.
ScienceWorld: AtomSD measures a substance at −31∘C and selects the destination specified for temperatures below 50∘C , completing the task.
Constraint-Driven Revision
An earlier choice must be revised after detecting a violated constraint.
WebShop: AtomSD rejects an over-budget product, returns to search, and selects a cheaper item satisfying the requested attributes.
Subgoal-State Dependency
A later subgoal requires the state produced by an earlier operation.
ALFWorld: AtomSD acquires a potato, heats it, and only then places the heated object in the target receptacle.
Feedback-Guided API Recovery
Uninformative API responses require adapting the query strategy based on returned feedback.
StableToolBench: After an empty response, AtomSD revises its pagination strategy, verifies the returned results, and reports nine products grounded in the API output.
Appendix
Table 6: Representative cases of recurring procedural patterns observed in evaluation trajectories.
Intent
When
Do
Avoid
Collect ingredients
Required ingredients have been identified from the recipe.
Move each required ingredient into a suitable common container.
Leaving ingredients outside the container or mixing prematurely.
Mix the contents
All required ingredients are in one accessible, stable container, and the recipe calls for mixing.
Mix the container and observe whether the target substance is produced.
Mixing incomplete or separated ingredients, or ignoring mixing errors.
Acknowledge the product
The desired substance has been created and its location is known.
Focus on the created substance and confirm task completion.
Focusing on the wrong object or continuing without acknowledging the goal.
Appendix
Table 7: Paraphrases of Procedural Atoms extracted from successful training trajectories.
Method
Score
Outcome
GRPO
33
Incomplete
OPSD
25
Incomplete
GRPO+OPSD
33
Incomplete
SDPO
8
Incomplete
Skill-SD
33
Incomplete
SDAR
33
Incomplete
Appendix
Table 8: Episode-level comparison on the dough task.
Figure 8: Prompt template for procedural Atom extraction.
Reinforcement learning with verifiable rewards (RLVR), along with recent selfdistillation variants such as SDPO, evaluates each rollout against a verifier and updates the policy from that episode-level signal. However, the richer procedural information in the rollout is rarely retained or reused. Across episodes and epochs, the model repeatedly encounters related problems under a changing policy, producing cross-episode signals that episode-local updates cannot capture: which strategies consistently pass verification, which failure modes persist, which patterns recur. We propose Procedural Memory Distillation (PMD), which converts these crossepisode signals into reusable procedural memory and distills it into the policy's weights during training. This memory functions as a training scaffold, absorbed into the policy itself, yielding a memory-free model at inference. PMD organizes the memory at three levels of abstraction: raw trajectories, self-reflected strategies and lessons, and higher-level behavioral patterns that recur across problems, all extracted online from the model's own trajectories. A memory-conditioned self-teacher draws on the accumulated experience to supervise the student on its own rollouts, enabling student to progressively internalize procedural knowledge within its parameters. The central design principle is co-evolution: the policy generates rollouts that update the memory, and memory shapes the supervision that updates the policy. Empirically, across Qwen3-8B and OLMo3-Instruct-7B, PMD improves over SDPO by 3.8-5.5% on SCIKNOWEVAL and 7.9-13.6% on LIVECODEBENCH. Co-evolution powers these gains: freezing either the memory or the policy trails PMD by more than 10% across SCIKNOWEVAL domains.
Reinforcement learning (RL) has been widely used to train LLM agents for multi-turn interactive tasks, but its sample efficiency is severely limited by sparse rewards and long horizons. On-policy self-distillation (OPSD) alleviates this by providing dense token-level supervision from a privileged teacher that has access to ground-truth answers. However, such fixed privileged information cannot capture the diverse valid strategies in agent tasks, and naively combining OPSD with RL often leads to training collapse. To address these limitations, we introduce Skill-SD, a framework that turns the agent's own trajectories into dynamic training-only supervision. Completed trajectories are summarized into compact natural language skills that describe successful behaviors, mistakes, and workflows. These skills serve as dynamic privileged information conditioning only the teacher, while the student always acts under the plain task prompt and learns to internalize the guidance through distillation. To stabilize the training, we derive an importance-weighted reverse-KL loss to provide gradient-correct token-level distillation, and dynamically synchronize the teacher with the improving student. Experimental results on agentic benchmarks demonstrate that Skill-SD substantially outperforms the standard RL baseline, improving both vanilla GRPO (+14.0%/+10.9% on AppWorld/Sokoban) and vanilla OPD (+42.1%/+40.6%). Project page: https://k1xe.github.io/skill-sd/
Hao Wang, Guozhi Wang, Han Xiao +8
Hangzhou Institute for Advanced Study, University of Chinese Academy of Sciences · vivo AI Lab · The Chinese University of Hong Kong +2
Skill-based prompting has become a practical mechanism for improving LLM agents, yet existing methods often treat skills as summaries of the agent's own experience or of successful demonstrations. This creates a mismatch for weaker student agents. A failed trajectory may not reveal the missing knowledge or strategy, while a teacher trajectory may be too implicit to internalize. We propose SKILL-KD, a contrastive skill distillation framework that treats skills as an explicit distillation medium between agents of different capabilities. Given a student failure and the teacher trajectory on the same task, SKILL-KD distills their actionable discrepancy into a textual skill patch, evaluates it by re-running the student, and iteratively refines it when the student still fails. To prevent skill drift from repeated local updates, SKILL-KD maintains trace-linked edit histories and performs Drift-Aware Skill Consolidation to decide whether each patch is added, merged, or skipped. Across five agent benchmarks and two student settings, SKILL-KD consistently improves frozen student agents over fixed-model adaptation baselines.
Qiming Shi, Yibo Dou, Jiawen Zhu +5
State Key Lab of CAD&CG, Zhejiang University · School of Software and Microelectronics, Peking University · School of Software Technology, Zhejiang University +1