TAIA: Large Language Models are Out-of-Distribution Data Learners

Summary

The paper introduces TAIA, a novel method for enhancing the performance of large language models (LLMs) in data-scarce domains with domain-mismatched data. The authors identify that during fine-tuning, only attention parameters significantly contribute to downstream task performance when training and test sets' distributions do not align. The proposed TAIA method retains updates to attention parameters while discarding updates to feed-forward network parameters, effectively improving performance on out-of-distribution (OOD) tasks. Extensive experiments demonstrate TAIA's superior performance across various datasets and model configurations, showcasing its robustness and generalizability.

Strengths

- The TAIA method is novel to selectively retain beneficial parameter updates during fine-tuning. - The paper provides a thorough analysis of the roles of self-attention and feed-forward networks in LLMs. - The comprehensive experiments validate the proposed method.

Weaknesses

- The presentation could be clear. For example, figure 2 could be more clear with a better caption. Currently, the caption is not informative for me to understand the method and its effectiveness. In addition, Figure 3 also made me confused. It seems that the base model performance pretty well. If this is the case, why should we use other methods?

Questions

See weakness

Rating

6

Confidence

3

Soundness

3

Presentation

2

Contribution

3

Limitations

yes

Summary

This paper proposes the hypotheses that under some circumstances LLMs are better trained by ft all parameters and then dropping the learnt FFNN layers. The proposed method is named TAIA from training all parameters but inferring with only Attention. A comparison of TAIA's the performance together with another variants like TOA (train only attention) is carried out. More over other ablated version using the FF components instead of the attention were also evaluated. Several model are considered during the study such as Llama2 LLama3 and Qwen. Several datasets (Alpaca , CoT-collection, ...) are used to train the models with the proposed technique and baselines. Extensive comparison are also done with other OOD techniques such as L2, EC and self-distill, among others. Several dimensions are taken into account such as varying data sizes. The results are positive, showing the new proposed techniques improves performance when models are trained OOD. Finally an honest limitation section with some experiments show conditions where the proposed method will not work better

Strengths

The paper performs several experiments varying many factors to validate there method. Several LLMS were used in the experiments, which helps into drawing generic conclusions. Many datasets are considered for evaluating 2 aspects : reasoning and knowledge. Ablation studies are carried out with other baselines varying TAIA is proposed. Models are compared with other OOD techniques such as Self-distillation, L2, or WEC. Representation analysis is carried out as well as quantitative analysis. The discussions are very interesting with clear research questions for which corresponding experiments were executed. Limitations include experiments to highlight some of the circumstances when the proposed technique doesn't work well. Finally, the extensive appendices specify many details on experiment, additional results and detailed explanations that help the reading of the paper.

Weaknesses

It would have nice to see what happens with even larger scale models. In some of the experiments additional rounds or repetitions would have helped to understand the variability of the results. Some deeper theoretical analysis could help to explain the unexpected and satisfactory results of the proposed technique.

Questions

It is not very clear how the model is trained and used when we have in-domain and out-of-domain data. It looks like both are merged together loosing the differentiation. In such a case, do you think differentiation of the different sources would further improve the performance? maybe some additional experiments would be interesting. How would the FF updates be done in such a case in order to embed any missing in-domain knowledge into the model ?

Rating

7

Confidence

3

Soundness

4

Presentation

3

Contribution

3

Limitations

Authors do an excellent work addressing and studying their proposal limitations.

Summary

This paper proposes a simple but effective strategy to improve fine-tuning of LLMs: fine-tune as normal but at inference time, only use the fine-tuned attention parameters but use the pre-trained MLP parameters. The intuition for this is that knowledge is generally stored in the MLP layers, and fine-tuning introduces harmful modifications to the pre-trained knowledge.

Strengths

The results are strong, showing that TAIA improves over both vanilla fine-tuning and the baseline of only training attention, across a number of test datasets and models. It also improves over continual learning strategies, which have the same intuition of wanting to preserve aspects of the pre-trained model. There is extensive analysis showing, among other results, that TAIA works both for LoRA and full fine-tuning, it is also better against red-teaming,

Weaknesses

My main concern is from Figure 4a, which seems to show that the gap between Vanilla and TAIA decreases as the training dataset gets larger. This also made me realize the paper does not discuss how much training data is used in the main experiments. Figure 4 also does not disclose which model is being evaluated. I think it is important to address exactly how the results in Figure 4a compare with the CoT-Collection results in Table 1, where TAIA seems to be much better than Vanilla. Relatedly, I would suggest discussing more about dataset size in Section 4.2.

