PolicyMem: Geometric Policy Memory for LLM Governance

• Closed-loop LLM governance observes unsafe behavior, acts by producing a safe yet helpful response, and verifies whether the intervention resolved the violation.Source: §1 · p. 1
• Learning-based guards encode policy behavior in model parameters and taxonomies, while programmable frameworks require manual translation of policies into prompts, rules, or workflows.Source: §1 · pp. 1–2
• Neither paradigm provides persistent policy representation that can share policy evidence consistently across governance stages.Source: §1 · p. 2
• PolicyMem externalizes natural-language policies as low-rank subspaces in shared geometric memory, queried by query-response representations.The resulting policy-evidence profile drives safety verdicts and is reused for attribution and post-intervention verification.Source: §1 · p. 2
• Operational policy memory is introduced as a reusable policy representation shared across the stages of closed-loop governance.Source: §1 · p. 2
• Across five widely used benchmarks, PolicyMem achieves state-of-the-art unsafe behavior detection while enabling policy attribution, response rewriting, and post-intervention verification.Source: §1 · p. 2

2.1 LLM Safety Governance

• Learned safety guards couple policy behavior to trained models, taxonomies, or prompting interfaces, while programmable guardrails require manual policy operationalization.Source: §2.1 · p. 2

2.2 Memory for LLM Agents

• LLM-agent memory systems externalize and retrieve selected experiences, facts, or reflections to support later reasoning and action.Examples include episodic experiences and reflections in Generative Agents, and structured notes and semantic links in A-MEM.Source: §2.2 · pp. 2–3
• Unlike contextual memory, governance requires enforcement-oriented policy state with a decision-causal interface.Source: §2.2 · p. 3

3.2 Geometric Policy Memory

• PolicyMem writes configured policies as low-rank subspaces and reads them through projection energy, creating reusable operational memory slots.Each slot pairs a natural-language policy address with a compact geometric representation queried through a shared read operation.Source: §3 · p. 3, §3.2 · p. 4

3 Method

• The method frames governance as an observe-act-verify process in which detection and attribution inform rewriting, and verification checks the revised response.The shared assessor and policy memory support both the initial assessment and post-rewrite verification.Source: §3 · p. 3, §3.1 · p. 3, §3.3 · p. 4

3.1 Problem Formulation

• The governor assesses the original response, maps implicated policies into targeted feedback, and reassesses each rewrite until safe or the attempt budget is exhausted.A remaining unsafe candidate after the bounded attempts results in a fallback refusal.Source: §3.1 · p. 3

3.3 Memory-Grounded Policy Assessment

• A shared policy assessor reads query-response case representations against configured policy slots and produces policy-indexed evidence for the safety decision.The verdict stage receives the evidence profile rather than query-response text or policy descriptions.Source: §3.3 · pp. 4–5

3.4 Policy Attribution and Coverage

• Ranking evidence coordinates supports multi-label policy attribution, while masking response spans measures their contribution to policy evidence.Theorem 1 states that max-energy coverage lower-bounds target-policy energy up to nearest-projector mismatch.Source: §3.4 · p. 5

3.5 Policy-Guided Rewriting and Verification

• The rewriter uses implicated-policy feedback to correct unsafe responses while preserving compliant and helpful information where possible.The same assessor and configured memory verify each rewrite; residual implicated policies guide another attempt, up to the attempt limit.Source: §3.5 · p. 5

3.6 Training PolicyMem

• Training first learns the policy memory and verdict interface, then freezes the assessor while training the rewriter with memory-guided fine-tuning.Rewriter learning includes rejection-sampling distillation from judge-approved candidates and a subsequent memory-gated preference optimization stage.Source: §3.6 · pp. 5–6

4.1 Experimental Setup

• Evaluation spans five benchmarks and four tasks, comparing 0.5B and 3B specialist and co-trained governors with programmable, learned-guard, and agent-memory baselines.Detection uses Safe-F1 and Unsafe-F1, attribution uses mAP, rewriting uses LLM-judge scores, and generalization uses AUROC.Source: §4.1 · pp. 6–7

4.2 Unsafe Behavior Detection

• The 3B specialist improves average Safe-F1 and Unsafe-F1 over the strongest baseline by 2.6% and 2.8%, respectively.The co-trained 3B governor remains within 1% of specialists on both averages while using one assessor and shared memory across five taxonomies.Source: Table 1 · p. 7, §4.2 · p. 7
• Relative to the strongest retrieval- or agent-memory baseline, the 3B specialist improves average Safe-F1 by 28.1% and Unsafe-F1 by 32.7%.Mem0 and A-Mem reduce average Unsafe-F1 by 24.6% and 20.9%, respectively, compared with the raw 3B backbone.Source: §4.2 · p. 7

