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

Large language models (LLMs) can be influenced by harmful or irrelevant context, which can significantly harm model performance on downstream tasks. This motivates principled designs in which LLM systems include built-in mechanisms to guard against such garbage in, garbage out'' scenarios. We propose a novel approach to limit the degree to which harmful context can degrade model performance. First, we define a baseline safe’’ behavior for the model – the model’s performance given no context at all (zero-shot). Next, we apply distribution-free risk control (DFRC) to control the extent to which the user-provided context can decay performance below this safe zero-shot baseline. We achieve this by leveraging dynamic early exit prediction, ignoring later attention heads that attend the most to the unsafe inputs. Finally, we propose modifications to DFRC that allow it to both control risk for harmful inputs and leverage performance and efficiency gains on helpful inputs. We present both theoretical and empirical results across 9 tasks spanning in-context learning and open-ended question answering, showing that our approach can effectively control risk for harmful context and simultaneously achieve substantial computational efficiency gains with helpful context.

💡 Research Summary

The paper tackles the safety problem that large language models (LLMs) can be severely misled by harmful or low‑quality user‑provided context, leading to degraded performance or unsafe outputs. To address this, the authors propose a principled framework that treats the model’s zero‑shot performance—i.e., performance with no context—as a “safe baseline.” Using this baseline, they apply distribution‑free risk control (DFRC) to bound the probability that context‑augmented inference falls below the safe level by more than a user‑specified tolerance ε.

The technical core consists of two intertwined components. First, they equip the LLM with a dynamic early‑exit mechanism. For each layer ℓ, the model computes a confidence score Cℓ (the maximum class probability) and compares it to a threshold λ. If Cℓ ≥ λ, the model stops processing further layers and emits a prediction based on the current hidden state. If no layer reaches λ, the model discards the supplied context entirely and falls back to the zero‑shot prediction from the final layer. This design directly mitigates the “overthinking” phenomenon observed in prior work, where deeper layers amplify the influence of harmful demonstrations and cause performance to drop.

Second, they introduce a novel context‑aware loss ℓc(λ; x, y, c) = ℓ(ŷλ(x, c), y) – ℓ(ŷ0(x), y), where ŷλ is the early‑exit prediction with context and ŷ0 is the zero‑shot prediction. Positive values indicate that the context hurt performance (risk), while negative values indicate a beneficial context (gain). Because ℓc is not monotonic in λ, standard risk‑control methods that rely on monotonicity cannot be used. The authors therefore adopt the Learn‑then‑Test (LTT) framework, which can handle non‑monotonic losses, and they further devise a loss‑transformation that maps the potentially unbounded ℓc into the