MultiRisk: Multiple Risk Control via Iterative Score Thresholding
Sunay Joshi, Yan Sun, Hamed Hassani, Edgar Dobriban
TL;DR
The paper tackles multi-criterion test-time filtering for generative systems by formulating a sequential risk-control problem: minimize an objective risk while enforcing multiple constraint risks via thresholds on score signals. It introduces two dynamic-programming algorithms, multirisk-base and multirisk, with the latter providing distribution-free guarantees through symmetric risk reductions and conformal-style safeguards. The authors establish finite-sample upper and lower bounds, concentration results, and near-optimality of the objective under mild regularity, and demonstrate the approach on a three-constraint LLM alignment task (PKU-SafeRLHF), achieving effective control of safety and uncertainty with improved objective performance. The work advances practical, provably reliable multi-risk control for test-time outputs and has implications for safer, more controllable AI deployments in real-world applications.
Abstract
As generative AI systems are increasingly deployed in real-world applications, regulating multiple dimensions of model behavior has become essential. We focus on test-time filtering: a lightweight mechanism for behavior control that compares performance scores to estimated thresholds, and modifies outputs when these bounds are violated. We formalize the problem of enforcing multiple risk constraints with user-defined priorities, and introduce two efficient dynamic programming algorithms that leverage this sequential structure. The first, MULTIRISK-BASE, provides a direct finite-sample procedure for selecting thresholds, while the second, MULTIRISK, leverages data exchangeability to guarantee simultaneous control of the risks. Under mild assumptions, we show that MULTIRISK achieves nearly tight control of all constraint risks. The analysis requires an intricate iterative argument, upper bounding the risks by introducing several forms of intermediate symmetrized risk functions, and carefully lower bounding the risks by recursively counting jumps in symmetrized risk functions between appropriate risk levels. We evaluate our framework on a three-constraint Large Language Model alignment task using the PKU-SafeRLHF dataset, where the goal is to maximize helpfulness subject to multiple safety constraints, and where scores are generated by a Large Language Model judge and a perplexity filter. Our experimental results show that our algorithm can control each individual risk at close to the target level.