On the Loss of Context Awareness in General Instruction Fine-tuning
Authors: Yihan Wang
Andrew Bai
Nanyun Peng
Cho-Jui Hsieh
Venue: NeurIPS 2025 poster ()
Keywords: Instruction fine-tuning
context awareness
catastrophic forgetting
Primary area: deep_learning
TL;DR: We identify and mitigate the loss of context awareness after instruction fine-tuning.

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 cause forgetting in 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. Surprisingly, we discovered that the loss of context awareness 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. The bias can be traced to training samples where the assistant response minimally relies on the user-provided instruction. 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. 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.

On the Loss of Context Awareness in General Instruction Fine-tuning
Yihan Wang" UCLA vangylhan6.70gmail.cc com com
Andrew Bai* UCLA andrewbai@ucla.edd
Nanyun Peng UCLA violetpeng@cs ucla. edu
Cho-Jui Hsieh
UCLA
chohsieh@cs ucla. edu
Abstract
Pre-trained Large Language Models (LLMs) require post-training methods such as supervised fine-tuning (SFT) instruction-respons pairs to enable instruction following. However, this process can cause forgetting in capabilities learned during pre-training. In this paper, investigate the loss of context awareness after SFT, where context awareness is defined as the ability to extract and understand information from user- rervidro context. Surprisingly, we discovered that the loss of context awareness occurs in instruction fine-tuned LLMs when the chat template is applied input prompts. We identify that the performance decline is associated with a bias toward different roles learned during conversational instruction fine- tning. The bias can be traced training samples where the assistant response minimally relies on the user-provided instruction. Based on these observations, we propose metric identify context-dependent examples from general instruction fine-tuning datasets. We then apply conditional instruction fine-tuning with a context- depeedeecyp indicator, enabling the model to preserve context awareness after SFT. Experiments on four context-dependen downstream tasks and three pre-trained LLMs of different sizes show that our method effectively mitigates the loss of context awareness without compromising general instraction-following capabilities.
Introduction
Large language models (LLMs) pretrained large-scale datasets acquire diverse language skills dur- ing pretraining. To enhance these models' ability to follow general instructions, further fine-tuning is typically required. This includes supervised instruction fine-tuning (SFT) [23, 20] and reinforcement learning from human feedback (RLHF) [6]. However, several studies have demonstrated additional fine-tuning can potentially harm existing capabilities learned during pretraining [17, 10].
Although some studies suggest that performance degradation can be mitigated or even eliminated through improved instruction fine-tuning methods [19, 3, 11], in this paper, demonstrate that instruction fine-tuning specifically leads to the worsening of a model's context awareness in series of open-source models. We define context awareness as model's ability accurately retrieve, process, and interpret specific information from user-provided context. Context awareness is highly
"Equal contribution
Code available at httpp////gggttt...////aaaaaaaaaa
39th Conference on Neural Information Processing Systems (NcurIPS 2025).

---

relevant to the intrinsic hallucinations of LLMs [13] and crucial to the truthfulness of chat models [20]. It also important for many real-workh use cases, including retrieval augmented generalization [15, 14, 26], in-context learning [1], and contextual question
[21, 8).
We first demonstrate the loss of context awareness through evaluations of several popular, open-source (NIH) task. We show that while many pretrained models
LLMs using the Beedle-in-a-Hassteee demonstrate near-perfect performance on NIH, their performance deteriorates consistently after SFT, regardless of context window sizes, chat templates, architectures, or model sizes, We show that this decline is correlated with the application of chat templates, which, however, are widely used and essential in building conversational LLM assistants. When these chat templates are removed from instruction fine-tuned models, NIH performance not only recovers but some instances surpasses that of their pretrained, non-fine non-fiineturn fine-tuned counterparts.
These observations suggest that the deterioration in NIH performance does not indicate catastrophic loss of context-retrieval capabilities during instruction fine-tuning. Instead, the chat template appears mask these underlying abilities by introducing systematic biases into the models' behavior. Through analysis, we observe differences in attention allocation patterns in input tokens when comparing instruction fine-tuned models with and without applying chat templates. Specifically, we examine how the attention allocation shifts when tokens are marked as "user tokens" by the chat template. As illustrated Figure 3, the application of chat templates leads to a redistribution of attention values: attention scores decrease for user input tokens while increasing for assistant tokens. We further validate the relationship between attention score reallocation and context awareness through targeted intervention experiments. By manually increasing attention values assigned to user tokens, we partially restored the performance on simple context-relevant tasks.
These findings motivate us to develop a fine-tuning strategy that mitigates attention bias acquired during instruction fine-tuning. Our approach stems from the intuition that the bias is learned from certain patterns in the training dataset, where for some examples the model does not need the context generate correct answers. Therefore, we develop quantitative metric assess the context-dependency of conversational instruction and response pairs based on attention allocation patterns. We discover that context- dependtte examples are notably sparse in commonly used open- source instruction fine-tuning datasets. To help the model distinguish examples with and without context-dependency during instruction finetuning, we add an indicator token to the identified context- dependent instructions. After fine-tuning the model with this enhanced dataset, it learns to allocate increased attention to user tokens when the indicator is present.
We evaluate our method on three open-source pretrained language models and several context- dependent and general tasks. Empirical results demonstrate that models fine-tuned using our method consistently achieve superior performance on context-dependent tasks compared to standard fine- tuning while maintaining similar performance on general tasks.
Our contributions are summarized as follows:
We identify that the context awareness of LLMS deteriorates after supervised instruction fine-tuning with chat templates applied, compared to pretrained models.
We pinpoint that the worsened context awareness is associated with attention allocation bias on tokens marked as from different roles.
We propose a quantitative metric to identify context-dependent instruction-response, pairs from general instruction fine-tuning datasets. By inserting an indicator into identified instructions during SFT, we mitigate the loss of context awareness in instruct models when the indicator is added during inference, while preserving general performance.
2 Loss of Context Awareness after Instruction Fine-tuning
We conduct preliminary studies to understand the loss of context awareness after instruction fine- tuning and its root cause. In Section 2.1, we present evidence that context awareness consistently deteriorates in instruct models and identified the main culprit as the roles indicated with chat templates. In Section 2.2, we analyze correlational relation between instruct models and decreased attention allocated to user role tokens. Section 2.3, we find causal relation between the decreased attention allocated to user role tokens (after instruction fine-tuning) and worsened context awareness.

---

2.1 Evaluating Context Awareness Through Needle-inna-Haasttac (NIH) Testing
Input prompt
[optional User template] You are helpful AI assistant that answers question using only
the provided document:
"July hard liquor, cigarettes, heroin,
common they' re
all concentrated
of addictive predecessore. thing Francisco sandwich
Dolores eunny day.
things describe addictive are."
Question:
is the thing [Optional template] thing
Francisco?
Target output
sandwich
sunny day
Figure 1:
An example the Needle-in-a-haystee (NIH) test used in our work. [Optional User template] and [Optional Assistant template] and assistant role indicators used in instruction fine-tuned models. The inserted needle highlighted in yellow.
We demonstrate the loss of context awareness after instruction fine-tuning with the needle-in-al haystack (NIH) test. The NIH task provides a fairer comparison between pretrained models and instruct models in terms of context retrieval performance, since it relies less on instruction-following capabilities. We remove the newlines between context and needle the original NIH test to increase difficulty and better discriminate among different models. An example of the NIH prompt is shown in Figure We rerun the evaluation with different prompt templates for more robust evaluation. More details can be found in Appendix A.2.
The NIH test evaluates the performance of language models in extracting given sentence (the needle) from irrelevant context, The needle can be inserted at different locations in contexts of varying lengths. We report the recall error:
Dataset.
recall
Σ WEK
1(w E output)
K
err 1
recall
where is the set keywords in the targeted output and output is the output the LLM. For all NIH evaluations, calculate the recall on the first 100 generated tokens. We average the recall eтror across 400 NIH tests with different insertion locations and context lengths within the model's context window. An example of the NIH prompt in experiments is shown in Figure 1. When the chat template is applied to the prompt, the whole input prompt is partitioned into the user instruction input and model response, indicated by special role markers in the chat template (e.g., - < | and <I iassistant |>). More details about the NIH tests can be found in Appendix A.2.
Models. We evaluate NIH on eight open-source language models from five model families. For each model, we compare the performance of the pretrained version (not instruction fine-tuned) and the official instruction-finttune stronger closed-souse source models as their pretrained versions are unavailable. The context window lengths of these models range from 4K to 32K.
version released by the model provider. Here, we do not consider
NIH performance drops after instruction finetuning on most models. We report the evaluation results on NIH in Figure 2. Given the significantly improved instruction -orllwtinnnfolooinng through instruction finetuning, we would expect that performance would always increase. However, when comparing the pretrained model (green bar) with the instruetion-iinetueedd model (red bar), the NIH error increases most models after instruction fine-tuning, which implies negative effects from worse context awareness after finetuning. The only outlier is Llama-3.1-8) which highlights the nuanced dual impact of instruction fine-tuning on different models: improvement instruction following and potential worsening of context sensitivity.
3

---

