Adaptive Memory as an Inference-Time Guardrail for Small Language Models
Abstract
The operational Safety of LLM-based agents is a growing concern as they transition from being just text generators to autonomous systems that can do planning, invoke external tools, and execute long-horizon tasks. Existing solutions primarily rely on different post-training techniques (e.g., SFT, RLHF), to teach the models from attempting unsafe behavior, which is generally expensive and don’t adapt after deployment. In this paper, we explore whether we can leverage self-evolving memory as an adaptive guardrail against the unsafe behavior in agents. Motivated by ExpRAG and ReMem [1] self-evolving memory adaptation technique, we investigate if we can reduce the unsafe behavior at inference time by selectively leveraging from memory. For evaluation, we simulate a task stream from Agent-Safety bench [2] which consists of 2000 tasks with 8 unsafe categories (for example, leak sensitive information) and safety score feedback coming from a safety judge (Shield Agent). We evaluated our approach across three different models - Qwen2.5-7B-Instruct, Hermes3-Llama3.1-8B, and Hermes2-Pro-Mistral-7B, served via vLLM on GH200 GPU. From Figures 1 and 2, we can find that overall unsafe rate of the model has reduced compared to the baseline (no memory) with some negative impact on the fulfillable tasks. As a future task, we plan to extend the evaluation to more datasets such as Agent Security Bench, AgentHarm, and PrivacyLens. We also plan to look at how we can improve memory by balancing over over-refusal and retrieval overhead