4.3 Policy Attribution and Localization

• PolicyMem outperforms prompted Qwen baselines in single-label and multi-label attribution, with particularly pronounced gains in the multi-label setting.The policy-indexed projection-energy profile retains distinguishable evidence when multiple policies are implicated; Figure 3 localizes contributions to spans s2 and s0, which contain actionable harmful methods.Source: Table 2 · p. 7, §4.3 · p. 7

4.2 Unsafe Behavior Detection

• The 0.5B specialist improves over the strongest baseline by 1.4% in Safe-F1 and 1.1% in Unsafe-F1, while the co-trained 0.5B variant also leads on both averages.The paper interprets this as reduced reliance on backbone scale while preserving strong detection performance.Source: §4.2 · p. 7

4.4 Safe Rewriting

• The 3B rewriter leads on safety, helpfulness, and overall scores across five benchmarks, with competitive information preservation despite training only on BeaverTails and WildGuardMix.Strong results on Aegis 2.0, BingoGuard, and SafeRLHF demonstrate transfer beyond the rewrite-training taxonomies; the case study shows harmful instructions removed and the response redirected safely.Source: Table 3 · p. 8, §4.4 · p. 8

4.5 Generalization Study

• Macro AUROC is 0.958 in distribution and 0.930 under leave-one-taxonomy-out evaluation, where target-taxonomy policy slots are removed.The study reports consistently strong performance across all five unseen taxonomies, using threshold-free coverage over the remaining memory.Source: Figure 4 · p. 8, §4.5 · p. 8

• PolicyMem reuses geometric policy memory across detection, attribution, rewriting, and verification, and experiments report state-of-the-art unsafe-behavior detection with accurate attribution and effective safe rewriting.Source: §5 · p. 8

• The study leaves multilingual and multimodal governance, broader deployment settings, and robustness under real-world distribution shift for future work.Source: Limitations · p. 9

A.1 Conditional Policy Coverage

• A.1 Conditional Policy Coverage: When a target policy projector is close to a configured slot, max-energy reading retains its target-policy energy up to the projector mismatch.The paper connects this guarantee to cross-taxonomy coverage in its experiments.Source: §A.1 · p. 11

A.2 Additional Policy-Slot Geometry

• A.2 Additional Policy-Slot Geometry: Identical projectors yield identical energy reads, so deterministic permutation-equivariant attribution assigns equal scores to exact aliases.Replacing an exact-alias class with one representative preserves max-energy coverage.Source: §A.2 · pp. 11–12
• A.2 Additional Policy-Slot Geometry: Consolidating nearby slots preserves max-energy coverage exactly for aliases and within additive δcan for approximate representatives.The stated stability guarantee is based on projector-distance consolidation.Source: §A.2 · p. 12

A.3 Multi-Policy Attribution

• A.3 Multi-Policy Attribution: Under the theorem’s block-incoherence condition, every implicated policy has greater energy than every non-implicated policy, enabling exact top-m recovery.The interpretation ties recovery to stronger policy-specific signal relative to inter-slot coherence and residual leakage.Source: §A.3 · pp. 12–13
• A.3 Multi-Policy Attribution: The recovery condition connects learned subspace geometry directly to multi-policy attribution.The result states that the top-m energy slots recover the implicated set exactly.Source: §A.3 · p. 13

B.1 Dataset Details

• The five benchmarks cover single-label and multi-label governance settings, with training sets ranging from approximately 13K to 136K examples.They are WildGuardMix, BingoGuard, BeaverTails, Aegis 2.0, and PKU-SafeRLHF.Source: §B.1 · p. 13
• BeaverTails uses 14 harm policies and averages 1.68 violated policies per unsafe case; Aegis 2.0 uses 20 slots and averages 1.60.Both datasets use multi-label unsafe-case annotations.Source: §B.1 · p. 13
• WildGuardMix assigns one policy to each unsafe case, while BingoGuard uses single-label annotations across 11 safety categories.WildGuardMix combines vanilla and adversarial prompts with human harmfulness annotations.Source: §B.1 · p. 13
• PKU-SafeRLHF has 136K training examples and 19 policy categories, with an average of 2.02 violated policies per unsafe case.It is the largest dataset in the study and has the highest multi-label density.Source: §B.1 · p. 13

