Uncertainty Aware Learning for Language Model Alignment (2024)

Abstract: As instruction-tuned large language models (LLMs) evolve, aligning pretrained foundation models presents increasing challenges. Existing alignment strategies, which typically leverage diverse and high-quality data sources, often overlook the intrinsic uncertainty of tasks, learning all data samples equally. This may lead to suboptimal data efficiency and model performance. In response, we propose uncertainty-aware learning (UAL) to improve the model alignment of different task scenarios, by introducing the sample uncertainty (elicited from more capable LLMs). We implement UAL by a simple fashion – adaptively setting the label smoothing value of training according to the uncertainty of individual samples. Analysis shows that our UAL indeed facilitates better token clustering in the feature space, validating our hypothesis. Extensive experiments on widely used benchmarks demonstrate that our UAL significantly and consistently outperforms standard supervised fine-tuning. Notably, LLMs aligned in a mixed scenario have achieved an average improvement of 10.62% on high-entropy tasks (i.e., AlpacaEval leaderboard), and 1.81% on complex low-entropy tasks (i.e., MetaMath and GSM8K).

## introduction large language models (llms), represented by @xcite , claude, and llama2 @xcite , have recently achieved significant success in a series of natural language generation and understanding tasks @xcite @xcite . the emergence of alignment methods has further enhanced the capabilities of llms, for example, the ability to follow human instructions or to achieve better zero-shot performance. the current popular alignment approaches including rlhf @xcite , consider supervised fine-tuning (sft) an essential part that contributes significantly. it is undeniable that the current sft paradigm achieves considerable success. some previous studies find that high-quality, complex, and diverse data contribute significantly to alignment @xcite . however, the data can exhibit different levels of uncertainty. for data points in highly technical, scientific, or specialized settings, context may have little or no ambiguity and a limited set of correct answers. in contrast, other dialogues might feature varied and dynamic social contexts with idiomatic language uses. appendix a presents a pair of such examples. nevertheless, the common sft paradigm applies the same level of supervision to all samples in the training set, overlooking the intrinsic uncertainty of the data. furthermore, due to the phenomenon of catastrophic forgetting @xcite @xcite , sft-aligned models may perform worse than their foundational models. the deepseek technical report reveals that the sft model often underperforms compared to its base model on several benchmarks @xcite , while llms generally show diminished performance on general tasks when forging their agent capability @xcite . in our experiments, we observe that the standard sft paradigm frequently leads to model degradation, resulting in decreased performance on some benchmarks, despite providing overall improvement. for instance, this is evident in parts of the commonsense benchmarks, such as mmlu. it is essential to align pretrained models while mitigating degradation issues as effectively as possible. therefore, we propose: uncertainty hypothesis. in an ideal paradigm, to further enhance alignment performance, the model should attend to samples differently based on their properties during alignment. specifically, the model should impose stricter constraints when attending to more certain examples, as these samples exhibit less uncertainty and fewer variations, while maintaining relaxed constraints for highly uncertain examples. to design our uncertainty-aware learning (ual) method and address potential model degradation caused by sft, we propose measuring uncertainty through a coarsely-grained approach that incorporates an autonomous judge (e.g., gpt-4) to assess the uncertainty of each sample in the training set. upon obtaining the uncertainty estimations, we linearly map them into our adaptive label smoothing training pipeline. further details of our algorithm are presented in section 2. ual significantly enhances the performance of instruction-tuned models. we apply this method across various model architectures and alignment datasets, observing that the ual paradigm consistently outperforms the vanilla sft paradigm on prominent benchmarks, including mmlu and truthfulqa @xcite , among others. as figure 2 shows, ual also helps the model improve in different scenarios. moreover, compared to its foundational model, as figure 1 presents, the sftaligned model brings same-class tokens closer together in the feature space, and ual enhances this tendency even further. this strongly correlates with low entropy high entropy performance (%) sft +10.62 +1.81 ual he* le* +3.81 -4.43 +13.73 -18.03 figure 2: mistral-7b (jiang et al., 2023) aligned with ual on the mixed scenario dataset has an improvement of 10.62% on the high entropy scenario and 1.81% on the low entropy scenario, compared with sft. evaluation results from models solely sft aligned with high-entropy data (i.e., he aligned with sharegpt-6k) or low-entropy data (i.e., le aligned with metamath-6k) are also provided for reference. the superior performance of ual-aligned models across benchmarks and scenarios, offering an important explanation for ual's outperformance over conventional sft. ## experiments in this section, we will introduce the experimental setup and present evaluation results from comprehensive benchmarks, as well as improvements observed in mixed scenarios, to demonstrate the superiority of our method. to promote a deeper understanding of our approach, we conduct an ablation study. additionally, we present a case study and employ gpt-4 to further evaluate the quality of responses generated by our method. ## one alternative view of uncertainty-aware learning feature clustering has been an interesting view in machine learning and deep learning @xcite . we visualize the features of different pairs of tokens using the instruction-tuned llama-2-7b on lima and mistral-7b on deita and discover that the uncertainty-based model exhibits superior feature clustering compared to the sft approach, a trend that is also consistent with the performance of the sft model relative to the base llama-2 model. to visualize a pair of randomly selected tokens, we perform two steps: 1) conduct model inference on hundreds of text inputs and collect the features from the layer preceding the token classifier head; 2) gather the features corresponding to the two tokens and apply principal component analysis (pca @xcite to reduce the dimensionality to a 2d space for visualization. to ensure the generality of our conclusions, we selected 100 pairs of randomly chosen tokens and used the silhouette coefficient to quantify their convergence degree, with the results displayed in additionally, due to the catastrophic forgetting phenomenon @xcite @xcite , some studies have mentioned that when aligning large language models, the deepseek technical report reveals that the sft model often underperforms compared to its base model on several benchmarks @xcite ; decreased performance is observed on general tasks when forging the agent capability of llms @xcite . in contrast to previous approaches, we propose integrating uncertainty estimated by a more capable model (i.e., @xcite ) into alignment to achieve considerably better performance. uncertainty-aware learning. estimating uncertainty to establish prediction confidence is vital for deep learning, and it has been studied thoroughly @xcite @xcite . uncertainty has played an important role in bayesian neural networks @xcite , semisupervised learning, and self-learning @xcite . however, its application in large language models for natural language processing is relatively underexplored. previous work has focused on improving neural machine translation @xcite and code generation @xcite , but these efforts have been limited to token-level uncertainty or perplexity (ppl). with the emergence of large-scale language models, some researchers have begun to elicit uncertainty from these models and have found that they can express their uncertainty well @xcite . this insight provides inspiration for subsequent research. ## conclusion in our study, we introduce a novel alignment approach called uncertainty-aware learning (ual). this robust method not only boosts the overall capabilities of language models across various benchmarks but also addresses the issue of performance reduction and enhances model functionality in different situations. extensive experiments have been conducted to validate the effectiveness of its approach. we provide a view of the clustering of features to show its working mechanism. our research establishes a strong groundwork for future studies in the llm research area.