🤖 machine learning

Training ML Models with Predictable Failures

This paper analyzes the bias in extrapolating ML model failure rates from limited evaluation sets, identifying conditions where such forecasts under-predict risks, and proposes a "forecastability loss" fine-tuning objective that significantly reduces forecast error while maintaining task performance and safety.

Original authors: Will Schwarzer, Scott Niekum

Published 2026-05-15

📖 6 min read🧠 Deep dive

CC BY 4.0

Original authors: Will Schwarzer, Scott Niekum

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

Imagine you are a safety inspector for a new self-driving car. You have a small test track with 100 cars to check for crashes. But the car will actually be deployed on a highway with 10 million cars.

The problem is that the "worst-case" crash you are worried about is so rare that it might not happen even once in your 100-car test. If you don't see a crash in your small test, you might wrongly assume the car is perfectly safe for the highway.

This paper proposes a clever way to fix that. It suggests two things:

Better Guessing: We can use math (specifically a branch called Extreme Value Theory) to look at the "near-misses" in our small test and mathematically predict how bad the worst crash would be on the huge highway.
Better Training: We can actually teach the AI model to behave in a way that makes this math prediction accurate, without making the car worse at its actual job (driving).

Here is a breakdown of the paper's ideas using simple analogies:

1. The Problem: The "Invisible Monster"

Imagine you are trying to predict the highest wave a surfer will ever ride. You watch them surf for one hour (your "evaluation set"). You see some big waves, but nothing catastrophic.

The Reality: The surfer is going to surf for 10 years (the "deployment set"). Somewhere in those 10 years, a "monster wave" will appear that is 100 feet tall.
The Failure: Because you only watched for one hour, you didn't see the monster wave. If you just guess based on what you saw, you might predict the max wave is 10 feet. When the 100-foot wave hits, you are in big trouble.

2. The Old Solution: "Extrapolating the Tail"

Researchers previously developed a method (by Jones et al.) to fix this. They look at the biggest waves you did see in your one-hour test. They assume the waves follow a specific mathematical shape (like a smooth curve) and draw a line to guess how high the waves would get if you watched for 10 years.

The Catch: This math trick only works if the waves actually follow that smooth shape. If the surfer suddenly encounters a totally different kind of wave (a "rare mode") that wasn't in your one-hour test, the math line will point way too low. It will under-predict the danger.

3. The Paper's Insight: "Teaching the Model to be Predictable"

The authors realized that the problem isn't just the math; it's the model itself. The AI model might be "messy" or "unpredictable" in how it fails.

The Analogy: Imagine the AI is a student taking a test. Sometimes it makes small mistakes, and sometimes it makes a huge, weird mistake that the teacher didn't expect. The teacher (the safety inspector) tries to guess the worst mistake the student will make on the final exam.
The Innovation: The authors created a new way to train the student. They didn't just tell the student "don't make mistakes." They told the student: "Make your mistakes follow a smooth, predictable pattern so I can accurately guess your worst-case scenario."

They call this new training goal the "Forecastability Loss."

4. How It Works (The "Magic" Trick)

The paper shows that you can tweak the AI's training so that:

It stays good at its job: The AI doesn't get worse at solving the actual task (like driving or answering questions).
It becomes "well-behaved": The AI's worst failures start to look like a smooth curve rather than a jagged, unpredictable mess.
The math works: Because the failures are now smooth and predictable, the "extrapolation math" can accurately predict the worst-case scenario, even if that scenario hasn't happened yet.

5. The Results: Two Experiments

The authors tested this idea in two very different worlds:

