Beyond Accuracy Optimization: Computer Vision Losses for Large Language Model Fine-Tuning (2024)

Abstract: Large Language Models (LLMs) have demonstrated impressive performance across various tasks. However, current training approaches combine standard cross-entropy loss with extensive data, human feedback, or ad hoc methods to enhance performance. These solutions are often not scalable or feasible due to their associated costs, complexity, or resource requirements. This study investigates the use of established semantic segmentation loss functions in natural language generation to create a versatile, practical, and scalable solution for fine-tuning different architectures. We evaluate their effectiveness in solving Math Word Problems and question answering across different models of varying sizes. For the analyzed tasks, we found that the traditional Cross-Entropy loss represents a sub-optimal choice, while models trained to minimize alternative (task-dependent) losses, such as Focal or Lovász, achieve a mean improvement of +36% on exact match without requiring additional data or human feedback. These findings suggest a promising pathway for more efficient and accessible training processes.

## introduction generative language models have shown impressive capabilities across various scenarios @xcite . recent advancements in large language models have made this even more evident @xcite . however, the performance of these models is influenced by three main factors: model size, amount of training data, and training strategy @xcite . increasing model size requires more computational resources while training on vast data collections is essential to achieve competitive results when increasing the size. additional training refinements have been introduced recently, some of which involve human experts providing feedback to enhance model performance, as in reinforcement learning from human feedback (rlhf) @xcite where human preferences are then used to align the model outputs. despite the improved performance, developing these models requires massive amounts of resources, power, time, and therefore significant costs, making them accessible only to very few leading companies. the increasing costs and resource requirements have already led to the development of solutions aimed at democratizing the training of these models. an example of this is the use of parameter efficient fine-tuning (peft) @xcite like low-rank adaptation (lora) @xcite or derived strategies, often combined with quantization techniques, to enable lightweight fine-tuning of these computationally expensive architectures. some works (e.g., zephyr @xcite ) circumvent the need for human feedback by distilling the knowledge of larger and more powerful models (e.g., gpt-4), leveraging the so-called ai feedback (aif). however, this may cause the propagation of potential biases from the larger model, or the generation of unfactual content, potentially resulting in a less accurate representation of reality @xcite . although direct preference optimization (dpo) @xcite and similar methods address the instability in reinforcement learning training @xcite and the computational requirements, their effectiveness is limited by the quality of the prior supervised fine-tuning stage @xcite . supervised fine-tuning still remains the most memory-efficient alternative, however the collection of fine-tuning and instruction-tuning datasets presents similar challenges, as it requires additional annotation costs that may not always be accessible, especially when dealing with multiple datasets to annotate @xcite . while costs can be mitigated by relying on weak annotation, similar to possible issues of using aif, the quality of annotations is not always guaranteed. the need for larger models and additional data pertains not only to general-purpose systems (e.g., chatgpt) but also to task-specific models for more complex scenarios, as in the case of math word problems (mwp). specifically, state-of-the-art solutions often employ more complex training procedures, involving multiple training stages (e.g., supervised fine-tuning, instruction tuning, preference-based tuning) @xcite @xcite to achieve improved performance. however, similar to general-purpose models, this leads to more expensive solutions. additionally, these approaches are also less portable since they are tailored to one specific task. to tackle the above-mentioned challenges, in this work we raise some concerns about the standard practice of cross-entropy (ce) loss optimization, which is the usual language modeling objective, and we show that a more accurate selection of the loss function to optimize can be incredibly beneficial for model training. specifically, by using a loss function tailored to the task under analysis and leveraging lora, we effectively fine-tune llms with small amounts of data. the underlying idea is that for certain language tasks, it is desirable to optimize not only for the correctness of the output but also for the structural adherence of the generated text to a specific format or syntax. this is particularly relevant for tasks involving formal languages or well-defined procedures, such as mathematical reasoning, where the intermediate steps and reason-ing process must follow a strict structure. consequently, we hypothesize that accounting for these characteristics by employing a more appropriate loss function could lead to improved performance. these challenges in natural language generation (nlg) are analogous to delineating boundaries in image segmentation, where each pixel (tokens in nlg) contributes to forming an accurate and coherent shape (textual output in nlg along with its syntax). in both fields, the presence of class imbalances (such as the distribution of output tokens in nlg) and well-defined yet complex shapes (the syntax and structure of the output text in nlg) are common challenges. our solution does not involve either the implementation of complex training procedures (e.g., luo et al. (2023) ), further pre-training (e.g., azerbayev et al. (2024) ), distilling knowledge from larger models (e.g., @xcite ) or collecting human preferences @xcite . in contrast, our approach focuses on selecting a more suitable loss function based on the task at hand, achieving improvements over the standard cross-entropy loss through a single training stage, as shown in figure 1 . in this work, we focus on mathematical reasoning and closed-question answering, which are common benchmark tasks @xcite @xcite that have a clear and well-defined expected output. furthermore, we choose these tasks since we claim that, for both of them, the role of human preferences is secondary since it is difficult to "prefer" one output with respect to another, especially in mathematical reasoning where there could be more than one procedure to get the final solution. our results show that accurately choosing the right loss function (combined with cross-entropy) can improve performance on the analyzed tasks using the same amount of data without adding complexity to the training process. the source code to reproduce the experiments is available for research purposes at https://github.com/darthreca/ segmentation-losses-nlp . ## related works the evolution of large language models has been driven by various innovative training methods. this section provides an overview of the existing approaches for training llms, highlighting the challenges and benefits of each approach. additionally, we explore the development of alternative loss functions beyond cross-entropy in both natural language processing and computer vision fields. ## methodology in this section, we formally introduce the loss functions we employ, shortlisted from the classification presented in @xcite , and explain their rationale. we describe our approach when employing them for causal generation. for the sake of readability, all loss formulations are reported in appendix a. ## experimental results we perform an extensive experimental evaluation on two tasks for a total of four datasets, five models, and five loss functions. in the following, we summarize the main results reporting the average across models. detailed results are available in appendix e. ## conclusion and future work in our work, we applied semantic segmentation losses to improve the fine-tuning of llms for mathematical reasoning and closed-ended questionanswering. our results demonstrate that using appropriate loss functions during fine-tuning can boost performance without the need for additional data or human feedback. as shown in our experiments, we found that the most effective loss functions are, on average, the lovász loss and focal loss for mwp and closed-ended qa tasks, respectively. however, we would like to highlight that our experiments indicate no definitive loss function exists to solve every problem. therefore, a thorough understanding of the specific problem and data is crucial to achieving the best results. in practice, this suggests a promising pathway for more efficient and accessible training processes. future work will focus on designing new task-specific loss functions and extending our approach to other tasks. pnrr -m4 c2, invest 1.3 -d.d. 1551.11-10-2022, pe00000004). this study was also partially carried out within the fair (future artificial intelligence research) and received funding from next-generationeu (italian pnrr -m4 c2, invest 1.3 -d.d. 1555.11-10-2022, pe00000013). this manuscript reflects only the authors' views and opinions, neither the european union nor the european commission can be considered responsible for them.