B.2 Baseline Details

• The baselines span learned safety judges and guards, programmable NeMo Guardrails, retrieval-based moderation, and general-purpose textual or structured memory systems.Examples include MD-Judge, Granite Guardian, Class-RAG, Mem0, and A-Mem.Source: §B.2 · pp. 13–14

B.3.1 Implementation Details

• The verdict module uses a permutation-invariant seven-dimensional summary of clamped policy energies, independent of the number or ordering of policy slots.The summary includes the second-largest energy and normalized Shannon entropy.Source: §B.3.1 · p. 14
• The detector is optimized for one epoch, while a separate LoRA-adapted Qwen2.5-3B/0.5B rewriter is trained through teacher rewrites and memory-gated preference optimization.The rewriter uses Claude Sonnet 4.5 as its Stage-1 teacher.Source: §B.3.1 · p. 15
• Policy anchors remain fixed for 100 optimization steps, then refresh every 100 steps; a shared compiler maps them to 256 × 8 matrices orthonormalized by reduced QR.Gradients pass through the compiler, not through detached anchors or their periodic text re-encoding.Source: §B.3.1 · p. 15
• At inference, configured subspaces can remain fixed; the shared assessor reuses policy energies for verdict, attribution, and reassessment of modified responses.Safe-rewriting evaluation uses one corrective pass followed by one policy-memory verification.Source: §B.3.1 · p. 15

B.3.2 Training Regimes

• The specialist regime trains one governor per benchmark, whereas the co-trained regime uses pooled data from all five benchmarks and their union 77-policy memory.Both regimes use the same modeling and optimization recipe.Source: §B.3.2 · p. 15

B.3.3 Evaluation Protocol

• Unsafe-behavior detection is evaluated on each official test split using Safe-F1 and Unsafe-F1, with the co-trained model applied to all five datasets.The co-trained model does not select a dataset-specific governor at inference.Source: §B.3.3 · p. 15
• Policy attribution is measured by mAP, while safe rewriting is scored on safety, information preservation, safe helpfulness, and a joint overall pass rate.The rewriter is trained on BeaverTails and WildGuardMix but evaluated on all five benchmarks.Source: §B.3.3 · pp. 15–16

C.1 Additional Safety Detection Performance

• 0.929 to 0.996 AUROC across five datasets, with macro-average AUROC 0.958, shows that unsafe responses have higher geometric coverage than safe responses.The authors caution that this separability does not establish a single calibrated threshold that transfers uniformly across taxonomies.Source: §C.1 · p. 16

C.2 Visualization of Per-Slot Attribution

• 0.618 mean energy for gold-policy slots versus 0.098 for other slots—a 6.3× gap—shows policy-specific concentration in the pooled 77-policy profile.The visualization samples 400 unsafe cases, with rows as cases and columns as policy slots.Source: §C.2 · p. 16

C.3 Rewriting-Verification Loop Dynamics

• 99.1% of cases are certified safe after the first rewrite, and at least 99.9% within three rounds; 0.1% remain flagged and trigger fallback refusal.Similar convergence is reported on Aegis 2.0, BingoGuard, and SafeRLHF, which were unseen during rewriter training.Source: §C.3 · p. 16

• 0.88 Unsafe-F1 at 28 ms per case for the 3B governor is the highest average, while the 0.5B variant reaches 0.87 at 19 ms and both outperform evaluated 7B-8B guards and the prompted 72B LLM.Source: §C.4 · p. 17
• The fixed, non-autoregressive policy-evidence readout is approximately 3× faster than Llama Guard 3 and 19× faster than WildGuard, while also providing per-policy attribution without another model call.Source: §C.4 · p. 18
• PolicyMem stores co-trained policy memory in 0.6 MB versus 418-493 MB for retrieval- and agent-memory baselines, a 697-822× reduction.It stores one compact low-rank subspace per policy, so deployment memory scales with the policy set rather than the number of training examples.Source: Table 5 · p. 18, §C.4 · p. 18

• Detection is robust across rs ∈{1, 2, 4, 8, 16} and dg ∈{128, 256, 512, 1024}; it drops at rs=32 when NT rs = 2464 exceeds dg=256, while rs=8 and dg=256 is stable and efficient.Source: Figure 10 · p. 19, §C.5 · p. 19
• 0.8799 Unsafe-F1 for co-trained PolicyMem versus 0.7765 for the anchor-based MLP is a 10.34-point absolute improvement with geometric subspaces and projection-energy read.Source: §C.5 · p. 19