Differential Smoothing Mitigates Sharpening and Improves LLM Reasoning
Abstract
It is widely recognized that reinforcement learning (RL) fine-tuning of large language models often leads to \textit{diversity collapse}, where outputs lack variety. Prior work has proposed a range of heuristics to counteract this effect, but these methods are ad hoc: they frequently trade off correctness for diversity, their effectiveness varies across tasks, and in some cases they even contradict one another. In this work, we place these observations on a rigorous foundation. We first provide a formal proof of why RL fine-tuning exhibits diversity collapse via a selection and reinforcement bias. Next, we make a key observation that any reward modification to address diversity collapse only needs to be applied on the correct trajectories. Building directly on this analysis, we introduce a principled method---\textit{differential smoothing}---that provably improves both correctness and diversity, outperforming vanilla RL as well as widely used entropy-based heuristics. Our theory precisely characterizes when existing heuristics help and why they fail, while showing that differential smoothing is universally superior. Extensive experiments with models from 1B to 7B parameters, across domains including CountDown and real-world mathematical reasoning, demonstrate consistent gains. Differential smoothing improves both Pass@1 and Pass@k, with up to 6.7\% improvements on AIME24 dataset.
Lay Summary
Modern AI assistants like ChatGPT are usually improved after their initial training through a process that rewards them for producing correct answers. While this boosts accuracy, it has a subtle side effect: the model becomes repetitive, generating the same kind of answer over and over. This is a real problem because hard tasks—such as challenging math—often require several different attempts before the right idea emerges. Many quick fixes have been proposed, but they are inconsistent: a fix that helps on one type of task can hurt on another, and different proposals often contradict each other. In our paper, we mathematically pin down why this loss of variety happens—the training is biased toward answers the model already favors, reinforcing them while pushing equally good alternatives aside. Building on this insight, we introduce a method called "differential smoothing," which gives the model different feedback depending on whether its attempt was right or wrong. This simple distinction keeps both accuracy and variety high at the same time. We tested our method on AI models of various sizes and on tasks ranging from number puzzles to high-school competition math. It consistently produced more accurate single answers and a higher success rate when the model was given several attempts—improving results by up to 11.4% on one of the toughest math benchmarks we evaluated.