Source-linked AI summary
A-MemGuard: A Proactive Defense Framework for LLM-Based Agent Memory
Qianshan Wei, Tengchao Yang, Yaochen Wang, Xinfeng Li, Lijun Li, Zhenfei Yin, Yi Zhan, Thorsten Holz, Zhiqiang Lin, XiaoFeng Wang
TL;DR
LLM-agent memory enables adaptation but can be poisoned by records that appear harmless in isolation and later reinforce incorrect behavior. A-MemGuard adds consensus-based validation and a dual-memory lesson structure without modifying the core architecture, substantially reducing attack success rates while preserving benign-task utility. The evaluations cover direct and indirect attacks across diverse scenarios and report strong overall performance.
Problem
Memory-injection attacks exploit context-dependent records and self-reinforcing error cycles, creating a security risk for memory-dependent LLM agents.
Method
A-MemGuard combines consensus-based validation of parallel reasoning paths with a separate lesson memory that stores detected failures for future correction.
Results
A-MemGuard reduced Attack Success Rate by over 97% in EHRAgent scenarios and more than 60% for indirect attacks while maintaining the highest benign-task accuracy among defense baselines.
Takeaways & Limitations
A-MemGuard provides a proactive memory defense that detects contextual anomalies and enables agents to learn from experience across diverse attack scenarios.
Takeaways & Limitations
The evaluation assumes benign users cannot directly inspect or modify memory, and the consensus premise relies on separability between benign and malicious reasoning structures.
Abstract
from arXiv · showhide
Large Language Model (LLM) agents use memory to learn from past interactions, enabling autonomous planning and decision-making in complex environments. However, this reliance on memory introduces a critical security risk: an adversary can inject seemingly harmless records into an agent's memory to manipulate its future behavior. This vulnerability is characterized by two core aspects: First, the malicious effect of injected records is only activated within a specific context, making them hard to detect when individual memory entries are audited in isolation. Second, once triggered, the manipulation can initiate a self-reinforcing error cycle: the corrupted outcome is stored as precedent, which not only amplifies the initial error but also progressively lowers the threshold for similar attacks in the future. To address these challenges, we introduce A-MemGuard (Agent-Memory Guard), the first proactive defense framework for LLM agent memory. The core idea of our work is the insight that memory itself must become both self-checking and self-correcting. Without modifying the agent's core architecture, A-MemGuard combines two mechanisms: (1) consensus-based validation, which detects anomalies by comparing reasoning paths derived from multiple related memories and (2) a dual-memory structure, where detected failures are distilled into ``lessons'' stored separately and consulted before future actions, breaking error cycles and enabling adaptation. Comprehensive evaluations on multiple benchmarks show that A-MemGuard effectively cuts attack success rates by over 95% while incurring a minimal utility cost. This work shifts LLM memory security from static filtering to a proactive, experience-driven model where defenses strengthen over time. Our code is available in https://github.com/TangciuYueng/AMemGuard
1 INTRODUCTION
A-MemGuard addresses memory-injection attacks whose harm is context-dependent and can propagate through stored precedents. It combines consensus validation with dual-memory lessons, substantially reducing attacks while preserving benign-task utility.
- Threat: Memory-injection attacks are difficult to detect because records may appear harmless alone but become harmful in specific contexts, potentially creating self-reinforcing error cycles.The attack can induce an incorrect decision that is later stored as a trusted precedent.
- Contribution: The work presents A-MemGuard as a non-invasive defense framework for securing agent memory without modifying the agent’s core architecture.It is described as the first framework addressing both context-dependent attacks and self-reinforcing error cycles.
- Method: A-MemGuard uses consensus-based validation to compare parallel reasoning paths derived from multiple related memories and flag deviations from the group consensus.This targets context-aware anomalies that isolated memory checks may miss.
- Method: A dual-memory structure stores detected flawed reasoning as corrective lessons, enabling the agent to reference past failures and avoid similar decisions.The lessons are kept separately from the primary memory and consulted before future actions.
- Results: Over 97% ASR reduction was achieved in EHRAgent scenarios, while indirect attacks saw more than 60% ASR reduction.The evaluation covered direct poisoning, indirect self-reinforcing attacks, and multi-agent scalability.
- Results: A-MemGuard maintained the highest accuracy on benign tasks compared with the evaluated defense baselines.The framework was evaluated across diverse agent models, tasks, and attack vectors.
2 RELATED WORK
Prior work studies memory-enabled agents and memory attacks, including poisoning and indirect injection, while existing defenses include filtering, alignment, and perplexity detection.
- LLM Agents with Memory: LLM agents use memory to store past experiences, supporting learning, adaptation, long-term planning, and multi-agent collaboration.The related work situates memory as a component of agents operating in dynamic environments.
- Existing Attacks against Memory: Existing memory attacks include malicious-record poisoning, knowledge-base backdoors, indirect injection, self-reinforcing error cycles, and data exfiltration.AgentPoison and MINJA exemplify poisoning and indirect injection respectively.
- Existing Attacks against Memory: Existing defenses include prompt filtering, alignment, and perplexity detection.These approaches are listed as prior defenses against memory-related threats.
3 PRELIMINARY
The paper models agents whose actions depend on retrieved memories and defines memory poisoning as context-triggered manipulation by a limited set of adversarial records. A validation layer must reduce attack impact while preserving useful memories.
- Memory-Augmented Agent Architecture: An LLM agent receives a query, retrieves K relevant records from dynamic memory, and uses them with its policy to generate an action plan.The core policy is a fixed-parameter pretrained LLM, while memory stores prior interactions or knowledge.
- Memory-Augmented Agent Architecture: Because retrieved memories directly inform the action plan, their integrity is a critical security dependency and attack target.The architecture treats memory as a potential single point of failure.
- Threat Model: The threat model assumes mostly benign memory containing a small fraction of adversarial records crafted to appear innocuous in isolation.The adversary can operate through direct or indirect injection pathways and has no ability to modify the agent architecture.
- Threat Model: A poisoning attack triggers a malicious action only for a trigger query and immediate context, leaving benign-query behavior largely unaffected.This context dependence makes malicious entries difficult to distinguish from legitimate records when inspected alone.
- Problem Formulation: The validation function produces a sanitized memory subset before retrieved records inform the agent policy.It must both filter malicious records and preserve useful records to maximize task success.
- A-MemGuard Overview: A-MemGuard intercepts the memory-to-action pipeline with online consensus validation and a dual-memory structure for long-term self-correction.The architecture is presented as a security layer over the existing agent.
4 METHODOLOGY: A-MEMGUARD FRAMEWORK
A-MemGuard intercepts the memory-to-action pipeline with consensus-based anomaly detection and a dual-memory mechanism for learning from detected failures. It compares parallel reasoning paths, filters anomalous memories, archives flawed paths as lessons, and consults similar lessons before action execution.
- A-MemGuard provides an external security layer with consensus validation and dual-memory self-correction, without modifying the agent’s core architecture.
- Consensus Validation via Reasoning Path Analysis: For each retrieved memory, the agent generates a structured reasoning path from the query, memory, and immediate conversational context.The process produces K parallel structured paths for immediate comparison.
- Consensus Validation via Reasoning Path Analysis: A divergence score measures each path’s deviation from the consensus, and paths exceeding threshold τ are marked anomalous.The validated memory set retains only memories whose reasoning paths are non-anomalous.
- Self-Taught Correction via Dual-Memory Structure: Detected anomalous paths become negative lessons and are archived in a dedicated lesson memory alongside the agent’s primary memory.The lesson memory stores flawed logical structures for direct comparison with newly proposed reasoning paths.
- Self-Taught Correction via Dual-Memory Structure: Before execution, A-MemGuard compares the proposed action path with relevant stored lessons and triggers deliberation when structurally similar lessons exist.The defended policy uses lesson memory when relevant lessons are found and otherwise follows the policy conditioned on validated memory.
- Self-Taught Correction via Dual-Memory Structure: The self-corrective loop turns detected threats into an adaptive defense that progressively hardens the agent’s security posture.
5 EXPERIMENTS
A-MemGuard is evaluated against direct and indirect memory attacks, benign-task utility, multi-agent misinformation, component ablations, hyperparameter sensitivity, and reasoning-path structure. Across these settings, it reduces attack success, preserves benign accuracy, scales to collaborative agents, and benefits from consensus-based detection and lesson memory.
- Experimental Setup: Experiments cover direct poisoning, indirect injection, multi-agent misinformation, benign-task utility, ablations, and top-k sensitivity across representative agent scenarios.The evaluation uses ReAct-StrategyQA, EHRAgent, MMLU, and MISINFOTASK, with GPT-4o-mini and LLaMA-3.1-8B backbones plus DPR and REALM retrieval.
- Direct Injection Defense: ASR-r fell from 100.0 to as low as 2.13 on EHRAgent, while ReAct-StrategyQA reached 0.00 for a LLaMA-3-8B agent.The reductions held across GPT-4o-mini and LLaMA-3.1-8B backbones and both DPR and REALM retrieval systems.
- Indirect Injection Defense: More than 60% ASR reductions were achieved for both GPT-4o-mini and LLaMA-3.1-8B against indirect injection, with final average ASRs of 0.256 and 0.233.The defense identifies anomalous reasoning paths before they are stored and reinforced, unlike several ineffective baseline defenses.
- Utility Cost: A-MemGuard consistently achieved the highest benign-task ACC among the evaluated defense mechanisms.The results indicate a low performance cost alongside defensive improvements.
- Multi-Agent Scalability: In the multi-agent setting, A-MemGuard achieved task success rate 0.950 and the lowest Final Score of 2.150, versus 3.200 for the unprotected baseline.The Final Score aggregates various error penalties.
- Ablation Study: Removing consensus and lesson mechanisms nearly doubled ASR-t from 36.17 to 71.27, while the full model outperformed the ablated variants.Removing lesson memory reduced ASR-a but weakened overall defense because final deliberation against past failures was omitted.
- Sensitivity Analysis: Increasing main-memory top-k from 2 to 8 lowered all ASR metrics and improved benign ACC, whereas lesson-memory top-k=6 gave the best end-to-end ASR.Retrieving too many lessons increased ASR-a, suggesting distracting noise can weaken final decisions.
- Reasoning-Path Analysis: Knowledge-graph analysis found benign reasoning paths structurally stable while malicious memories produced clear structural outliers.This supports comparing multiple reasoning paths to detect deviations that isolated audits would miss.
6 CONCLUSION
The conclusion presents A-MemGuard as a proactive framework for securing LLM agent memory through consensus validation and dual-memory learning. Evaluations show reduced attack success across diverse scenarios while maintaining high benign-task utility.
- A-MemGuard combines consensus-based validation with a dual-memory structure so agents detect contextual anomalies and learn from experience.The conclusion describes the framework as the first proactive defense framework designed to secure LLM agent memory.
- Extensive evaluations demonstrate substantially reduced attack success rates across diverse scenarios while maintaining the highest utility on benign tasks.
7 ETHICS STATEMENT
The ethics statement frames A-MemGuard as defensive security research and limits experimentation to public benchmarks and open-source models. The study reports no new data collection, human subjects, or personal data involvement.
- Experiments are confined to public benchmarks and open-source models, excluding deployed or proprietary systems.The stated purpose is reproducibility while preventing real-world harm.
- The study reports no new data collection, human subjects, or personal data involvement.
8 REPRODUCIBILITY STATEMENT
The reproducibility statement commits to releasing the framework, baselines, and experimental scripts publicly. Evaluation uses accessible public benchmarks and specified model backbones, while the appendix describes alternative divergence-scoring implementations and their trade-offs.
- The framework source code, baseline implementations, and experimental scripts are intended for public release upon publication.
- Experiments use publicly available ReAct-StrategyQA, EHRAgent, and MMLU benchmarks with GPT-4o-mini and LLaMA-3.1-8B backbones.
- The generic path-divergence score has three explored implementations trading off performance, computational cost, and complexity.The LLM-based implementation was selected for the main-paper experiments because of its reported performance and ability to handle nuanced inconsistencies without manual threshold tuning.
A.1 INSTANTIATION 1: LLM-BASED DIRECT DECISION-MAKING (MAIN METHOD)
A-MemGuard’s main validation method uses an LLM judge to compare multiple reasoning paths against a synthesized consensus, avoiding manually tuned divergence thresholds. Alternative embedding-distance and DBSCAN validators expose difficult security–utility trade-offs, while structured reasoning improves path separability.
- LLM-Based Direct Decision-Making: The LLM judge synthesizes a consensus plan from K reasoning paths, then classifies each path as consistent or inconsistent through pairwise checks.The output is a binary decision with a brief justification, eliminating the need to tune a divergence threshold.
- Embedding Distance: Embedding-distance validation classifies a reasoning path as anomalous when its cosine distance from the semantic centroid exceeds threshold τ.This approach uses sentence embeddings to quantify deviation from the group’s central tendency.
- Embedding Distance: τ = 0.5 yields task-level ASR 0.915, while τ = 0.1 lowers ASR-t to 0.723 but leaves ASR-r at 1.0 and harms benign-task accuracy.The results show sensitivity to threshold selection and failure to prevent malicious-memory retrieval.
- Density-Based Clustering: DBSCAN can reduce retrieval ASR to zero with strict eps values, but benign-task accuracy falls to 0.383.Lenient eps values preserve accuracy while allowing attacks to succeed at high rates, creating an impractical security–utility trade-off.
- Baseline Methods: The LLM Auditor baseline sanitizes retrieved context, while Distil Classifier and Perplexity Filter baselines classify or filter memory entries using textual content and perplexity.The Perplexity Filter applies a batch-aware Median Absolute Deviation threshold with k=1.0 after an initial check.
- Representation Separability: Structured reasoning paths improve separability between benign and malicious memories, with overlapping graph edges remaining below 1% across scenarios.The method transforms raw records into representations whose embeddings and structural relationships better distinguish malicious behavior.
D.2 COSINE SIMILARITY DISTRIBUTION ANALYSIS
The analysis shows that raw memory records are semantically clustered in ways that obscure benign–malicious differences, whereas structured reasoning paths disperse similarity and reveal separable reasoning structures. Knowledge-graph visualizations provide a complementary structural view of this separation.
- Cosine Similarity Distributions: Raw memory records have tightly clustered, high cosine similarity, making benign and malicious memories semantically difficult to distinguish.The raw distribution is concentrated near a high-similarity peak, whereas structured processing changes the distribution’s shape.
- Cosine Similarity Distributions: Structured reasoning paths produce a more dispersed cosine-similarity distribution with a lower primary peak, amplifying latent logical and semantic differences.This semantic diffusion supports anomaly detection by making anomalous paths more distinguishable as outliers.
- Knowledge Graph Visualization: Knowledge graphs encode benign-only edges in green, malicious-only edges in orange, and shared edges in bright red.The graphs cover six representative application scenarios.
- Knowledge Graph Visualization: Across the visualized scenarios, bright-red overlap edges are scarce relative to distinct benign and malicious edges.This visual pattern supports low structural overlap and the separability of reasoning paths derived from different memory types.
E TOKEN COST ANALYSIS
On EHRAgent, A-MemGuard achieves the strongest reported defense performance while using fewer tokens than the Auditor baseline. The comparison frames security gains against the computational cost of different defenses.
- Performance–Token Trade-off: A-MemGuard’s token cost is lower than the Auditor baseline while its reported performance is higher.The figure compares performance against token consumption on EHRAgent.
- Performance–Token Trade-off: A-MemGuard reaches nearly 0.9 performance (1 - ASR) with 7.8K tokens, outperforming the Auditor’s approximately 0.68 at 8.0K tokens.Performance is measured as 1 - Attack Success Rate, so higher values indicate stronger defense.
- Performance–Token Trade-off: The No-defense baseline uses approximately 3.6K tokens but achieves only 0.5 performance (1 - ASR).The comparison illustrates the cost of adding defense mechanisms relative to the undefended setting.
G SELF-TAUGHT CORRECTION IMPLEMENTATION DETAILS
Self-Taught Correction turns detected failures into lessons that are retrieved by query and action similarity, then injected as explicit warnings before future decisions. Case studies show this mechanism rejecting or counteracting adversarial reasoning in QA tasks.
- Self-Taught Correction: Self-Taught Correction lets the agent identify malicious memories, learn from failures, and avoid repeating similar mistakes.Its implementation has three stages: Lesson Generation, Context-Aware Retrieval, and Preventive Prompt Injection.
- Lesson Generation: Detected malicious memories are converted into concise lessons and stored with the problematic memory entry for later retrieval.The system first generates a reasoning chain for each malicious memory, then synthesizes the chain into a lesson.
- Context-Aware Retrieval: Lesson retrieval first uses query similarity, then searches for memories with similar actions and collects their associated lessons.Action-centric retrieval can connect different questions that produced similar potentially flawed operational steps.
- Preventive Prompt Injection: Retrieved lessons are presented as explicitly framed warnings that instruct the model to avoid repeating past counterexamples, pitfalls, or erroneous logic.They are used as cautionary instructions rather than direct examples.
- Case Studies: In a QA case, the defense identifies an adversarial ASCII-encryption instruction as an inconsistent reasoning chain and discards it, producing the correct answer.Without defense, the embedded instruction causes the model to follow the malicious behavior.
- Case Studies: Other case studies show adversarial signals inverting the Miami answer or suppressing reasoning about Picts and Old English.The examples illustrate manipulation of reasoning and answer generation in distinct QA scenarios.