The Password Game (Language Model):
- The Setup: An AI is given a secret password and told never to say it. The test is to see if it accidentally leaks the password.
- The Problem: Most prompts are safe, but a few "adversarial" prompts are designed to trick the AI into leaking the password. These bad prompts are so rare they might not show up in a small test.
- The Fix: They trained the AI using their new method. The AI became much better at predicting when it might leak the password, and it actually leaked the password far less often than before, all while staying good at normal conversation.
The Gridworld (Robotics/RL):
- The Setup: A robot navigates a grid. Most grids are safe, but some have hidden traps.
- The Problem: The robot might fall into a trap it has never seen before.
- The Fix: They trained the robot to have "predictable failures." The robot learned to avoid the traps better, and the safety inspector could accurately predict the worst-case scenario for the robot's future performance.

6. The Key Takeaway

The paper argues that we shouldn't just hope our safety tests are big enough to catch every disaster. Instead, we should train our AI models to be "safety-compliant" in a specific way: they should fail in a way that is easy for us to measure and predict.

By doing this, we can use a small test set to accurately forecast the safety of a massive, real-world deployment. It's like teaching a surfer to ride waves in a way that allows a scientist to accurately predict the size of the next monster wave, even if the scientist has only watched for a few minutes.

What the paper does NOT claim:

It does not claim this solves all safety problems.
It does not claim this works for every single type of AI failure (though it worked well in their two tests).
It does not claim this is ready for immediate use in life-or-death medical or military systems; it is a "proof-of-concept" showing the math and method work in controlled experiments.

Technical Summary: Training ML Models with Predictable Failures

Problem Statement

Pre-deployment safety assessment for machine learning systems relies on estimating failure rates at deployment scale. However, catastrophic failures are often so rare that no practical evaluation set contains a single example, creating a gap between observed performance and future risk. While Jones et al. [21] proposed using Extreme Value Theory (EVT) to extrapolate failure rates from the largest $k$ observed scores in an evaluation set to deployment scale, this method relies on two critical assumptions: (A1) the failure tail is asymptotically log-linear (Gumbel domain), and (A2) the evaluation set is representative of the deployment distribution.

The paper identifies a specific failure regime where these assumptions are violated: when the deployment set contains a rare, high-failure mode (e.g., adversarial prompts or hazardous layouts) that is absent from the smaller evaluation set. In this scenario, the EVT-based forecaster systematically under-predicts the worst-case risk. Furthermore, the paper demonstrates that even when assumptions hold, the standard estimator possesses a built-in bias toward over-prediction, which, while safety-favorable, can be offset by the "occupancy" failure mode described above.

Methodology

The authors propose forecastability fine-tuning, a method to adjust a pre-trained model so that its failure tail conforms to the structural assumptions required for accurate EVT extrapolation.

1. Theoretical Decomposition

The core theoretical contribution is a finite- $k$ decomposition of the forecast error for the Gumbel-tail estimator. The error $\Delta$ is decomposed into:

Rank Term ( $T_\theta$ ): A bias inherent to the estimator based on the top- $k$ count and deployment ratio. This term is generally positive (over-prediction) and independent of the model's specific parameters, except for a scaling factor.
Curvature Term ( $C_\theta$ ): Proportional to the second derivative of the tail-quantile curve. This term depends on the model and can be trained. It reflects whether the tail is "well-behaved" (log-linear) or curved.
Occupancy Term ( $G_\theta$ ): A non-zero term that appears when a rare high-risk component is present in the deployment set but absent from the evaluation set. This term causes under-prediction.
Higher-Order Remainder: Terms related to smoothness violations (e.g., bounded support or heavy tails).

The authors show that the forecast error is differentiable with respect to model parameters when the risk score is differentiable. This allows the error itself to serve as a loss function.

2. The Forecastability Loss

The proposed objective, forecastability loss, minimizes the squared difference between the OLS-predicted worst-case score and the realized worst-case score across extrapolated ranks.
$\mathcal{L}_{\text{forecast}} = \sum_{j \in \mathcal{J}} w_j \left( \hat{Q}_\theta(y_j) - Y^{(j)}_\theta \right)^2$
where $\mathcal{J}$ represents ranks deeper than the fit set, and weights $w_j$ follow a log-uniform prior over deployment sizes.

