Reference Trustable Decoding: A Training-Free Augmentation Paradigm for Large Language Models

We sincerely appreciate your willingness to thoroughly read our paper. The time you’ve invested in our work is truly invaluable to us. Your recognition of our method’s clear motivation, extensive experiments, and ease of optimization is an honor. Additionally, we thank you for pointing out the shortcomings in our paper. We highly value your feedback, and these suggestions will undoubtedly contribute to refining our work.

Weakness

1. How RTD behave when auto-regressive decoding: Regarding RTD, it follows the traditional autoregressive decoding pattern, where the generation of subsequent tokens is handled in the same manner as the first token. Specifically, after generating a token, it is appended to the input sequence, and the new last hidden states are used to perform another round of RTD for the next token. We will provide more detailed explanations in future versions of the paper to clarify these points.

2. What is the benchmark based on RAG and Open Book Question Answering: For benchmark based on RAG and Open Book Question Answering, we add two components to the original QA setup to test whether RTD could inject new knowledge into LLMs:

   1. Source of New Knowledge: We chose Wikipedia articles, and the filter the relevant pages out using a RAG system.
   2. Process of Knowledge Injection: Since RTD operates at the token level, and Wikipedia articles do not directly contain options like A, B, C, and D, we need a process where the model generates an analysis of the question. During this generation process, RTD injects new knowledge. So we adopted a CoT decoding approach for this purpose.

   We will provide more detailed descriptions in the revised version.

3. The unstable performance of RTD in some dataset: In the NLU experiments, whether it’s fine-tuning, ICL, or our RTD approach, the core goal is not merely injecting more knowledge into the model but rather guiding the model to better utilize its existing knowledge. Essentially, any optimization method is constrained by the upper limit of the model’s knowledge capacity. Breaking through this limit is not feasible given the current computational budget. Even if these methods are orthogonal, when the model has already touched or approaching this upper limit, further improvement in capability becomes restricted. In some cases, fluctuations due to errors may occur. Our primary experiments have demonstrated two key points:

   1. RTD alone significantly enhances model performance
   2. RTD has the potential to complement methods like ICL However, when ICL is already sufficient to reach the theoretical upper limit, combining it with RTD may unavoidably introduce performance variations due to error factors.

   We will add more detailed analysis in the next version.

4. Detailed description to baseline: Thank you for your suggestion. We will provide a more detailed discussion in future version. For most benchmarks, specific example data is already prepared for constructing ICL. For instance, in the case of the MMLU benchmark, the dev subset provides five sample questions for each question category. This is why we chose to use a 5-shot setup. For benchmarks that do not have corresponding sections explicitly designed for ICL, we randomly sample five questions from the provided training data to create a 5-shot setup. Both the test group and the control group use the same set of questions to ensure experimental fairness.

5. Discussion about the Last Hidden Space: For last hidden space, our modification target lies within the LM_Head, and the Last Hidden State serves as the input to the LM_Head. The LM_Head represents the final layer of the entire LLM network, which means that our method essentially operates on a relatively shallow portion of a deep neural network. This design choice provides us with an opportunity to employ richer methods compared to a single gradient descent approach. Additionally, it allows us to maintain a “training free” approach. In fact, the core of our paper lies in exploring the optimization of the model using the Last Hidden State. Our ultimate goal is not only to bypass the LM_Head but also to delve further into the properties of this specific region. Notably, in some newer models like llama3-8b, the LM_head constitutes a substantial portion of the model parameters (6.6% or 0.53B out of 8.03B), yet it has received relatively little attention.

Finally, we would like to express our gratitude for the time you’ve spent reviewing our paper again. Your suggestions have been highly relevant to the content and have significantly contributed to improving the paper. If you have any further questions or additional feedback, please feel free to reach out, and we will promptly respond.

Dear reviwer tphZ,

We sincerely appreciate your timly reply.

We propose RTD, an approach that provide a new balance between performance and efficiency. We conducted supplementary tests on the OBQA benchmark, with LLaMA-3-8B, and $20,000$ RTD entries, comparing different methods in terms of generation speed and memory footprint. Notably, RTD exhibits significant advantages over ICL and performs on par with the baseline in the aspact of generation speed. Analyzing memory usage, RTD’s additional memory footprint is only $\frac{1}{2}$ of ICL’s, resulting in reduced bandwidth bottlenecks and computational load. For other datasets, especially those with longer contexts, RTD’s efficiency advantage becomes even more pronounced, as ICL’s extra overhead depends on context length, whereas RTD remains context-agnostic.

Considering both performance and efficiency analysis, we can conclude that RTD offers a satisfactory solution for various task types due to its excellent cost-effectiveness, which makes it always worth trying. If RTD alone doesn’t fully meet your performance requirements and you’re willing to sacrifice some efficiency, consider using the ICL method to assess performance. Alternatively, combining ICL with RTD may yield a further performance boost without significantly compromising efficiency in most tasks. Finally, if you have abundant computational resources and performance remains unsatisfactory, fine-tuning could be explored. A prefered workflow with RTD can be look like:

prompt tuning -> RTD -> ICL -> ICL+RTD -> finetune

