Teaching LLMs How to Learn with Contextual Fine-Tuning

Weakness 2: Writing

In addition to this major weakness, there are some aspects of exposition I'd try to improve. The first paragraph of the intro seems unnecessarily "hype-y." I actually think you could cut this paragraph and the paper would be fine. Also in the abstract and intro, I would say it's not clear to me that this method is necessarily "learning to learn": that might be one explanation, but there may be others. In the conclusion, the two paragraphs seem to repeat each other a bit. The authors might consider cutting the first one and just keeping the second.

Thank you sincerely for carefully reading our manuscript and for providing such thoughtful and constructive feedback.

We will correct the typos, revise the abstract, the introduction and the conclusion as you suggest.

I also think it might be interesting for the authors, at some point, to speculate a bit on the mechanisms that make the technique work, or suggest some future experiments that might elucidate these mechanisms.

In our response to the reviewer ULmj, we referenced the work "Why Think Step by Step? Reasoning Emerges from the Locality of Experience" by Prystawski et al. (2023). This study suggests that chain-of-thought (CoT) prompting works because local reasoning steps are embedded in the pretraining data, effectively simulating an internal thought process during training that manifests at inference time.

Drawing inspiration from this, we hypothesize that our Contextual Fine-Tuning (CFT) method may work by leveraging similar mechanisms. The contextual prompts could be guiding the model to process information more effectively during training, influencing its internal representations and gradient updates.

We believe that a future study examining how CFT influences the model's learning dynamics—as a function of the pretraining data—could shed light on the mechanisms behind the effectiveness of contextual prompts.

Thank you again for your insightful suggestion.

References:

[1] Ben Prystawski and Noah D. Goodman. 2023. Why think step-by-step? Reasoning emerges from the locality of experience. CoRR abs/2304.03843 (2023)

Weakness 1: Synthetic experiment and Question 1

Here's why I don't think Section 4 is relevant. In the synthetic function setting, the prepended prompts contain customized relevant information for the subsequent prediction tasks, helpfully factoring the functions that the system is learning to compute. This contrasts with the real-world CFT case, where—if I understand correctly—the prepended prompt has no customized connection at all to the upcoming text. I am happy to be corrected if I'm mistaken on this difference, but if I'm reading correctly, I think the synthetic task just doesn't shed light on the real-world task. That said, I don't think the synthetic task is particularly necessary for the argument, so I recommend simply cutting this section completely.

My main question is whether my reading of the difference between the synthetic vs. the real-world setting is correct. If the real-world prompts contain customized information about the upcoming text (contrary to my reading), then I would say Section 4 is relevant after all.

Thank you for your thoughtful comment. First, you are correct that, in the real-world setting, the prepended contextual prompts do not have a customized connection to the upcoming text. However, we'll highlight that we have newer experiments in Table 1 in the general comments that continue to indicate improved performance against CPT. Both original contextual fine-tuning and fine-tuning with contextual prompts that are dependent on the upcoming text (AutoDep-CFT) outperform the base Chat model and CPT across most tasks. i.e. even in experimental settings where real-world prompts contain customized information, CFT continues to outperform CPT. We hope the additional experiment, that we will fold into the manuscript, alleviates concern about the potential disconnect with respect to Section 4.

That said, your comment also serves as a reminder to better motivate our synthetic experiment and our response to the reviewer TQCm to their question might be helpful.

Our current hypothesis for why our approach works is that gradients under prompts that contain semantic content relevant for learning serve to regularize the process of learning via fine-tuning. However testing this hypothesis directly is challenging since (a) different LLMs might interpret semantic information in a prompt differently (as a function of scale) and (b) it requires knowing which neurons are responsible for representing the inferred semantic information in the prompt -- an open problem in mechanistic interpretability.

The primary objective of the synthetic experiment was to analyze how contextual prompts affect the gradients of transformer models during training in a controlled setting where we can describe the semantic information that is necessary for learning explicitly via text. The advantage of this is that it enables us to not worry about how the transformer encodes semantic information (thus enabling the study of this phenomenon on much smaller models) and consequently better understand what properties of the gradient enable this.

To expand on this further, the sequence of tokens we use in the synthetic data, by design, $(x_1,f(x_1),x_2,f(x_2),\ldots,x_k,f(x_k))$ encode the semantic information necessary for learning this synthetic class of problem, facilitated by conditioning on the prompts.

Weakness 3: Including computational cost

It might be beneficial to include computational cost etc. for efficiency evaluation.

We have added the following paragraph in Appendix D.2 in blue, which will be included in the revised version of the paper.