protrained model
Instruct model wi chat tomp late Instruct model chat template
Average recall enror on
Uama-2-7b
Liemn-2-133b
Liema-3.1-80 Lleme-3-8b
gemmma.2-ge gemma--.227 22hb Mistral-78-0.22
Mistral-78.a
Mistral-78
Figure 2: Average recall error (I recall) on NIH for different model series (lower better). We report the performance of official instruction-tunce models (both with and without chat templates) and their pretrained counterparts from five model families, with sizes ranging from 7B 27B. Some errors too small to be visible in the figure. Detailed numerical values can be found in Appendix B.I.
The performance drop is associated with chat templates. To determine whether the performance difference mainly comes from a different input format or fine-tuned model weights, we remove the chat templates (i.e., the role indicators with blue bars in Figure 2. The NIH error without applying chat templates (blue bar) is significantly lower than with templates (red bar). These results indicate that context retrieval capabilities are not climinated by instruction fine-tuning, but are instead impacted by biases associated with the presence of chat templates.
instruction-tuned models) and visualize the NIH errors
The aforementioned phenomenon is consistent across models with varying context window lengths, model families, chat templates, and small medium model sizes. Although we are unable conduct experiments on extremely large models, context awareness in medium-sized models remains relevant, as they are widely adopted in cost-sensitive settings such as edge devices and small businesses.
2.2 Attention Allocation Bias Across Different Roles
Based on our observations, performance on NIH drops significantly when the chat template is applied. We hypothesize that the performance deterioration stems from the bias in instruction data and the bias embedded different roles marked by the chat template. When the model generates a response, balances information from the input context and internal knowledge stored within its weights. It pays attention to user tokens to maintain consistency with the user ever-provided context while attending to previously generated assistant response tokens maintain consistency with its output. If the model learns to assign lower importance to user-provided context and higher importance to its internal knowledge during SFT, it may develop bias that causes it weigh user tokens less. To support our hypothesis, we analyze the attention allocation between user and assistant tokens, both with and without chat templates.
Experimental settings. We prepare two inputs for each NIH test case: one with the chat template and one without. The prompt formats follow the input prompt shown in Figure 1. collect attention weights from each layer, focusing on the last token (which generates the next answer token) and its attention to input tokens. We separately sum the attention weights for user and assistant tokens. When calculating the attention weight allocation with chat templates, exclude the attention weights on chat template tokens and renormalize the attention weights across the user, assistant, and BOS tokens. We report the attention allocation from arbitrary middle layer (e.g., Layer 15) representative context retrieval head that allocates the highest attention to user tokens without chat templates. Further discussion head and layer selection can be found in the Appendix B.3.
Less attention on user tokens with chat templates. We visualize the changes in attention allo- cation, both with and without chat templates in Figure 3. When chat templates are applied mark tokens as from different roles, attention allocated to user tokens decreases while attention to assistant tokens increases for all models. This indicates that the models learn to assign lower attention to user tokens compared to the baseline level (where the chat template is not applied). In our experiments, we collect attention allocation data for context lengths less than 4,000. Although several models (e.g. Llama-3 and Gemma-2) achieve perfect NIH accuracy under 4,000 context length, the decrease in

---

