PolicyMem: Geometric Policy Memory for LLM Governance

Yuanchen Bei; Zhengzhang Chen; Yanjun Zhao; Haoyu Wang; Hanghang Tong; Haifeng Chen

cs.CLSep 12, 2026

PolicyMem: Geometric Policy Memory for LLM Governance

Authors: Yuanchen Bei, Zhengzhang Chen, Yanjun Zhao, Haoyu Wang, Hanghang Tong, Haifeng Chen

Organizations: University of Illinois Urbana-Champaign · NEC Laboratories America

Abstract

As large language models (LLMs) are increasingly deployed in real-world high-stakes applications, effective governance has become essential. Existing safeguards largely follow two paradigms: learning-based guards provide strong semantic discrimination but couple policy behavior to trained models and taxonomies, while programmable frameworks offer flexible control but require substantial manual prompt and workflow engineering. Neither externalizes policies as reusable operational states, making it difficult to consistently reuse policy evidence across detection, intervention, and verification. In this paper, we introduce PolicyMem, a geometric policy memory that externalizes natural-language policies as reusable geometric memory objects represented by low-rank subspaces in a shared representation space. A memory writer compiles natural-language policies into policy memory slots, and query-response pairs read the policy memory through projection energy. The resulting policy-evidence profile directly mediates the safety verdict and is reused for policy attribution and post-intervention verification. Coupled with a response rewriter, PolicyMem enables a detect-rewrite-verify loop for LLM governance. Across five widely used benchmarks, PolicyMem achieves state-of-the-art unsafe behavior detection while enabling effective policy attribution, rewriting, and post-intervention verification through the shared policy memory.

Figures & tables

Figure 1: Comparison between PolicyMem and existing LLM governance paradigms.

Figure 2: Overview of the PolicyMem governor. Detection and verification use the policy assessor $\mathcal{A}_{\Theta}$ and geometric memory, while $\mathcal{R}_{\Psi}$ performs policy memory-guided rewriting.

Model	Size	Safe-F1						Unsafe-F1
Model	Size	BT	WGM	Aeg	Bin	SRL	Avg.	BT	WGM	Aeg	Bin	SRL	Avg.
Qwen-2.5-3B	3B	0.4362	0.3935	0.1818	0.3771	0.6060	0.3989	0.7560	0.3378	0.6753	0.6449	0.7876	0.6403
Qwen-2.5-72B	72B	0.7753	0.9109	0.5997	0.9358	0.8456	0.8135	0.8369	0.6685	0.7487	0.9187	0.8905	0.8127
NeMo Guardrails	-	0.7679	0.9443	0.8390	0.9188	0.8734	0.8687	0.7889	0.7285	0.8232	0.8716	0.8834	0.8191
Nemotron-Safety-Reasoning	4B	0.8043	0.9519	0.8783	0.9108	0.9310	0.8953	0.8109	0.7244	0.8624	0.8421	0.9384	0.8356
NeMoGuard-ContentSafety	8B	0.7862	0.9528	0.8878	0.8962	0.9199	0.8886	0.7842	0.7682	0.8780	0.8207	0.9248	0.8352

Table 1: Safe and unsafe response detection across five policy taxonomies. Best and second-best results are highlighted in bold and underlined , respectively.

Model	Single-Label		Multi-Label			Avg.
Model	WGM	Bin	BT	Aeg	SRL	Avg.
Qwen2.5-3B	0.3158	0.4090	0.3271	0.1479	0.3209	0.3041
Qwen3-4B	0.4180	0.6162	0.4640	0.4830	0.5701	0.5103
Qwen3-8B	0.5630	0.6928	0.5115	0.5177	0.6511	0.5872
PolicyMem	0.8097	0.8826	0.8648	0.8633	0.8954	0.8632

Table 2: Policy attribution of unsafe responses (mAP).

Rewriter		Safety	Information	Helpfulness	Overall
BT	Qwen2.5-3B	2.6636	3.7842	2.1991	0.4253
	Qwen3-4B	2.9879	3.7357	2.3923	0.5309
	Qwen3-8B	3.5678	3.6555	2.5755	0.6982
	PolicyMem (0.5B)	3.8171	3.6469	2.7459	0.7790
	PolicyMem (3B)	3.9071	3.7432	3.1533	0.8707
WGM	Qwen2.5-3B	1.5708	3.6652	1.8933	0.1416

Table 3: LLM-as-a-judge evaluation of safe rewriting.

Figure 3: Case study on interpretable localization and safe rewriting.

Figure 4: Max-energy coverage separates safe and unsafe responses (a) in distribution and (b) under leave-one-taxonomy-out evaluation.

