LightWM: Training-Free Hierarchical Working Memory for Small Language Model Agents
Abstract
Small language models (SLMs) are attractive for low-cost agent deployment, and their growing capabilities make them promising for procedure-centric workloads that repeatedly execute specialized task families with similar workflows. However, SLM agents still struggle in long-horizon interaction because limited context budgets make it difficult to reliably retain and reuse decision-relevant state across many steps. Existing working-memory methods attempt to mitigate context growth, but their reliance on unstructured natural-language summarization can discard critical facts, introduce state drift, and compound errors in SLM execution. We present LightWM, a training-free hierarchical working-memory framework that decomposes procedure-centric tasks into subgoals and organizes memory into task-level global memory and subtask-level local memory, where local memory directly conditions SLM action selection and is updated from new observations through structured updates. To instantiate such memories without training, a one-time offline LLM-based induction pipeline builds reusable schemas per task family from a few successful traces, requiring no SLM parameter updates or online LLM calls. On ALFWorld valid_unseen, Qwen3-4B reaches 0.910 success, whereas representative prompting and prior working-memory baselines under the same setting remain below 0.320.
Lay Summary
Small language models are attractive for AI assistants because they are cheaper and easier to deploy than very large models. However, they often struggle with long tasks that require remembering important information across many interaction steps. Existing approaches usually rely on compressing past interactions into free-form text summaries, but small models can easily lose key details or become confused by these summaries over time. In this work, we introduce a training-free memory framework that helps small language model agents keep track of important task information more reliably. Instead of storing the entire interaction history, our system organizes memory into structured task-level and subtask-level information, allowing the model to focus only on the information needed for the current goal. The memory structure is built automatically from a small number of example task traces and can then be reused without additional training. We evaluate our approach on a benchmark for long-horizon household tasks. Our method substantially improves the success rate of small language model agents while also reducing the amount of context they need to process. These results suggest that better memory organization can make smaller and cheaper AI agentsmore practical for real-world applications.