Entropy-Based Adaptive Weighting for Self-Training

TL;DR

The paper addresses improving mathematical reasoning in LLMs by reweighting self-generated training data according to model uncertainty. It introduces EAST, which computes entropy over clustered reasoning-path outputs for each question and maps this uncertainty through an exponential-style function with a tunable exponent a to obtain training weights, then applies these weights to base losses like SFT, DPO, or KTO. Empirical results on GSM8K and MATH show EAST yields consistent improvements over vanilla self-training and local-uncertainty baselines, with notable gains on GSM8K and modest gains on MATH, and positive effects observed under iterative training. The approach demonstrates that focusing on uncertain, diverse answers can better guide learning and improve reasoning capabilities in LLMs.

Abstract

The mathematical problem-solving capabilities of large language models have become a focal point of research, with growing interests in leveraging self-generated reasoning paths as a promising way to refine and enhance these models. These paths capture step-by-step logical processes while requiring only the correct answer for supervision. The self-training method has been shown to be effective in reasoning tasks while eliminating the need for external models and manual annotations. However, optimizing the use of self-generated data for model training remains an open challenge. In this work, we propose Entropy-Based Adaptive Weighting for Self-Training (EAST), an adaptive weighting strategy designed to prioritize uncertain data during self-training. Specifically, EAST employs a mapping function with a tunable parameter that controls the sharpness of the weighting, assigning higher weights to data where the model exhibits greater uncertainty. This approach guides the model to focus on more informative and challenging examples, thereby enhancing its reasoning ability. We evaluate our approach on GSM8K and MATH benchmarks. Empirical results show that, while the vanilla method yields virtually no improvement (0%) on MATH, EAST achieves around a 1% gain over backbone model. On GSM8K, EAST attains a further 1-2% performance boost compared to the vanilla method.

Figures (5)

• Figure 1: Comparison between the traditional self-training pipeline and EAST. The LLM generates $n$ responses per question, clustered by final answers. Questions with all incorrect answers are discarded. Self-training fine-tunes uniformly on the rest, while EAST assigns higher weights to questions with diverse (uncertain) answers and lower weights to consistent (confident) ones.
• Figure 2: The framework of EAST. For each training question, the LLM generates $n$ responses, clustered by final answers. Entropy value is computed from the cluster distribution, transformed via mapping function, and integrated as weight into the loss objective.
• Figure 3: Performance(accuracy (%)) of various exponent parameters $a$ on GSM8K and MATH datasets using LLaMA-3.2-1B.
• Figure 4: Comparison of iterative learning performance (accuracy (%)) between vanilla SFT and EAST on LLaMA-3.2-1B.
• Figure 5: The figure illustrates the distribution of training data in entropy-based, accuracy-based, and rejected-based values. Each point represents a training example ($x_i$), with coordinates ($H(x_i), 1-A(x_i)$) for entropy-based and accuracy-based values, and color indicating the rejected-based value ($R(x_i)$). The accompanying table reports the performance (accuracy(%)) of three weighting strategies on the GSM8K and MATH datasets.