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.

Large language models can often get confused and give bad answers if users accidentally feed them incorrect or misleading prompts. To fix this "garbage in, garbage out" problem, we develop a built-in safety mechanism with provable guarantees that acts like an emergency brake on model output. As the AI processes a user's prompt step-by-step, the system constantly checks its confidence; if it detects that the extra information is harming the model's output in some way, it stops early or simply ignores the user's input, relying instead on its own pre-trained knowledge. Testing across various language modeling tasks proved that our approach successfully prevents the AI from being derailed by bad information. As an added bonus, because the AI can stop processing early when it reaches a reliable answer, this method also makes the model run much faster and saves compute.