From left to right both the cost and expected performance increase. And this workflow demonstrates how RTD seek a new balance.

In our upcoming revisions, we will include more comprehensive efficiency tests and further analyze the trade-off between efficiency and performance of variouse methods, emphasizing the high efficiency of our proposed method.

Best regards,

Authors

We sincerely appreciate the time you’ve spent on our paper. Thank you for recognizing our approach, experiments, and writing. The following is our response regarding the weaknesses, questions, and limitations, you raised. We aim to clarify any potential misunderstandings and indicate the direction of modifications for our paper.

Weakness

1. Comparison with other methods: In the second section of our paper, we discussed the comparison between our approach and other methods in terms of efficiency and capability. However, we plan to supplement a small table to provide a concise and quick overview. The table would look like this:

   Methods	Training Cost	Inference Cost	Capability
   No Method	Low	Low	Low
   RTD (Ours)	Minimum	Minimum	Moderate
   PEFT	High	Low	Moderate
   ICL	Low	High	Moderate
   FFT	Gigantic	Low	Best

   And we've also tested the inference efficiency of RTD and it's counterparts.

   Methods	Speed (it/s)
   RTD	$23.6$
   PEFT	$25.1$
   ICL	$7.9$

2. Real-world application: In our experiments, we utilized highly challenging and widely used datasets that are derived from real-world scenarios, such as MMLU which covers 57 subjects across STEM, the humanities, the social sciences, and tests both world knowledge and problem-solving ability. This makes the benchmark more challenging and more similar to how we evaluate humans. We will provide more detailed case studies in future versions of the paper.

3. Additional generative benchmarks: In section 4.2, we tested different LLM generation tasks, including reasoning with context and style transfer tasks. For example, we used a Retrieval-Augmented Generation (RAG) benchmark to evaluate knowledge injection and the ability to provide step-by-step reasoning using Chain-of-Thought prompting. Additionally, we conducted style transfer experiments using the Tiny-Shakespeare dataset to measure the model's ability to adapt to different language styles. In our revised version, we will include additional examples of RTD-generated instances to support the effectiveness of the RTD method.

4. Solutions of large datastores: While we acknowledge the importance of compressing large reference datastores, this particular challenge falls outside the primary scope of our current research. Our main focus was on introducing the reference trustable framework (RTD), which aims to better integrate and utilize reference datastores. Besides, in Sections 3.4 and 4.3 of our paper, we discussed a resource-saving approach by reducing the number of utilized heads, which can partially solve this. We see this limitation as an opportunity for future research rather than a weakness of our current work. We believe RTD provides a solid foundation for further exploration into effective compression techniques. We appreciate your insight and will consider discussing potential approaches in future works.

Questions

1. Larger and more complex tasks: We would like to emphasize that the dataset used in our experiments in Section 4, such as MMLU, is already a large-scale and challenging benchmark commonly used for testing. Its complexity and difficulty level are significant.

2. Integrated into existing systems: This aspect was indeed underrepresented in our paper, and we agree that it deserves further elaboration. For instance, when using LLM for text screening, a small set of screening cases can be used to generate a Reference Datastore. By combining this with RTD, we can replace the original ICL portion, thereby increasing screening efficiency and reducing overall computational load. Additionally, analyzing screening results can help reverse-engineer potential issues in the Reference Datastore, highlighting the “Trustable” feature of RTD.

3. Combine methods In Table 2, we have conducted experiments combining RTD with other methods, demonstrating that generally, the more methods we combine, the better the overall performance tends to be. However, such improvements often come with additional computational resource requirements, necessitating a trade-off between performance and cost. The best integration strategy needs to be determined based on the specific use case. The core contribution of RTD lies in providing a method to augment model capabilities without imposing excessive additional computation, thus opening up new possibilities for model enhancement under limited resources.

Limitations

1. Large datastores: See weakness #4

2. Generative testings: See weakness #3

3. Settings generalization: RTD has been tested on two types of tasks: language understanding and language generation, across a total of seven different benchmarks and six models. We believe this extensive evaluation effectively demonstrates RTD's ability to generalize to various domains and tasks.

4. Responsible AI development: We propose RTD as a model enhancement method, which will have similar societal effects as other existing methods and won’t introduce any additional issues for discussion. We will emphasize this point more clearly in future versions of the paper.

Dear Reviewer,

We sincerely appreciate your thorough review and valuable feedback. We have taken your comments seriously and have already conducted additional experiments and provided further explanatory details based on your recommendations.

We've further measured our RTD methods on efficiency, both time and memory comsumption wise, here is our latest result.

Considering both performance and efficiency analysis, we can conclude that RTD offers a satisfactory solution for various task types due to its excellent cost-effectiveness, which makes it always worth trying. A prefered workflow with RTD can be like:

prompt tuning -> RTD -> ICL -> ICL+RTD -> finetune

To ensure that we address the issues you raised effectively, we kindly request your prompt response. Your feedback is crucial for enhancing our work. We understand that the review process is demanding, and we greatly appreciate your time and effort.

Thank you once a gain for your support and help.

Best regards,

Authors

