Learning to Extract Context for Context-Aware LLM Inference

2025-12-12

AI Summary

To address the problem of unsafe or overly evasive LLM responses caused by ambiguous user prompts, this paper proposes an endogenous context-aware method that automatically extracts implicit contextual elements—such as user intent, prior knowledge, and risk signals—from the original prompt to guide safe and precise response generation. The core contribution is the design of the first reinforcement learning–based autoencoding context generator, enabling end-to-end two-stage reasoning: context extraction followed by conditional response generation—thereby departing from conventional single-stage response paradigms. Evaluated on SafetyInstruct, the method reduces harmful responses by 5.6%; on XSTest and WildJailbreak, it improves the safety–compliance harmonic mean by 6.2%. These gains demonstrate enhanced fine-grained semantic understanding and improved risk–utility trade-off capability.

Technology Category

• Natural Language Processing: Safety and Robustness
• Machine Learning: Large Multimodal Models (LMMs)
• Search and Optimization: Learning to Search

Application Category

• Search and Retrieval-Augmented AI: Search Tool Learning with LLM: Teaching LLMs to invoke search and make use of retrieved information
• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• User Modeling, Personalization and Recommendation: Large Language Models (LLM) for user modeling and recommendation

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.

Problem

• Extracts contextual cues from ambiguous user prompts
• Reduces harmful outputs by interpreting user intentions and risks
• Improves safety and reliability in LLM response generation

Innovation

• Reinforcement learning extracts contextual cues from prompts
• Autoencoder-like generator infers user intentions and risks
• Context-guided response improves safety and reduces refusals