精打细算的训练:Cost-Aware Learning 如何让 LLM 训练省出 30% 的算力
精打细算的训练:Cost-Aware Learning 如何让 LLM 训练省出 30% 的算力
> 作者:小凯 | 来源:arXiv:2604.28020v1 [cs.LG] | 机构:Google Research / Tel Aviv University
---
一、食堂经济学
想象你在食堂吃饭,预算有限。你面前有两个窗口:
- 窗口 A:5 块钱的素菜,能提供 50 大卡热量
- 窗口 B:20 块钱的荤菜,能提供 80 大卡热量
这就是成本感知学习(Cost-Aware Learning)的核心直觉:不同数据点的"性价比"不同,训练时应该优先选性价比高的。
---
二、AI 训练中的"隐形账单"
传统随机梯度下降(SGD)的理论有个隐藏假设:所有梯度步骤的计算成本相同。这在简单分类任务里大致成立——每张图片的前向传播都差不多快。
但在现代大模型训练中,这个假设彻底崩了:
- RLHF 中的长短差异:一个 1000 token 的推理轨迹和一个 50 token 的简洁回答,前向传播的 FLOPs 差了近 20 倍
- 多模态数据:处理一张高清图 vs 一段文本,成本天差地别
- 代码生成:编译失败的长代码和一次通过的短代码,模型都要处理,但价值可能完全不同
---
三、重要性采样:老统计学的冷饭,新场景的热锅
Cost-Aware Learning 的解法并不新——它来自 1978 年的重要性采样(Importance Sampling)。但这篇论文的妙处在于,把古老的统计工具精准地嫁接到现代大模型训练流程里。
3.1 核心公式
假设每个样本 i 有一个已知的成本 cᵢ(比如 token 数量、FLOPs 估计)。传统 SGD 均匀采样,期望成本是 E[c] = Σ cᵢ / n。
Cost-Aware SGD 改为按逆成本概率采样:
pᵢ ∝ 1 / cᵢ
然后,为了保持估计的无偏性,用重要性权重 wᵢ = cᵢ / c̄ 来校正梯度。最终每个样本 i 对总目标的贡献是:
(1 / (n * pᵢ)) * ∇fᵢ(θ) = (cᵢ / c̄) * ∇fᵢ(θ)
这保证了虽然采样频率变了,但期望梯度仍然正确指向总体目标的最小值。
3.2 直观理解
用食堂类比:你不再每道菜都吃一口(均匀采样),而是多吃性价比高的菜(逆成本采样)。为了不让自己"营养不良"(偏差),你在计算总热量摄入时,给多吃的那几道菜打一个折扣(重要性权重校正)。
最终你摄入的热量估计仍然是准确的,但花的钱更少。
---
四、从 SGD 到 GRPO:进入 RL 的世界
论文把 Cost-Aware Learning 扩展到了 GRPO(Group Relative Policy Optimization),这是 DeepSeek-R1 等推理模型使用的核心 RL 算法。
4.1 为什么 RL 更需要成本感知
在监督学习(SFT)中,成本差异主要来自输入长度。但在 RL 中,成本差异被放大了:
- 生成阶段:模型要生成一条推理链,长度未知。长的 CoT 可能比短的贵 5-10 倍
- 评估阶段:每个生成的回答都要用验证器或奖励模型打分,这又是一次前向传播
- 组内差异:GRPO 的一个核心设计是"组采样"——对同一个问题生成 N 个回答。有的回答 50 token,有的 800 token,成本极度不均
4.2 Cost-Aware GRPO
论文提出的修改极其简洁:
1. 成本定义:每个样本(一个问题-回答对)的成本 = prompt token 数 + response token 数 2. 采样策略:在组内和跨组都引入逆成本采样 3. 权重校正:重要性权重重新中心化(re-centering)以维持训练稳定性
具体来说,权重被重新中心化到 1(即减去均值再加 1),这保持了梯度范数的一致性,避免了某些低成本样本因权重过大而主导梯度。
这个修改的核心洞察是:强化学习天然有"探索-利用"的权衡,成本感知把"计算预算"也纳入了这个权衡。 模型不仅要探索哪些动作能获得高奖励,还要探索哪些动作能以更低的计算成本获得高奖励。
---
五、实验:30% 的节省从何而来
5.1 合成验证
论文首先在合成凸优化问题上验证了理论:
- 设定:1000 个样本,成本均匀分布在 [1, 10] 之间
- 目标:达到 ε = 0.01 的精度
- 结果:Cost-Aware SGD 的总计算成本比均匀采样降低 ~40%
5.2 LLM 后训练
在真实的 LLM 训练场景中(基于 Verl 框架,vLLM 推理,FlashAttention 优化):
| 设置 | 基线 GRPO | Cost-Aware GRPO | 节省 |
|---|---|---|---|
| 1.5B 模型 | 24 GPU 小时 | ~17 GPU 小时 | ~30% |
| 8B 模型 | 55 GPU 天 | ~38 GPU 天 | ~30% |
1. 精度不打折:Cost-Aware GRPO 的验证集损失和基线持平,没有因为"省钱"而"降质" 2. 长尾样本不丢失:逆成本采样确保低成本样本被更频繁地看到,但重要性权重保证它们不会过度影响模型。高成本样本虽然采样频率低,但仍有机会被选中 3. 稳定性:重新中心化的权重设计避免了训练初期的梯度爆炸
---
六、数据预筛选:另一种省钱思路
除了训练时的动态采样,论文还探讨了训练前的静态数据筛选:
> "给定一个固定数据集和一个成本预算,如何选子集使得训练误差最小?"
这引入了一个有趣的权衡:
- 无偏但昂贵:用完整数据集,期望成本 E[c]
- 有偏但便宜:筛选掉高成本样本,引入偏差,但显著降低总成本
对 LLM 训练的现实意义:如果你的训练预算固定,也许不应该盲目追求"数据越多越好",而是精挑细选"性价比最高"的数据子集。
---
七、一句话总结
Cost-Aware Learning 的核心启示是:训练大模型和经营食堂一样,光看好不好吃(损失函数)不够,还得看贵不贵(计算成本)。 当 AI 训练进入"万亿参数 × 万亿 token"的时代,计算成本已经和数据质量同等重要。这篇论文把经济学里最朴素的"性价比"思维,用严谨的统计语言带进了梯度下降的每一步——让我们在精打细算中,把模型训得更好。
---
参考链接
- 论文原文:https://arxiv.org/abs/2604.28020
🌟 智谱 GLM-5 已上线
我正在智谱大模型开放平台 BigModel.cn 上打造 AI 应用,智谱新一代旗舰模型 GLM-5 已上线,在推理、代码、智能体综合能力达到开源模型 SOTA 水平。
🎁 领取 2000万 Tokens