On the Loss of Context-awareness in General Instruction Fine-tuning

TL;DR

This work shows that supervised instruction fine-tuning with chat templates can degrade context-awareness in open-source LLMs, not by destroying retrieval ability but by biasing attention away from user-provided context. The authors diagnose this via NIH testing and attention analysis, revealing a template-induced shift toward assistant tokens. They propose context-dependency indicators and conditional finetuning to preserve context reliance when needed, demonstrating improved context-aware performance on NIH, SQuAD, and QuAC while maintaining general instruction-following on MT-Bench and DROP. The approach provides a practical pathway to retain context understanding in instruction-tuned LLMs, with potential benefits for retrieval-augmented generation and contextual QA in real-world deployments.

Abstract

Pre-trained Large Language Models (LLMs) require post-training methods such as supervised fine-tuning (SFT) on instruction-response pairs to enable instruction following. However, this process can potentially harm existing capabilities learned during pre-training. In this paper, we investigate the loss of context awareness after SFT, where context awareness is defined as the ability to extract and understand information from user-provided context and respond accordingly. We identify and demonstrate that the loss of context awareness, particularly in open-source models, occurs in instruction fine-tuned LLMs when the chat template is applied to input prompts. We identify that the performance decline is associated with a bias toward different roles learned during conversational instruction fine-tuning. We demonstrate this correlation by visualizing changes in attention allocation after the chat template is applied and manually steering the attention heads. The bias can be learned from training examples that align with the model's internal knowledge and rely less on the user-provided context to generate correct responses. Based on these observations, we propose a metric to identify context-dependent examples from general instruction fine-tuning datasets. We then apply conditional instruction fine-tuning with a context-dependency indicator, enabling the model to preserve context awareness after SFT. Empirical experiments on four context-dependent downstream tasks and three pre-trained LLMs of different sizes show that our method effectively mitigates the loss of context awareness without compromising general instruction-following capabilities.

Figures (8)

• Figure 1: An example of the Needle-in-a-haystack (NIH) test used in our work. [Optional User template] and [Optional Assistant template] are user and assistant role indicators used in instruction finetuned models. The inserted needle is highlighted in yellow.
• Figure 2: Average recall error (1 - recall) on NIH for different model series. We report the performance of official instruction-tuned models (both with and without chat templates) and their corresponding pretrained models from five model families, with sizes ranging from 7B to 27B. Note: Errors on some models are too small to be visible in the figure. Detailed numerical values can be found in Appendix \ref{['apd:NIH_performance']}
• Figure 3: We visualize the changes in attention allocation between user tokens and assistant tokens after applying chat templates. The attention allocation is calculated when the model generates the first answer token in its response. For cases where a chat template is applied, we normalize the attention values on user tokens, assistant tokens, and the BOS token so that their combined attention scores sum to 1. The attention weights are averaged across 400 tests with context lengths ranging from 200 to 4,000 and needle depths from 0% to 100%. Detailed absolute attention scores for each component can be found in Appendix \ref{['sec:full_result_attention_score_diff']}.
• Figure 4: We visualize the full attention allocation on user tokens, assistant tokens, and BOS token with and without applying the chat templates. The attention allocation is calculated when the model is generating the first answer token in its response. For cases where the chat template is applied, we normalize the attention values on user tokens, assistant tokens, and the BOS token such that attention scores allocated to these three sum up to 1. The attention weights are averaged across 400 tests with context lengths ranging from 200 to 4000 and needle depths from 0% to 100%.
• Figure 5: We visualize the disagreement heatmap of $\hat{S}$ selection when the context-dependency score $S_M({\mathbf{Y}}_m)$ is calculated across different layers. We select as $\hat{S}$ the 10% of conversation turns with the highest context-dependency scores on each layer. The disagreement is measured by the number of non-overlapping conversation turns in $\hat{S}$ selected by any two layers.