3. Training Refinements

To prevent the model from simply worsening its performance to make the tail "predictable" (e.g., by increasing risk scores to match an over-predicted line), the authors introduce an improving-only gradient mask. This mask ensures that gradient updates only occur if they reduce the risk score on the specific task, effectively coupling the forecast objective to the primary safety objective.

Additionally, to avoid overfitting to specific partition splits of the data, the method employs meta-multi-task fine-tuning. At each training step, the task pool is randomly partitioned into a fit set and a deployment set. A partition-invariant cache of top-scoring tasks is maintained to make this re-partitioning computationally tractable without re-scoring the entire dataset at every step.

Key Contributions

Finite- $k$ Error Decomposition: A theoretical breakdown of the Gumbel-tail forecast error that isolates model-dependent components (curvature, occupancy) from estimator-dependent components (rank bias). This identifies exactly which parts of the error are trainable.
Forecastability Loss: A differentiable objective that directly minimizes the forecast error by shaping the model's failure tail to be well-behaved (log-linear) and representative of the deployment distribution.
Empirical Validation: Demonstration that fine-tuning for forecastability substantially reduces held-out forecast error while preserving (or improving) primary task performance and achieving safety comparable to supervised baselines.

Experimental Results

The method was evaluated in two settings:

Language Model (Password Game): A model is instructed to keep a secret password but is tested against adversarial prompts designed to extract it.
- Result: Forecastability fine-tuning reduced the worst-case leak probability by $\approx 40$ decades compared to the pre-trained baseline, matching the safety of Supervised Fine-Tuning (SFT). Crucially, it preserved model capability (measured by WildChat KL drift) significantly better than SFT (two orders of magnitude lower drift). It also improved forecast precision by $\approx 60\times$ over the baseline, whereas SFT and post-hoc calibration plateaued below $20\times$ .
Reinforcement Learning (Gridworld): A policy navigates a gridworld with rare, densely hazardous layouts.
- Result: Forecastability training improved capability (return) by $1.27\times$ and reduced worst-case regret by $3.5\times$ , outperforming SFT. In terms of forecast precision, the method achieved a $52\times$ improvement with post-hoc calibration, significantly outperforming SFT ( $26\times$ ) and calibration alone ( $6\times$ ).

In both cases, the "improving-only" mask ensured that reductions in forecast error correlated with genuine reductions in worst-case risk (leak probability or regret).

Significance and Claims

The paper claims to address a "missing layer" in pre-deployment safety assessment: the ability to not only extrapolate risk scores from evaluation to deployment scale but to train models such that this extrapolation is reliable.

Safety-Favorable Bias: The authors note that the standard estimator has a built-in bias toward over-prediction, which is desirable for safety. However, this bias is negated when the evaluation set misses a rare failure mode. Forecastability training aims to eliminate this specific failure mode.
Complement to Calibration: The results suggest that while post-hoc affine calibration can correct systematic biases, it cannot fix structural mismatches in the tail shape or occupancy gaps. Forecastability training shapes the underlying tail, making subsequent calibration more effective.
Scalability: The authors acknowledge their experiments are small-scale (0.6B and 8B models, synthetic mixtures). They claim the method is a proof-of-concept that generalizes to other distributions but note that scaling to frontier models and natural rare modes is the necessary next step.
Differentiation from Existing Methods: The paper distinguishes its approach from distributionally robust training (which minimizes worst-case loss) and importance sampling (which requires sampling failures). Forecastability loss specifically targets the shape of the risk distribution to make it amenable to EVT extrapolation, rather than just minimizing the risk itself.

The authors conclude that pre-deployment safety assessment increasingly depends on extrapolating tail behavior, and their results show that the model itself need not be a fixed input to this process; a small amount of fine-tuning can produce a model whose tails are described much more accurately by standard forecasting methods.