Appendix figures & tables8 assets

Supplementary material from the paper’s appendix.

Appendix

Dataset	Policies	Train	Dev	Test	Label Type
BeaverTails	14	26,643	543	3,021	Multi-label
WildGuardMix	13	37,176	758	1,709	Single-label
Aegis2	20	14,431	687	813	Multi-label
BingoGuard	11	13,270	1,658	1,658	Single-label
SafeRLHF	19	136,397	7,230	16,419	Multi-label

Appendix

Table 4: Dataset statistics.

Panel of Figure 5 — Figure 5: The instruction of the LLM-as-a-Judge metric for safe rewriting.

Figure 6: Coverage-based safe detection. The coverage score

\kappa=\max_{p}e_{p}

is shown for safe and unsafe responses on each held-out test set; dashed lines denote class medians.

Figure 7: Visualization of per-slot policy attribution. Projection energies $e_{p}$ are shown for 400 randomly sampled unsafe cases.

Figure 8: Dynamics of the bounded rewriting–verification loop. Left: cumulative fraction of initially flagged responses certified safe by the assessor after each rewrite round. Right: distribution of the round at which each response first becomes assessor-safe.

Figure 9: Efficiency comparison in detection quality, model size, and inference latency.

Memory	Storage
RAR Buonocore and Parimbelli (2025)	418 MB
Class-RAG Chen et al. (2024)	463 MB
Mem0 Chhikara et al. (2025)	442 MB
A-Mem Xu et al. (2026)	493 MB
PolicyMem	0.6 MB

Appendix

Table 5: Memory storage cost on co-trained corpus.

Figure 10: Parameter study on the subspace rank and governance-space representation dimension.

LLM Guardrails

Large language model guardrails detect or constrain harmful outputs, prompt attacks, and unsafe agent actions before they cause harm.

See all 195 papers

SAFE-Mem: Secure Agentic Memory with Fine-Grained Authority and Governance-Aware Retrieval
Oct 16, 2026Haoming Li, Dingyi Kang, Xu Hu +8Persistent Memory for Language Models Agent Memory Poisoning
The University of Texas at Dallas · Emory University · Virginia Tech +3
PolicyBank: Evolving Policy Understanding for LLM Agents
Apr 16, 2026Jihye Choi, Jinsung Yoon, Long T. Le +2AI Agent Benchmarks AI Agent Governance
Google Cloud · University of Wisconsin-Madison
Memory Provenance Laundering in LLM Agents: A Non-Amplification Firewall for Persistent Memory
Jul 31, 2026Jinghan Xu, Yiyong Xiao, Wanru Shao +2Agent Memory Poisoning Tool Access Control for LLM Agents

Oct 16, 2026J/K move · Enter open · S save

SAFE-Mem: Secure Agentic Memory with Fine-Grained Authority and Governance-Aware Retrieval

Haoming Li, Dingyi Kang, Xu Hu, Jinxin Yang+7

The University of Texas at Dallas · Emory University · Virginia Tech+3

Persistent Memory for Language Models Agent Memory Poisoning Graph-Based Agent Memory

Open paper AskPDF arXiv

Large language model (LLM) agents increasingly rely on persistent memory to support long-term reasoning and personalized interactions. However, information collected from external sources can introduce security risks when untrusted claims (i.e., atomic statements) are later treated as authorized preferences, policies, or permissions. Existing defenses typically rely on content-based filtering or coarse-grained provenance tracking, which either fail to prevent unauthorized memory use or block legitimate information, sacrificing memory utility. To address these limitations, we propose SAFE-Mem, a fine-grained memory governance framework that improves security while preserving legitimate memory use. First, SAFE-Mem annotates each claim with three governance labels (i.e., kind, authority, and scope) to capture its provenance and permitted use. Second, governance-aware retrieval matches query requirements with claim permissions to block unauthorized influence while retaining legitimate information. Third, a governance-enabled memory graph connects source sentences, claims, and entities to support provenance tracking and efficient retrieval. Across four LLM backbones on MPBench and LoCoMo, SAFE-Mem attains the lowest attack success among defenses that retain legitimate memory, at least 77% below the strongest of them, and the highest answer accuracy of all evaluated defenses, improving it by up to 15 points. The code is available at https://github.com/LVenum/SAFE-Mem.

Abstract

LLM Guardrails

Explore similar work

SAFE-Mem: Secure Agentic Memory with Fine-Grained Authority and Governance-Aware Retrieval

PolicyBank: Evolving Policy Understanding for LLM Agents

Memory Provenance Laundering in LLM Agents: A Non-Amplification Firewall for Persistent Memory