ProbGuard: Calibrated Safety Risk Estimation from LLM Output Distributions

2026-08-11

AI Summary

This work addresses a critical limitation in current safety mechanisms for large language models, which typically reduce safety evaluation to deterministic classification and thereby overlook the inherent uncertainty in early generation stages and probabilistic information within output distributions. To overcome this, the authors propose ProbGuard—a architecture-agnostic, fully probabilistic safety guardrail that, for the first time, models safety risk as an unsafe probability embedded in the generative dynamics. By leveraging early output distribution signals and employing Monte Carlo sampling to estimate and calibrate risk, ProbGuard enables early termination of unsafe generations. Experimental results across nine model–dataset combinations demonstrate that ProbGuard achieves state-of-the-art calibration performance, reducing average Brier score and expected calibration error (ECE) by 79.6% and 71.9%, respectively, and suppresses success rates of six types of jailbreak attacks to below 1% using only the first ten generation steps.

Abstract

Recent research on Large Language Model (LLM) safety has widely adopted guardrails to identify unsafe LLM outputs. Existing guardrails typically formulate safety assessment as a deterministic classification task, mapping a discrete token sequence to a discrete safety label. However, this paradigm has two limitations: First, safety assessment is inherently an uncertain problem, particularly during the early generation state. Second, relying solely on discrete token sequences discards the rich probabilistic information embedded in the LLM output distribution. To address these limitations, we propose the first completely probabilistic architecture-agnostic guardrail \textsc{ProbGuard} to leverage the LLM early output distributional signals for estimating and calibrating the safety probability, thereby enabling early stopping of unsafe ongoing outputs. Specifically, given an LLM's generated prefix distribution, we formulate the safety risk as the unsafe probability of its continued generation dynamics and estimate this risk by Monte-Carlo sampling. Through post-training on the distributional signals and calibrated safety risk, \textsc{ProbGuard} achieves the best calibration performance across all nine model--dataset combination settings, reducing the average Brier score and ECE by 79.6\% and 71.9\%, respectively, over the best baseline. \textsc{ProbGuard} further limits the attack success rate to at most 1\% across six representative jailbreak attacks after observing the LLM early output distributions from only the first ten decoding steps.

Problem

• LLM safety
• safety risk estimation
• output distribution
• uncertainty
• early stopping

Innovation

• probabilistic safety estimation
• LLM output distribution
• early stopping
• calibrated risk
• guardrail