NAISS
SUPR
NAISS Projects
SUPR
Risk Control for LLM Agents: Uncertainty Signals, Risk-Aware Post-Training and Memory Leakage
Dnr:

NAISS 2026/3-970

Type:

NAISS Medium

Principal Investigator:

Panagiotis Papadimitratos

Affiliation:

Kungliga Tekniska högskolan

Start Date:

2026-10-27

End Date:

2027-05-01

Primary Classification:

10210: Artificial Intelligence

Secondary Classification:

10211: Security, Privacy and Cryptography

Tertiary Classification:

10208: Natural Language Processing

Webpage:

Allocation

Abstract

LLM agents that call tools over many turns are deployed as email, banking and computer-use assistants, where a mistake is an action with side effects, not a bad sentence. AgentDojo (prompt injection), AgentHarm and OS-Harm (harmful behavior) show that agents can be hijacked or misused and that risk builds up over a trajectory. Agents with memory add another risk: private data stored in one session can surface in another. Existing defenses trade benign-task utility for security and rarely give a statistical guarantee. We ask four questions about uncertainty as a control signal for agent risk. RQ1: can uncertainty signals prevent unsafe behavior? At each step the agent samples K candidate actions and scores each by dependence on the user's task (MELON-style masked re-execution) and disagreement among samples. A threshold calibrated on held-out cases gives a conformal usability guarantee, and uncertified steps go to an independent monitor model. RQ2: do models contain risk-aware signals? On logged trajectories, we test whether token surprisal and hidden-state probes separate risky from benign steps, beyond sample disagreement. RQ3: how can risk awareness be trained from few samples without over-refusal? We train 7-8B agents with GRPO and rewards combining task success, safety and calibrated uncertainty (cf. RLCR), using full data and 25 or 100 labelled cases, and measure over-refusal on benign tasks. RQ4: how can memory leakage be quantified and prevented? We add retrieval memory to AgentDojo-style agents, plant private canary records, measure how often they reach unauthorized parties, and test uncertainty-gated memory use with a conformal bound on leakage. The 12-month plan has seven work packages and four investigators. Investigator 1 covers RQ1 and RQ2. WP0 (months 1-2) replicates undefended AgentDojo results on four open-weight models (70B to about 235B parameters) and measures GPU cost. WP1 (months 3-4) reproduces MELON as baseline and source of scoring code. WP2 (months 5-9) builds the full pipeline and compares it with MELON and AgentDojo's defenses on the utility-versus-attack-success frontier. WP3 (months 10-12) tests transfer to AgentHarm, OS-Harm and an adaptive attacker. WP4 (months 3-9) replays logged trajectories to extract surprisal and hidden states and trains probes. Investigator 2 covers RQ3 and RQ4. WP5 (months 3-12) runs 27 GRPO runs (five reward variants at full data, three at limited data) and compares them with the pipeline. WP6 (months 8-12) builds the memory-augmented agents and leakage metric and tests the defenses.