AI Safety in Critical Clinical Applications

Module: Healthcare AI | Difficulty: Advanced

Failure Mode Probability

Safety Metric

where is the reliability of safety layer .

AI Safety Framework Layers

import torch
import torch.nn as nn
import numpy as np

class SafetyWrapper:
    def __init__(self, model, uncertainty_threshold=0.3,
                 confidence_threshold=0.7):
        self.model = model
        self.uncertainty_threshold = uncertainty_threshold
        self.confidence_threshold = confidence_threshold
        self.prediction_history = []

    def predict_with_safety(self, x):
        self.model.eval()
        with torch.no_grad():
            prediction = self.model(x)

        uncertainty = self._estimate_uncertainty(x)
        is_valid_input = self._validate_input(x)
        confidence = torch.softmax(prediction, dim=1).max().item()

        safe = (is_valid_input and
                uncertainty < self.uncertainty_threshold and
                confidence > self.confidence_threshold)

        self.prediction_history.append({
            'prediction': prediction,
            'uncertainty': uncertainty,
            'confidence': confidence,
            'safe': safe
        })

        return {
            'prediction': prediction,
            'uncertainty': uncertainty,
            'confidence': confidence,
            'safe': safe,
            'requires_human_review': not safe
        }

    def _estimate_uncertainty(self, x, n_samples=10):
        self.model.train()
        predictions = []
        for _ in range(n_samples):
            with torch.no_grad():
                pred = self.model(x)
                predictions.append(torch.softmax(pred, dim=1))
        predictions = torch.stack(predictions)
        uncertainty = predictions.std(dim=0).mean().item()
        return uncertainty

    def _validate_input(self, x):
        if torch.isnan(x).any():
            return False
        if torch.isinf(x).any():
            return False
        if x.abs().max() > 100:
            return False
        return True

class DriftDetector:
    def __init__(self, reference_distribution, threshold=0.05):
        self.reference = reference_distribution
        self.threshold = threshold
        self.current_samples = []

    def add_sample(self, sample):
        self.current_samples.append(sample)
        if len(self.current_samples) > 1000:
            self.current_samples.pop(0)

    def detect_drift(self):
        if len(self.current_samples) < 100:
            return False
        current = np.array(self.current_samples[-100:])
        reference = np.array(self.reference)
        ks_stat = np.abs(np.mean(current) - np.mean(reference))
        return ks_stat > self.threshold

simple_model = nn.Linear(50, 10)
safety_wrapper = SafetyWrapper(simple_model)

for i in range(5):
    x = torch.randn(1, 50)
    result = safety_wrapper.predict_with_safety(x)
    print(f'Sample {i}: safe={result["safe"]}, '
          f'confidence={result["confidence"]:.3f}, '
          f'uncertainty={result["uncertainty"]:.3f}')

detector = DriftDetector(reference_distribution=np.random.randn(1000))
detector.current_samples = list(np.random.randn(100))
print(f'Drift detected: {detector.detect_drift()}')

Research Insight: Safety in medical AI requires defense in depth: multiple independent safety layers that catch different failure modes. No single safety mechanism is sufficient. The most dangerous failures are silent ones where the model is confidently wrong. Uncertainty quantification methods (MC Dropout, ensemble disagreement) can detect these cases but add computational overhead. The balance between safety and latency is a critical design decision.