TS -Align: A Teacher-Student Collaborative Framework for Scalable Iterative Finetuning of Large Language Models (2024)

Abstract: Mainstream approaches to aligning large language models (LLMs) heavily rely on human preference data, particularly when models require periodic updates. The standard process for iterative alignment of LLMs involves collecting new human feedback for each update. However, the data collection process is costly and challenging to scale. To address this issue, we introduce the “TS-Align” framework, which fine-tunes a policy model using pairwise feedback data automatically mined from its outputs. This automatic mining process is efficiently accomplished through the collaboration between a large-scale teacher model and a small-scale student model. The policy fine-tuning process can be iteratively repeated using on-policy generations within our proposed teacher-student collaborative framework. Through extensive experiments, we demonstrate that our final aligned policy outperforms the base policy model with an average win rate of 69.7% across seven conversational or instruction-following datasets. Furthermore, we show that the ranking capability of the teacher is effectively distilled into the student through our pipeline, resulting in a small-scale yet effective reward model for policy model alignment.

## introduction general-purpose conversational ai assistants, such as gpt-4 @xcite and @xcite , are empowered by aligning large pretrained language models with humanpreferred behaviors @xcite @xcite . these aligned llms showcase exceptional capabilities in instruction following @xcite , natural conversation @xcite , safety @xcite , reasoning (wei et al., 2022b; kojima et al., * corresponding author 2022), among others. commonly-used llm alignment techniques include instruction tuning @xcite , reinforcement learning from human feedback (rlhf) @xcite , and direct preference optimization (dpo) @xcite . while recent research has focused significantly on the development of more sophisticated alignment techniques @xcite @xcite @xcite , it is worth noting that llm alignment is not a one-time process and the model requires continuous refinement to adapt to evolving user needs and changing linguistic patterns. the standard practice for iterative alignment of the llms is to gather new human preference data for every subsequent update to the model. for instance, @xcite performs five iterations of rlhf finetuning on the base sft llama-2 model. for each iteration, they update the reward model with newly collected human preference data. this process poses challenges regarding scalability and resource requirements. to alleviate the issue, existing research adopts self-evolution @xcite @xcite or external model supervision @xcite @xcite . the effectiveness of self-evolution is highly dependent on the quality of the base model as it operates without the introduction of external supervision or knowledge during refinement. for instance, in their study, @xcite utilize a sophisticated 70b llama-2 model to demonstrate the potential of their iterative self-rewarding procedure. when employing external model supervision, it is crucial to utilize a robust model that can effectively generalize to new data. typically, these models are substantially large to avoid reward overoptimization @xcite . despite being reliable, labeling abundant data with a large-scale model is still very costly and time-consuming. in this paper, we aim to balance reliability and efficiency in the data labeling process during the iterative fine-tuning of the policy model. to achieve this, we propose ts-align, a teacher-student collaborative framework that leverages the reliability of the large-scale teacher model without requiring it to process all the candidates. specifically, ts-align uses a base supervised fine-tuned policy model to generate response candidates for a diverse set of instruction prompts sampled from public instruction-tuning datasets. a small-scale student reward model (rm) provides coarse-grained annotations, allowing for the quick processing of abundant unlabeled data and the selection of preference pairs from the candidates. next, the strong teacher helps re-rank the selected pairs reliably. the policy model is then fine-tuned on the re-ranked preference data using dpo. this process is repeated in several iterations. given that the student rm, with its smaller parameter size, is not as robust as the teacher model, we iteratively update the student using an adapter-based multi-task training setup @xcite . this training uses the same model-labeled preference data to enhance the student's reliability, which can be perceived as distilling new knowledge from the large teacher model to the small student rm. our contributions are three-fold: (1) we introduce "ts-align", an efficient and reliable pipeline for the iterative alignment of large language models. this approach circumvents the need for costly human annotations by employing a teacher-student model collaboration to automatically extract preference data from the policy model's own outputs. ## the ts-align pipeline this section details ts-align, with standardized notations in ## related work iterative llm alignment can be broadly divided into two categories: the first focuses on self-evolution without relying on an external reward model @xcite @xcite . for example, @xcite proposes self-rewarding language models, where the process begins by bootstrapping instructions from the policy model, which then creates candidate responses based on these instructions. the model employs "llmas-a-judge" prompting @xcite to evaluate and reward its own outputs. this approach allows the model to align itself through directed preference optimization using the selfcurated data. @xcite introduces instruction back-translation. this involves using the policy model to generate new instructions from text spans within the clueweb corpus. the model then produces responses given the newly generated instructions. the resulting instruction-response pairs serve as a basis for further fine-tuning the policy model, enhancing its alignment through continuous refinement. however, these approaches heavily rely on the scale of the llms as the "llm-as-a-judge" may not work well on smaller language models. additionally, the self-rewarding mechanism tends to bias towards their generations. the second category, in contrast, relies on an external rm to guide the alignment process @xcite @xcite @xcite . @xcite uses human annotations of policy generations during each alignment iteration and employs rejection sampling to guide the policy model to produce human-favored outputs. the rest adopt a similar pipeline to ours, using an external reward model to annotate policy model generations and derive pseudo-labeled preference data for alignment. the key difference between ts-align and other approaches is the teacher-student collaboration mechanism, which enables reliable and efficient annotation of large-scale preference data for policy model alignment. our approach is also more practically feasible under conditions of limited budget and resources. proaches propose to curate preference data through ai feedback @xcite @xcite , which is an efficient way to obtain large-scale preference data than using human annotators. ## conclusion we introduce ts-align, a teacher-student collaborative framework designed to balance reliability and efficiency in the data labeling process for iterative fine-tuning of policy models. by leveraging the strengths of a large-scale teacher model without requiring it to process all candidates, ts-align combines the efficiency of a smaller student reward model with the reliability of a robust teacher model. this iterative alignment process results in a highly aligned policy model with an impressive average win rate of 69.7% over the base policy, as judged by gpt-4. human evaluations also confirm the effectiveness of ts-align. additionally, we demonstrate that the teacher's knowledge is effectively distilled into the student, and the final student reward model, after iterative alignment, can be transferred to align other base policy models.