Change of Attention Allocation an User Tokens
Change of Attention Allocation an Assistant Tokens
Percentage
0
4.352
2353
2.50
Lema-2-78.Chat Lleme-2-13B.CM 2-att
Lieme-3-00 instruct
Lime--311-80 instruct
Mistral-Te-Inutnucii -e-nnntncctvvv 3ottruuccc
7
genuna 2-98
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. The attention weights are averaged across 400 tests with context lengths ranging from 200 4,000 and needle depths from 0% 100%. More detailed scores can be found in Appendix B.2.
user attention remains noticeable. Note that comparison is only reasonable within variants of each model and not between different ones.
2.3 Attention Steering to Compensate for Attention Bias
In the previous section, we observed a trend of decreased attention allocated user tokens when chat templates were applied to instructional models, associated with performance decline the NIH task. To establish a more robust causal relationship. we further verify our hypothesis by manually steering attention toward user tokens to compensate for the attention bias.
Post-hoc attention steering of user tokens. To compensate for the attention bias observed in tokens.
instruction models, manually steer the attention
Specifically, we modify the self-attention weights in each transformer layer:
Au(x,y) - - aat((x))) i U y) otherwise,
(I)
where x and y two tokens in the input sequence, more emphasis on user tokens), U is the subset of all user tokens, and Z is the normalization constant that renormalizes the altered attention scores across all tokens. Aut(x, y) is the original attention weight from token x to token y.
(0, 1) is the steering strength (lower for
We adopt the same attention steering implementation Zhang al. [28]. They steered the attention of pretrained language models to emphasize user-specified portions user instructions, enabling models to follow instructions without explicit instruction fine-tuning. In setting, we increase models on the entire user input prompt, which conse-
the attention weights instruction-fine quently decreases weights other tokens (chat template role tokens, BOS/EOS tokens, and partially generated model responses). We steer on all heads with intervention factor a 0.95.
Post-hoc attention steering partially recovers the NIH performance but produces side effects. We report the performance attention steering in Table 1. Attention steering requires customized attention calculations for different heads and layers, which limits the use of several existing efficient attention implementations. Therefore, we are only able to apply attention steering to two Llama- 2 models with 4,000-token context windows. We report the recall on NIH task as well the performances on two additional contextual QA tasks: QuAC [5] and DROP [8]. Detailed descriptions and metrics of these tasks can be found in Section 4. Unlike NIH and QuAC, which retrieve exact sentences from the context, DROP requires the model not only to understand and retrieve relevant information from the context but also to apply discrete mathematical operations to retrieved information. As shown in the table, attention steering can boost performance on simple context retrieval tasks such as NIH and QuAC. However, on DROP, which requires more complex combination of different capabilities, performance with attention steering is negatively impacted. The
5

---

Table 1: NIH recall and QuAC/DROP containing score with attention steering. "Baseline" and "+ Attention Steering" are evaluated with chat templates. "w/o chat template" shows the NIH performance without the chat template for reference (same as Figure 2).
Task
Capabilities
Model Name
Baseline
Attention Steering
w/o chat template
Llama-2-7B-Chan
0.9917
0.9932
0.9975 0.9578
NIH
sentence retrieval
Llama-2-13B-
0.9207
0.9225
sentence retrieval
Llama 1-7ma-2-BB-a 2-7B-Chir
22.20
24.00
QuAC
reading comprehension
1lama-2-13B-hhar
18.60
20.00
context retrieval math operation
44.22
43,46
DROP
Llama-2-13B-n
46.20
45.11
deteriorating performance on DROP suggests that intervening in attention scores to emphasize user tokens, while improving context awareness, might impair other capabilities of the model.
In the next section, we introduce a fine-tuning strategy to better mitigate the loss context awareness with fewer side effects based on all of our aforementioned observations.
3 Instruction Fine-tuning with Context-dependenccy Indicators
Our method is based on the intuition that by explicitly marking context-dependent data samples with special indicator during instruction fine-tuning, the model learns to associate the indicator with paying more attention to the user-provided context. After fine-tuning whenever the indicator is appended to user instruction, the conditional generation allocates more attention the user-provided content and responds to the instruction with more context awareness. The main technical challenge is identifying context-dependent data samples from the instruction dataset.
3.1 Identifying Context-dependent Instructions
training sample in the instruction fine-tuning dataset is conversation between the user and model assistant, which may consist of multiple Instruction-respon pairs. Formally, denote user instruc- tions as X, assistant responses as Y, and a conversation of n total turns as C [X ,I...,,,,,, Yn].
Identifying context-dependent instructions with reference language model. We identify the context-dependent instructions by calculating the attention allocation on user tokens. We start by preparing a seed instruction fine-tuned model M, which can be the same weaker pretrained model fine-tuned with the original instruction fine-tuning dataset. We then define the context-dependency score for the m turn response Ym given its instruction X., and conversation history:
MM
m) = Y
1
max Σ YEY
Σ XEX,U...UX..
Att,, (M,
(2)
REW
where H is the set of attention heads in model M and Att(y, x) is the attention weight from token y x. Intuitively, the score $ (M(Y,)) measures the sum of attention scores allocated all user instructions in prior turns X, U...U Xm. averaged over response tokens E Y As different heads learn different capabilities, we calculate the score on the most representative head for context retrieval on each layer, specifically the head that allocates the highest attention weight to user tokens. We compute the score on a single middle layer for practical efficiency, as we find the relative scores s.M to be insensitive layer choice. We defer the detailed discussion of layer and head selection to Appendix B.5.
3.2 Instruction Fine-tuning with Context-dependenncy Indicators
A threshold 3 E (0, can be selected after the context-dependency score is obtained for each instruction-response pair. A conversation turn (Xm: Y,) with 8M (Y,)) > /3 is considered context- dependent. We append special token [IND] to the user instruction X,, if it context-dependent, In
our implementation, the special token [IND] is added as an additional special token to the vocabulary to avoid conflicts with existing ones.

---

After conditional instruction fine-tuning, the user can specify whether to add this indicator to their query, depending on whether the model response should rely more on user- user-provided context.
4
Experiments
We validate the effectiveness our method using three open-source pretrained models trained on three instruction fine-tuning datasets and benchmarked set of context-dependent and general tasks.
4.1 Experiment Settings
Models We evaluate method on three open-source pretrained large language models: TinyLlama- 1.1B [27]. Llama-2-7B [22], and Llama-3-88 [9]. TinyL.ama-l 1...amm.-1.1 is a 1.1B Ilama model pretrained on 3 trillion tokens with context window length of 2048. Llama-2-7B and Llama-3-8B have context windows of 4096 and 8192 tokens, respectively. Due to limited computational resources in academic labs, we can only fine-tune models with up to 8B parameters. We also truncate the training examples to 4096 tokens. Detailed hyperparameters can be found in Appendix A.I.I.
Instruction Fine-tuning Datasets We experiment with three popular open-source instruction fine-tuning datasets: ShareGPT, adopted by Vicuna [4), UltraChat-200k [7], and --7 [25]. For ShareGPT, we follow the same preprocessing process as Chiang et al. [4]. We remove refusal responses from ShareGPT and M-ZardLM-70K to prevent the fine-tuned models from becoming oversensitive and frequently refusing to respond. For all three datasets, we remove model responses from incomplete conversation chunks that lack user input instructions. Statistics of the processed datasets presented in Table 2.
Table 2: Statistics of instruction fine-tuning tinettuning datasets our experiments. We report the statistics after performing preprocessing detailed in Section 4.1. Average length is measured in the number of tokens with Tamy tokenization.
Datasets
Avg. conversation length
conversations
# instructions
ShareGPT
1,567.68
93,645
331,722
UltraChat-20Ok
1,437.33
207.865
657.794
WizardLM-70K
484.00
57,523
57,523
Context awareness benchmarks. In addition to NIH, we report the performance on three closed- book QA tasks benchmark context awareness: SQuAD [21], QuAC [5], and DROP [8]. SQuAD is a reading comprehension benchmark where the answer to each question can be found in the context. We evaluate only the answerable subset of questions in SQuAD 1.0. QuAC is similar SQuAD, but its questions are more open- endene and the lengths of the answers are longer. While NIH, SQuAD, and QuAC only require direct retrieval from context, DROP requires more complicated reasoning based on the given context, and its answers require discrete operations on the retrieved context such as addition, sorting, or counting.
As instruction fine-tuned models are not specifically trained on QA tasks to provide concise answers, their responses are generally more verbose. Therefore, we report the containment score, defined as whether the model response contains the ground-truth answer with keyword string matching, rather than the Fl score to exclude the effects of different models' response styles. Prompt templates for QA tasks are listed in Appendix A.3.
For Needle-in-a-haystade
we report the recall defined in Section 2.1, which is also the default metric used in Dubey et al. [9]. We set the maximum NIH context length to 1,000 for models fine- tuned on WizardL_M-70k due to its shorter instruction lengths. For models fine-tuned ShareGPT and Utra we set the maximum NIH context length to the maximum context window considered in fine-tuning, which is 2,000 for Liny and 4,000 for Llama-2 / Llama-3. The prompt template used in NIH is the same as Section 2.1 except that we remove the response prefix and keep only the user input prompt.
General instruction-folloving benchmarks. To validate that our method maintains strong per- formance general instructioo-follooing tasks, we evaluate the fine-tuned models on MT-Bench
7

---

[29] where the response quality is rated by GPT-4 judge based on helpfulness, relevance, accuracy. depth, creativity, and level of detail. We report the average rating across the MT-Bennc test cases.
Table 3: Comparing vanilla instruction finetuning with finetuning with context-relevan indicators (+ indicator). For "+ Indicator" models, [IND] added in all evaluations. As a reference, we also list the performances evaluated on official Llama-2 and Llama-3 instruct models, which are finetuned with closed-soure, datasets. NIH, SQuAD, and Qu are simple context-dependent tasks, while DROP and MT Bench require more complex capabilities.
Context-dependent tasks
Complex-skin tasks
SFT datasct
Pretrained Model
Method
NIHI
SQuAD
QuAC
DROP
MT-Bench
Vanilla
0,9846
59.73
15,50
27.39
3.7250
TinyLlama-1..
Indicator
0.9921
62.05
17.40
27.84
3.7375
ShareGPT (Vicuna)
Vanilla 0.3378
76.78
23.60
33,90
6,4875
Llama-2-76
Indicator
0.7007
79.09
24.20
33.90
5.7375
Vanilla 0.8957
83,06
24.80
42.15
7.4375
Llama-3-8h
Indicator
0.9404
84.80
24.50
43.17
7.1625
Vanilla 1.0000
73.03
22.70
30.96
3,9000
TinyLlama-I.I I.IB
Indicator
1.0000
74.47
23.10
30.96
4.1125
UltraChat-200K
Vanilla 0.9850
83.81
24.20
37.91
5.7125
Llama-2-
Indicator
0.9725
85.76
26.10
37.58
5.8125
Vanilla 1.0000
85.12
25.50
50.99
7.2375
Llama-3-8b
Indicator
1,0000
86.28
26.40
50.22
6,8500
Vanilla 0,9250
60.51
13,80
27.53
4.2750
Timaaa 1.1B
Indicator
0.9925
63.39
14.60
28.36
4,3000
WizardLM-70K
Vanilla 0.7375
82.89
23.70
34.07
5.7750
1-ma-2-7h
Indicator
0.9254
83.13
25.30
34.44
6,2250
Vanilla 0.9846
88.25
24.60
46.87
7.1125
Llama-3-8h
Indicator
0.9871
88.53
26.00
47.85
7.5250
Llama-2-77-ce
0.8264"
83.28
22.20
44.22
6,9375
(Closed-source datasets)
Llama-3 3-8ttunsttt
1.0"
86.96
27.40
46.54
8.0750
Here NIH is evaluated without the response prefix used in Section 2.1 and the maximum context length is set to 4096 for fair comparison. Therefore, the exact numbers differ from Figure 2.
4.2 Instruction Fine-tuning with Context-dependency Indicators
Settings and hyperparameters. We adopt a TinyLlama model fine-tuned on the original ShareGPT (Vicuna) dataset as the seed model M and compute the context-dependency score middle layer (15 all of our main experiments) for faster computation. We set the threshold for context-awareness as 3 = 0.6 for all experiments reported in Table 3. An ablation study on the choice of threshold value can be found in Appendix B.6.
Sparsity of context-dependern instructions. We compute the context- sentexxt-deendenncy scores on all three instruction fine-tuning datasets (sce Table 10) and find that context-dependeet instructions are consistently scarce in all datasets. Note that the scarcity is an intrinsic property of the datasets overlooked by the original curators of the data. This observation supports hypothesis that the model learns the bias to weigh user tokens less importantly from the instruction dataset.
4.3 Experiment Results
Conditional finetuning improves performance on context-dependent tasks. Table 3 shows that "+ Indicator" (ours) outperforms "vanilla" fine-tunin tine-tuning consistently across different models and SFT datasets on NIH, SQuAD, and QuAC. The benchmarks isolate and measure context awareness performance. For the "vanilla" fine-tuning setting, we train and evaluate the model without the indicator token. For the "+ Indicator" setting, we add the indicator token to the selected subset of prompts in fine-tuning and all queries for evaluation. The results confirmed that models learn focus more the user-provided context when the indicator token is present in the prompt.

---

Sparsity of context-dependee instructions impacts performance on general tasks. The purpose of evaluating on general tasks is to demonstrate that our method minimally impacts other capabilities than context awareness. 'abble 3 presents the evaluation results on MT--eecc as an assessment of general cntructtio-followinng capabiliites and DROP to assess more complex reasoning capabilities inadditionl to purely contextual retrieval. Models fine-tuned with the indicator perform comparably or sometimes better. Comparing the performance between the three SFT datasets, ShareGPT suffers from the most negative impact with the indicator, while W-Z improves. This can be explained by the net number and thus diversity of the context-dependent subset each dataset. Table 10 shows that WizardLM-70K has the highest proportion of context dettext-dependent samples, while ShareGPT has the least.
Related Work
Instruction fine-tuning and chat templates. To enable instruction-follovingg pretrained LLMs usually require supervised fine-tuning on cnstructioo-followwiit datasets (SFT) [23, 20], and optionally reinforcement learning with human feedback (RLHF) [6]. Instruction fine-tuning datasets consist of user instruction and target model response pairs, which can collected from modified NLP tasks [23, 19), human annotations [20, 4] or synthesized data from existing LLMs [7, 25]. Instruction fine-tuning usually converts training examples into a dialog format with a chat template, which typically consists of user, assistant, and system role indicators. However, these role indicators and role partition in the conversation are not sufficiently presented and learned during pretraining, making them prone to bias during fine-tuning.
Context awareness and hallucinations in LLMs Context awareness is crucial for mitigating hallucinations where the response is not consistent with the provided context, such as the "closed- domain" hallucination in Ouyang al. [20] and intrinsic hallucination in Huang et al. [13]. Several existing works aim to understand and mitigate these intrinsic hallucinations. Liu et al. [18] studies
the failure of context retrieval when the relevant information is in the middle of the provided context, showing the existence of positional bias in context retrieval. Follow-up work [12] proposes to calibrate this positional bias and mitigate the issue. our work, we study a new role bias in popular open-source instruction- finetuned models, where the context receives less attention when marked as user tokens by the chat template. Some other papers such as An et al. [2] synthesize examples targeted toward specific tasks (e.g., contextual QA) to increase the performance tasks. Instead, our method is more general in terms of task and input format compared. This makes our method more applicable different types of user queries that require greater attention to user-provided context.
context-dependent
Side effects of instruction finetuning Neural networks are known to catastrophically forget existing knowledge or capabilities when sequentially trained on new tasks or domains [16]. It is commonly believed that traditional catastrophic forgetting on pretraining stage capabilities can be significantly mitigated by finetuning the model on a diversified mixture of prompts [19, 3, 11]. We discovered that, contrary to popular belief, context awareness can deteriorate after instruction fine-tuning.
6 Conclusion
This work highlights the detrimental effects supervised instruction fine-tuning on the context awareness of pretrained language models, even in scenarios involving short context lengths. We have identified that the decline in context awareness is closely linked to attention allocation biases within chat templates, which are learned during conversational instruction finetuning. Our proposed method utilizes conditional supervised fine-tuning with an indicator marking samples with context-relevant training samples. Our method effectively maintains contextual understanding while benefiting from supervised instruction.
Limitation Due to computational resource limitations, our methods were only validated on smaller models. We hope to extend the experiments to full finetuning or larger models when more resources become available. Additionally, our technique of associating context-dependent user instructions with the indicator token may also encode other unintended styles from the selected subset of instructions. This issue is particularly aggravated when the subset of context-relevant samples is small and

---

significantly different from the remaining dataset. One future direction is to better disentangle the context-dependencc signal within the selected corpus. Another future direction is expanding the evaluation of context-dependent benchmarks (e.g., RAG) to discover the extent of the harm caused by the decreased context awareness.
Acknowledgments and Disclosure of Funding
This work is partially supported by NSF 2048280, 2325121, 2244760, 2331966 and ONR N00014- 23-002300.P00001.
References
[1] Rishabh Agarwal, Avi Singh. Lei M. Zhang, Bernd Bohnet, Luis Rosias. Stephanie Chan, Biao Zhang, Ankesh Anand, Zaheer Abbas, Azade Nova, John D. Co-Reyes, Eric Chu. Feryal Behbahani, Aleksandra Faust, and Hugo Larochelle. Many-shot in-context learning, 2024. URL attps://arxiy org/ abs/2404. 11018.
[2] Shengnan An, Zexiong Ma. Zeqi Lin, Nanning Zheng, and -uang Lou. Make your llm fully utilize the context, 2024. URL https://aara org/ /2044 16811
[3] Yuntao Bai, Andy Jones, Kamal Ndousse, Amanda Askell, Anna Chen, Nova DasSarma, Dawn Drain, Stanislav Fort, Deep Ganguli, Tom Henighan, et al. Training helpful and harmless assistant with reinforcement learning from human feedback. arXiv preprint arXiv:2204.05800 2022.
[4] Wei-Lin Chiang, Zhuohan Li, Zi Lin, Ying Sheng, Zhanghao Wu, Hao Zhang, Lianmin Zheng, Siyuan Zhuang, Yonghao Zhuang, Joseph E. Gonzalez, lon Stoica, and Eric P. Xing. Vicuna: An open-source chatbot impressing gpt-4 with 90%* chatgpt quality, March 2023. URL
[5] Eunsol Choi, He He, Mohit lyyer, Mark Yatskar, Wen tau Yih, Yejin Choi, Percy Liang, and Luke Zettlemoyer. Quac Question answering in context, 2018. URL https: org/ abs/1808.0700
[6] Paul F Christiano, Jan Leike, Tom Brown, Miljan Martic, Shane Legg, and Dario Amodei. Deep reinforcement learning from human preferences. Advances in neural information processing systems, 30, 2017.
[7] Ning Ding, Yulin Chen, Bokai Xu, Yujia Qin, Zhi Zheng, Shengding Hu, Zhiyuan Liu, Maosong Sun. and Bowen Zhou. Enhancing chat language models by scaling high-quality instructional conversations.
arXiv preprint arXiv:2305. 14233, 2023.
[8] Dheeru Dua, Yizhong Wang, Pradeep Dasigi, Gabriel Stanovsky, Sameer Singh, and Matt Gard- ner. DROP: A reading comprehension benchmark requiring discrete reasoning over paragraphs. In Jill Burstein, Christy Doran, and Thamar Solorio (eds.), Proceedings of the 2019 Conference of the North American Chapter of the Association for Computational Linguistics: Human Language Technologies, Volume (Long and Short Papers), pp. 2368-231 2378, Minneapolis, Minnesota, June 2019, Association for Computational Linguisties. doi: 10.1 11119////111111246 URL httpp::/aclaattholl
org/oogg// N11//1199111 1246.
[9] Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang. Angela Fan, Anirudh Goyal. Anthony Hartshorn, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurelien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Roziere, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olano, Diego Perino, Dieuwke Hupkes, Egor Lakomkin, Ehab AlBadawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Gregoire Mialon, Guan Pang, Guillem Cucurell, Hailey
10

---

Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet. Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadeokar, Jeet Shah, Jelmer van der I Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang, Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jia, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, Khalid El-Arini, Krithika Iyer, Kshitiz Malik, Kuenley Chiu, Kunal Bhalla, Lauren Rantala-Yeary, Laurens van der Maaten, Lawrence Chen, Liang Tan, Liz Jenkins, Louis Martin, Lovish Madaan, Lubo Malo, Lukas Blecher, Lukas Landzaat, Luke de Oliveira, Madeline Muzzi, Mahesh Pasupuleti, Mannat Singh, Manohar Paluri, Marcin Kardas, Mathew Oldham, Mathieu Rita, Maya Pavlova, Melanie Kambadur, Mike Lewis, Min Si, Mitesh Kumar Singh, Mona Hassan, Naman Goyal, Narjes Torabi, Nikolay Bashlykov, Nikolay Bogoychev, Niladri Chatterji, Olivier Duchenne, Onur Çelebi, Patrick Alrassy, Pengchuan Zhang. Pengwei Li, Petar Vasic, Peter Weng, Prajjwal Bhargava, Pratik Dubal, Praveen Krishnan, Punit Singh Koura, Puxin Xu, Qing He, Qingxiao Dong. Ragavan Srinivasan, Raj Ganapathy, Ramon Calderer, Ricardo Silveira Cabral, Robert Stojnic, Roberta Railcanu, Rohit Girdhar, Rohit Patel, Romain Sauvestre, Romnie Polidoro, Roshan Sumbaly, Ross Taylor, Ruan Silva, Rui Hou, Rui Wang, Saghar Hosseini, Sahana Chennabasappa, Sanjay Singh, Sean Bell, Seohyun Sonia Kim, Sergey Edunov, Shaoliang Nie, Sharan Narang, Sharath Raparthy, Sheng Shen, Shengye Wan, Shruti Bhosale, Shun Zhang, Simon Vandenhende, Soumya Batra. Spencer Whitman, Sten Sootla, Stephane Collot, Suchin Gururangan, Sydney Borodinsky, Tamar Herman, Tara Fowler, Tarek Sheasha, Thomas Georgiou, Thomas Scialom, Tobias Speckbacher, Todor Mihaylov, Tong Xiao, Ujjwal Karn, Vedanuj Goswami, Vibhor Gupta, Vignesh Ramanathan, Viktor Kerkez, Vincent Gonguet, Virginie Do. Vish Vogeti, Vladan Petrovic, Weiwei Chu, Wenhan Xiong, Wenyin Fu, Whitney Meers, Xavier Martinet, Xiaodong Wang. Xiaoqing Ellen Tan, Xinfeng Xie, Xuchao Jia, Xuewei Wang. Yaelle Goldschlag, Yashesh Gaur, Yasmine Babaci, Yi Wen, Yiwen Song, Yuchen Zhang, Yue Li, Yuning Mao, Zacharie Delpierre Coudert, Zheng Yan, Zhengxing Chen, Zoe Papakipos, Aaditya Singh, Aaron Grattafiori, Abha Jain, Adam Kelsey, Adam Shajnfeld, Adithya Gangidi, Adolfo Victoria, Ahuva Goldstand, Ajay Menon, Ajay Sharma, Alex Boesenberg, Alex Vaughan, Alexei Baevski, Allie Feinstein, Amanda Kallet, Amit Sangani, Anam Yunus, Andrei Lupu, Andres Alvarado, Andrew Caples, Andrew Gu, Andrew Ho, Andrew Poulton, Andrew Ryan, Ankit Ramchandani, Annie Franco, Aparajita Saraf, Arkabandhu Chowdhury, Ashley Gabriel, Ashwin Bharambe, Assaf Eisenman, Azadeh Yazdan, Beau James, Ben Maurer, Benjamin Leonhardi, Bernie Huang, Beth Loyd, Beto De Paola, Bhargavi Paranjape, Bing Liu, Bo Wu, Boyu Ni, Braden Hancock, Bram Wasti, Brandon Spence, Brani Stojkovic, Brian Gamido, Britt Montalvo, Carl Parker, Carly Burton, Catalina Mejia, Changhan Wang, Changkyu Kim, Chao Zhou, Chester Hu, Ching-Hsiang Chu, Chris Cai, Chris Tindal, Christoph Feichtenhofer, Damon Civin, Dana Beaty, Daniel Kreymer, Daniel Li, Danny Wyatt, David Adkins, David Xu, Davide Testuggine, Delia David, Devi Parikh, Diana Liskovich, Didem Foss, Dingkang Wang. Duc Le, Dustin Holland, Edward Dowling. Eissa Jamil, Elaine Montgomery. Elconora Presani, Emily Hahn, Emily Wood, Erik Brinkman, Esteban Arcaute, Evan Dunbar, Evan Smothers, Fei Sun, Felix Kreuk, Feng Tian, Firat Ozgenel, Francesco Caggioni, Francisco Guzmán, Frank Kanayet, Frank Seide, Gabriela Medina Florez, Gabriella Schwarz, Gada Badeer, Georgia Swee, Gil Halpern, Govind Thattai, Grant Herman, Grigory Sizov, Guangyi, Zhang, Guna Lakshminarayanan, Hamid Shojanazeri, Han Zou, Hannah Wang, Hanwen Zha, Haroun Habeeb, Harrison Rudolph, Helen Suk, Henry Aspegren, Hunter Goldman, Ibrahim Damlaj, Igor Molybog, Igor Tufanov, Irina-Elena Veliche Itai Gat, Jake Weissman, James Geboski, James Kohli, Japhet Asher, Jean-Baptiste Gaya, Jeff Marcus, Jeff Tang. Jennifer Chan, Jenny Zhen, Jeremy Reizenstein, Jeremy Teboul. Jessica Zhong, Jian Jin, Jingyi Yang, Joe Cummings, Jon Carvill, Jon Shepard, Jonathan McPhie, Jonathan Torres, Josh Ginsburg, Junjic Wang, Kai Wu. Kam Hou U, Karan Saxena, Karthik Prasad, Kartikay Khandelwal, Katayoun Zand, Kathy Matosich, Kaushik Veeraraghavan, Kelly Michelena, Keqian Li, Kun Huang, Kunal Chawla, Kushal Lakhotia, Kyle Huang, Lailin Chen, Lakshya Garg, Lavender A, Leandro Silva, Lee Bell, Lei Zhang, Liangpeng Guo, Licheng Yu, Liron Moshkovich, Luca Wehrstedt, Madian Khabsa, Manav Avalani, Manish Bhatt, Maria Tsimpoukelli, Martynas Mankus, Matan Hasson, Matthew Lennie, Matthias Reso, Maxim Groshev, Maxim Naumov, Maya Lathi, Meghan Keneally, Michael L. Seltzer, Michal Valko, Michelle Restrepo, Mihir Patel, Mik Vyatskov, Mikayel Samvelyan, Mike Clark, Mike Macey, Mike Wang, Miquel Jubert Hermoso, Mo Metanat, Mohammad Rastegari, Munish Bansal,
11

---

Nandhini Santhanam, Natascha Parks, Natasha White, Navyata Bawa, Nayan Singhal, Nick Egebo, Nicolas Usunier, Nikolay Pavlovich Laptev, Ning Dong, Ning Zhang. Norman Cheng, Oleg Chernoguz, Olivia Hart, Omkar Salpekar, Ozlem Kalinli, Parkin Kent, Parth Parekh, Paul Saab, Pavan Balaji, Pedro Rittner, Philip Bontrager, Pierre Roux, Piotr Dollar, Polina Zvyagina, Prashant Ratanchandani, Pritish Yuvraj, Qian Liang, Rachad Alao, Rachel Rodriguez, Rafi Ayub, Raghotham Murthy, Raghu Nayani, Rahul Mitra, Raymond Li, Rebekkah Hogan, Robin Battey, Rocky Wang. Rohan Maheswari, Russ Howes, Ruty Rinott, Sai Jayesh Bondu, Samyak Datta, Sara Chugh, Sara Hunt, Sargun Dhillon, Sasha Sidorov, Satadru Pan. Saurabh Verma, Seiji Yamamoto, Sharadh Ramaswamy, Shaun Lindsay, Shaun Lindsay, Sheng Feng. Shenghao Lin. Shengxin Cindy Zha, Shiva Shankar, Shuqiang Zhang. Shuqiang Zhang, Sinong Wang. Sneha Agarwal, Soji Sajuyigbe, Soumith Chintala. Stephanie Max, Stephen Chen, Steve Kehoe, Steve Satterfield, Sudarshan Govindaprasad, Sumit Gupta, Sungmin Cho, Sunny Virk Suraj Subramanian, Sy Choudhury, Sydney Goldman, Tal Remez, Tamar Glaser, Tamara Best, Thilo Kohler, Thomas Robinson, Tianhe Li, Tianjun Zhang, Tim Matthews, Timothy Chou, Tzook Shaked, Varun Vontimitta, Victoria Ajayi, Victoria Montanez, Vijai Mohan, Vinay Satish Kumar, Vishal Mangla, Vitor Albiero, Vlad Ionescu, Vlad Poenaru, Vlad Tiberiu Mihailescu, Vladimir Ivanov, Wei Li. Wenchen Wang, Wenwen Jiang, Wes Bouaziz, Will Constable, Xiaocheng Tang, Xiaofang Wang, Xiaojian Wu, Xiaolan Wang, Xide Xia, Xilun Wu, Xinbo Gao, Yanjun Chen, Ye Hu, Ye Jia, Ye Qi, Yenda Li, Yilin Zhang, Ying Zhang, Yossi Adi, Youngjin Nam, Yu, Wang, Yuchen Hao, Yundi Qian, Yuzi He. Zach Rait, Zachary Vito, Zef Rosnbrick, Zhaoduo Wen, Zhenyu Yang, and Zhiwei Zhao. The llama herd of models, 2024. URL https://asee /arxiv..arivv.//aaa org/ abs/2407 21783.
[10] Tingchen Fu, Deng Cai, Lemao Liu, Shuming Shi, and Rui Yan. theeereet Pushing the limits of instruction tuning via alignment tax reduction. In Lun- Wei Ku, Andre Martins, and Vivek Srikumar (eds.). Findings of the Association for Computational Linguistics ACL 2024, pp. 2967-2985, Bangkok, Thailand and virtual meeting, August 2024. Association for Computational Linguistics. doi: 10.1 (0.186333v1/0024 2 URL https:// acl aclanthology org/2024. 2024. f findings-acl. 175.
[11] Amelia Glaese, Nat McAleese, Maja Trebacz, John Aslanides, Vlad Firoiu, Timo Ewalds, Maribeth Rauh, Laura Weidinger, Martin Chadwick, Phoebe Thacker, et al. Improving alignment of dialogue agents via targeted human judgements. arXiv preprint arXiv:2209.1433 14375, 2022.
[12] Cheng-Yo Yu Hsieh, Yung-Sung Chuang, Chun-Liang Li, Zifeng Wang, Long Le, Abhishek Kumar, James Glass, Alexander Ratner, Chen-Yu Lee, Ranjay Krishna, et al. Found in the middle: Calibrating positional attention bias improves long context utilization. In Findings of the Association for Computational Linguistics ACL 2024, pp. 14982-14995, 2024.
[13] Lei Huang, Weijiang Yu, Weitao Ma, Weihong Zhong, Zhangyin Feng. Haotian Wang, Qiang- long Chen, Weihua Peng, Xiaocheng Feng, Bing Qin, et al. A survey hallucination in large language models: Principles, taxonomy, challenges, and open questions. ACM Transactions on Information Systems, 2023.
[14] Gautier Izacard, Patrick Lewis, Maria Lomeli, Lucas Hosseini, Fabio Petroni, Timo Schick, Jane Dwivedi- Yu, Armand Joulin, Sebastian Riedel, and Edouard Grave Atlas: Few-shot learning with retrieval augmented language models. Journal of Machine Learning Research, 24(251): 1-43, 2023.
[15] Urvashi Khandelwal, Angela Fan, Dan Jurafsky, Luke Zettlemoyer, and Mike Lewis. Nearest neighbor machine translation. arXiv preprint arXiv:2010.00710 2020.
[16] James Kirkpatrick, Razvan Pascanu, Neil Rabinowitz, Joel Veness, Guillaume Desjardins, Andrei A Rusu, Kieran Milan, John Quan, Tiago Ramalho, Agnieszka Overcoming catastrophic forgetting in neural networks. Proceedings of the national academy of sciences, 1 114(13):3521-322 2017.
et al.
[17] Yong Lin, Hangyu Lin, Wei Xiong, Shizhe Diao, Jianmeng Liu, Jipeng Zhang, Rui Pan, Haoxiang Wang. Wenbin Hu, Hanning Zhang, Hanze Dong, Renjie Pi, Han Zhao, Nan Jiang, Heng Ji, Yuan Yao, and Tong Zhang. Mitigating the alignment tax of rlhf 2024. URL. https:/ / /aaxxvvvvvv. org/ abs/23/9 06256.
12

---

[18] Nelson F. Liu, Kevin Lin, John Hewitt, Ashwin Paranjape, Michele Bevilacqua, Fabio Petroni, and Percy Liang. Lost in the middle: How language models use long contexts. Transactions of the Association for Computational Linguistics, 12:157-173 2024. doi: 10.1 (0.1162/tacl_a00k a0063..
URL. htt.https//acllathologg.. hwwpeehhooooo.ogg/2/eeee2 www//024. taco-1...
[19] Shayne Longpre, Le Hou, Tu Vu, Albert Webson, Hyung Won Chung, Yi Tay, Denny Zhou, Quoc V Le, Barret Zoph, Jason Wei, et al. The flan collection: Designing data and methods for effective instruction tuning. In International Conference on Machine Learning, pp. 22631- 22648. PMLR, 2023.
[20] Long Ouyang, Jeffrey Wu, Xu Jiang, Diogo Almeida, Carroll Wainwright, Pamela Mishkin, Chong Zhang, Sandhini Agarwal, Katarina Slama, Alex Ray, et al. Training language models to follow instructions with human feedback. Advances in neural information processing systems, 35:27730-27744, 2022.
[21] Pranav Rajpurkar, Jian Zhang, Konstantin Lopyrev, and Percy Liang. Squad: 100,000+ questions for machine comprehension text, 2016. URL https: http://aaxvvvvoooo abs/ 1606. 05250.
[22] Hugo Touvron, Louis Martin, Kevin Stone, Peter Albert, Amjad Almahairi, Yasmine Babaei, Nikolay Bashlykov, Soumya Batra, Prajjwal Bhargava, Shruti Bhosale, et al. Llama 2: Open foundation and tnn-tuned chat models. arXiv preprint arXiv:2307.09288, 2023.
[23] Jason Wei, Maarten Bosma, Vincent Y Zhao, Kelvin Guu, Adams Wei Yu, Brian Lester, Nan Du, Andrew M Dai, and Quoc V Le. Finetuned language models zero-shot learners. arXiv preprint 11Xiv:2109.01652 2021.
[24] Wenhao Wu, Yizhong Wang, Guangxuan Xiao, Hao Peng, and Yao Fu. Retrieval head mechanis- tically explains long-context factuality, 2024. URL https://arxiv oraiiiivorr//bbaa arggabs//444. 15574.
[25] Can Xu, Qingfeng Sun, Kai Zheng, Xiubo Geng, Pu Zhao, Jiazhan Feng, Chongyang Tao, and Daxin Jiang. Wizardlm: Empowering large language models to follow complex instructions. arXiv preprint arXiv:2304. 12244, 2023.
[26] Peng Xu, Wei Ping. Xianchao Wu, Lawrence McAfee, Chen Zhu, Zihan Liu, Sandeep Subrama- nian, Evelina Bakhturina, Mohammad Shoeybi, and Bryan Catanzaro. Retrieval meets long context large language models. arXiv preprint arXiv:2310.03025, 2023.
[27] Peiyuan Zhang, Guangtao Zeng, Tianduo Wang, and Wei Lu. Tinyllama: An open-source small language model, 2024.
[28] Qingru Zhang, Chandan Singh, Liyuan Liu, Xiaodong Liu, Bin Yu, Jianfeng Gao, and Tuo Zhao. Tell your model where to attend: Post-hoc attention steering for LLMs. In The Twelfth International Conference on Learning Representations, 2024. URL https://openrevee/ net/forum? (ddforrommmxxx200000000ee
[29] Lianmin Zheng, Wei-Lin Chiang, Ying Sheng, Siyuan Zhuang, Zhanghao Wu, Yonghao Zhuang, Zi Lin, Zhuohan Li, Dacheng Li, Eric Xing, Hao Zhang, Joseph E. Gonzalez, and Ion Stoica. Judging alm-as---uudge with menhench and chatbot arena, 2023. URL https: /arxivvvvv org/ abs/2306. 25662230600688.
13

---

A Appendix
A.1 Experimental Details
A.1.1 Instruction Fine-tuning
We adopted the fine-tuning recipes from the Huggingface hignment-handbookt for Llama-2 and Llama-3 QLoRA tuning. For the TinyLlama model, we used the fine-tuning recipe provided by the author?. See detail configurations in Table 4. We fine-tuned the models for one epoch ShareGPT and UltraChat- 200K,, and two epochs on WizardLM-708 due to its smaller training set. We used the TinyLlama chat template for all instruct models fine-tuned in Table 3. All experiments conducted on A6000 GPUs on local server. Execution time of training runs usually range from a couple hours to at most 1-2 days.
Table 4: Fine- Fine-tuning hyperparameters configuration
Models
Fine-tune config
Learning rate
Batch size
Precision
TinyLlama
Full fine-tune
2c-5
128
bfl6
1l3ma-2/3
QLoRA with rank = 16, alpha =16
2c-4
64
bf16
A.2
NIH Evaluation Details
For all NIH evaluations, we average the recall error across 400 tests. Specifically, we evaluate on 20 context lengths uniformly distributed between 200 and the maximum context length, and 20 needle insertion depths uniformly located between 0% and 100%.
To ensure that our NIH evaluation is not sensitive to differences in prompt templates, we run evaluations with different prompt templates on small-scale experiments and report the mean and standard deviation. For evaluations larger models or larger context window sizes, report the evaluation results using only one prompt template, We illustrate the mean errors in Figure 1. Complete results with standard deviations can be found in Table 7.
A.3 Contextual QA Evaluation Details
We list the prompts used in contextual QA tasks in Table 5 and Table 6. For contextual QA tasks, we generate answers of up to 100 tokens and truncate them at the end of the first complete sentence. For NIH tests, we generate answers of up to 50 tokens.
As UltraChat-200K constructs its data with fixed set of prompt templates similar to our default ones used evaluation (the templates used for ShareGPT and WizardLM models in Table 6 and 5), we evaluate U finnttundd models with simpler template to exclude the impact of overfitting on finetuning prompt templates.
Table 5: Prompt templates used for SQuAD and DROP in Table and Table when the model is finetuned on different instruction finetuning datasets.
Instruct Finetuning Dataset
Template for SQuAD and DROP
ShareGPT & WizardLM-70K
(contex()laAnswee) the question according the abuve passage: [question] [context) (question)
UltraChat-200K
14

---

Table 6: Prompt templates used for QuAC in Table 1 and Table 3 when the model is finetuned on different instruction finetuning datasets.
Instruct Finetuning Dataset
Template for QuAC
ShareGPT WizardL M-0KK
(context /nnAwwwr the question with pieces from the the above passage: [question]
UltrCChat-20KK
[context] {question)
B Additional Experiment Results
B.1 Full NIH Results on Open-source Official Models
In Figure 2, only report the NIH performances when the response prefix is added, for fair comparison. In Table 7, we show the exact numbers for Figure 2 as well as additional evaluation results without the response prefix. When the response prefix is removed, the performance drop NIH is even more significant compared to results without chat templates.
The standard deviations shown in the table are explained in Section A.2.
Table 7: NIH performance with and without chat templates on different models. The mean and standard deviations were calculated using different prompt templates listed in Table
w/o chat template w/ response prefix
w/ chat template
Model Name
Context window
w/ response prefix
w/o response prefix
Llama-2-7
98.96 94 = 0.33% 99.40 + 0.48%
4K
Llama-2-77 2-7b-chat
99.17 + 0.28 %
92.45 + 1.24%
Llama- Llaba-2-13b
98.89 = 0.34%
4K
1lmm-22-3h-chhn
96.13 + 0.64%
91.50 0.76%
92.78 + 0.244
Llama-3-8h
99.62 0.26%
8K
99.92 + 0.09%
96.10 16 + 0.58%
95.89 + 0.62%
Llama-3.1
86.89%
128K
Llama-3.1 3.1.-...iii-ttuuee
98.17%
95.64%
94.75%
mistral-v0.2
100%
32K
mistral-vo..2innsrruut
99.00%
94.14%
93.92%
mistral-v0.3
100%
32K
mistral-vO.3-innstuuct
99.32%
84.71%
72.00%
gemma-2-9b
100
0%
8K
gemma-2-9b-i
98.75 + 0%%
98.03:t 0%
98.72 + 0.54%
gemma-2-276
100%
8K
gemma-2-27b-it
100%
99.64%
99.25%
B.2 Full Results for Figure
In Figure 3, we only show the changes in attention allocation with and without chat templates. In Figure 4, we show the absolute numbers of attention allocation for each part of the input prompts. When the chat template is added, we normalize the attention weights on user tokens, response tokens, and the BOS token only, with the sum of attention allocation being 1.
B.3 Probing Context Retrieval Heads
In Sections 2.2, 2.3, and 3.1, we mentioned identifying a representative context retrieval head on each layer that allocates the largest attention to user tokens. Because different attention heads can have very different functionalities, and context retrieval heads can sparse among heads [24], we believe that selecting a representative context retrieval head for visualization, attention steering, and data selection is both necessary and important. Specifically, given an input sequence X1, X2, ..., X with tokens, select the context retrieval head h from each layer that allocates the highest
15

---

Table 8: Four different prompt templates used in the NIH evaluation. In Table 7, report the mean and standard deviation across different prompt templates for small models and small context windows. (context) represents the context with the needle inserted.
Prompt templates in NIH evaluation
You are helpful AI assistant that answers question using only the provided document:
(context)
Question: [retrieval question]
You are a helpful Al assistant that answers question using only the provided context:
[context]
Question:
veetrievll_question
Document:
[context)
Answer the question accoriding to provided document: [retrieval_question)
Context:
{context]
Answer the question accoriding to the provided context: (retrieval_guestion)
divence witted of
Thorm
003 Thves
Percentage 66
Ulmaa22 2-7--2770000 L/ama22-385 136 Chat
Llama-3 3 90 90-nstruct
Ulma- 2.2 1-9--trttr
Matrall 7e-tstttt-tss
Mhtrral 7istrri-7-
gemma.2 2010022-2200 2-98-2
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 4000 and needle depths from 0% to 100%.
attention to user tokens when generating the first answer token: ft"
arg max (x),,.(x_1,xx,)
bE
where x_1 and x; are the last token and ith token in the sequence of tokens on each layer, respectively. U is subset of all user tokens.
In Section 2.2, we select the retrieval head on layer 15 for each input prompt and visualize the average attention change the selected head across different prompts. In Section 2.3, we select one retrieval head on each layer using a sampled NIH prompt and steer all identified retrieval heads. In Section 3.1, select the retrieval head on layer 15 for each input prompt and calculate the context-dependency score. In Section B.5, we provide a further discussion on the agreement between different layers.
B.4 Attention Distribution Analysis
This section compares the attention distributions of inference- time attention steering and training-time conditional SFT to illustrate why the latter method better for recovering context-awareness. We analyze the attention allocated to user versus assistant responses in Llama-2, fine-tuned on ShareGPT. The attention steering method is applied with a 0.95, while the conditional SFT model was trained using a context-relevance filtering threshold of B 0.6.
16

---

Fig presents the attention distribution averaged across all heads (left) and the distribution for the retrieval head specifically (right). We observe that while attention steering boosts user attention more than conditional SFT when averaged across all heads, the opposite is true for the retrieval head, This discrepancy highlights the coarse-grained nature of attention steering; it uniformly steers all heads, which can unnecessarily boost attention in ways that do not improve context-awarenes:. Conversely, the conditional SFT method learns to selectively increase user attention on the critical retrieval head, thus avoiding the unintended performance regressions associated with indiscriminate attention manipulation.
Average of Heads
Retrieval Head
0.010
User Attention
0.00010
Assistant Attention
Attention delta
0.005
Attention delta
0.00005
0.000
0.00000
-0.005
-0.00005
User Attention
-0.010
0.00010
Assistant Attention
Attention Steering
Conditional SFT
Attention Steering
Conditional SFT
Figure 5: We visualize the attention distribution shift the attention steering and conditional SFT variants relative to the vanilla instruction fine-tuned model.
B.5 Agreement Between Different Layers
In Figure 6, we calculate and visualize the disagreement heatmap in S selection when the context- dependency score is calculated across different layers. We use the same TinyLlama model, fine-tuned on the vanilla ShareGPT dataset, the seed model M. Specifically, we first caleulate the context- dependency scores for each conversation turn 500 randomly sampled examples from the ShareGPT dataset across different layers. We then select the top 10% of conversation turns with the highest context-dependency scores each layer as the subset St We compute the disagreement between two layers and I' by calculating the ratio of non-overlapping conversation turns their respective subsets S; and Sr. We can see from the figure that the disagreement among the middle layers is low, indicating that safely choose an arbitrary layer for the context- context-dependency score calculation.
B.6 Ablation Study for Different Threshold /
Table 9: Ablation study with different threshold B, which is used in Section 3.
Threshold /
SQuAD
QuAC
DROP
MT-Bench
1.0 (Vanilla)
0.5918
0.1130
0.2739
3.725
0.5
0.6207
0.1270
0.2872
4.075
0.6
0.6144
0.1290
0.2784
3.825
0.7
0.6160
0.1290
0.2786
3.675
We use 0.6 in all our main experiments. To evaluate the sensitivity the threshold B, select S with different thresholds and prepare the final modified instruction finetuning dataset. We fine-tune a TinyLlama-1.11 model on these three datasets and evaluate it on three contextual QA tasks and MT-Bench. As shown in Table 9, all three models outperform vanilla finetuning on the contextual QA tasks. However, performance MT-Bench shows decreasing trend when the threshold increases from 0.5 to 0.7, potentially due more drastic difference between $ and the unselected subset.
17

---

01
layer
0.000.020
00000000000000000000000000000010
layer
0.02
0.020.020.00
3.040.040..0 0.040.640.0000.000
000
layer
layer_6
layer
0.05 45g.::00........................
8
080005
00000000000000000000000000000000
0.07 0.070.000.000
layer_10
layer_11 00000300000000000000000000000050
layer_ _____12____ 0...00.......0..0700000000 01111110000077000000000000000000 0.000.000000000000
layer_1_ 13 01..
layer_14 14 2..00......000000000005000100010
layer 15
0.030.030 04000033300.00000000000000000000
layer 0.020.030.0400000000 12
2.050.060.000 00 01000000000000000000000000000000
layer 17
layer 0.060.060 011.0000000000000000000000110000
0...0.070.006.0000 wwt.e1/1900000100000000000000000 0.090 0.08
layer 20
layer_21 21 0.08 080 08 0.0.0.0000 w.0.000000000000 0.000000000000000000...00..00000 0 0 0.07 0 0.000000 0.00 0.09 0.02
0.066.09
05 0.06 080 07 03
0.000000000000 0 0.00
layer 0
layer _
layer 3
layer 5
layer 7
layer 9
layer 10
layer 11
layer 12
layer 13
layer 14
layer 15
layer 16
layer 17
layer 18
layer 19
layer_ 20
layer 21
layer
layer
layer
layer
Figure 6: We visualize the disagreement heatmap of $ selection when the context-dependency score (ym) is calculated across different layers. We select $ the 10% of conversation turns with the highest context-dependengg scores on each layer. The disagreement is measured by the number of non-overlapping conversation turns in Š selected by any two layers.
Table 10: Ratio of identified context-dependent instructions in each instruction finetuning dataset. Total number of instructions in each dataset can be found in Table 2
Dataset
0.5
0.6
ShareGPT (Vicuna)
0.14
0.10
0.07
0.04
UltraCHat-200K
0.22
0.18
0.14
0.11
WizardL.M-709
0.34
0.23
0.13
0.06
B.7 Distribution of Instruction Lengths
Here we visualize the changes in the distribution of instruction lengths between the original instruction finetuning dataset and the selected context-dependen subset S. Although higher context-dependency to some extent correlated with longer instruction lengths, there is still a large number of short instructions showing high context dependency and selected for inclusion in S.
B.8 Proportion of context-dependent samples filtered with different thresholds
We report the ratio of identified context-dependen instructions with different threshold values / in Table 10.
18

---

Original dataset 5
Selected context--eeendenn subset 5
Number of Instructions
500
1000
1500
2000
2500
3000
3500
4000
Instruction Length
Figure 7: Change of instruction lengths between the original and the selected subset from ShareGPT dataset.
Original dataset
Selected context-dependent subset 5
Number of Instructions
500
1000
1500
2000
2500
2000
3500
4000
Instruction Length
Figure 8: Change of instruction lengths between the original and the selected subset from UltraChat- 200K dataset.
Original dataset
Selected context-dependent subset 5
Number of Instructions
30
500
1000
1500
2000
2500
3000
3500
4000
Instruction Length
Figure 9: Change of instruction lengths between the original and the selected subset from WizardLM- 70K dataset.
19

---

NeurIPS Paper Checklist
Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made the paper and important assumptions and limitations. A No or NA answer to this question will not be perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Section 6
Guidelines:
The answer NA means that the paper has no limitation while the answer No means that paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create a separate "Limitations" section their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on a few datasets or with a few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-text system might not be used reliably to provide closed captions for online lectures because it fails to handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimesss
While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, a worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community, Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory assumptions and proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and complete (and correct) proof?
Answer: [NA]
20

---

Justification: [NA]
Guidelines:
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but they appear in the supplemental material, the authors are encouraged to provide short proof sketch provide intuition.
Inversely, any informal proof provided in the core of the paper should complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental result reproducibility
Question: Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that affects the main claims and/or conclusions of the paper (regardless whether the code and data are provided or not)?
Answer: [Yes]
Justification: Section 4 and Appendix A.I.1
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is a specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g.. in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) the contribution is primarily a new algorithm, the paper should make clear how to reproduce that algorithm.
(b) the contribution is primarily a new model architecture, the paper should describe the architccture clearly and fully.
(c) If the contribution is new model (e.g., a large language model), then there should either way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome to describe the particular way they provide for reprodueibility. the case of closed-source models, may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
21

---

Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurIPS code and data submission guidelines (https://nips.c cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (c.g., for a new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https: /niiss wwtppcccpppppaeeeccccecccccceeee / ssissssiiciiyyy for more details.
The authors should provide instructions on data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments are reproducible, they should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental setting/details
Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary understand the results?
Question:
Answer: [Yes]
Justification: Section and Appendix A.1.1 Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, appendix, as supplemental material.
7. Experiment statistical significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [No]
Justification:
The temperature when evaluating LLMs set 0 so the output is completely deterministic.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the enror bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
22

---

It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality of errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables figures symmetric error bars that would yield results that are out range (e.g. negative error rates).
error bars are reported in tables or plots, The authors should explain the text how they were calculated and reference the corresponding figures or tables in the text.
8. Experiments compute resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed to reproduce the experiments?
Question:
Answer: [Yes]
Justification: Appendix A.1.1
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate type of compute workers CPU or GPU, internal cluster, or cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported in the paper (e.g., preliminary or failed experiments that didn't make into the paper).
9. Code of ethics
Question: Does the research conducted the paper conform, in every respect, with the NeurlPS Code of Ethics https: h/t//nnurip.../ wtweerrps.ccppplliii/oooooccccee httpp//liic.cccccccueeeeeeeeeeee
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the authors have not reviewed the NeurlPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- cration due to laws or regulations their jurisdiction).
10. Broader impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: This paper presents work aimed at advancing the field of machine learning. We investigate and mitigate language models' loss of context awareness after supervised fine-tuning (SFT). Our work can potentially benefit many real-world applications, such as generation, in-context learning, and contextual question-anen answering. We have not identified any potential negative societal consequences.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact why the paper does not address societal impact.
23

---

Examples of negative societal impacts include potential malicious unintended uses (e.g., disinformation, generating fake profiles, surveillance), faireess considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and not tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional unintentional) misuse of the technology.
If there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms monitor how system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data or models that have a high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper poses no such risks.
Released models that have high risk for misuse or dual-use should be released with necessary safeguards to allow for controlled use of the model, for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards is challenging, and many papers do not require this, but we encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g., code, data, models), used in the paper, properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Justification: [NA]
Guidelines:
The answer NA means that the paper does use existing assets. The authors should cite the original paper that produced the code package dataset.
The authors should state which version of the asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g., website). the copyright and terms of service of that source should be provided.
If assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswithcode. thcode. com/ datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.

---

For existing datasets that are re-packaged, both the original license and the license of the derived asset (if has changed) should be provided.
If this information is not available online, the authors are encouraged to reach out to the asset's s creators.
13. New assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer:
[NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model as part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember to anonymize your assets (if applicable). You can either create anonymized URL or include anonymized zip file.
14. Crowdsourcing and research with human subjects
Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According the NeurIPS Code of Bthies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional review board (IRB) approvals or equivalent for research with human subjects
Question: Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the paper does not involve crowdsourcing nor research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurlPS Code of Ethics and the guidelines for their institution.
For initial submissions, not include any information that would break anonymity (if applicable), such the institution conducting the review.
25

---

16. Declaration of LLM usage
Question: Does the paper describe the usage of LLMs if it is an important, original, or non-standard component of the core methods in this research? Note that if the LLM is used only for writing, editing, or formatting purposes and does not impact the core methodology. scientific rigorousness, originality of the research, declaration is not required.
Answer: [NA]
Justification: [NA]
Guidelines:
The answer NA means that the core method development in this research does not involve LLMs as any important, original, non-standard components.
Please refer to L.M policy (https: http://neurisscpee Satrrripssccccoofffeeeeee/////// for what should or should not be described.
26


Reviews:
[{"review_id": "8QJhpgu9tS", "reviewer": "Reviewer_3Mu4", "summary": "The paper shows that instruction tuning, particularly with chat templates, weakens the model’s ability to remain context-aware. The paper further discusses that during conversational finetuning, models become biased toward specific roles, as seen through attention patterns. It was found that user tokens get less attention with templates. While steering more attention to user tokens improves context awareness (like on the NIH dataset), it hurts performance on others like DROP.  To address this, the paper introduces a new conditioning method that measures how much a model’s response depends on earlier user instructions. This context dependency is indicated by a score which is generated via reference LLM. If a conversation turn is context-dependent, a special token [IND] is added to the user instruction. This approach helps the model perform better on both context-aware and complex tasks.", "questions": "None", "limitations": "Yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 4, "significance": 4, "originality": 4, "strengths_and_weaknesses": "Strength\n\n - I liked the idea and flow in which the methods are derived in a coherent, step-by-step manner that makes the paper easy to follow.\n- The paper is well-written, making the proposed method accessible to readers\n- The paper uses a solid and relevant set of models to support its findings.\n- Proposes two ways to mitigate context awareness in LLMs. The claims are supported with empirical evidence for most part of study.\n\n\nWeaknesses\n\n- The paper uses only the NIH dataset to show the negative impact of chat templates. Evaluating at least top 1 or 2 models on one more context aware dataset will strengthen the claim. It would also help if performance on complex skills is not added to understand how models perform on these before and after instruction tuning.\n\n- The attention steering approach appears limited to a specific type of attention mechanism. A discussion on how the method could be extended to other attention implementations would improve understanding of its generalizability.\n\n- The paper would benefit from a direct comparison, ideally with graphs or figures, showing how attention to different prompt components differs between the two methods: attention steering and conditional instruction fine-tuning.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and for appreciating the clarity of our presentation, the relevance of our model choices, and the coherence of our proposed methods. Please see the responses to the mentioned weaknesses below:\n\n> W1: The paper uses only the NIH dataset to show the negative impact of chat templates. Evaluating at least top 1 or 2 models on one more context aware dataset will strengthen the claim. It would also help if performance on complex skills is not added to understand how models perform on these before and after instruction tuning.\n\nA1: In response to the reviewer’s suggestion, we explored simpler QA-style tasks to isolate the impact of chat templates. We found that the pretrained (non-instruction-tuned) models were unable to answer even simple questions, despite adding hints to the prompts. Nevertheless, we conducted a comparison between instruction-tuned models with and without chat templates on LongBench, a benchmark specifically designed to evaluate context awareness with relatively straightforward questions that do not require complex reasoning. Performance was measured using a bag-of-words recall metric, and the results are as follows:\n\n| | Llama-3-8B-Instruct | Llama-3.1-8B-Instruct |  \n| --- | --- | --- |\n| w/ chat template | 0.3504 |0.5176 |\n| w/o chat template | **0.3662** | **0.5267** |\n\nWe observe that removing the chat template consistently improves performance, which further supports our claim that the use of role tokens introduces a bias that reduces context awareness.\n\n> W2: The attention steering approach appears limited to a specific type of attention mechanism. A discussion on how the method could be extended to other attention implementations would improve understanding of its generalizability.\n\nA2: While our attention steering experiments were conducted using standard scaled dot-product attention, the method is not inherently tied to this mechanism. In the revision, we will add a discussion section describing how the same idea could be adapted to architectures with alternative attention mechanisms (e.g., multi-query or linear attention) since it only relies on reweighting cross-token dependencies.\n\n> W3: The paper would benefit from a direct comparison, ideally with graphs or figures, showing how attention to different prompt components differs between the two methods: attention steering and conditional instruction fine-tuning.\n\nA3: We thank the reviewer for suggesting a more direct comparison between attention steering and conditional fine-tuning. We will add a new figure to the supplementary material that visualizes the attention distribution differences between these two methods across prompt components. The attention scores of the user (assistant) tokens for the steering are merely multiplied (divided) by a constant factor. This further highlights why conditional fine-tuning provides a more flexible trade-off across tasks."}, {"review_id": "oVb8u9Clfa", "reviewer": "Reviewer_uYwu", "summary": "The author identifies the loss of context-awarness measured by NIH test after fine-tuning a pre-trained model, and shows that the observation is related to the use of chat template. The author further shows evidence that the root cause can be the shift of attention weight toward assistant tokens when chat template is applied during SFT. Based on the observation, the author proposes a fine-tuning techniques that identifies and adds a context-awareness indicator to the fine-tuning data, and shows general improvement after applying the technique on various combinations of models and tasks.", "questions": "1. It would be interesting to include results when different models (potentially the original pre-trained model) are used to determine the context-aware indicator. The results would give further insights on how the capabilities of indicator model affects the recovery of context-awareness.", "limitations": "Yes.", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 4, "strengths_and_weaknesses": "Strength:\n1. The paper is very well-written and well-motivated.\n2. The paper presents concise and convincing study of the loss of context awareness and its root cause.\n3. The method developed based on manual modification is simple yet effective.\n\nWeakness:\n1. It seems that the loss of context-awareness is not catastrophic. Therefore, it is unclear to the reviewer whether paying the compute cost to generate context-aware indicators is worthwhile when the user perform SFT in general scenarios.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and for recognizing the clarity, motivation, and effectiveness of our study and proposed method. Please see the response to the question below:\n\n> Q1: It would be interesting to include results when different models (potentially the original pre-trained model) are used to determine the context-aware indicator. The results would give further insights on how the capabilities of indicator model affects the recovery of context-awareness.\n\nA1: We thank the reviewer for pointing out this aspect, as it highlights an important benefit of our approach. Our current implementation uses a different model to compute the context-aware indicator than the one being fine-tuned. Specifically, the indicator is generated using a TinyLlama model fine-tuned on the ShareGPT data (Page 7, Line 261). This indicator is then reused across all experiments with different pretrained models. This design addresses the reviewer’s concern in the weakness section: the indicator only needs to be computed once per dataset, and the model used for this computation can be significantly smaller, minimizing the computational cost."}, {"review_id": "Z068Yka8LW", "reviewer": "Reviewer_YvtK", "summary": "This paper uncovers a “role-bias” in instruction-tuned LLMs: when a chat template is used, user tokens receive disproportionately low attention. The authors document the phenomenon on NIH question-answering tasks across Llama, Mistral, and Gemma families. They then propose a mitigation strategy: measure the attention mass on user tokens, and—if it exceeds a tunable threshold—prepend a special marker token to the user instruction during training. Fine-tuning with this marker improves context utilization on NIH, Closed-Book QA, and MT-Bench.", "questions": "1. Can the authors provide a comparison experiment to inference time reweighting methods like PASTA? Reweight the attention towards user instructions and see if the proposed addition of special token is better and worth the training efforts. Doesn't need to be large scale experiments but some baselines would significantly strength the work.\n\n2. If best, the authors are encouraged to provide a experiment that scales over 8k tokens given that some major LLM usage scenes nowadays are long context (long COT, multi-turn model environment interactions). The experiments don't need to be on reasoning models or agentic tool call models but some tasks that require longer context will be evidence that the proposed method can generalize to longer inputs. \n\nQ1 : Baseline request\nPlease include a comparison with attention-reweighting techniques applied only at inference, such as PASTA. An experiment that amplifies attention on user instructions—without retraining—would reveal whether the proposed special-token fine-tuning justifies its additional cost.\n\nQ2: Long-context evaluation\nModern LLM workloads often exceed 8 k tokens (e.g., lengthy chain-of-thought prompts or multi-turn tool-use settings).  The experiments don't need to be on reasoning models or on agentic tool-call tasks but demonstrating the method on a task that truly stresses long contexts would strengthen the claim that it generalizes beyond inputs of moderate-lengths.", "limitations": "yes", "rating": 5, "confidence": 4, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 4, "clarity": 4, "significance": 3, "originality": 3, "strengths_and_weaknesses": "Strengths\n1. Comprehensive empirical evidence supports both the diagnosis of role-bias and the proposed fix. The experiments span three model families, three instruction-tuning corpora, and multiple evaluation suites, showing consistent gains on context-dependent tasks.\n\n2. The manuscript is clear and well-structured, making the core insight and methodology easy to follow.\n\nWeakness\n1. Task coverage – NIH and short-span QA test ≤ 8 K tokens; real-world RAG are not studied.\nIt would be beneficial to add one or two tasks that require long context (>8k) to showcase the proposed solutions can scale up to longer context. \n\n2. Dependency on a seed model & threshold: \nContext-dependency score depends on head selection and a heuristic β. In practition, if we want to scale up training and include more datasets, I can imagine that this threshold needs to be tuned.", "was_revised": true, "final_justification": "This paper uncovers a “role-bias” and proposes a simple yet effective approach to improve user instruction following. I had some misunderstandings about the paper in my original review. But the authors pointed them out and addressed almost all of my comments. \nThe only remaining experiment I would like to see is how they evaluate on more complex or long-horizon tasks (for example multi-turn, or agentic tasks instead of single turn tasks) but that doesn't diminish the values of this paper. I would recommend acceptance.", "rebuttal": "Thank you for your detailed review and for appreciating the clarity of our manuscript as well as the breadth and consistency of our empirical evidence across models, datasets, and evaluation suites. Please see the responses to the questions below:\n\n> Q1: Baseline request Please include a comparison with attention-reweighting techniques applied only at inference, such as PASTA. An experiment that amplifies attention on user instructions—without retraining—would reveal whether the proposed special-token fine-tuning justifies its additional cost.\n\nA1: We believe there may be a misunderstanding, as Section 2.3 of the paper is entirely dedicated to inference-time attention steering methods (specifically PASTA). To reiterate our findings, such attention steering can partially mitigate the reduced context-awareness. However, stronger steering magnitudes introduce negative side effects, which limit its effectiveness compared to the fine-tuning approach (see Page 5 Line 171).\n\n> Q2: Long-context evaluation Modern LLM workloads often exceed 8 k tokens (e.g., lengthy chain-of-thought prompts or multi-turn tool-use settings). The experiments don't need to be on reasoning models or on agentic tool-call tasks but demonstrating the method on a task that truly stresses long contexts would strengthen the claim that it generalizes beyond inputs of moderate-lengths.\n\nA2: We agree with the reviewer that it is important to verify whether the identified issue also appears in long-context settings. As shown in Fig. 2 (Table 8), our NIH benchmark does evaluate models at their full context lengths, which range from 4K up to 128K tokens (e.g., 32K for Mistral models and 128K for Llama‑3.1). Due to computational constraints, we were unable to fine-tune models beyond 8K context for additional complex tasks. Nevertheless, since the observed reduction in context-awareness persists even in models with extended context windows, we expect our proposed method to remain effective in these long-context scenarios.\n\n> W2: Context-dependency score depends on head selection and a heuristic β. In practition, if we want to scale up training and include more datasets, I can imagine that this threshold needs to be tuned.\n\nWe appreciate the reviewer’s concern regarding the dependence on hyperparameters when applying our method to new settings. As described on Page 6 Line 207, the attention head is selected by identifying the head that assigns the highest weight to user tokens. Importantly, the seed model used for head selection does not need to match the fine-tuned model, as demonstrated by our experiments. For the heuristic choice of $\\beta$, we provide a sensitivity analysis in Appendix B.5, which shows that the performance of our method is robust and not overly sensitive to the specific value of this threshold."}, {"review_id": "VVETnR7cmu", "reviewer": "Reviewer_wJSG", "summary": "This paper investigates the degradation of context awareness in LLMs after instruction fine-tuning. The authors show that instruction-tuned models exhibit a bias that reduces attention to user-provided context, particularly when prompted using chat templates. The authors trace this issue to instruction datasets where many responses do not rely on context using the Needle-in-a-Haystack (NIH) task and analysis on attention. They propose to use a token indicating context-dependent instructions during fine-tuning, encouraging the model to reallocate attention appropriately.", "questions": "* Can the proposed indicator token method generalize effectively beyond simple retrieval tasks, particularly to tasks requiring complex reasoning or operations?\n* Can you analyse further to answer the second point of limitations (above mentioned)?\n* Would the proposed method be scalable and effective for larger models or larger datasets?", "limitations": "yes", "rating": 4, "confidence": 3, "soundness": null, "presentation": null, "contribution": null, "strengths": "", "weaknesses": "", "quality": 3, "clarity": 3, "significance": 3, "originality": 3, "strengths_and_weaknesses": "### Strengths\n* The paper identifies a specific and underexplored issue in instruction fine-tuning: loss of context awareness due to role-based attention bias introduced by chat templates.\n* The analysis of attention on NIH tasks shows how context sensitivity deteriorates after SFT.\n* The proposed solution, appending an indicator token to context-dependent instructions, is practical and easy to apply in real-world fine-tuning pipelines.\n\n\n### Limitations\n* The effectiveness of the indicator token approach is clearly shown on the NIH task; however, improvements are marginal or inconsistent on more complex tasks such as DROP and MT-Bench, raising concerns about the generalizability of the proposed method.\n* While the paper identifies bias in training data as a root cause, it lacks comprehensive ablation studies. Comparing models fine-tuned on balanced datasets, specifically ones with a higher proportion of context-dependent examples, would provide stronger evidence regarding the influence of data versus model architecture or training procedures.", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your thoughtful review and for recognizing our identification of role-based attention bias and the practicality of our proposed indicator-token solution. Please see the responses to the questions below:\n\n> Q1: Can the proposed indicator token method generalize effectively beyond simple retrieval tasks, particularly to tasks requiring complex reasoning or operations?\n\nA1: Thank you for the suggestion. We conducted an additional evaluation of the fine-tuned Llama-3-8B models on the math_qa dataset, and the resulting accuracies are reported in the table below. Consistent with our findings on DROP, fine-tuning improves the performance of the Wizard and UltraChat models, while it slightly degrades the performance of the ShareGPT models. These results further support our discussion in Section 4.3 that the sparsity of context-dependent instructions can affect performance on complex tasks when using our method.\n\n| | **ShareGPT** | **ShareGPT (ours)** | **UltraChat** | **UltraChat (ours)** | **Wizard** | **Wizard (ours)** |\n| --- | --- | --- | --- | --- | --- | --- |\n| math_qa | $28.87 \\pm  0.68$ | $26.19 \\pm 0.66$ | $10.73 \\pm  0.46$ | $11.66 \\pm  0.48$ | $25.77 \\pm  0.65$ | $26.7 \\pm 0.66$ |\n\n> Q2: Can you analyse further to answer the second point of limitations (above mentioned)?\n\nA2: While we do not have experiments on datasets with an extremely high proportion of context-dependent samples, our work already includes experiments on three fine-tuning datasets that differ in their density of context-dependent examples. As shown in Table 10, the proportion of context-dependent samples increases from ShareGPT to WizardLM. Correspondingly, performance on MT-Bench and DROP improves more significantly on datasets with a higher density of context-dependent conversations (e.g., WizardLM-70K), which also exhibit a more diverse distribution within their context-dependent subset. The additional math_qa results presented above (A1) show a consistent trend. Therefore, although we have not tested datasets where context-dependent samples approach 100%—which is impractical given the broad diversity of use cases—we expect our method’s effectiveness to hold, and potentially improve, as the proportion of context-dependent samples increases.\n\n> Q3: Would the proposed method be scalable and effective for larger models or larger datasets?\n\nA3: Since the issue we identify stems from biases in the post-training data, we expect that larger models would be similarly affected by the same attention bias. Consequently, our method should remain effective for larger models as well. For larger datasets, the impact primarily depends on the data distribution: as long as context-dependent samples do not dominate the dataset, the same bias will persist, and our method should continue to be effective. In other words, the scalability of our approach is determined by the nature of the data rather than the model or data size."}]

Author rebuttal:

Meta-review: The paper investigates how instruction fine-tuning with chat templates reduces LLMs’ context awareness and identifies role-bias in attention patterns as a root cause. The authors propose conditional fine-tuning with a context-dependency indicator token, demonstrating consistent improvements across multiple models and tasks without sacrificing general instruction-following ability. Strengths highlighted by reviewers include the clarity of exposition, thorough empirical analysis across datasets and model families, and a simple yet effective mitigation strategy. Weaknesses noted were limited evaluations on long-context and complex reasoning tasks, reliance on a heuristic threshold, and modest gains in some cases. However, during rebuttal the authors addressed these concerns with additional experiments (e.g., LongBench, math_qa) and clarifications on generalizability, scalability, and computational cost, which reviewers acknowledged as satisfactory even if some open questions remain. Overall, the contribution is original, timely, and practically relevant, making a compelling case for acceptance.
Decision: Accept (poster)