Learning to Extract Context for Context-Aware LLM Inference

TL;DR

This paper introduces a framework that extracts contextual cues from user prompts to improve the safety and reliability of large language model responses, addressing ambiguity and risk factors.

Contribution

It proposes a reinforcement learning-based context generator that guides LLM responses by leveraging broader contextual information from prompts.

Findings

Reduces harmful responses by 5.6% on SafetyInstruct dataset

Improves safety and compliance metrics on XSTest and WildJailbreak

Enhances LLM inference safety and reliability through context extraction

Abstract

User prompts to large language models (LLMs) are often ambiguous or under-specified, and subtle contextual cues shaped by user intentions, prior knowledge, and risk factors strongly influence what constitutes an appropriate response. Misinterpreting intent or risks may lead to unsafe outputs, while overly cautious interpretations can cause unnecessary refusal of benign requests. In this paper, we question the conventional framework in which LLMs generate immediate responses to requests without considering broader contextual factors. User requests are situated within broader contexts such as intentions, knowledge, and prior experience, which strongly influence what constitutes an appropriate answer. We propose a framework that extracts and leverages such contextual information from the user prompt itself. Specifically, a reinforcement learning based context generator, designed in an autoencoder-like fashion, is trained to infer contextual signals grounded in the prompt and use them to guide response generation. This approach is particularly important for safety tasks, where ambiguous requests may bypass safeguards while benign but confusing requests can trigger unnecessary refusals. Experiments show that our method reduces harmful responses by an average of 5.6% on the SafetyInstruct dataset across multiple foundation models and improves the harmonic mean of attack success rate and compliance on benign prompts by 6.2% on XSTest and WildJailbreak. These results demonstrate the effectiveness of context extraction for safer and more reliable LLM inferences.