Search papers, labs, and topics across Lattice.
This paper introduces Cost-Aware Learning, a framework for minimizing the total cost of training finite-sum objectives where sampling different component functions incurs varying costs. They propose Cost-Aware Stochastic Gradient Descent with cost complexity analysis and a subset selection algorithm. Applied to reinforcement learning with language models, their Cost-Aware GRPO algorithm reduces token usage in policy optimization by up to 30% while maintaining or improving accuracy on 1.5B and 8B parameter models.
Cut your LLM fine-tuning costs by 30% without sacrificing accuracy by intelligently sampling training data based on cost.
We consider the problem of Cost-Aware Learning, where sampling different component functions of a finite-sum objective incurs different costs. The objective is to reach a target error while minimizing the total cost. First, we propose the Cost-Aware Stochastic Gradient Descent algorithm for convex functions, and derive its cost complexity to attain an error of $\epsilon$. Furthermore, we establish a lower bound for this setting and provide a subset selection algorithm to further reduce the cost of training. We apply our theoretical insights to reinforcement learning with language models, where the computational cost of policy gradients varies with sequence length. To this end, we introduce Cost-Aware GRPO, an algorithm designed to reduce the cost of policy optimization while preserving performance. Empirical results on 1.5B and 8B LLMs demonstrate that our approach reduces the tokens used in policy optimization by up to about 30% while matching or exceeding baseline accuracy.
Google Research