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

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:

Q1: Can the proposed indicator token method generalize effectively beyond simple retrieval tasks, particularly to tasks requiring complex reasoning or operations?

A1: 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.

Q2: Can you analyse further to answer the second point of limitations (above mentioned)?

A2: 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.

Q3: Would the proposed method be scalable and effective for larger models or larger datasets?

A3: 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.

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:

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.

A1: 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).

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.

A2: 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.

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.

We 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.

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:

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.

A1: 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.

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:

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.

A1: 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:

We 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.

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.

A2: 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.

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.

A3: 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.

Here is a gentle reminder: The Author-Reviewer discussions will end by Aug 8, 11.59pm AoE. Please also be aware that there is a "Mandatory Acknowledgement" for the reviewers.