To assess the efficiency of CFT, we carefully measured the computational resources required for our experiments and compared the overhead introduced by incorporating contextual prompts. Below are the details of our computational setup and findings. We utilized the Fully Sharded Data Parallel (FSDP) training to efficiently distribute the model across multiple GPUs. Training was performed using the bf16 (Brain Floating Point) data format. We implemented Flash Attention 2. All training was conducted with 8 NVIDIA A100 GPUs. With the above configuration, we achieved a training speed of approximately 55,188 tokens per second, measured using the Llama tokenizer. The fine-tuning required a total of approximately 111.11 GPU-hours to complete. Incorporating contextual prompts increased the total training time by approximately 0.89 GPU-hours, resulting in a total of 112 GPU-hours. Each contextual prompt added only about 0.8% to the length of each training example on average. This slight increase in input length led to less than a 1% increase in total training time.

Weakness 4: More information about OpenMedText

While OpenMedText is a comprehensive dataset proposed in a research paper (not dataset/benchmark papers), more information regarding statistics, potential biases, quality issues etc. in the dataset are not thoroughly discussed.

In addition to the total number of tokens, we added a detailed breakdown of the number of journals in each category.

We acknowledge that OpenMedText may have inherent limitations and potential biases, which we have added the following paragraph in Appendix C.4, which will be included in the revised version of the paper.

For the textbook data, since the textbooks were originally in PDF format, we used an Optical Character Recognition (OCR) API to extract the text. Despite careful processing, OCR can introduce typos or parsing errors, especially with complex formatting or specialized terminology. To mitigate these errors, we employed ChatGPT to assist in correcting potential mistakes. While this approach improved the overall quality, some errors may persist. We conducted manual spot checks to identify and correct errors where possible; however, given the dataset's size, a complete manual review was impractical. Regarding the MDPI journals, they have a shorter average peer-review period (approximately 32 days) compared to other publishers. While this expedites the dissemination of research, it may affect the depth and rigor of the review process. The shorter review time could lead to variations in article quality, with some papers potentially not meeting the highest standards of scientific rigor. Additionally, relying primarily on MDPI journals may introduce a source bias. We acknowledge that including journals from a wider range of publishers could enhance the dataset's balance and representativeness.

Thank you very much for your suggestion to provide more details.

References:

[2] Daixuan Cheng, Shaohan Huang, and Furu Wei. Adapting large language models via reading comprehension. arXiv:2309.09530, 2023.

Weakness 2: Limited exploration of prompt generation

While the authors provide thoughtful prompts based on educational theories in Appendix B1, it seems to be very limited exploration of prompt optimization or automated prompt generation methods, as the prompt template seems very various and task-specific.

We appreciate your recognition of the thoughtful prompts based on educational theories presented in Appendix B1 and we understand your concern about the limited exploration of prompt optimization and the task-specific nature of our prompt templates.

Our primary objective in this paper was to explore the capability of incorporating contextual prompts during the training phase of language models, specifically through Contextual Fine-Tuning (CFT) rather than identifying the optimal set of prompts that, with CFT, would yield the highest improvements.

Posing the identification of optimal prompts as an optimization problem is difficult. Prompts influence the entire learning trajectory of the model, affecting model weights and internal representations over many training rendering gradient-based and few-shot learning based methods for prompt optimization computationally infeasible.

Please also see Table 1 in the general comments. We attempted an initial exploration of automated prompt generation to address your concerns. The results demonstrate that contextual fine-tuning with automatically generated prompts outperforms continued pre-training and is on par with our original contextual fine-tuning using hand-crafted prompts. This suggests that automatic prompt generation, even without explicit optimization, can produce effective contextual prompts. It also indicates that the content of the contextual prompts does not need to be unique or manually designed to enhance the model's performance.

Our findings indicate that such methods can be a viable alternative to manually crafting prompts, potentially simplifying the fine-tuning process and making it more accessible. We believe this contributes to the broader understanding of how automated approaches can be employed in prompt design and optimization. We will include these new results and discussions in the revised version of our paper to provide a more comprehensive exploration of prompt generation methods. We are grateful for your feedback, which has helped us enhance our work and consider new avenues for research.

Thank you for your valuable feedback. We are happy to address your comments.

Weakness 1: Baseline comparison

For the experimental settings, the main experiments focus on comparison between CFT and CPT. To demonstrate the effectiveness, shall the comparisons also include other ICL methods, also RAG-based methods?

