Technical Summary: Controlling the Risk of Corrupted Contexts for Language Models via Early-Exiting

Problem Statement

Large Language Models (LLMs) are increasingly adapted to specific tasks via in-context learning (ICL) and prompt tuning, which allow models to learn from user-provided demonstrations without expensive fine-tuning. However, this flexibility introduces a critical safety vulnerability: "garbage in, garbage out." User-provided context can be incorrect, misleading, or maliciously crafted, causing the model to degrade in performance or produce unsafe outputs.

Recent research indicates that LLMs may "overthink" on harmful inputs, where performance peaks at intermediate layers but degrades in deeper layers as the model overfits to the corrupted context. Existing safety mechanisms often lack principled ways to detect when context is harmful and revert to a safer baseline without sacrificing the efficiency gains of early exiting on helpful inputs.

Methodology

The authors propose a framework that combines Dynamic Early Exiting with Distribution-Free Risk Control (DFRC) to mitigate the risks of corrupted context while maintaining efficiency on helpful context.

1. Safe Context-Aware Predictor

The core mechanism defines a "safe" baseline using the model's zero-shot performance (performance with no context). The proposed predictor, $\bar{p}_\lambda$, operates as follows:

• Early Exit: The model processes input through layers sequentially. At each layer $l$, a confidence score $C_l$ is calculated. If $C_l$ exceeds a threshold $\lambda$, the model exits early and returns the prediction from that layer.
• Fallback to Zero-Shot: If the model processes all layers without any confidence score exceeding $\lambda$, it ignores the user-provided context entirely and reverts to the zero-shot prediction (the prediction made without context).
• Goal: This design allows the model to leverage helpful context for efficiency and accuracy but forces a fallback to the known-safe zero-shot behavior if the context causes the model to become uncertain or "overthink."

2. Context-Aware Prediction Risk

To select the optimal threshold $\lambda$, the authors define a novel context-aware loss function ($\ell_c$):
$$\ell_c(\lambda; x, y, c) = \ell(\bar{y}_\lambda(x, c), y) - \ell(\hat{y}(x), y)$$
Where:

• $\bar{y}_\lambda(x, c)$ is the prediction from the safe context-aware model.
• $\hat{y}(x)$ is the prediction from the full zero-shot model.
• $\ell$ is a standard loss function (e.g., 0-1 loss for classification).

This formulation measures overthinking:

• Positive Loss: Indicates the context-aware model performed worse than the zero-shot baseline (harmful context).
• Negative Loss: Indicates the context-aware model performed better than the zero-shot baseline (helpful context).

3. Risk Control via Learn-Then-Test (LTT)

The authors utilize the Learn-Then-Test (LTT) framework to find a threshold $\hat{\lambda}$ that guarantees the expected risk $R(\hat{\lambda})$ remains below a user-specified tolerance $\epsilon$.

• Challenge: Standard LTT requires loss functions to be bounded within $[0, 1]$. However, the context-aware loss $\ell_c$ can be negative (ranging from $[-1, 1]$), and clipping negative values to zero (a common workaround) discards valuable information about helpful context, leading to overly conservative behavior.
• Solution: The authors propose a Domain-Preserving Risk Transformation. They linearly scale the loss and the risk level $\epsilon$ to map the interval $[a, b]$ to $[0, 1]$ before applying LTT, then reverse the transformation. This preserves the distinction between helpful (negative loss) and harmful (positive loss) contexts, allowing the system to balance safety with efficiency.

Key Contributions

1. Novel Early-Exit Formulation for Safety: A predictor that dynamically exits on confidence but defaults to a zero-shot baseline if confidence is never achieved, preventing the model from relying on corrupted context.
2. Context-Aware Loss for Overthinking: A loss function explicitly designed to measure performance degradation relative to the zero-shot baseline, enabling the detection of harmful context.
3. Adapted Risk Control Framework: A modification of the LTT framework using risk transformation to handle unbounded (specifically, negative) losses, allowing for simultaneous risk control on harmful inputs and efficiency gains on helpful inputs.

Experimental Results

The framework was evaluated across 9 tasks (spanning sentiment analysis, hate speech detection, semantic classification, and open-ended QA) and 5 distinct models (including LLaMA-2, LLaMA-3, and T5).

• Risk Control: The approach successfully controlled the risk of overthinking across all models and datasets. Empirical risk remained below the user-specified $\epsilon$ threshold even with mixed distributions of helpful and harmful context.
• Safety vs. Efficiency Trade-off: The method identified thresholds where the model avoided harmful context (performing better than the full model on harmful inputs) while retaining performance gains from helpful context.
• Efficiency Gains: Compared to prior approaches using loss clipping, the proposed risk transformation approach achieved significantly higher computational efficiency.
  • On in-context learning tasks, it evaluated 53% fewer layers on average at $\epsilon=0.05$.
  • On open-ended QA (SQuAD v2.0), it evaluated up to 81.8% fewer layers.
• Zero-Shot Fallback: In open-ended QA, the system defaulted to zero-shot predictions 3.3x more frequently on unanswerable (misleading) questions compared to answerable ones, demonstrating its ability to discriminate context quality.

Significance and Claims

The authors claim this is the first work to establish a principled framework that controls the risk of potentially harmful user-provided context in LLMs.

• Safety: The approach provides robust safety guarantees by anchoring risk control to the zero-shot baseline, which has undergone extensive pre-deployment safety testing. It effectively prevents "overthinking" on harmful inputs.
• Efficiency: Unlike static pruning or conservative clipping methods, this dynamic approach leverages early exiting to achieve substantial computational savings when context is helpful.
• Principled Design: By using distribution-free risk control, the method offers statistical guarantees on risk without requiring monotonicity assumptions or specific distributional knowledge of the user inputs.

The paper concludes that while specific design choices (e.g., exiting only in the second half of layers) may need adaptation for different tasks, the core mechanism of detecting and intervening on harmful outputs via intermediate representations remains a viable and necessary approach for safe LLM deployment.