We sincerely appreciate your thorough reading of our paper. The time you’ve invested in our work is truly invaluable to us. Your endorsement of the efficiency, orthogonality, and practical significance of our proposed RTD method means a great deal, and we are genuinely grateful for your recognition.

Additionally, we thank you for pointing out the limitations and constraints of RTD. Such constructive feedback is immensely helpful in improving our paper.

Weakness

1. Lack of comparison with other retrieval-based decoding methods: The reason we didn’t compare with retrieval methods lies in two factors. First, other retrieval-based methods typically rely on an gigantic reference datastore. For instance, in the original paper of $k$NN-LM, over 1 billion entries were used, which significantly exceeds the data volume typically employed in our RTD method (usually less than 1 million). This disparity makes direct comparisons a lower reference value, especially given RTD’s focus on efficiency. Second, running a complete $k$NN-LM requires substantial hardware resources, often involving massive CPU memory and potentially distributed cluster deployments. Unfortunately, we lack the necessary resources to fully replicate knn-LM’s performance.

2. Lack of generalization analysis: Thank you for your suggestion. We want to emphasize that RTD’s core objective is to efficiently augment LLMs within specific domains. However, exploring cross-domain generalization is indeed a promising avenue. We plan to supplement our paper with cross-domain tests within the NLU field, given the structural similarities across different benchmarks. Preliminary exploration suggests that our RTD method exhibits some cross-domain generalization capability, particularly when task formats are similar. As soon as detailed testing is complete, we’ll promptly share the results. Our first batch of result are as follows, we tested the reference datastore generate from different benchmarks on OBQA.

3. Requirement on additional hyperparameter searches: Our method introduces a certain number of hyperparameters, and determining the optimal hyperparameters does indeed pose new challenges for RTD. However, in Section 4.3 of our paper, we extensively discuss the impact of hyperparameters on performance. Our experiments reveal that RTD is relatively insensitive to hyperparameters, and performance trends are quite evident as hyperparameters vary. Therefore, while identifying an optimal set of hyperparameters still requires some experimentation, finding a satisfactory set is relatively simple and straightforward.

4. Lacks association with trustable decoding: Thank you for your suggestion. Our proposed RTD method fundamentally targets the LM_Head within the entire LLM network. Since the LM_Head is the final layer of the LLM, our method essentially operates on a relatively shallow portion of a deep neural network. This design choice allows us to employ more interpretable methods compared to traditional gradient-based optimization. For example, if during RTD deployment, we observe unexpected model behavior or undesirable outputs, we can rectify these issues by inspecting the data within the Reference Datastore. The data in the Reference Datastore is traceable, meaning we can identify which specific example led to an error and easily remove it. In contrast, achieving a similar process with gradient-based methods, such as gradient descent, would be extremely challenging, if possible at all. We will supplement relevant content in the next version to address this issue and further clarify the association with "Trustable Decoding."

5. Typos: We are sorry that there are some typos in our paper, we will correct them immediately. Thank you for your careful reading.

Questions

• The explainaion of setting $\lambda$ to 1: In RTD, we introduce a hyperparameter $λ$ to blend the results from RTD and those from the base model. The rationale behind this choice is that we typically don’t want to completely discard information from the base model. By doing so, we avoid the risk of RTD inadvertently causing a regression in overall system performance, ensuring greater robustness. In the specific case you mentioned, which pertains to NLU tasks, the relatively constrained format allows RTD to perform well on the task itself while enhancing performance. In such scenarios, we believe that additional information from the base model is unnecessary. Furthermore, setting $λ$ to $1$ avoids costly vocabulary operations, conserving computational resources, especially crucial for large-vocabulary modern models like llama3-8b, where the LM_head constitutes 6.6% of the model parameters (0.53B out of 8.03B). We appreciate your perspective, and indeed, the setting of $λ$ still worth exploration for this task. We will conduct an additional experiment to present the influence of $λ$. We will provide more detailed discussion in the next version. Our result of tuning $\lambda$ on LLaMA3-8B, OBQA are as follows:

  $\lambda$	$1$	$0.8$	$0.6$	$0.4$	$0.2$	$0$
  score	$71.4$	$68.0$	$67.0$	$66.8$	$66.6$	$53.3$

Allow us to express our gratitude once again for your diligent reading of our paper. Your suggestions are valuable and can definitely help us to improve our paper. If you have any further insights or questions, please share them with us.

Dear Reviewer,

We sincerely appreciate your thorough review and valuable feedback. We have taken your comments seriously and have already conducted additional experiments and provided further explanatory details based on your recommendations.

We've further measured our RTD methods on efficiency, both time and memory comsumption wise, here is our latest result.

Considering both performance and efficiency analysis, we can conclude that RTD offers a satisfactory solution for various task types due to its excellent cost-effectiveness, which makes it always worth trying. A prefered workflow with RTD can be like:

prompt tuning -> RTD -> ICL -> ICL+RTD -> finetune

To ensure that we address the issues you raised effectively, we kindly request your prompt response. Your feedback is crucial for enhancing our work. We understand that the review process is demanding, and we greatly appreciate your time and effort.

Thank you once a gain for your support and help.

Best regards,

Authors