Questions

While I understand the intuition for TAIA, it also seems that it would introduce some distribution shift, as the fine-tuned attention parameters were trained to operate with the fine-tuned MLP parameters, but those are not present at inference time. Why does this not cause the model to perform badly? Since you are doing fine-tuning, I wonder about the choice of using the chat models rather than base (no instruction tuning) models. Why use the chat models? Could you also do experiments on the base models? Figure 3: What dataset is this accuracy reported on? In Section 4.6, I wonder if you can also measure similarity to the pre-trained model. This would measure whether TAIA is essentially regularizing the model towards the pre-trained model (similar to what L2 or EWC would do, but apparently better). I wonder if Section 5 should be called "Analysis" rather than "Discussion", as I usually think of Discussion as not introducing more results. Table 4: Why is there a "-" in Infer Mode for the first row? I'm not sure what the first row represents. Other comments: * Line 22: Sentence should be rewritten as "which" currently attaches to NLP and not LLM

Rating

7

Confidence

4

Soundness

3

Presentation

3

Contribution

3

Limitations

Limitations are appropriate.

Summary

The paper proposes a novel inference-time intervention method that trains all model parameters but retains only the self-attention updates for inference. This approach, named "Training All parameters but Inferring with only Attention" (TAIA), optimizes performance across a variety of downstream and closed-book tasks, as validated by extensive experiments. The paper confirms the reproducibility of TAIA in terms of fine-tuning methods and dataset scales. TAIA also maintains the few-shot adaptation ability of base models and withstands multi-level adversarial attacks, consistently improving performance in vertical domains like healthcare even with increasing OOD data.

Strengths

1. This paper introduced an innovative inference-time intervention strategy termed "Training All parameters but Inferring with only Attention" (TAIA), which trains all model parameters while retaining solely the refined self-attention updates for inference. 2. Through extensive experimental validation on a variety of datasets and with LLMs of varying parameter sizes, the authors showcased that TAIA consistently outperforms both fully fine-tuned models and base models under most conditions, achieving substantial performance enhancement. 3. The authors substantiate the robustness and generalization capabilities of TAIA across different fine-tuning methodologies and dataset scales. TAIA also preserves the base models' few-shot adaptation prowess and defends against multi-tiered adversarial assaults, markedly elevating performance in niche domains such as healthcare, even amidst escalating OOD data.

Weaknesses

1. The authors posit that the parameters of Feed-Forward Networks (FFNs) encapsulate a substantial amount of pre-trained knowledge. However, these parameters are prone to shifts during Supervised Fine-Tuning (SFT), which may result in suboptimal performance on Out-Of-Distribution (OOD) data. But there is a notable absence of empirical evidence to substantiate this claim and merely citing other studies is insufficient. 2. The authors train the FFN parameters during the training process but discards them during inference. Could this lead to an inconsistency between the training and inference? 3. The authors lack implementation details on how to apply TAIA to the full finetune in response to RQ1. 4. The author should include a comparative experiment, utilizing only the parameters of the FFN while keeping the attention parameters fixed during the inference process. This will highlight the validity of TAIA. 5. The method section primarily contains text without any mathematical formulation, making it challenging to follow. Including formulations on how the model learns from different types of data and how it only infers with attention during inference would improve clarity and make it easier to evaluate the method's contribution. 6. Adding some theoretical proof to support your claims would strengthen the paper. 7. The overall results indicate a serious overfitting issue, as the base models have shown competitive results, such as 53.76 (base) vs. 53.24 (TAIA). Could you provide more insights into this observation?

Questions

Refer to the above.

Rating

4

Confidence

4

Soundness

2

Presentation

2

Contribution

2

Limitations

Yes

Thanks for valuable reviews

We hope this message finds you well. First and foremost, we would like to express our sincere gratitude for your efforts in overseeing the review process. We deeply appreciate the time and attention that you have devoted to evaluating my submission. We are writing to kindly inquire if it would be possible to engage in the discussion. We have responded to your questions in our rebuttal and we believe that a collaborative discussion with you could be highly beneficial in ensuring a balanced and fair assessment of our work. We understand the demands on your time, and we truly appreciate your consideration of this request. We are hoping to receive your thoughtful feedback. Thank you once again for your continued support and valuable reviews about our work.