Training Agents to Fail Predictably: Fine-Tuning for Forecastable Tail Risk

원문 초록

Before deployment, an AI system might be evaluated on thousands or even millions of inputs, but after deployment it may encounter billions, so its worst failures may not appear during evaluation. Jones et al. (2025) proposed forecasting how risky the riskiest deployment inputs will be by scoring how close the system comes to failure on each evaluation input and extrapolating from the highest scores. We show that the error of this forecast splits into a noise term, whose distribution is the same for every model up to scale, and terms that depend on the shape of the tail of the model's risk scores. These model-dependent terms can make the forecast under-predict badly, both in theory, whenever the evaluation set lacks a rare kind of input that deployment contains, and in practice: on real user prompts marked as toxic, the forecast under-predicted the probability that a response begins with an offensive word for all eleven language models we tested. Because the shape of this tail is a property of the model, we propose fine-tuning the model itself to make the forecast accurate, using the forecast error as the training loss. In a model organism of jailbreaking and in an RL gridworld with rare traps, our method reduced the held-out squared forecast error by factors of 14 to 60, lowered worst-case risk about as much as direct risk minimization (or more), and preserved the model's primary behavior. For settings in which the model should change as little as possible, we show that training with an alternative safety constraint and an added loss term improved the forecast with far smaller changes to the model. Our method also improved the forecast in RL from sampled returns, in language models of up to 8B parameters, and (with minimal-change training) on the toxic user prompts.