Fast on the Easy, Deep on the Hard: Efficient Reasoning via Powered Length Penalty

TL;DR

This paper introduces a novel length penalty method to improve reasoning efficiency in large language models, reducing output length for simple tasks while maintaining or improving accuracy on complex benchmarks.

Contribution

It proposes a new reward function with a length penalty that adapts to problem complexity, enhancing reasoning efficiency and performance across diverse datasets.

Findings

Shortened outputs on GSM8K and MATH500 with maintained or improved accuracy.

Achieved higher accuracy on the challenging AIME2024 dataset.

Demonstrated effectiveness of the method across three benchmark datasets.

Abstract

Large language models (LLMs) have demonstrated significant advancements in reasoning capabilities, performing well on various challenging benchmarks. Techniques like Chain-of-Thought prompting have been introduced to further improve reasoning. However, these approaches frequently generate longer outputs, which in turn increase computational latency. Although some methods use reinforcement learning to shorten reasoning, they often apply uniform penalties without considering the problem's complexity, leading to suboptimal outcomes. In this study, we seek to enhance the efficiency of LLM reasoning by promoting conciseness for simpler problems while preserving sufficient reasoning for more complex ones for accuracy, thus improving the model's overall performance. Specifically, we manage the model's reasoning efficiency by dividing the reward function and including a novel penalty for output length. Our approach has yielded impressive outcomes in benchmark evaluations across three datasets: GSM8K, MATH500, and AIME2024. For the comparatively simpler datasets GSM8K and MATH500, our method has effectively shortened output lengths while preserving or enhancing accuracy. On the more demanding AIME2024 dataset, our approach has resulted in improved accuracy.