Cost-Aware Learning

2026-04-30

AI Summary

This work addresses the challenge of heterogeneous sampling costs across component functions in finite-sum optimization by proposing the first cost-aware stochastic gradient descent framework and establishing its theoretical lower bound. The core contributions include a novel cost model based on sequence length, a subset selection strategy designed to minimize total sampling overhead, and the development of the Cost-Aware GRPO algorithm for policy optimization in large language models. Experimental results on 1.5B and 8B parameter models demonstrate that the proposed method achieves comparable or superior accuracy to baseline approaches while reducing the number of tokens required for policy optimization by approximately 30%, thereby substantially lowering computational costs.

Abstract

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 $ε$. 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.

Problem

• Cost-Aware Learning
• finite-sum optimization
• sampling cost
• cost minimization
• target error

Innovation

• Cost-Aware Learning
• Stochastic Gradient Descent
• Cost Complexity
• Subset Selection
• Policy Optimization