用策略梯度优化序列生成,突破基础模型限制。
Post-Training with Policy Gradients: Optimality and the Base Model Barrier
- 基于结果与过程奖励的策略梯度方法优化序列输出
- 在非平凡似然下可达到近似完美生成,查询次数近乎最优
- 过程奖励能避免维度灾难,适合长序列生成任务
我们研究了后训练阶段的线性自回归模型,结合结果奖励与过程奖励。给定上下文 $\boldsymbol{x}$,模型需生成长度为 $N$、满足 $γ$ 间隔条件的响应序列 $\boldsymbol{y} \in Y^N$,该条件是标准可分性的序列扩展。我们证明:当基础模型在测试样本上具有非平凡似然 $α$ 时,一种策略梯度(PG)变体可在几乎最小最大意义下以 $\tilde{O}((α^{-1} + \varepsilon^{-1})/γ^2)$ 次奖励查询实现似然 $1 - \varepsilon$。然而,超越基础模型支持域存在障碍:仅用结果奖励时,后训练总体期望误差受基础模型的似然分位数(LQ)控制,且任何 PG 变体可能需要指数级于 $N$ 的奖励查询,无论预训练方式如何。为克服此障碍,我们引入过程奖励模型,证明在该设置下 PG 变体可通过依赖词级别 LQ 避免 $N$ 的维度灾难。此外,我们还证明在间隔条件下,自适应学习率的 SGD 可实现近似最优统计学习误差;自适应学习率的 PG 在在线学习中可实现近似最优误分类数,且计算高效,可能具有独立意义。
原文摘要 · Abstract (English)
We study post-training linear autoregressive models with outcome and process rewards. Given a context $\boldsymbol{x}$, the model must predict the response $\boldsymbol{y} \in Y^N$, a sequence of length $N$ that satisfies a $γ$ margin condition, an extension of the standard separability to sequences. We prove that on test samples where the base model achieves a non-trivial likelihood $α$, a variant of policy gradient (PG) can achieve likelihood $1 - \varepsilon$ with an essentially minimax optimal number of reward queries $\tilde{O}((α^{-1} + \varepsilon^{-1})/γ^2)$. However, a barrier arises for going beyond the support of the base model. We prove that the overall expected error after post-training with outcome rewards is governed by a property of the base model called the Likelihood Quantile (LQ), and that variants of PG, while minimax optimal, may require a number of reward queries exponential in $N$ to go beyond this support, regardless of the pre-training algorithm. To overcome this barrier, we study post-training with a process reward model, and demonstrate how PG variants in this setting avoid the curse of dimensionality in $N$ via dependence on a token-level LQ. Along the way, we prove that under the margin condition, SGD with adaptive learning rate (LR) achieves a near optimal test error for statistical learning, and PG with adaptive LR achieves a near optimal number of mistakes for online learning while being computationally efficient whenever possible, both of which may be of independent interest.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。