We agree that including evaluation of other In-Context Learning (ICL) methods and Retrieval-Augmented Generation (RAG)-based models would provide additional insights. However, ICL is only viable for models with very large context length and RAG based methods keep track of the entire content at inference time. These methodological differences make apples-apples direct comparisons challenging. Furthermore, both of these methods could also be used on a model that has been fine-tuned with CFT and so we see these methods as complementary and are looking to explore their use with CFT trained models in future work.

To expand our set of baselines however, we have included AdaptLLM [2] as an additional baseline to provide a more comprehensive evaluation. AdaptLLM is better aligned methodologically with our approach, as it enhances domain-specific knowledge during training by converting specialized corpora into a reading comprehension format for fine-tuning. We evaluated our method against AdaptLLM on several medical benchmarks. The results are presented in the table below:

Table. Comparison against AdaptLLM

We conclude that:

1. Our CFT method consistently outperforms AdaptLLM across all tasks.
2. AdaptLLM did not perform as well as anticipated on our dataset. We speculate that this may be due to differences in the datasets used. The original AdaptLLM paper utilizes PubMed abstracts to create reading comprehension tasks. Abstracts typically provide concise summaries of articles, making it easier to generate meaningful question-answer pairs. In contrast, our dataset consists of full-text articles from MDPI journals and textbooks, where not every paragraph contains information that readily lends itself to question-answer generation. This may limit the effectiveness of AdaptLLM's methodology when applied to our dataset.
3. As AdaptLLM is the most recent work closely related to our approach, these results suggest that our CFT method provides superior performance in domain-specific fine-tuning.

We hope that this additional comparison provides further confidence in the value of our method relative to existing work.

Weakness: Design of the contextual prompts

I also find that the design of the contextual FT prompts used in the real scenario is insufficiently justified. Currently, these prompts are associated with various educational theories. However, to me this appears to be insufficient justification because it is unclear -- and unlikely-- that any parallels can be drawn between the human learning process and the way that large language models use facts.

