Search papers, labs, and topics across Lattice.
This paper introduces ProbGuard, a novel probabilistic guardrail for assessing the safety of Large Language Model (LLM) outputs by leveraging early output distribution signals. By reformulating safety risk estimation as a probabilistic problem rather than a deterministic classification, ProbGuard significantly improves calibration performance, achieving a 79.6% reduction in average Brier score and a 71.9% reduction in expected calibration error across various model-dataset combinations. Additionally, it effectively limits the success rate of jailbreak attacks to 1% after analyzing only the first ten decoding steps, showcasing its potential for enhancing LLM safety in real-time applications.
Early output distributions can be harnessed to drastically improve LLM safety assessments, cutting calibration errors by nearly 80%.
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.
