Keywords: Large Language Models (LLMs), Attention, Prompt, alignment, XAI
TL;DR: We introduce MSP, a framework that automatically identifies salient prompt tokens and reinforces model attention to them throughout decoding.
Abstract: The efficacy of Large Language Models (LLMs) is heavily dependent on the quality of user-provided prompts. Consequently, many optimization methods focus on augmenting prompts with extensive details to provide comprehensive context. However, these methods often produce verbose and information-saturated prompts, which inadvertently causes LLMs to lose focus on the most critical instructions. This phenomenon, known as attention dilution, significantly constrains model performance on tasks requiring comprehension of long contexts. To address this issue, we propose Salience Aware Mark-Steered Prompting (MSP), a novel framework designed to mitigate attention dilution by explicitly steering the model's focus toward the most critical information within the prompt. MSP consists of two stages: first, Gradient-Guided Mask Search (GGMS) automatically identifies the most influential tokens. Second, Mark-Steered Decoding (MSD) persistently guides the model by amplifying the influence of these key tokens at every step of the generation process, improving the model's alignment with core user instructions. We evaluate the effectiveness of MSP on five widely used benchmarks with three representative LLMs of multiple scales. The experimental results show that MSP yields consistent performance gains over state-of-the-art baselines, and its strong performance across diverse tasks and models highlights its robustness and generalizability. Our implementation is provided in the supplementary material.
Supplementary Material: zip
Primary Area: alignment, fairness, safety, privacy, and societal considerations
Submission Number: 3397
Loading