Furthermore, if I understand the paper correctly, it seems that there is no actual task supervision corresponding to these contextual prompts (i.e. for the critical analysis prompt: "Critically analyze the upcoming information. Look for underlying assumptions, evaluate arguments, and consider different perspectives.”, there is no actual ground-truth supervision given to the model on what the underlying assumptions/arguments in the provided information are). As a result, the mechanisms behind how these contextual prompts actually improve performance are quite unclear, and as I mentioned previously the synthetic data setup is not convincing in its relation to the real data setup.

Thank you for your insightful comments. They indicate we could have been clearer in our manuscript about the hypothesized mechanisms behind how the prompts change the process of learning. There are two core questions you ask, why pick these prompts, and why do these prompts work?

re: your first question -- Our current hypothesis is that there exists prompts whose semantic information can regularize the gradients during learning and improve generalization. This leads to a natural question -- how do you pick these prompts? In general, we do not anticipate the existence of a unique answer to this question and so we chose prompts that were compact yet general purpose. We settled on a variety of prompts summarized from the educational theories because they satisfied our desiderata of being both compact, yet general purpose. We have also experimented with an extension of our approach wherein the prompts are generated by another LLM. Table 1 in the general comments indicate that such prompts can also enhance model performance, outperforming standard continued pretraining.

What are the implications of this choice and does it imply that LLMs learn using mechanisms similar to humans? We do not believe we have provided sufficient evidence to reach this conclusion.

re: your second question -- do these prompts work? In a nutshell, having a comprehensive understanding of why this phenomena occurs might require a deep dive into the pretraining data, which unfortunately is rarely made available for open source models.

However, we conjecture that the rationale for why our method works has to do with explanatory text that exists in the training corpora.

To ground this conjecture we point the reader to experimental evidence for a different phenomena, chain of thought prompting, A recent work Why Think Step by Step? Reasoning Emerges from the Locality of Experience (Prystawski et al., 2023) [1] suggests that chain of thought prompting works because there are local steps embedded in pretraining corpora that simulate at training time the internal thought process that we have come to expect at test time from CoT prompting. While our goal in this work is to demonstrate the value from CFT, we believe a similar, future study along this lines studying why CFT works, as a function of the pretraining data would be valuable to the community to test this conjecture and shed light on the mechanisms behind why prompts provide useful supervisory signal during learning.

Thank you very much for your detailed review and insightful comments. We would like to address your concerns.

Weakness: Synthetic setting and Question 2

In the synthetic setting, the added context tokens actually depend on the specific inputs and functions used in the example. On the other hand, in the real-world setting, the additional context tokens are sourced from the educational prompts and randomly sampled without considering he contents of a particular example. To summarize, it appears to me that contextual fine-tuning in the synthetic experiments actually provides a significant source of additional supervision (reminiscent of COT approaches) whereas this is absent in real-world instantiation of contextual FT.

I would like it if the authors could further justify how the synthetic data setup should be viewed as comparable to the real-world setup. In particular, why are the additional tokens added in the synthetic settings input dependent while the contextual prompts used in real settings are randomly sampled independently of the contents of the document/example?

Our current hypothesis for why our approach works is that gradients under prompts that contain semantic content relevant for learning serve to regularize the process of learning via fine-tuning. However testing this hypothesis directly is challenging since (a) different LLMs might interpret semantic information in a prompt differently (as a function of scale) and (b) it requires knowing which neurons are responsible for representing the inferred semantic information in the prompt -- an open problem in mechanistic interpretability.

To that end the primary objective of the synthetic experiment was to analyze how contextual prompts affect the gradients of transformer models during training in a controlled setting where we can describe the semantic information that is necessary for learning explicitly via text. The advantage of this is that it enables us to not worry about how the transformer encodes semantic information (thus enabling the study of this phenomenon on much smaller models) and consequently better understand what properties of the gradient enable this. To expand on this further, the sequence of tokens we use in the synthetic data, by design, $(x_1,f(x_1),x_2,f(x_2),\ldots,x_k,f(x_k))$ encode the semantic information necessary for learning this synthetic class of problem, facilitated by conditioning on the prompts.

Our empirical results, presented in Appendix G, show that contextual fine-tuning is more effective for instruction-tuned and chat models compared to non-chat models. This observation suggests that models capable of following instructions are better at leveraging contextual prompts during fine-tuning, even when the prompts are not customized to each example. Our intention with the synthetic experiment was to provide insight into the potential mechanisms by which contextual prompts can enhance learning, acknowledging that direct analysis of gradients in large-scale language models is infeasible.

Question 1: Additional ablations on contextual prompts

I would like to see some additional ablations about the contents of the context prompts that go beyond the negative CFT setup. In particular, it would be nice if authors considered prompts consisting of various writing styles (but which are not correlated to educational theories or contradiction of the documents). This could help relate this work to prior research on the personas hypothesis etc.

Thank you for the suggestion.

To your suggestion of ablations, we have created a variant of our method where instead of CFT with custom education inspired prompts, we experiment with an automated prompt generation system that uses an auxiliary LLM to create prompts for each sampled paragraph. Please see Table 1 in the general comments.

Our findings indicate that models fine-tuned with these automatically generated prompts denoted as AutoDep-CFT show performance improvements over the baseline models without contextual prompts and are comparable to those fine-tuned with our original contextual prompts. This suggests that incorporating prompts with varied writing styles—even those generated automatically without specific alignment to educational theories—can enhance the model's performance.

We are running a few more experiments that we hope can better address this question and will report back soon.

References:

[1] Ben Prystawski and Noah D. Goodman. 2023. Why think step-by-step? Reasoning emerges from the locality of experience. CoRR abs/2304.03843 (2023)

Thank you to the reviewer for their time and feedback. Please see our response below:

The reviewer recommends evaluating the impact of domain-specific fine-tuning on the model's general and instruction-following capabilities. We evaluate the OpenMedText fine-tuned Llama-2-13B model on general benchmarks and instruction-following benchmarks, and we find that the capabilities of the base model are largely retained. We present results from MMLU, MMLU-Pro, and IFEval which provide coverage of these capabilities.

Table. Llama-2-13B (Accuracy)

We do not observe catastrophic forgetting as a result of fine-tuning, the general knowledge of the model is only slightly diminished. CPT is slightly more robust than CFT to knowledge degradation, however, the performance difference is small and we emphasize CFT's stronger in-domain performance as demonstrated by Table 1 in the manuscript.

Table. Medical Benchmarks (Manuscript)

To conclude, we find that CFT up-holds a model's general capabilities while providing significant boosts for in-domain performance.

Thank you for your insightful suggestion to evaluate our method on general benchmarks.

Results:

Table 1. Evaluation of CFT with auto-generated contextual prompts that are dependent on the upcoming text on medical benchmarks

We will refer to this table in more detail in our individual responses to each reviewer.

The results from Table 1 indicate that AutoDep-CFT, which uses automatically generated, content-dependent prompts, achieves an average accuracy of 46.31%, outperforming both the baseline Chat model (45.05%) and the Chat CPT model (45.96%). While it does not surpass the Chat (CFT) model with manually designed prompts (47.81%), these findings indicate that auto-generated, context-dependent prompts can effectively enhance model performance across medical benchmarks.

We will incorporate the table and update the manuscript. Thank you again for your thoughtful feedback.