UFT: Unifying Fine-Tuning of SFT and RLHF/DPO/UNA through a Generalized Implicit Reward Function (2410.21438v1)

Abstract: By pretraining on trillions of tokens, an LLM gains the capability of text generation. However, to enhance its utility and reduce potential harm, SFT and alignment are applied sequentially to the pretrained model. Due to the differing nature and objective functions of SFT and alignment, catastrophic forgetting has become a significant issue. To address this, we introduce Unified Fine-Tuning (UFT), which integrates SFT and alignment into a single training stage using the same objective and loss functions through an implicit reward function. Our experimental results demonstrate that UFT outperforms SFT on instruction-tuning data alone. Moreover, when combining instruction-tuning data with alignment data, UFT effectively prevents catastrophic forgetting across these two stages and shows a clear advantage over sequentially applying SFT and alignment. This is evident in the significant improvements observed in the \textbf{ifeval} task for instruction-following and the \textbf{truthful-qa} task for factuality. The proposed general fine-tuning framework UFT establishes an effective and efficient pretraining-UFT paradigm for LLM training.

Overview of UFT: Unifying Fine-Tuning of SFT and RLHF/DPO/UNA through a Generalized Implicit Reward Function

The paper presents a novel approach for fine-tuning LLMs with the introduction of Unified Fine-Tuning (UFT). This method aims to mitigate the prevalent issue of catastrophic forgetting that occurs when supervised fine-tuning (SFT) and various alignment techniques such as Reinforcement Learning from Human Feedback (RLHF), Direct Preference Optimization (DPO), and Unified Alignment (UNA) are sequentially applied. The authors propose a unified methodology that integrates these processes into a single stage using a generalized implicit reward function, demonstrating superiority in performance across a variety of tasks.

Methodology

UFT combines SFT and alignment into a single training framework, utilizing a shared objective and loss function through an implicit reward model. This contrasts with traditional approaches where SFT and alignment are distinct stages, leading to the phenomenon known as catastrophic forgetting. The authors leverage UNA's ability to process various types of feedback, including pairwise, binary, and score-based feedback. They extend these capabilities to encompass the goals of SFT, thereby ensuring compatibility and synergy between instruction-tuning data and alignment data.

The paper presents a mathematical formulation whereby both SFT and UNA aim to maximize the likelihood of response generation in instruction-tuning data. Through experimentation, it is shown that UFT not only outperforms SFT on instruction-tuning datasets but also effectively prevents catastrophic forgetting when both instruction and alignment data are considered.

Experimental Results

The experimental evaluation demonstrates that UFT consistently surpasses traditional SFT across several tasks, particularly in instruction-following (ifeval) and factuality (truthful-qa) assessments. This is attributed to UFT's dual focus on maximizing reward scores while minimizing divergence from the pretrained model. The results illustrate UFT's effectiveness in maintaining the alignment and instructional capabilities of LLMs, contrasting with the performance declines observed in sequential training paradigms.

When analyzing the impact of instruction-tuning and alignment data distributions, UFT shows optimized performance for both types of data, underscoring the necessity of a balanced dataset for enhancing LLM capabilities. Furthermore, the integration with UNA ensures that UFT can handle the complexities associated with various feedback types, establishing a robust framework for future applications.

Implications and Future Developments

The implications of this research are significant, particularly in the field of natural language processing, due to the potential improvements in both the generation capabilities and ethical alignment of LLMs. By harmonizing SFT and alignment processes, UFT promises to enhance the efficiency and effectiveness of LLM fine-tuning, potentially influencing future methodologies in AI alignment and LLM training.

Future developments may involve exploring the integration of additional feedback mechanisms and optimizing the balance between instruction-tuning and alignment data. Further studies could also investigate UFT's adaptability to different LLM architectures and its applicability across various ethical and instructional contexts in AI.

In conclusion, the proposed UFT methodology signifies a considerable advancement in the fine-tuning landscape for LLMs, addressing key challenges through an innovative, unified approach. The paper contributes a foundational framework that promises to refine and enhance LLM training paradigms, fostering advancements in AI alignment and LLM utility.