Table of Contents
Fetching ...

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

Zhichao Wang, Bin Bi, Zixu Zhu, Xiangbo Mao, Jun Wang, Shiyu Wang

TL;DR

This paper addresses the problem of catastrophic forgetting and alignment tax that arise when neural LLMs are fine-tuned in separate SFT and alignment stages. It introduces Unified Fine-Tuning (UFT), a framework that unifies SFT and alignment under a generalized implicit reward, treating instruction-tuning data as alignment data and optimizing with a shared objective to maximize rewarded quality while staying close to pretrained behavior. Empirically, UFT outperforms pure SFT on instruction-tuning data and, when mixing instruction-tuning with alignment data, surpasses sequential SFT+alignment approaches on multiple tasks and leaderboards, with notable gains in instruction-following (ifeval) and factuality (truthful). The work demonstrates that a unified, data-distribution-aware fine-tuning approach can achieve strong performance with reduced forgetting, offering a practical parallel to pretraining and guiding future deployments in diverse linguistic and real-world settings.

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.

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

TL;DR

This paper addresses the problem of catastrophic forgetting and alignment tax that arise when neural LLMs are fine-tuned in separate SFT and alignment stages. It introduces Unified Fine-Tuning (UFT), a framework that unifies SFT and alignment under a generalized implicit reward, treating instruction-tuning data as alignment data and optimizing with a shared objective to maximize rewarded quality while staying close to pretrained behavior. Empirically, UFT outperforms pure SFT on instruction-tuning data and, when mixing instruction-tuning with alignment data, surpasses sequential SFT+alignment approaches on multiple tasks and leaderboards, with notable gains in instruction-following (ifeval) and factuality (truthful). The work demonstrates that a unified, data-distribution-aware fine-tuning approach can achieve strong performance with reduced forgetting, offering a practical parallel to pretraining and guiding future deployments in diverse linguistic and real-world settings.

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.

Paper Structure

This paper contains 15 sections, 7 equations, 2 figures, 13 tables.

Figures (2)

  • Figure 1: UFT integrates SFT and alignment through a generalized implicit reward function. It likens pre-training and fine-tuning of LLMs to Chinese proveb "Read ten thousand books, travel ten thousand miles". In pre-training, the LLM processes vast amounts of data without feedback, gaining broad language understanding. In fine-tuning, it generates responses to prompts and receives feedback, refining its abilities and improving performance on specific tasks.
  • Figure 2: Subfigure (A) refers to SFT, Subfigure (B) refers to alignment including RLHF, DPO and UNA and Subfigure (C) refers to UFT. Traditionally, the fine-tuning process begins with SFT followed by alignment. However, the proposed UFT method integrates both SFT and alignment into a single, cohesive process.