autopasta
LLMSteer: Improving Long-Context LLM Inference by Steering Attention on Reused Contexts
Gu, Zhuohan, Yao, Jiayi, Du, Kuntai, Jiang, Junchen
Large Language Models (LLMs) have demonstrated remarkable capabilities in complex tasks such as question answering, summarization, and reasoning (llm [a,b,c]). To enhance their reliability, LLMs are often augmented with domain-specific or user-specific knowledge that extends beyond their inherent training data (Lewis et al. [2020], Jiang et al. [2023], Chen et al. [2024]). However, incorporating these supplemental contexts, which can exceed thousands of tokens (Jin et al. [2024], Gao et al. [2023]), presents two challenges: (1) models often struggle to comprehend long context (e.g., lost-in-the-middle problem (Liu et al. [2023a], Junqing et al. [2023])) and (2) processing long context incurs substantial runtime costs (Liu et al. [2024], Lin et al. [2024], Zhong et al. [2024]). Since the Key-Value (KV) cache of the same context text chunks is often reused multiple times (Liu et al. [2023b], Yao et al. [2024], Jin et al. [2024]), many recent systems adopt prefix caching (Jin et al. [2024], Liu et al. [2023b], Qin et al. [2024]), which stores the KV caches for the frequently reused contexts such that LLMs no longer need to prefill these contexts repeatedly. However, the model persists in losing track of key information from the context as its KV pairs remain unchanged. So, is there a way to simultaneously achieve high efficiency and high quality without fine-tuning models?
Model Tells Itself Where to Attend: Faithfulness Meets Automatic Attention Steering
Zhang, Qingru, Yu, Xiaodong, Singh, Chandan, Liu, Xiaodong, Liu, Liyuan, Gao, Jianfeng, Zhao, Tuo, Roth, Dan, Cheng, Hao
Large language models (LLMs) have demonstrated remarkable performance across various real-world tasks. However, they often struggle to fully comprehend and effectively utilize their input contexts, resulting in responses that are unfaithful or hallucinated. This difficulty increases for contexts that are long or contain distracting information, which can divert LLMs from fully capturing essential evidence. To address this issue, many works use prompting to help LLMs utilize contextual information more faithfully. For instance, iterative prompting highlights key information in two steps that first ask the LLM to identify important pieces of context and then derive answers accordingly. However, prompting methods are constrained to highlighting key information implicitly in token space, which is often insufficient to fully steer the model's attention. To improve model faithfulness more reliably, we propose AutoPASTA, a method that automatically identifies key contextual information and explicitly highlights it by steering an LLM's attention scores. Like prompting, AutoPASTA is applied at inference time and does not require changing any model parameters. Our experiments on open-book QA demonstrate that AutoPASTA effectively enables models to grasp essential contextual information, leading to substantially improved model faithfulness and performance, e.g., an average improvement of 7.95% for LLAMA3-70B-Instruct. Code will be publicly available at https://github.com/QingruZhang/AutoPASTA .