Introduction
Reinforcement learning (RL) has become a central component of LLM post-training, producing substantial performance improvements in mathematical reasoning, code generation, and other verifiable domains [5 ] (Shao et al. 2024; Yu et al. 2025). Yet, an increase in average performance does not by itself reveal how the model has improved. Pretrained LLMs contain extensive latent capabilities, and RL can affect them through two distinct mechanisms. It may uncover successful behavior that is rarely expressed by the base policy, or it may concentrate probability on solutions that are already readily available. The latter phenomenon, referred to as sharpening , improves single-trial success but often reduces the range of problems solvable by the model under repeated trials (Yue et al., 2025). Further capability progress thus depends on effectively controlling sharpening.
Recent studies confirm this exact limitation: RL post-training dynamics are indeed dominated by sharpening (Yue et al., 2025; Song et al., 2025). The impressive single-trial gains frequently diminish or reverse in multi-trial settings, as the generated reasoning paths are confined to a narrow subset of the base model’s original distribution. Counteracting this over-concentration is notoriously difficult. For example, simply injecting randomness or diversity at the token level rarely translates into meaningful diversity in high-level reasoning strategies [6 ] (Walder and Karkhanis 2025).
Multi-trial metrics provide a practical lens for quantifying the extent of sharpening (Yue et al., 2025). For binary rewards, pass@k measures the probability that at least one of k k k k k k k k k k k k
We derive the closed-form optimal policies for KL-regularized pass@k and max@k. Our analysis reveals that @k optimization acts as prompt-adaptive pressure: it restricts optimization on tasks the model already solves reliably and redirects pressure toward tasks on which successful behavior is rare. We show that this behavior is achievable with ordinary @1 objectives via prompt-adaptive regularization or reward transformation. Leveraging this connection, we introduce @k Adaptive Pressure for Policy Optimization (KAPPO), a framework for turning existing @1 policy-optimization methods into @k optimizers. After a one-time reference calibration, KAPPO introduces prompt-adaptive pressure while preserving all other characteristics of the base method. Our contributions are:
We derive closed-form optimal policies for KL-regularized pass@k and max@k objectives.
We prove that @k optimization redistributes optimization pressure according to task difficulty via prompt-adaptive regularization or reward transformation, prioritizing hard tasks.
We introduce KAPPO, a framework that turns existing @1 methods into @k optimizers. In KAPPO, k k k k k k
We instantiate KAPPO across several policy optimization methods, matching the performance of a direct multi-completion @k-optimization baseline at a fraction of its rollouts.
Related Work
Exploration mechanisms in RLVR. Beyond the sharpening evidence discussed above, RL has been shown to consistently reduce generation diversity [11 ] , including on problems the model has not yet solved (Song et al., 2025). Application of standard exploration strategies, such as entropy regularization [7 , 2 ] (Mnih et al. 2016) and noisy reward perturbations [14 ] , has shown limited or sensitive improvements relative to methods designed specifically for multi-attempt reasoning [6 ] (Walder & Karkhanis 2025). More structured approaches are guided by epistemic uncertainty via policy ensembles [18 ] , or introduce exploration bonuses based on final outcomes (Song et al., 2025), hidden representations [19 ] (Tuyls et al. 2026), or sequence-level diversity [9 ] . These methods rely on different proxies for novelty and have been evaluated in heterogeneous settings; no established exploration mechanism for LLM post-training has yet emerged.
Pass@k and max@k optimization. Originally introduced as an evaluation metric for code generation [4 ] , pass@k has recently been adopted as a training objective. Walder & Karkhanis (2025) derive unbiased policy-gradient estimators, called PKPO, for binary pass@k and continuous max@k rewards. They show that increasing k k k [12 ] , and encourage “hedging the bets” given limited expressivity, a principle associated with exploration. Subsequent work studies analytical advantage design [6 ] , token-level mechanisms for preventing probability over-concentration (Peng et al., 2025), improved policy-gradient estimators for order-statistic objectives (Parmas et al., 2026), and a pass@k interpretation of maximum likelihood objectives (Tajwar et al., 2026). These approaches primarily analyze training gradients or construct multi-completion learning signals. We instead derive the corresponding KL-regularized optimal policies and show that the effect of @k optimization can be reproduced through prompt-adaptive @1 objectives. This provides a general route for turning existing @1 methods into @k variants without constructing expensive @k-dependent multi-completion learning signals.
Preliminaries
KL-regularized reward optimization. Let x ∼ D x \sim\mathcal{D} x ∼ D y ∼ π ( ⋅ ∣ x ) y \sim\pi(\cdot\mid x) y ∼ π ( ⋅ ∣ x ) R ( y , x ) R(y,x) R ( y , x )
J ( π ) : = E x ∼ D [ E y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) ] − β K L ( π ( ⋅ ∣ x ) ∥ π r e f ( ⋅ ∣ x ) ) ] . (1) \mathcal{J}(\pi) := \mathbb{E}_{x\sim\mathcal{D}}\left[\mathbb{E}_{y\sim\pi(\cdot\mid x)}[R(y,x)]-\beta\,\mathrm{KL}\left(\pi(\cdot\mid x)\|\pi_{\mathrm{ref}}(\cdot\mid x)\right)\right]. \tag*{(1)} J ( π ) := E x ∼ D [ E y ∼ π ( ⋅ ∣ x ) [ R ( y , x )] − β KL ( π ( ⋅ ∣ x ) ∥ π ref ( ⋅ ∣ x ) ) ] . (1) where π r e f \pi_{\mathrm{ref}} π ref β > 0 \beta>0 β > 0 [13 ] . Beyond potential empirical advantages, KL regularization is a prerequisite for non-degenerate optimal policies. Indeed, the optimal policy is given in closed form:
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) / β ) . (2) \pi^*(y\mid x)=\frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(R(y,x)/\beta\right). \tag*{(2)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( R ( y , x ) / β ) . (2) Methods such as REINFORCE (Williams, 1992), RLOO [1 ] , PPO (Schulman et al., 2017), and GRPO (Shao et al., 2024) approximate gradients of Eq. (1) . In contrast, methods like DPO (Rafailov et al., 2023), A∗ ^* ∗ [3 ] , and QRPO [17 ] directly fit a parameterized policy π θ \pi_\theta π θ π ∗ \pi^* π ∗ (2) . Our analysis and contributions are agnostic to the choice of optimizer, separating @k capabilities from other advancements.
Pass@k and max@k. Let y 1 , … , y k ∼ i . i . d . π ( ⋅ ∣ x ) y_1,\ldots,y_k \overset{\mathrm{i.i.d.}}{\sim}\pi(\cdot\mid x) y 1 , … , y k ∼ i.i.d. π ( ⋅ ∣ x )
max@k ( π ; x ) = E [ max i = 1 , … , k R ( y i , x ) ] , pass@k ( π ; x ) = P [ ∃ i ∈ { 1 , … , k } : R ( y i , x ) = 1 ] . \operatorname{max@k}(\pi;x)=\mathbb{E}\left[\max_{i=1,\ldots,k}R(y_i,x)\right],\qquad\operatorname{pass@k}(\pi;x)=\mathbb{P}\left[\exists\,i\in\{1,\ldots,k\}:R(y_i,x)=1\right]. max@k ( π ; x ) = E [ i = 1 , … , k max R ( y i , x ) ] , pass@k ( π ; x ) = P [ ∃ i ∈ { 1 , … , k } : R ( y i , x ) = 1 ] . Note that for binary rewards R ( y , x ) ∈ { 0 , 1 } R(y,x)\in\{0,1\} R ( y , x ) ∈ { 0 , 1 } k = 1 k=1 k = 1 E y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) ] \mathbb{E}_{y\sim\pi(\cdot\mid x)}[R(y,x)] E y ∼ π ( ⋅ ∣ x ) [ R ( y , x )] J \mathcal{J} J (1) . Let V k ( π ; x ) V_k(\pi;x) V k ( π ; x ) V k ( π ; x ) V_k(\pi;x) V k ( π ; x ) J \mathcal{J} J
J k ( π ) = E x ∼ D [ V k ( π ; x ) − β K L ( π ( ⋅ ∣ x ) ∥ π r e f ( ⋅ ∣ x ) ) ] . (3) \mathcal{J}_k(\pi)=\mathbb{E}_{x\sim\mathcal{D}}\left[V_k(\pi;x)-\beta\,\mathrm{KL}\left(\pi(\cdot\mid x)\|\pi_{\mathrm{ref}}(\cdot\mid x)\right)\right]. \tag*{(3)} J k ( π ) = E x ∼ D [ V k ( π ; x ) − β KL ( π ( ⋅ ∣ x ) ∥ π ref ( ⋅ ∣ x ) ) ] . (3) @k Optimal Policies and Adaptive Pressure
Binary Rewards: pass@k
Proofs of all results in this subsection, including boundary cases, are provided in Appendix § 13 . We first consider binary rewards, i.e., we set V k ( π ; x ) = pass@k ( π ; x ) V_k(\pi;x)=\operatorname{pass@k}(\pi;x) V k ( π ; x ) = pass@k ( π ; x ) (3) for J k \mathcal{J}_k J k J k \mathcal{J}_k J k
Proposition 4.1: Closed-form optimal policy for the pass@k objective Let R ( y , x ) ∈ { 0 , 1 } R(y,x)\in\{0,1\} R ( y , x ) ∈ { 0 , 1 } J k \mathcal{J}_k J k
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) β e f f ( x ) ) , (4) \pi^*(y\mid x)=\frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{R(y,x)}{\beta_{\mathrm{eff}}(x)}\right), \tag*{(4)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β eff ( x ) R ( y , x ) ) , (4) where
β e f f ( x ) = β k ( 1 − p ∗ ( x ) ) k − 1 \beta_{\mathrm{eff}}(x)=\frac{\beta}{k(1-p^*(x))^{k-1}} β eff ( x ) = k ( 1 − p ∗ ( x ) ) k − 1 β and
Z ( x ) = 1 − p r e f ( x ) + p r e f ( x ) exp ( 1 β e f f ( x ) ) . (5) Z(x)=1-p_{\mathrm{ref}}(x)+p_{\mathrm{ref}}(x)\exp\left(\frac{1}{\beta_{\mathrm{eff}}(x)}\right). \tag*{(5)} Z ( x ) = 1 − p ref ( x ) + p ref ( x ) exp ( β eff ( x ) 1 ) . (5) Here, p r e f ( x ) = P y ∼ π r e f ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] p_{\mathrm{ref}}(x)=\mathbb{P}_{y\sim\pi_{\mathrm{ref}}(\cdot\mid x)}[R(y,x)=1] p ref ( x ) = P y ∼ π ref ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] p ∗ ( x ) = P y ∼ π ∗ ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] p^*(x)=\mathbb{P}_{y\sim\pi^*(\cdot\mid x)}[R(y,x)=1] p ∗ ( x ) = P y ∼ π ∗ ( ⋅ ∣ x ) [ R ( y , x ) = 1 ]
Compared with the optimal policy (Eq. (2) ) for the single-reward objective (Eq. (1) ), the constant regularization strength β \beta β β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) p ∗ ( x ) p^*(x) p ∗ ( x ) π ∗ \pi^* π ∗ (4) over successful completions, which yields a scalar equation relating p ∗ ( x ) p^*(x) p ∗ ( x ) p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x )
Proposition 4.2: Scalar fixed-point characterization. For 0 < p r e f ( x ) < 1 0 < p_{\mathrm{ref}}(x) < 1 0 < p ref ( x ) < 1 p ∗ ( x ) p^*(x) p ∗ ( x ) ( 0 , 1 ) (0,1) ( 0 , 1 )
k ( 1 − p ∗ ( x ) ) k − 1 = β log p ∗ ( x ) ( 1 − p r e f ( x ) ) p r e f ( x ) ( 1 − p ∗ ( x ) ) (6) k\left(1-p^*(x)\right)^{k-1}=\beta\log\frac{p^*(x)\left(1-p_{\mathrm{ref}}(x)\right)}{p_{\mathrm{ref}}(x)\left(1-p^*(x)\right)} \tag*{(6)} k ( 1 − p ∗ ( x ) ) k − 1 = β log p ref ( x ) ( 1 − p ∗ ( x ) ) p ∗ ( x ) ( 1 − p ref ( x ) ) (6) The reference success probability p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x ) 1 p ∗ ( x ) p^*(x) p ∗ ( x ) (6) . For an error analysis, see Appendix § 22 . Proposition 4.1 suggests two equivalent ways to recover the pass@k optimal policy with single-reward optimization.
Corollary 4.3: pass@k as prompt-adaptive regularization. Optimizing the ordinary single-reward objective in Eq. (1) using the prompt-adaptive regularization strength β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) (5) results exactly in the pass@k optimal policy provided in Eq. (4) .
This equivalence holds because both objectives decompose over prompts and β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) y y y x x x (4) depends on the reward and KL coefficient only through their ratio, we obtain a complementary reward-side interpretation.
Corollary 4.4: pass@k as prompt-adaptive reward scaling. Optimizing the ordinary single-reward objective J \mathcal{J} J (1) with constant regularization strength β \beta β
R ~ k ( y , x ) = β β e f f ( x ) R ( y , x ) = k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) (7) \widetilde{R}_k(y,x)=\frac{\beta}{\beta_{\mathrm{eff}}(x)}R(y,x)=k\left(1-p^*(x)\right)^{k-1}R(y,x) \tag*{(7)} R k ( y , x ) = β eff ( x ) β R ( y , x ) = k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) (7) results exactly in the pass@k optimal policy provided in Eq. (4) .
Two views of adaptive pressure. Corollary 4.3 and Corollary 4.4 show that pass@k optimization can be expressed as ordinary single-reward optimization with either prompt-adaptive regularization or, equivalently, prompt-adaptive reward scaling. Taking 1 − p r e f ( x ) 1-p_{\mathrm{ref}}(x) 1 − p ref ( x ) x x x p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x ) p ∗ ( x ) p^*(x) p ∗ ( x ) Figure 1 a). Equivalently, under the reward-side view, a successful completion receives a smaller reward multiplier on easy tasks and a larger multiplier on difficult ones.
Figure 1. Adaptive pressure induced by @k objectives. (a) For binary rewards, increasing k k k k = 4 k = 4 k = 4 [ 0 , 1 ] [0, 1] [ 0 , 1 ] β = 0.5 \beta= 0.5 β = 0.5
Intuition: reducing sharpening. Consider optimizing the usual single-reward objective (Eq. (1) ) on a binary-reward task. The reward gain depends only on the increase in success probability, regardless of initial performance—for example, moving from near zero to 0.1 on a hard task is worth as much as moving from 0.9 to 1 on an already-easy task. For an LLM, however, the latter is often much easier: the model already knows how to solve the task and only needs to concentrate probability on known successful behavior. In contrast, increasing a near-zero success probability requires acquiring entirely novel reasoning strategies. The adaptive pressure induced by pass@k objectives counteracts this tendency towards sharpening by imposing stronger effective regularization on already-easy tasks or, equivalently, through the prompt-adaptive reward scaling described above.
Continuous Rewards: max@k
Next, we consider real-valued rewards, i.e., we set V k ( π ; x ) = max @ k ( π ; x ) V_k(\pi; x) = \operatorname{max}@k(\pi; x) V k ( π ; x ) = max @ k ( π ; x ) (3) for J k \mathcal{J}_k J k J k \mathcal{J}_k J k k > 1 k > 1 k > 1
Proposition 4.5: Closed-form optimal policy for the max@k objective The optimal policy for the max@k objective J k \mathcal{J}_k J k
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( k β ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) , (8) \pi^*(y \mid x) = \frac{1}{Z(x)}\pi_{\mathrm{ref}}(y \mid x)\exp\left(\frac{k}{\beta}\int_{-\infty}^{R(y,x)} F_{\pi^*}(r \mid x)^{k-1}\,\mathrm{d}r\right), \tag*{(8)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β k ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) , (8) where F π ∗ ( r ∣ x ) : = P y ′ ∼ π ∗ ( ⋅ ∣ x ) [ R ( y ′ , x ) ≤ r ] F_{\pi^*}(r \mid x) := \mathbb{P}_{y'\sim\pi^*(\cdot\mid x)}[R(y',x) \le r] F π ∗ ( r ∣ x ) := P y ′ ∼ π ∗ ( ⋅ ∣ x ) [ R ( y ′ , x ) ≤ r ]
The optimal policy identified above remains implicit because the CDF F π ∗ F_{\pi^*} F π ∗ π ∗ \pi^* π ∗ F π ∗ F_{\pi^*} F π ∗ F π r e f F_{\pi_{\mathrm{ref}}} F π ref § 15 . All derivations and regularity conditions are explained in Appendix § 14 .
Following the derivation for pass@k, a natural question is whether we can also identify a task-dependent β e f f \beta_{\mathrm{eff}} β eff
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) β e f f ( R ( y , x ) , x ) ) for β e f f ( r , x ) = β k r ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u . (9) \pi^*(y \mid x) = \frac{1}{Z(x)}\pi_{\mathrm{ref}}(y \mid x)\exp\left(\frac{R(y,x)}{\beta_{\mathrm{eff}}(R(y,x),x)}\right) \quad\text{for } \beta_{\mathrm{eff}}(r,x) = \frac{\beta}{k}\,r\int_{-\infty}^{r} F_{\pi^*}(u \mid x)^{k-1}\,\mathrm{d}u. \tag*{(9)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β eff ( R ( y , x ) , x ) R ( y , x ) ) for β eff ( r , x ) = k β r ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u . (9) Note that, to ensure well-definedness, one can assume R ( y , x ) > 0 R(y, x) > 0 R ( y , x ) > 0 (1) and (3) are invariant to constant reward shifts.
(9) provides an interpretation analogous to the binary-reward case: max@k replaces constant regularization with a strength that adapts to task complexity. In the max@k case, however, the effective strength also varies with the reward of the individual completion. As an illustrative example, Figure 1 b uses two synthetic reference reward distributions on the [ 0 , 1 ] [0, 1] [ 0 , 1 ] β e f f ( R ( y , x ) , x ) \beta_{\mathrm{eff}}(R(y, x), x) β eff ( R ( y , x ) , x ) β e f f \beta_{\mathrm{eff}} β eff
Corollary 4.6. max@k as prompt-adaptive reward transformation. Optimizing the ordinary single-reward objective J \mathcal{J} J (1) with the transformed reward
R ~ ( y , x ) = S x ( R ( y , x ) ) for S x ( r ) = k ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u (10) \widetilde{R}(y,x)=S_x(R(y,x))\text{ for }S_x(r)=k\int_{-\infty}^{r}F_{\pi_*}(u\mid x)^{k-1}\,\mathrm{d}u \tag*{(10)} R ( y , x ) = S x ( R ( y , x )) for S x ( r ) = k ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u (10) results in the max@k optimal policy derived in Proposition 4.5.
The transformed-reward view shows the same prompt-adaptive effect. In the illustration in Figure 1 c, the same raw reward is assigned a larger transformed value on the difficult task than on the easy task, providing a complementary reward-space view of how max@k can be used to control sharpening.
Explicit closed-form optimal policy for win-rate reward. As discussed, the optimal policy in Eq. (8) remains implicit because it depends on the unknown CDF F π ∗ F_{\pi_*} F π ∗ F r e f ( r ∣ x ) = P y ′ ∼ π r e f ( ⋅ ∣ x ) [ R ( y ′ , x ) ≤ r ] F_{\mathrm{ref}}(r\mid x)=\mathbb{P}_{y'\sim\pi_{\mathrm{ref}}(\cdot\mid x)}[R(y',x)\le r] F ref ( r ∣ x ) = P y ′ ∼ π ref ( ⋅ ∣ x ) [ R ( y ′ , x ) ≤ r ]
Proposition 4.7. Explicit optimal policy for win-rate max@k. For a continuous reference reward distribution with CDF F r e f ( r ∣ x ) F_{\mathrm{ref}}(r\mid x) F ref ( r ∣ x ) J k \mathcal{J}_k J k F r e f ( R ( y , x ) ∣ x ) F_{\mathrm{ref}}(R(y,x)\mid x) F ref ( R ( y , x ) ∣ x )
π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) exp ( S x q ( R ( y , x ) ) β ) (11) \pi_q^*(y\mid x)=\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{S_x^q(R(y,x))}{\beta}\right) \tag*{(11)} π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) exp ( β S x q ( R ( y , x )) ) (11) where
S x q ( r ) = β log ( C + 1 β Ψ − 1 ( F r e f ( r ∣ x ) k ) ) , Ψ ( z ) = ∫ 0 z d u C + 1 β u k , and C s.t. Ψ ( 1 ) = 1. (12) S_x^q(r)=\beta\log\left(C+\frac{1}{\beta}\Psi^{-1}\left(F_{\mathrm{ref}}(r\mid x)^k\right)\right),\qquad\Psi(z)=\int_0^z\frac{\mathrm{d}u}{C+\frac{1}{\beta}u^k},\qquad\text{and }C\text{ s.t. }\Psi(1)=1. \tag*{(12)} S x q ( r ) = β log ( C + β 1 Ψ − 1 ( F ref ( r ∣ x ) k ) ) , Ψ ( z ) = ∫ 0 z C + β 1 u k d u , and C s.t. Ψ ( 1 ) = 1. (12) For a proof, see Appendix § 17 . The scalar C C C Ψ \Psi Ψ k k k β \beta β S x q S_x^q S x q F r e f ( r ∣ x ) F_{\mathrm{ref}}(r\mid x) F ref ( r ∣ x ) p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x ) § 22 for an error analysis). Thus, the adaptive-reward principle of Corollary 4.6 admits a directly computable realization for win-rate-based optimal policies.
KAPPO: @k Adaptive Pressure for Policy Optimization
We introduce @k Adaptive Pressure for Policy Optimization (KAPPO), a framework for adapting existing @1 policy-optimization methods to @k settings. The parameter k k k § 4 . The required prompt-dependent transformations and coefficients are estimated once before training at negligible cost and reused across runs; we call this stage reference calibration (see Appendix § 21 ). We instantiate both interfaces below to demonstrate KAPPO’s generality and effectiveness across several popular RLRV optimizers.
Regularization-Side Adaptation: @k-DPO
The regularization-side interface is particularly useful for methods that do not directly optimize scalar rewards, most prominently DPO ((Rafailov et al., 2023)). We apply it in the binary-reward setting,2 β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) (5) . We call the resulting method @k-DPO:
L @ k -DPO ( θ ) = − E ( x , y w , y l ) ∼ D p r e f log σ ( β e f f ( x ) [ log π θ ( y w ∣ x ) π r e f ( y w ∣ x ) − log π θ ( y l ∣ x ) π r e f ( y l ∣ x ) ] ) . (13) \mathcal{L}_{@k\text{-DPO}}(\theta)=-\mathbb{E}_{(x,y_w,y_l)\sim D_{\mathrm{pref}}}\log\sigma\left(\beta_{\mathrm{eff}}(x)\left[\log\frac{\pi_\theta(y_w\mid x)}{\pi_{\mathrm{ref}}(y_w\mid x)}-\log\frac{\pi_\theta(y_l\mid x)}{\pi_{\mathrm{ref}}(y_l\mid x)}\right]\right). \tag*{(13)} L @ k -DPO ( θ ) = − E ( x , y w , y l ) ∼ D pref log σ ( β eff ( x ) [ log π ref ( y w ∣ x ) π θ ( y w ∣ x ) − log π ref ( y l ∣ x ) π θ ( y l ∣ x ) ] ) . (13) Before training, we perform reference calibration by generating a fixed bank of M M M x x x p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x ) (6) for p ∗ ( x ) p^*(x) p ∗ ( x ) β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) β e f f \beta_{\mathrm{eff}} β eff
Reward-Side Adaptation: @k-A*PO and @k-GRPO
To demonstrate the reward-side interface, we consider two complementary instantiations. First, we choose A ∗ A^* A ∗ [3 ] , a method that uses only a single online completion per prompt. This choice isolates a central implication of our analysis: online tuples of k k k
@k-A*PO. A ∗ A^* A ∗
L A ∗ P O ( θ ) = E x ∼ D , y ∼ π o l d ( ⋅ ∣ x ) [ ( R ( y , x ) − β log Z ( x ) − β log π θ ( y ∣ x ) π r e f ( y ∣ x ) ) 2 ] , (14) \mathcal{L}_{A^*\mathrm{PO}}(\theta)=\mathbb{E}_{x\sim D,\,y\sim\pi_{\mathrm{old}}(\cdot\mid x)}\left[\left(R(y,x)-\beta\log Z(x)-\beta\log\frac{\pi_\theta(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}\right)^2\right], \tag*{(14)} L A ∗ PO ( θ ) = E x ∼ D , y ∼ π old ( ⋅ ∣ x ) [ ( R ( y , x ) − β log Z ( x ) − β log π ref ( y ∣ x ) π θ ( y ∣ x ) ) 2 ] , (14) with the partition function Z ( x ) Z(x) Z ( x ) R ( y , x ) R(y,x) R ( y , x )
For binary rewards, we substitute the adaptive reward β R ( y , x ) β e f f ( x ) \frac{\beta R(y,x)}{\beta_{\mathrm{eff}}(x)} β eff ( x ) β R ( y , x ) (14) . We obtain β e f f \beta_{\mathrm{eff}} β eff Z ( x ) Z(x) Z ( x )
For continuous rewards, we apply the win-rate-based adaptive reward transformation S x q S_x^q S x q (12) and set Z ( x ) = 1 Z(x)=1 Z ( x ) = 1 F r e f ( r ∣ x ) F_{\mathrm{ref}}(r\mid x) F ref ( r ∣ x ) { y j r e f } j = 1 M \{y_j^{\mathrm{ref}}\}_{j=1}^{M} { y j ref } j = 1 M F ^ r e f ( r ∣ x ) = 1 M ∑ j = 1 M I [ R ( y j r e f , x ) ≤ r ] \widehat{F}_{\mathrm{ref}}(r\mid x)=\frac{1}{M}\sum_{j=1}^{M}\mathbb{I}[R(y_j^{\mathrm{ref}},x)\le r] F ref ( r ∣ x ) = M 1 ∑ j = 1 M I [ R ( y j ref , x ) ≤ r ] x x x
@k-GRPO. Here, we construct @k-GRPO for the continuous-reward setting. For the corresponding binary-reward adaptation as well as additional properties of GRPO instantiations, see Appendix § 18 . Given a group of n n n R i = R ( y i , x ) R_i=R(y_i,x) R i = R ( y i , x ) S x q ( R i ) S_x^q(R_i) S x q ( R i )
L @ k -GRPO ( θ ) = − E x ∼ D , y 1 , … , y n ∼ π o l d ( ⋅ ∣ x ) [ 1 n ∑ i = 1 n 1 ∣ y i ∣ ∑ t u i , t c l i p ( θ ) ] , (15) \mathcal{L}_{@k\text{-GRPO}}(\theta)=-\mathbb{E}_{x\sim D,\,y_1,\ldots,y_n\sim\pi_{\mathrm{old}}(\cdot\mid x)}\left[\frac{1}{n}\sum_{i=1}^{n}\frac{1}{|y_i|}\sum_t u_{i,t}^{\mathrm{clip}}(\theta)\right], \tag*{(15)} L @ k -GRPO ( θ ) = − E x ∼ D , y 1 , … , y n ∼ π old ( ⋅ ∣ x ) [ n 1 i = 1 ∑ n ∣ y i ∣ 1 t ∑ u i , t clip ( θ ) ] , (15) u i , t c l i p ( θ ) = min ( π θ ( y i , t ∣ y i , < t , x ) π o l d ( y i , t ∣ y i , < t , x ) A ^ i k , clip ( π θ ( y i , t ∣ y i , < t , x ) π o l d ( y i , t ∣ y i , < t , x ) , 1 − ϵ , 1 + ϵ ) A ^ i k ) , A ^ i k = S x q ( R i ) − 1 n ∑ j = 1 n S x q ( R j ) std j = 1 , … , n [ S x q ( R j ) ] . u_{i,t}^{\mathrm{clip}}(\theta)=\min\left(\frac{\pi_\theta(y_{i,t}\mid y_{i,<t},x)}{\pi_{\mathrm{old}}(y_{i,t}\mid y_{i,<t},x)}\widehat{A}_i^k,\operatorname{clip}\left(\frac{\pi_\theta(y_{i,t}\mid y_{i,<t},x)}{\pi_{\mathrm{old}}(y_{i,t}\mid y_{i,<t},x)},1-\epsilon,1+\epsilon\right)\widehat{A}_i^k\right),\quad\widehat{A}_i^k=\frac{S_x^q(R_i)-\frac{1}{n}\sum_{j=1}^{n}S_x^q(R_j)}{\operatorname{std}_{j=1,\ldots,n}[S_x^q(R_j)]}. u i , t clip ( θ ) = min ( π old ( y i , t ∣ y i , < t , x ) π θ ( y i , t ∣ y i , < t , x ) A i k , clip ( π old ( y i , t ∣ y i , < t , x ) π θ ( y i , t ∣ y i , < t , x ) , 1 − ϵ , 1 + ϵ ) A i k ) , A i k = std j = 1 , … , n [ S x q ( R j )] S x q ( R i ) − n 1 ∑ j = 1 n S x q ( R j ) . @k-GRPO uses the same one-time reference calibration as the continuous @k-A*PO instantiation. Apart from the reward transformation S x q S_x^q S x q
Experiments
Setup
Tasks and scope. We evaluate KAPPO on mathematical reasoning with binary rewards (correctness) and code generation with continuous rewards (test case pass rate). We investigate all instantiations of KAPPO presented in Section 5, but note that KAPPO could be applied to other algorithms.
Model selection. Two principles guide our choice of initial checkpoints. First, we avoid models that have undergone reasoning-focused preference optimization or RL. This reflects the practical setting in which RL for reasoning is applied directly after supervised fine-tuning (SFT), and separates the effect of our training from reasoning behavior acquired via an undocumented prior curriculum. Second, we aim to minimize possible contamination of the training and evaluation datasets. These criteria lead us to OLMo-3-7B-Instruct-SFT and Qwen3-4B-Base. The former is an SFT-only, pre-RL checkpoint with a fully open training pipeline including explicit data decontamination procedures ((Team Olmo et al., 2025)). For a second model family, we use the Qwen pre-trained rather than a post-trained checkpoint to limit exposure to undocumented post-training data, although its pretraining corpus cannot be fully audited ((Yang et al., 2025)). Hyperparameters and checkpoints are selected on a validation split; complete experimental details are provided in Appendix § 19 .
Baselines. As elaborated in § 5 , KAPPO can be applied to a wide range of RL optimizers by adjusting the effective regularization or transforming the reward. Since applying KAPPO with k = 1 k = 1 k = 1 k = 1 k = 1 k = 1 A ⋆ A^\star A ⋆ n ≥ k n \ge k n ≥ k
Binary Rewards: Mathematical Reasoning
Data and evaluation. We design the train–validation–test split to mirror a common model development setting: training data and validation benchmarks are closely coupled during development, while the final test consists of newer and often more difficult problems, appearing after model release. To closely match this setting, we perform RL on the training split of MATH and use its test split for in-distribution validation ([8 ] ). Final evaluation is performed on AIME 2026 [16 ] , a markedly harder dataset whose problems predominantly exceed MATH’s highest difficulty tier. AIME 2026 also postdates both model initialization checkpoints, ruling out direct contamination of our test benchmark. A completion receives reward one if its extracted final answer is correct and zero otherwise, i.e., we consider binary rewards.
Controlling sharpening by adjusting k k k Figure 2 plots pass@k e v a l k_{\mathrm{eval}} k eval k e v a l k_{\mathrm{eval}} k eval exclusively at small k e v a l k_{\mathrm{eval}} k eval k e v a l k_{\mathrm{eval}} k eval A ⋆ A^\star A ⋆ k e v a l k_{\mathrm{eval}} k eval k > 1 k > 1 k > 1 k = 1 k = 1 k = 1 Figure 2 b exhibits the same effect for the regularization-side interface. Standard DPO produces gains concentrated at small evaluation budgets, whereas @k-DPO trained with k = 2 k = 2 k = 2 k e v a l k_{\mathrm{eval}} k eval
Figure 2. Figure 2 : pass@k e v a l k_{\mathrm{eval}} k eval k e v a l k_{\mathrm{eval}} k eval k > 1 k > 1 k > 1 ⋆ ^\star ⋆ k ∈ { 1 , 2 , 4 } k \in\{1, 2, 4\} k ∈ { 1 , 2 , 4 } k ∈ { 1 , 2 } k \in\{1, 2\} k ∈ { 1 , 2 }
@k adaptation counteracts excessive sharpening and shifts optimization pressure from already-easy tasks toward difficult tasks where successful behavior remains rare.
Agreement with dedicated multi-completion pass@k optimization. We next verify that a KAPPO adaptation of a single-completion optimizer can match the performance of a state-of-the-art dedicated multi-completion pass@k algorithm. To that end, we compare @k-APO with PKPO for k ∈ { 2 , 4 } k \in\{2,4\} k ∈ { 2 , 4 } n = k n = k n = k PO achieves validation performance comparable to PKPO (see Figure 3 ) while requiring only a fraction of the on-policy rollouts, reducing training cost. This supports our central prediction: @k optimization can be achieved without relying on online k k k
Figure 3. MATH Level 5 validation performance of the starting checkpoints, @k-APO checkpoints, and PKPO checkpoints. @k-A PO matches the performance of PKPO at a fraction of the on-policy training rollouts, requiring only a single-time reference calibration at ≤ 5 % \le5\% ≤ 5%
Continuous Rewards: Code Generation
Data and evaluation. We follow the same data selection strategy as in the binary-reward experiments: training and validation use closely related data, while final evaluation uses a newer and substantially harder problem source to approximate the distribution shift encountered after model deployment. We conduct our max@k experiments using Qwen3-4B-Base. Specifically, we train on the training split of LeetCodeDataset ((Xia et al., 2025)) and use its test split for validation. The two splits share the same source and task format, with a temporal cutoff on July 2024 ((Xia et al., 2025)). Each completion receives a continuous reward in [ 0 , 1 ] [0,1] [ 0 , 1 ] [10 ] ) released between January and April 2025, limiting the opportunity for their inclusion in Qwen3 pretraining (released April 28th). Nevertheless, since Qwen’s pretraining cutoff is not documented, contamination cannot be ruled out completely.
Controlling sharpening by adjusting k k k We observe the same qualitative effect as in the binary-reward experiments. Figure 4 compares the initial checkpoint with win-rate-based @k-A*PO and @k-GRPO trained with k ∈ { 1 , 2 , 4 } k \in\{1,2,4\} k ∈ { 1 , 2 , 4 } k = 1 k = 1 k = 1 k > 1 k > 1 k > 1 k = 1 k = 1 k = 1
Figure 4. max@k_eval on the hard subset of LiveCodeBench v6 for Qwen3-4B. Win-rate-based @k-A*PO (left) and @k-GRPO (right) mitigate excessive sharpening for k > 1 k > 1 k > 1
Conclusion
We derive closed-form optimal policies for KL-regularized pass@k and max@k optimization, proving that @k optimization redistributes pressure according to task difficulty. For binary rewards, this pressure can be expressed equivalently through prompt-adaptive reward scaling or regularization strength. For continuous rewards, it appears as an adaptive reward transformation or a completion-dependent regularization strength, with a fully explicit construction for win-rate-based rewards. Leveraging these insights, we introduce KAPPO, a framework that adapts existing @1 policy-optimization methods to @k settings without requiring online multi-completion rollouts. In particular, our @k-A⋆PO instantiation requires only one new completion per prompt, yet achieves performance comparable to a SoTA multi-completion baseline. Validated across mathematical reasoning and code generation, @k-A⋆PO, @k-DPO, and @k-GRPO successfully control sharpening, confirming that diverse reasoning can be effectively encouraged via prompt-adaptive optimization pressure.
Limitations and future work. Our reference calibration relies on a static bank of M M M p ∗ p^* p ∗ S x q S_x^q S x q S x q S_x^q S x q F π ∗ F_{\pi^*} F π ∗ § 15 ) to future work. Lastly, our win-rate solution assumes a continuous reference reward distribution; its application to discrete rewards with ties is an approximation.
Acknowledgements
This work was supported as part of the Swiss AI Initiative by a grant from the Swiss National Supercomputing Centre (CSCS) under project ID infra01 on Alps. It was partially funded by an unrestricted gift from Coefficient Giving and by the Swiss National Science Foundation under grant number 212111. We are grateful to Robert West for his supervision and support.
AI Use Statement
We used generative AI tools to check proofs, improve their exposition, edit the manuscript, and assist with literature search and reference checking. AI also supported the implementation of methods, training infrastructure, evaluation, and plotting, following author-specified designs and guidance. For data processing, AI assisted with deduplication and analysis of dataset diversity, task difficulty, and metadata. The central research ideas and essential proof steps were developed by the authors. All experimental decisions, including experimental design, hyperparameters, comparisons, and all interpretations of results were determined solely by the authors. The authors reviewed the AI-assisted code, verified mathematical claims and proofs, checked cited sources and their relevance, and checked figures against experimental outputs. We take responsibility for the final content of this work, including all AI-assisted text, claims, and artifacts.
Ethics Statement
We improve understanding of pass@k and max@k optimization and introduce a method to adapt existing reinforcement learning methods from @1 to @k. As a result, post-training of frontier models may be accelerated, which can result in all kinds of secondary effects, both economic and societal. Our experiments use only publicly available models and datasets in accordance with their licenses.
Reproducibility Statement
All considered datasets and models are publicly available. Instead of proposing a novel RL optimization method with many degrees of freedom to be specified, we identify a minimalist intervention to turn any standard @1 RL method into an @k method. We provide pseudocode for reference calibration with binary and continuous win-rate rewards in Appendix § 21 .§ 1 . As a result, all of our experiments can be easily reproduced. For completeness, we describe the exact experimental details in Appendix § 19 . Our entire code base is available in the supplementary materials.
A Variational Optimality Lemma
We first establish a general result used in the derivations of both optimal policies. Since all objectives considered in this work decompose over prompts, we fix a prompt x x x
S x = { y : π r e f ( y ∣ x ) > 0 } (16) S_x = \{y : \pi_{\mathrm{ref}}(y \mid x) > 0\} \tag*{(16)} S x = { y : π ref ( y ∣ x ) > 0 } (16) denote the support of the reference policy. Any policy assigning positive probability outside S x S_x S x S x S_x S x
Consider an objective of the form
J x ( π ) = U x ( π ) − β KL ( π ( ⋅ ∣ x ) ∥ π r e f ( ⋅ ∣ x ) ) , (17) \mathcal{J}_x(\pi) = U_x(\pi) - \beta\operatorname{KL}\left(\pi(\cdot\mid x)\Vert\pi_{\mathrm{ref}}(\cdot\mid x)\right), \tag*{(17)} J x ( π ) = U x ( π ) − β KL ( π ( ⋅ ∣ x ) ∥ π ref ( ⋅ ∣ x ) ) , (17) where U x U_x U x β > 0 \beta> 0 β > 0
For a feasible policy π \pi π
π ε ( y ∣ x ) = π ( y ∣ x ) + ε h ( y ∣ x ) , (18) \pi_\varepsilon(y \mid x) = \pi(y \mid x) + \varepsilon h(y \mid x), \tag*{(18)} π ε ( y ∣ x ) = π ( y ∣ x ) + ε h ( y ∣ x ) , (18) where
∑ y ∈ S x ∣ h ( y ∣ x ) ∣ < ∞ , ∑ y ∈ S x h ( y ∣ x ) = 0 , (19) \sum_{y\in S_x} \lvert h(y \mid x)\rvert< \infty,\qquad\sum_{y\in S_x} h(y \mid x) = 0, \tag*{(19)} y ∈ S x ∑ ∣ h ( y ∣ x )∣ < ∞ , y ∈ S x ∑ h ( y ∣ x ) = 0 , (19) and there exists ε 0 > 0 \varepsilon_0 > 0 ε 0 > 0 π ε ( y ∣ x ) ≥ 0 \pi_\varepsilon(y \mid x) \ge0 π ε ( y ∣ x ) ≥ 0 ε ∈ [ 0 , ε 0 ] \varepsilon\in[0,\varepsilon_0] ε ∈ [ 0 , ε 0 ] y ∈ S x y \in S_x y ∈ S x
We assume that the corresponding one-sided directional derivative of U x U_x U x
d d ε U x ( π + ε h ) ∣ ε = 0 + = ∑ y ∈ S x h ( y ∣ x ) A π ( y , x ) (20) \left.\frac{d}{d\varepsilon}U_x(\pi+\varepsilon h)\right|_{\varepsilon=0+} = \sum_{y\in S_x} h(y \mid x)A_\pi(y,x) \tag*{(20)} d ε d U x ( π + ε h ) ε = 0 + = y ∈ S x ∑ h ( y ∣ x ) A π ( y , x ) (20) for some score A π ( y , x ) A_\pi(y,x) A π ( y , x )
We additionally assume that an optimal policy exists and has full support on S x S_x S x ε \varepsilon ε O ( ε ) O(\varepsilon) O ( ε ) π ( y ∣ x ) = 0 \pi(y \mid x) = 0 π ( y ∣ x ) = 0 y ∈ S x y \in S_x y ∈ S x ε \varepsilon ε y y y
− β ε log ε = β ε log 1 ε . (21) -\beta\varepsilon\log\varepsilon= \beta\varepsilon\log\frac{1}{\varepsilon}. \tag*{(21)} − β ε log ε = β ε log ε 1 . (21) Although this quantity tends to zero as ε → 0 \varepsilon\to0 ε → 0 ε \varepsilon ε ε \varepsilon ε O ( ε ) O(\varepsilon) O ( ε ) S x S_x S x
Lemma A.1 (Variational form of a KL-regularized optimum). Under the assumptions above, any maximizer π ∗ \pi^* π ∗ (17) satisfies
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( A π ∗ ( y , x ) β ) , (22) \pi^*(y \mid x) = \frac{1}{Z(x)}\pi_{\mathrm{ref}}(y \mid x)\exp\left(\frac{A_{\pi^*}(y,x)}{\beta}\right), \tag*{(22)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β A π ∗ ( y , x ) ) , (22) where
Z ( x ) = ∑ y ∈ S x π r e f ( y ∣ x ) exp ( A π ∗ ( y , x ) β ) . (23) Z(x) = \sum_{y\in S_x}\pi_{\mathrm{ref}}(y \mid x)\exp\left(\frac{A_{\pi^*}(y,x)}{\beta}\right). \tag*{(23)} Z ( x ) = y ∈ S x ∑ π ref ( y ∣ x ) exp ( β A π ∗ ( y , x ) ) . (23) Conversely, any feasible policy satisfying Eq. (22) is the unique maximizer of Eq. (17) .
Proof. For compactness, define
B π ∗ ( y , x ) = A π ∗ ( y , x ) − β ( log π ∗ ( y ∣ x ) π r e f ( y ∣ x ) + 1 ) . (24) B_{\pi^*}(y,x) = A_{\pi^*}(y,x) - \beta\left(\log\frac{\pi^*(y \mid x)}{\pi_{\mathrm{ref}}(y \mid x)} + 1\right). \tag*{(24)} B π ∗ ( y , x ) = A π ∗ ( y , x ) − β ( log π ref ( y ∣ x ) π ∗ ( y ∣ x ) + 1 ) . (24) Consider any perturbation h h h π ∗ + ε h \pi^*+\varepsilon h π ∗ + ε h ε \varepsilon ε (20) to h h h − h -h − h ε = 0 \varepsilon=0 ε = 0 J x ( π ∗ + ε h ) \mathcal{J}_x(\pi^*+\varepsilon h) J x ( π ∗ + ε h )
d d ε J x ( π ∗ + ε h ) ∣ ε = 0 = 0. (25) \left.\frac{d}{d\varepsilon}\mathcal{J}_x(\pi^*+\varepsilon h)\right|_{\varepsilon=0}=0. \tag*{(25)} d ε d J x ( π ∗ + ε h ) ε = 0 = 0. (25) The directional derivative of the KL term is
d d ε K L ( π ∗ + ε h ∥ π r e f ) ∣ ε = 0 = ∑ y ∈ S x h ( y ∣ x ) ( log π ∗ ( y ∣ x ) π r e f ( y ∣ x ) + 1 ) . (26) \left.\frac{d}{d\varepsilon}\mathrm{KL}\left(\pi^*+\varepsilon h\middle\|\pi_{\mathrm{ref}}\right)\right|_{\varepsilon=0}
=\sum_{y\in\mathcal{S}_x}h(y\mid x)\left(\log\frac{\pi^*(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}+1\right). \tag*{(26)} d ε d KL ( π ∗ + ε h ∥ π ref ) ε = 0 = y ∈ S x ∑ h ( y ∣ x ) ( log π ref ( y ∣ x ) π ∗ ( y ∣ x ) + 1 ) . (26) Together with Eq. (20) , this gives
∑ y ∈ S x h ( y ∣ x ) B π ∗ ( y , x ) = 0. (27) \sum_{y\in\mathcal{S}_x}h(y\mid x)B_{\pi^*}(y,x)=0. \tag*{(27)} y ∈ S x ∑ h ( y ∣ x ) B π ∗ ( y , x ) = 0. (27) Now consider any two completions y , z ∈ S x y,z\in\mathcal{S}_x y , z ∈ S x
h = e y − e z , (28) h=e_y-e_z, \tag*{(28)} h = e y − e z , (28) where e y e_y e y y y y π ∗ \pi^* π ∗ ε \varepsilon ε (27) yields
B π ∗ ( y , x ) − B π ∗ ( z , x ) = 0. (29) B_{\pi^*}(y,x)-B_{\pi^*}(z,x)=0. \tag*{(29)} B π ∗ ( y , x ) − B π ∗ ( z , x ) = 0. (29) Because y y y z z z B π ∗ ( y , x ) B_{\pi^*}(y,x) B π ∗ ( y , x ) y y y S x \mathcal{S}_x S x c ( x ) c(x) c ( x )
A π ∗ ( y , x ) − β ( log π ∗ ( y ∣ x ) π r e f ( y ∣ x ) + 1 ) = c ( x ) (30) A_{\pi^*}(y,x)-\beta\left(\log\frac{\pi^*(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}+1\right)=c(x) \tag*{(30)} A π ∗ ( y , x ) − β ( log π ref ( y ∣ x ) π ∗ ( y ∣ x ) + 1 ) = c ( x ) (30) for every y ∈ S x y\in\mathcal{S}_x y ∈ S x
Rearranging Eq. (30) gives
π ∗ ( y ∣ x ) = π r e f ( y ∣ x ) exp ( A π ∗ ( y , x ) β − c ( x ) β − 1 ) . (31) \pi^*(y\mid x)=\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{A_{\pi^*}(y,x)}{\beta}-\frac{c(x)}{\beta}-1\right). \tag*{(31)} π ∗ ( y ∣ x ) = π ref ( y ∣ x ) exp ( β A π ∗ ( y , x ) − β c ( x ) − 1 ) . (31) The final two terms are independent of y y y ∑ y π ∗ ( y ∣ x ) = 1 \sum_y\pi^*(y\mid x)=1 ∑ y π ∗ ( y ∣ x ) = 1 (22) and (23) .
It remains to show that this stationary policy is the unique global maximizer. Let π \pi π
h ( y ∣ x ) = π ( y ∣ x ) − π ∗ ( y ∣ x ) . (32) h(y\mid x)=\pi(y\mid x)-\pi^*(y\mid x). \tag*{(32)} h ( y ∣ x ) = π ( y ∣ x ) − π ∗ ( y ∣ x ) . (32) This perturbation satisfies
∑ y ∣ h ( y ∣ x ) ∣ ≤ 2 , ∑ y h ( y ∣ x ) = 0 , (33) \sum_y\lvert h(y\mid x)\rvert\le2,\qquad\sum_y h(y\mid x)=0, \tag*{(33)} y ∑ ∣ h ( y ∣ x )∣ ≤ 2 , y ∑ h ( y ∣ x ) = 0 , (33) and
π ∗ + ε h = ( 1 − ε ) π ∗ + ε π (34) \pi^*+\varepsilon h=(1-\varepsilon)\pi^*+\varepsilon\pi \tag*{(34)} π ∗ + ε h = ( 1 − ε ) π ∗ + ε π (34) is a valid policy for every ε ∈ [ 0 , 1 ] \varepsilon\in[0,1] ε ∈ [ 0 , 1 ]
Define
g ( ε ) = U x ( π ∗ + ε h ) . (35) g(\varepsilon)=U_x(\pi^*+\varepsilon h). \tag*{(35)} g ( ε ) = U x ( π ∗ + ε h ) . (35) Because U x U_x U x g g g [ 0 , 1 ] [0,1] [ 0 , 1 ]
g ( 1 ) − g ( 0 ) ≤ g ′ ( 0 + ) . (36) g(1)-g(0)\le g'(0^+). \tag*{(36)} g ( 1 ) − g ( 0 ) ≤ g ′ ( 0 + ) . (36) Using Eq. (20) , this gives
U x ( π ) − U x ( π ∗ ) ≤ ∑ y ∈ S x ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) A π ∗ ( y , x ) . (37) U_x(\pi)-U_x(\pi^*)\le\sum_{y\in\mathcal{S}_x}\left(\pi(y\mid x)-\pi^*(y\mid x)\right)A_{\pi^*}(y,x). \tag*{(37)} U x ( π ) − U x ( π ∗ ) ≤ y ∈ S x ∑ ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) A π ∗ ( y , x ) . (37) The KL difference admits the exact decomposition
KL ( π ∥ π r e f ) − KL ( π ∗ ∥ π r e f ) = ∑ y ∈ S x ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) ( log π ∗ ( y ∣ x ) π r e f ( y ∣ x ) + 1 ) + KL ( π ∥ π ∗ ) . (38) \operatorname{KL}(\pi\|\pi_{\mathrm{ref}})-\operatorname{KL}(\pi^*\|\pi_{\mathrm{ref}})
= \sum_{y\in S_x}\left(\pi(y\mid x)-\pi^*(y\mid x)\right)\left(\log\frac{\pi^*(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}+1\right)+\operatorname{KL}(\pi\|\pi^*). \tag*{(38)} KL ( π ∥ π ref ) − KL ( π ∗ ∥ π ref ) = y ∈ S x ∑ ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) ( log π ref ( y ∣ x ) π ∗ ( y ∣ x ) + 1 ) + KL ( π ∥ π ∗ ) . (38) To verify this identity, expand
log π ( y ∣ x ) π r e f ( y ∣ x ) = log π ( y ∣ x ) π ∗ ( y ∣ x ) + log π ∗ ( y ∣ x ) π r e f ( y ∣ x ) (39) \log\frac{\pi(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)}
= \log\frac{\pi(y\mid x)}{\pi^*(y\mid x)}+\log\frac{\pi^*(y\mid x)}{\pi_{\mathrm{ref}}(y\mid x)} \tag*{(39)} log π ref ( y ∣ x ) π ( y ∣ x ) = log π ∗ ( y ∣ x ) π ( y ∣ x ) + log π ref ( y ∣ x ) π ∗ ( y ∣ x ) (39) inside the first KL term and use ∑ y ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) = 0 \sum_y(\pi(y\mid x)-\pi^*(y\mid x))=0 ∑ y ( π ( y ∣ x ) − π ∗ ( y ∣ x )) = 0
Combining Eqs. (37) and (38) yields
J x ( π ) − J x ( π ∗ ) ≤ ∑ y ∈ S x ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) B π ∗ ( y , x ) − β KL ( π ∥ π ∗ ) . (40) \begin{aligned}
J_x(\pi)-J_x(\pi^*) &\le\sum_{y\in S_x}\left(\pi(y\mid x)-\pi^*(y\mid x)\right)B_{\pi^*}(y,x)\\
&\quad-\beta\operatorname{KL}(\pi\|\pi^*). \tag*{(40)}
\end{aligned} J x ( π ) − J x ( π ∗ ) ≤ y ∈ S x ∑ ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) B π ∗ ( y , x ) − β KL ( π ∥ π ∗ ) . (40) By Eq. (30) , B π ∗ ( y , x ) = c ( x ) B_{\pi^*}(y,x)=c(x) B π ∗ ( y , x ) = c ( x ) y y y
∑ y ∈ S x ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) B π ∗ ( y , x ) = c ( x ) ∑ y ∈ S x ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) = 0. (41) \sum_{y\in S_x}\left(\pi(y\mid x)-\pi^*(y\mid x)\right)B_{\pi^*}(y,x)
=c(x)\sum_{y\in S_x}\left(\pi(y\mid x)-\pi^*(y\mid x)\right)=0. \tag*{(41)} y ∈ S x ∑ ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) B π ∗ ( y , x ) = c ( x ) y ∈ S x ∑ ( π ( y ∣ x ) − π ∗ ( y ∣ x ) ) = 0. (41) Consequently,
J x ( π ) − J x ( π ∗ ) ≤ − β KL ( π ∥ π ∗ ) ≤ 0. (42) J_x(\pi)-J_x(\pi^*)\le-\beta\operatorname{KL}(\pi\|\pi^*)\le0. \tag*{(42)} J x ( π ) − J x ( π ∗ ) ≤ − β KL ( π ∥ π ∗ ) ≤ 0. (42) Equality holds only if π = π ∗ \pi=\pi^* π = π ∗ π ∗ \pi^* π ∗
The score A π ( y , x ) A_\pi(y,x) A π ( y , x ) x x x y y y (20) because every permitted perturbation has zero total mass, and it is absorbed into Z ( x ) Z(x) Z ( x ) (22) . The following sections derive the relevant score for the pass@k and max@k reward functionals.
Binary-Rewards Derivation
We now apply Lemma A.1 to the pass@k objective. As before, the objective decomposes over prompts, so we fix a prompt x x x π \pi π
p π ( x ) = P y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] = ∑ y π ( y ∣ x ) R ( y , x ) . (43) p_\pi(x)=\mathbb{P}_{y\sim\pi(\cdot\mid x)}[R(y,x)=1]=\sum_y\pi(y\mid x)R(y,x). \tag*{(43)} p π ( x ) = P y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] = y ∑ π ( y ∣ x ) R ( y , x ) . (43) For binary rewards, the probability that at least one of k k k
U x ( π ) = 1 − ( 1 − p π ( x ) ) k . (44) U_x(\pi)=1-(1-p_\pi(x))^k. \tag*{(44)} U x ( π ) = 1 − ( 1 − p π ( x ) ) k . (44) The prompt-wise KL-regularized objective is therefore
J k , x ( π ) = 1 − ( 1 − p π ( x ) ) k − β KL ( π ( ⋅ ∣ x ) ∥ π r e f ( ⋅ ∣ x ) ) . (45) J_{k,x}(\pi)=1-(1-p_\pi(x))^k-\beta\operatorname{KL}\left(\pi(\cdot\mid x)\|\pi_{\mathrm{ref}}(\cdot\mid x)\right). \tag*{(45)} J k , x ( π ) = 1 − ( 1 − p π ( x ) ) k − β KL ( π ( ⋅ ∣ x ) ∥ π ref ( ⋅ ∣ x ) ) . (45) Optimal Policy
Proposition 4.1 (restated). For R ( y , x ) ∈ { 0 , 1 } R(y,x)\in\{0,1\} R ( y , x ) ∈ { 0 , 1 } J k , x J_{k,x} J k , x
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) β e f f ( x ) ) , \pi^*(y\mid x)=\frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{R(y,x)}{\beta_{\mathrm{eff}}(x)}\right), π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β eff ( x ) R ( y , x ) ) , where
β e f f ( x ) = β k ( 1 − p ∗ ( x ) ) k − 1 , Z ( x ) = 1 − p r e f ( x ) + p r e f ( x ) e 1 / β e f f ( x ) , \beta_{\mathrm{eff}}(x)=\frac{\beta}{k(1-p^*(x))^{k-1}},\qquad Z(x)=1-p_{\mathrm{ref}}(x)+p_{\mathrm{ref}}(x)e^{1/\beta_{\mathrm{eff}}(x)}, β eff ( x ) = k ( 1 − p ∗ ( x ) ) k − 1 β , Z ( x ) = 1 − p ref ( x ) + p ref ( x ) e 1/ β eff ( x ) , with p r e f ( x ) = p π r e f ( x ) p_{\mathrm{ref}}(x)=p_{\pi_{\mathrm{ref}}}(x) p ref ( x ) = p π ref ( x ) p ∗ ( x ) = p π ∗ ( x ) p^*(x)=p_{\pi^*}(x) p ∗ ( x ) = p π ∗ ( x )
Proof. We first consider 0 < p r e f ( x ) < 1 0 < p_{\mathrm{ref}}(x) < 1 0 < p ref ( x ) < 1 p ↦ 1 − ( 1 − p ) k p \mapsto1 - (1-p)^k p ↦ 1 − ( 1 − p ) k [ 0 , 1 ] [0,1] [ 0 , 1 ] k ≥ 1 k \ge1 k ≥ 1 k = 1 k = 1 k = 1 k ≥ 2 k \ge2 k ≥ 2
d 2 d p 2 [ 1 − ( 1 − p ) k ] = − k ( k − 1 ) ( 1 − p ) k − 2 ≤ 0. (46) \frac{\mathrm{d}^2}{\mathrm{d}p^2}\left[1-(1-p)^k\right] = -k(k-1)(1-p)^{k-2} \le0. \tag*{(46)} d p 2 d 2 [ 1 − ( 1 − p ) k ] = − k ( k − 1 ) ( 1 − p ) k − 2 ≤ 0. (46) Since p π ( x ) p_\pi(x) p π ( x ) π \pi π U x ( π ) U_x(\pi) U x ( π )
Consider a perturbation π ε = π + ε h \pi_\varepsilon= \pi+\varepsilon h π ε = π + ε h
p π ε ( x ) = p π ( x ) + ε ∑ y h ( y ∣ x ) R ( y , x ) . (47) p_{\pi_\varepsilon}(x) = p_\pi(x) + \varepsilon\sum_y h(y \mid x)R(y,x). \tag*{(47)} p π ε ( x ) = p π ( x ) + ε y ∑ h ( y ∣ x ) R ( y , x ) . (47) Consequently, the directional derivative of the reward functional is
d d ε U x ( π + ε h ) ∣ ε = 0 + = k ( 1 − p π ( x ) ) k − 1 ∑ y h ( y ∣ x ) R ( y , x ) . (48) \left.\frac{\mathrm{d}}{\mathrm{d}\varepsilon}U_x(\pi+\varepsilon h)\right|_{\varepsilon=0+}
= k\left(1-p_\pi(x)\right)^{k-1}\sum_y h(y \mid x)R(y,x). \tag*{(48)} d ε d U x ( π + ε h ) ε = 0 + = k ( 1 − p π ( x ) ) k − 1 y ∑ h ( y ∣ x ) R ( y , x ) . (48) Thus, the score in Eq. (20) is
A π ( y , x ) = k ( 1 − p π ( x ) ) k − 1 R ( y , x ) . (49) A_\pi(y,x) = k\left(1-p_\pi(x)\right)^{k-1}R(y,x). \tag*{(49)} A π ( y , x ) = k ( 1 − p π ( x ) ) k − 1 R ( y , x ) . (49) Applying Lemma A.1 gives
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( k β ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) ) , (50) \pi^*(y \mid x) = \frac{1}{Z(x)}\pi^{\mathrm{ref}}(y \mid x)\exp\left(\frac{k}{\beta}\left(1-p^*(x)\right)^{k-1}R(y,x)\right), \tag*{(50)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) ) , (50) where
p ∗ ( x ) = p π ∗ ( x ) . (51) p^*(x) = p_{\pi^*}(x). \tag*{(51)} p ∗ ( x ) = p π ∗ ( x ) . (51) Define
p r e f ( x ) = P y ∼ π r e f ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] . (52) p_{\mathrm{ref}}(x) = \mathbb{P}_{y\sim\pi_{\mathrm{ref}}(\cdot\mid x)}[R(y,x)=1]. \tag*{(52)} p ref ( x ) = P y ∼ π ref ( ⋅ ∣ x ) [ R ( y , x ) = 1 ] . (52) Since the reward is binary, the reference policy assigns total mass 1 − p r e f ( x ) 1-p_{\mathrm{ref}}(x) 1 − p ref ( x ) p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x ) (50) therefore gives
Z ( x ) = 1 − p r e f ( x ) + p r e f ( x ) exp ( k β ( 1 − p ∗ ( x ) ) k − 1 ) . (53) Z(x) = 1-p_{\mathrm{ref}}(x)+p_{\mathrm{ref}}(x)\exp\left(\frac{k}{\beta}\left(1-p^*(x)\right)^{k-1}\right). \tag*{(53)} Z ( x ) = 1 − p ref ( x ) + p ref ( x ) exp ( β k ( 1 − p ∗ ( x ) ) k − 1 ) . (53) Introducing
β e f f ( x ) = β k ( 1 − p ∗ ( x ) ) k − 1 , (54) \beta_{\mathrm{eff}}(x) = \frac{\beta}{k\left(1-p^*(x)\right)^{k-1}}, \tag*{(54)} β eff ( x ) = k ( 1 − p ∗ ( x ) ) k − 1 β , (54) we recover the form reported in Eq. (4) :
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) β e f f ( x ) ) . (55) \pi^*(y \mid x) = \frac{1}{Z(x)}\pi^{\mathrm{ref}}(y \mid x)\exp\left(\frac{R(y,x)}{\beta_{\mathrm{eff}}(x)}\right). \tag*{(55)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β eff ( x ) R ( y , x ) ) . (55) The solution remains implicit because its right-hand side depends on the success probability of the unknown optimal policy.
Scalar Equation and Uniqueness
It remains to determine p ∗ ( x ) p^*(x) p ∗ ( x )
Proposition 4.2 (restated). For 0 < p r e f ( x ) < 1 0 < p_{\mathrm{ref}}(x) < 1 0 < p ref ( x ) < 1 p ∗ ( x ) p^*(x) p ∗ ( x ) ( 0 , 1 ) (0,1) ( 0 , 1 )
k ( 1 − p ∗ ( x ) ) k − 1 = β log p ∗ ( x ) ( 1 − p r e f ( x ) ) p r e f ( x ) ( 1 − p ∗ ( x ) ) . k\left(1-p^*(x)\right)^{k-1} = \beta\log\frac{p^*(x)\left(1-p_{\mathrm{ref}}(x)\right)}{p_{\mathrm{ref}}(x)\left(1-p^*(x)\right)}. k ( 1 − p ∗ ( x ) ) k − 1 = β log p ref ( x ) ( 1 − p ∗ ( x ) ) p ∗ ( x ) ( 1 − p ref ( x ) ) . Proof. Summing Eq. (50) over successful completions gives
p ∗ ( x ) = p r e f ( x ) exp ( k β ( 1 − p ∗ ( x ) ) k − 1 ) 1 − p r e f ( x ) + p r e f ( x ) exp ( k β ( 1 − p ∗ ( x ) ) k − 1 ) . (56) p^*(x)=\frac{p_{\mathrm{ref}}(x)\exp\left(\frac{k}{\beta}(1-p^*(x))^{k-1}\right)}{1-p_{\mathrm{ref}}(x)+p_{\mathrm{ref}}(x)\exp\left(\frac{k}{\beta}(1-p^*(x))^{k-1}\right)}. \tag*{(56)} p ∗ ( x ) = 1 − p ref ( x ) + p ref ( x ) exp ( β k ( 1 − p ∗ ( x ) ) k − 1 ) p ref ( x ) exp ( β k ( 1 − p ∗ ( x ) ) k − 1 ) . (56) Taking the odds ratio and then the logarithm yields
k ( 1 − p ∗ ( x ) ) k − 1 = β log p ∗ ( x ) ( 1 − p r e f ( x ) ) p r e f ( x ) ( 1 − p ∗ ( x ) ) , (57) k(1-p^*(x))^{k-1}=\beta\log\frac{p^*(x)(1-p_{\mathrm{ref}}(x))}{p_{\mathrm{ref}}(x)(1-p^*(x))}, \tag*{(57)} k ( 1 − p ∗ ( x ) ) k − 1 = β log p ref ( x ) ( 1 − p ∗ ( x )) p ∗ ( x ) ( 1 − p ref ( x )) , (57) which is Eq. (6) in the main text.
To establish uniqueness, fix q = p r e f ( x ) ∈ ( 0 , 1 ) q=p_{\mathrm{ref}}(x)\in(0,1) q = p ref ( x ) ∈ ( 0 , 1 ) p ∈ ( 0 , 1 ) p\in(0,1) p ∈ ( 0 , 1 )
Φ q ( p ) = β log p ( 1 − q ) q ( 1 − p ) − k ( 1 − p ) k − 1 . (58) \Phi_q(p)=\beta\log\frac{p(1-q)}{q(1-p)}-k(1-p)^{k-1}. \tag*{(58)} Φ q ( p ) = β log q ( 1 − p ) p ( 1 − q ) − k ( 1 − p ) k − 1 . (58) Equation 57 is equivalent to Φ q ( p ) = 0 \Phi_q(p)=0 Φ q ( p ) = 0
lim p → 0 + Φ q ( p ) = − ∞ , lim p → 1 − Φ q ( p ) = + ∞ , (59) \lim_{p\to0^+}\Phi_q(p)=-\infty,\qquad\lim_{p\to1^-}\Phi_q(p)=+\infty, \tag*{(59)} p → 0 + lim Φ q ( p ) = − ∞ , p → 1 − lim Φ q ( p ) = + ∞ , (59) and
Φ q ′ ( p ) = β p ( 1 − p ) + k ( k − 1 ) ( 1 − p ) k − 2 > 0. (60) \Phi_q'(p)=\frac{\beta}{p(1-p)}+k(k-1)(1-p)^{k-2}>0. \tag*{(60)} Φ q ′ ( p ) = p ( 1 − p ) β + k ( k − 1 ) ( 1 − p ) k − 2 > 0. (60) Hence, Φ q \Phi_q Φ q − ∞ -\infty − ∞ + ∞ +\infty + ∞ ( 0 , 1 ) (0,1) ( 0 , 1 )
Φ q ( q ) = − k ( 1 − q ) k − 1 < 0 , (61) \Phi_q(q)=-k(1-q)^{k-1}<0, \tag*{(61)} Φ q ( q ) = − k ( 1 − q ) k − 1 < 0 , (61) which implies that the unique solution satisfies
p ∗ ( x ) > p r e f ( x ) . (62) p^*(x)>p_{\mathrm{ref}}(x). \tag*{(62)} p ∗ ( x ) > p ref ( x ) . (62) For k = 1 k=1 k = 1 (57) reduces to
1 = β log p ∗ ( x ) ( 1 − p r e f ( x ) ) p r e f ( x ) ( 1 − p ∗ ( x ) ) , (63) 1=\beta\log\frac{p^*(x)(1-p_{\mathrm{ref}}(x))}{p_{\mathrm{ref}}(x)(1-p^*(x))}, \tag*{(63)} 1 = β log p ref ( x ) ( 1 − p ∗ ( x )) p ∗ ( x ) ( 1 − p ref ( x )) , (63) whose solution is
p ∗ ( x ) = p r e f ( x ) e 1 / β 1 − p r e f ( x ) + p r e f ( x ) e 1 / β . (64) p^*(x)=\frac{p_{\mathrm{ref}}(x)e^{1/\beta}}{1-p_{\mathrm{ref}}(x)+p_{\mathrm{ref}}(x)e^{1/\beta}}. \tag*{(64)} p ∗ ( x ) = 1 − p ref ( x ) + p ref ( x ) e 1/ β p ref ( x ) e 1/ β . (64) In this case, β e f f ( x ) = β \beta_{\mathrm{eff}}(x)=\beta β eff ( x ) = β (50) reduces to the standard KL-regularized single-reward optimal policy.
Boundary Cases
The scalar equation above assumes 0 < p r e f ( x ) < 1 0<p_{\mathrm{ref}}(x)<1 0 < p ref ( x ) < 1
Boundary cases. If p r e f ( x ) ∈ { 0 , 1 } p_{\mathrm{ref}}(x)\in\{0,1\} p ref ( x ) ∈ { 0 , 1 } J k , x \mathcal{J}_{k,x} J k , x
π ∗ ( ⋅ ∣ x ) = π r e f ( ⋅ ∣ x ) , p ∗ ( x ) = p r e f ( x ) . \pi^*(\cdot\mid x)=\pi_{\mathrm{ref}}(\cdot\mid x),\qquad p^*(x)=p_{\mathrm{ref}}(x). π ∗ ( ⋅ ∣ x ) = π ref ( ⋅ ∣ x ) , p ∗ ( x ) = p ref ( x ) . Proof. If p r e f ( x ) = 0 p_{\mathrm{ref}}(x)=0 p ref ( x ) = 0
p ∗ ( x ) = 0 , π ∗ ( ⋅ ∣ x ) = π r e f ( ⋅ ∣ x ) . (65) p^*(x)=0,\qquad\pi^*(\cdot\mid x)=\pi_{\mathrm{ref}}(\cdot\mid x). \tag*{(65)} p ∗ ( x ) = 0 , π ∗ ( ⋅ ∣ x ) = π ref ( ⋅ ∣ x ) . (65) The reward term is identically zero over the feasible set, and the reference policy uniquely minimizes the KL term.
If p r e f ( x ) = 1 p_{\mathrm{ref}}(x)=1 p ref ( x ) = 1
p ∗ ( x ) = 1 , π ∗ ( ⋅ ∣ x ) = π r e f ( ⋅ ∣ x ) . (66) p^*(x)=1,\qquad\pi^*(\cdot\mid x)=\pi_{\mathrm{ref}}(\cdot\mid x). \tag*{(66)} p ∗ ( x ) = 1 , π ∗ ( ⋅ ∣ x ) = π ref ( ⋅ ∣ x ) . (66) The reference policy again uniquely minimizes the KL term.
For k > 1 k>1 k > 1 p r e f ( x ) = 1 p_{\mathrm{ref}}(x)=1 p ref ( x ) = 1 β e f f ( x ) = + ∞ \beta_{\mathrm{eff}}(x)=+\infty β eff ( x ) = + ∞ k = 1 k=1 k = 1 β \beta β
Objective-Level Equivalences
We finally prove the two objective-level interpretations stated in the main text.
Corollary 4.3 (restated). Replacing the constant KL coefficient β \beta β β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) (54) recovers the pass@k optimal policy, with β e f f ( x ) = + ∞ \beta_{\mathrm{eff}}(x)=+\infty β eff ( x ) = + ∞
Proof. For each prompt x x x β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x )
π r e g ∗ ( y ∣ x ) = 1 Z r e g ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) β e f f ( x ) ) . (67) \pi_{\mathrm{reg}}^*(y \mid x) = \frac{1}{Z_{\mathrm{reg}}(x)}\pi^{\mathrm{ref}}(y \mid x)\exp\left(\frac{R(y,x)}{\beta_{\mathrm{eff}}(x)}\right). \tag*{(67)} π reg ∗ ( y ∣ x ) = Z reg ( x ) 1 π ref ( y ∣ x ) exp ( β eff ( x ) R ( y , x ) ) . (67) This is exactly the pass@k optimal policy derived in Eq. (50) . Since the objective decomposes over prompts and β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) y y y x x x β e f f ( x ) = + ∞ \beta_{\mathrm{eff}}(x)=+\infty β eff ( x ) = + ∞ π r e f \pi_{\mathrm{ref}} π ref
Corollary 4.4 (restated). The ordinary single-reward objective with constant KL coefficient β \beta β
R ~ k ( y , x ) = β β e f f ( x ) R ( y , x ) = k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) \widetilde{R}_k(y,x)=\frac{\beta}{\beta_{\mathrm{eff}}(x)}R(y,x)=k(1-p^*(x))^{k-1}R(y,x) R k ( y , x ) = β eff ( x ) β R ( y , x ) = k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) has the same unique optimal policy as the pass@k objective.
Proof. Retain the constant KL coefficient β \beta β
R ~ k ( y , x ) = β β e f f ( x ) R ( y , x ) = k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) . (68) \widetilde{R}_k(y,x)=\frac{\beta}{\beta_{\mathrm{eff}}(x)}R(y,x)=k(1-p^*(x))^{k-1}R(y,x). \tag*{(68)} R k ( y , x ) = β eff ( x ) β R ( y , x ) = k ( 1 − p ∗ ( x ) ) k − 1 R ( y , x ) . (68) The corresponding single-reward optimal policy is
π r e w ∗ ( y ∣ x ) = 1 Z r e w ( x ) π r e f ( y ∣ x ) exp ( R ~ k ( y , x ) β ) (69) \pi_{\mathrm{rew}}^*(y \mid x)=\frac{1}{Z_{\mathrm{rew}}(x)}\pi^{\mathrm{ref}}(y \mid x)\exp\left(\frac{\widetilde{R}_k(y,x)}{\beta}\right) \tag*{(69)} π rew ∗ ( y ∣ x ) = Z rew ( x ) 1 π ref ( y ∣ x ) exp ( β R k ( y , x ) ) (69) = 1 Z r e w ( x ) π r e f ( y ∣ x ) exp ( R ( y , x ) β e f f ( x ) ) . (70) =\frac{1}{Z_{\mathrm{rew}}(x)}\pi^{\mathrm{ref}}(y \mid x)\exp\left(\frac{R(y,x)}{\beta_{\mathrm{eff}}(x)}\right). \tag*{(70)} = Z rew ( x ) 1 π ref ( y ∣ x ) exp ( β eff ( x ) R ( y , x ) ) . (70) This again coincides exactly with the pass@k optimal policy.
Continuous-Rewards Derivation
We now derive the optimal policy for the max@k objective with real-valued rewards. As before, the objective decomposes over prompts, so we fix a prompt x x x
R ( y , x ) ∈ [ a x , b x ] (71) R(y,x)\in[a_x,b_x] \tag*{(71)} R ( y , x ) ∈ [ a x , b x ] (71) for every y y y x x x
Boundedness guarantees that the relevant expectations and directional derivatives are finite and permits differentiation under the expectation. The result extends to unbounded rewards whenever the corresponding expectations, integrals, and directional derivatives are finite, the required interchange of differentiation and expectation is valid, and the optimum is attained.
No continuity assumption on the reward distribution is required in this section. In particular, the reward CDF may contain atoms.
For a policy π \pi π
U k , x ( π ) = E y 1 , … , y k ∼ i . i . d . π ( ⋅ ∣ x ) [ max i = 1 , … , k R ( y i , x ) ] . (72) U_{k,x}(\pi)=\mathbb{E}_{y_1,\ldots,y_k \mathrel{\overset{\mathrm{i.i.d.}}{\sim}\pi(\cdot\mid x)}\left[\max_{i=1,\ldots,k}R(y_i,x)\right].} \tag*{(72)} U k , x ( π ) = E y 1 , … , y k ∼ i.i.d. π ( ⋅ ∣ x ) [ m a x i = 1 , … , k R ( y i , x ) ] . (72) The prompt-wise KL-regularized objective is
J k , x ( π ) = U k , x ( π ) − β K L ( π ( ⋅ ∣ x ) ∥ π r e f ( ⋅ ∣ x ) ) . (73) \mathcal{J}_{k,x}(\pi)=U_{k,x}(\pi)-\beta\,\mathrm{KL}\left(\pi(\cdot\mid x)\Vert\pi_{\mathrm{ref}}(\cdot\mid x)\right). \tag*{(73)} J k , x ( π ) = U k , x ( π ) − β KL ( π ( ⋅ ∣ x ) ∥ π ref ( ⋅ ∣ x ) ) . (73) Optimal Policy
Proposition 4.5 (restated). For k > 1 k > 1 k > 1 β > 0 \beta> 0 β > 0 J k , x J_{k,x} J k , x
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( k β ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) . \pi^*(y \mid x) = \frac{1}{Z(x)}\pi_{\mathrm{ref}}(y \mid x)\exp\left(\frac{k}{\beta}\int_{-\infty}^{R(y,x)}F_{\pi^*}(r \mid x)^{k-1}\,\mathrm{d}r\right). π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β k ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) . where F π ∗ ( r ∣ x ) = P y ∼ π ∗ ( ⋅ ∣ x ) [ R ( y , x ) ≤ r ] F_{\pi^*}(r \mid x) = \mathbb{P}_{y \sim\pi^*(\cdot\mid x)}[R(y,x) \le r] F π ∗ ( r ∣ x ) = P y ∼ π ∗ ( ⋅ ∣ x ) [ R ( y , x ) ≤ r ] Z ( x ) Z(x) Z ( x )
Proof. Step 1: Concavity and existence. Let
F π ( r ∣ x ) = P y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) ≤ r ] (74) F_\pi(r \mid x) = \mathbb{P}_{y \sim\pi(\cdot\mid x)}[R(y,x) \le r] \tag*{(74)} F π ( r ∣ x ) = P y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) ≤ r ] (74) denote the reward CDF under π \pi π k k k a x a_x a x
U k , x ( π ) = a x + ∫ a x b x [ 1 − F π ( r ∣ x ) k ] d r . (75) U_{k,x}(\pi) = a_x + \int_{a_x}^{b_x}\left[1-F_\pi(r \mid x)^k\right]\,\mathrm{d}r. \tag*{(75)} U k , x ( π ) = a x + ∫ a x b x [ 1 − F π ( r ∣ x ) k ] d r . (75) For each fixed r r r F π ( r ∣ x ) F_\pi(r \mid x) F π ( r ∣ x ) π \pi π
u ⟼ 1 − u k (76) u \longmapsto1-u^k \tag*{(76)} u ⟼ 1 − u k (76) is concave on [ 0 , 1 ] [0,1] [ 0 , 1 ] U k , x ( π ) U_{k,x}(\pi) U k , x ( π ) π \pi π (73) has at most one maximizer.
For a finite completion space, existence follows from compactness of the probability simplex. For a countably infinite space, boundedness of the reward allows a maximizing sequence to be chosen with uniformly bounded KL divergence from π r e f \pi_{\mathrm{ref}} π ref § 12 .
Step 2: Directional derivative. Consider a perturbation
π ϵ = π + ϵ h (77) \pi_\epsilon= \pi+ \epsilon h \tag*{(77)} π ϵ = π + ϵ h (77) satisfying the conditions of Appendix § 12 . Expanding the product measure in U k , x ( π ϵ ) U_{k,x}(\pi_\epsilon) U k , x ( π ϵ ) ϵ = 0 \epsilon= 0 ϵ = 0 k k k
d d ϵ U k , x ( π + ϵ h ) ∣ ϵ = 0 + = k ∑ y h ( y ∣ x ) E y 2 , … , y k ∼ i . i . d . π ( ⋅ ∣ x ) [ max { R ( y , x ) , R ( y 2 , x ) , … , R ( y k , x ) } ] . (78) \left.\frac{\mathrm{d}}{\mathrm{d}\epsilon}U_{k,x}(\pi+\epsilon h)\right|_{\epsilon=0^+}
= k\sum_y h(y \mid x)\mathbb{E}_{y_2,\ldots,y_k \stackrel{\mathrm{i.i.d.}}{\sim}\pi(\cdot\mid x)}
\left[\max\{R(y,x),R(y_2,x),\ldots,R(y_k,x)\}\right]. \tag*{(78)} d ϵ d U k , x ( π + ϵ h ) ϵ = 0 + = k y ∑ h ( y ∣ x ) E y 2 , … , y k ∼ i.i.d. π ( ⋅ ∣ x ) [ max { R ( y , x ) , R ( y 2 , x ) , … , R ( y k , x )} ] . (78) Thus, one valid score in Eq. (20) is
Γ π ( y , x ) = k E y 2 , … , y k ∼ i . i . d . π ( ⋅ ∣ x ) [ max { R ( y , x ) , R ( y 2 , x ) , … , R ( y k , x ) } ] . (79) \Gamma_\pi(y,x) = k\mathbb{E}_{y_2,\ldots,y_k \stackrel{\mathrm{i.i.d.}}{\sim}\pi(\cdot\mid x)}
\left[\max\{R(y,x),R(y_2,x),\ldots,R(y_k,x)\}\right]. \tag*{(79)} Γ π ( y , x ) = k E y 2 , … , y k ∼ i.i.d. π ( ⋅ ∣ x ) [ max { R ( y , x ) , R ( y 2 , x ) , … , R ( y k , x )} ] . (79) To isolate the part that depends on y y y
M k − 1 = max i = 2 , … , k R ( y i , x ) . (80) M_{k-1} = \max_{i=2,\ldots,k}R(y_i,x). \tag*{(80)} M k − 1 = i = 2 , … , k max R ( y i , x ) . (80) Using
max { a , b } = b + ( a − b ) + , ( u ) + = max { u , 0 } , (81) \max\{a,b\}=b+(a-b)_+,\qquad(u)_+=\max\{u,0\}, \tag*{(81)} max { a , b } = b + ( a − b ) + , ( u ) + = max { u , 0 } , (81) we obtain
Γ π ( y , x ) = k E π [ M k − 1 ] + k E π [ ( R ( y , x ) − M k − 1 ) + ] . (82) \Gamma_\pi(y,x) = k\mathbb{E}_\pi[M_{k-1}] + k\mathbb{E}_\pi\left[(R(y,x)-M_{k-1})_+\right]. \tag*{(82)} Γ π ( y , x ) = k E π [ M k − 1 ] + k E π [ ( R ( y , x ) − M k − 1 ) + ] . (82) The first term depends on π \pi π x x x y y y a a a M M M
( a − M ) + = ∫ − ∞ a I [ M ≤ r ] d r . (83) (a-M)_{+}=\int_{-\infty}^{a}\mathbb{I}[M\le r]\,\mathrm{d}r. \tag*{(83)} ( a − M ) + = ∫ − ∞ a I [ M ≤ r ] d r . (83) Taking expectations and applying Tonelli’s theorem gives
E [ ( a − M ) + ] = ∫ − ∞ a P [ M ≤ r ] d r . (84) \mathbb{E}[(a-M)_{+}]=\int_{-\infty}^{a}\mathbb{P}[M\le r]\,\mathrm{d}r. \tag*{(84)} E [( a − M ) + ] = ∫ − ∞ a P [ M ≤ r ] d r . (84) This identity holds for arbitrary reward distributions and does not require their CDFs to be continuous.
Because M k − 1 M_{k-1} M k − 1 k − 1 k-1 k − 1 π \pi π
P [ M k − 1 ≤ r ] = F π ( r ∣ x ) k − 1 . (85) \mathbb{P}[M_{k-1}\le r]=F_{\pi}(r\mid x)^{k-1}. \tag*{(85)} P [ M k − 1 ≤ r ] = F π ( r ∣ x ) k − 1 . (85) Substituting this into Eq. (84) , with a = R ( y , x ) a=R(y,x) a = R ( y , x )
E π [ ( R ( y , x ) − M k − 1 ) + ] = ∫ − ∞ R ( y , x ) F π ( r ∣ x ) k − 1 d r . (86) \mathbb{E}_{\pi}\left[(R(y,x)-M_{k-1})_{+}\right]=\int_{-\infty}^{R(y,x)}F_{\pi}(r\mid x)^{k-1}\,\mathrm{d}r. \tag*{(86)} E π [ ( R ( y , x ) − M k − 1 ) + ] = ∫ − ∞ R ( y , x ) F π ( r ∣ x ) k − 1 d r . (86) We may therefore choose the completion-dependent score
A π ( y , x ) = k ∫ − ∞ R ( y , x ) F π ( r ∣ x ) k − 1 d r . (87) A_{\pi}(y,x)=k\int_{-\infty}^{R(y,x)}F_{\pi}(r\mid x)^{k-1}\,\mathrm{d}r. \tag*{(87)} A π ( y , x ) = k ∫ − ∞ R ( y , x ) F π ( r ∣ x ) k − 1 d r . (87) Step 3: Optimal policy. Applying Lemma A.1 with the score in Eq. (87) gives
π ∗ ( y ∣ x ) = 1 Z ( x ) π r e f ( y ∣ x ) exp ( k β ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) , (88) \pi^{*}(y\mid x)=\frac{1}{Z(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{k}{\beta}\int_{-\infty}^{R(y,x)}F_{\pi^{*}}(r\mid x)^{k-1}\,\mathrm{d}r\right), \tag*{(88)} π ∗ ( y ∣ x ) = Z ( x ) 1 π ref ( y ∣ x ) exp ( β k ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) , (88) where
Z ( x ) = ∑ y π r e f ( y ∣ x ) exp ( k β ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) . (89) Z(x)=\sum_y\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{k}{\beta}\int_{-\infty}^{R(y,x)}F_{\pi^{*}}(r\mid x)^{k-1}\,\mathrm{d}r\right). \tag*{(89)} Z ( x ) = y ∑ π ref ( y ∣ x ) exp ( β k ∫ − ∞ R ( y , x ) F π ∗ ( r ∣ x ) k − 1 d r ) . (89) This is (8) from the main text. By the converse statement of Lemma A.1, any feasible policy satisfying this characterization is the unique global maximizer.
The characterization remains implicit because the CDF F π ∗ F_{\pi^{*}} F π ∗ F π ∗ F_{\pi^{*}} F π ∗ § 15 .
Objective-Level Equivalence
Corollary 4.6 (restated). For k > 1 k>1 k > 1
S x ( r ) = k ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u . S_x(r)=k\int_{-\infty}^{r}F_{\pi^{*}}(u\mid x)^{k-1}\,\mathrm{d}u. S x ( r ) = k ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u . The ordinary single-reward objective with constant KL coefficient β \beta β R ~ ( y , x ) = S x ( R ( y , x ) ) \widetilde{R}(y,x)=S_x(R(y,x)) R ( y , x ) = S x ( R ( y , x ))
Proof. Treating S x S_x S x
π r e w ∗ ( y ∣ x ) = 1 Z r e w ( x ) π r e f ( y ∣ x ) exp ( S x ( R ( y , x ) ) β ) . \pi_{\mathrm{rew}}^{*}(y\mid x)=\frac{1}{Z_{\mathrm{rew}}(x)}\pi_{\mathrm{ref}}(y\mid x)\exp\left(\frac{S_x(R(y,x))}{\beta}\right). π rew ∗ ( y ∣ x ) = Z rew ( x ) 1 π ref ( y ∣ x ) exp ( β S x ( R ( y , x )) ) . Substituting the definition of S x S_x S x (88) . Normalization therefore yields Z r e w ( x ) = Z ( x ) Z_{\mathrm{rew}}(x)=Z(x) Z rew ( x ) = Z ( x ) π r e w ∗ = π ∗ \pi_{\mathrm{rew}}^{*}=\pi^{*} π rew ∗ = π ∗
Numerical Reconstruction of the Raw-Reward Optimal CDF
The optimal policy in Eq. (88) depends on the unknown reward CDF F π ∗ F_{\pi^*} F π ∗ x x x
Define the optimal-policy score
H ( r ) = k ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u (90) H(r) = k \int_{-\infty}^{r} F_{\pi^*}(u \mid x)^{k-1}\,du \tag*{(90)} H ( r ) = k ∫ − ∞ r F π ∗ ( u ∣ x ) k − 1 d u (90) and let
c ( x ) = 1 Z ( x ) . (91) c(x) = \frac{1}{Z(x)}. \tag*{(91)} c ( x ) = Z ( x ) 1 . (91) Because the likelihood ratio in Eq. (88) depends on a completion only through its reward, the reward distributions under the optimal and reference policies satisfy
d P π ∗ ( r ∣ x ) = c ( x ) exp ( H ( r ) β ) d P r e f ( r ∣ x ) . (92) dP_{\pi^*}(r \mid x) = c(x) \exp\left(\frac{H(r)}{\beta}\right)dP_{\mathrm{ref}}(r \mid x). \tag*{(92)} d P π ∗ ( r ∣ x ) = c ( x ) exp ( β H ( r ) ) d P ref ( r ∣ x ) . (92) Thus, once the scalar c ( x ) c(x) c ( x ) c ( x ) c(x) c ( x )
Continuous Reference Reward Distribution
Suppose first that the reference reward distribution has density f r e f ( r ∣ x ) f_{\mathrm{ref}}(r \mid x) f ref ( r ∣ x ) [ a x , b x ] [a_x,b_x] [ a x , b x ] f π ∗ ( r ∣ x ) f_{\pi^*}(r \mid x) f π ∗ ( r ∣ x )
f π ∗ ( r ∣ x ) = c ( x ) f r e f ( r ∣ x ) exp ( H ( r ) β ) . (93) f_{\pi^*}(r \mid x) = c(x)f_{\mathrm{ref}}(r \mid x)\exp\left(\frac{H(r)}{\beta}\right). \tag*{(93)} f π ∗ ( r ∣ x ) = c ( x ) f ref ( r ∣ x ) exp ( β H ( r ) ) . (93) Since
d d r F π ∗ ( r ∣ x ) = f π ∗ ( r ∣ x ) , d d r H ( r ) = k F π ∗ ( r ∣ x ) k − 1 , (94) \frac{d}{dr}F_{\pi^*}(r \mid x) = f_{\pi^*}(r \mid x), \qquad\frac{d}{dr}H(r) = kF_{\pi^*}(r \mid x)^{k-1}, \tag*{(94)} d r d F π ∗ ( r ∣ x ) = f π ∗ ( r ∣ x ) , d r d H ( r ) = k F π ∗ ( r ∣ x ) k − 1 , (94) the unknown functions satisfy the coupled system
d d r F π ∗ ( r ∣ x ) = c ( x ) f r e f ( r ∣ x ) exp ( H ( r ) β ) , d d r H ( r ) = k F π ∗ ( r ∣ x ) k − 1 , (95) \begin{aligned}
\frac{d}{dr}F_{\pi^*}(r \mid x) &= c(x)f_{\mathrm{ref}}(r \mid x)\exp\left(\frac{H(r)}{\beta}\right), \\
\frac{d}{dr}H(r) &= kF_{\pi^*}(r \mid x)^{k-1}, \tag*{(95)}
\end{aligned} d r d F π ∗ ( r ∣ x ) d r d H ( r ) = c ( x ) f ref ( r ∣ x ) exp ( β H ( r ) ) , = k F π ∗ ( r ∣ x ) k − 1 , (95) with initial conditions
F π ∗ ( a x ∣ x ) = 0 , H ( a x ) = 0. (96) F_{\pi^*}(a_x \mid x) = 0, \qquad H(a_x) = 0. \tag*{(96)} F π ∗ ( a x ∣ x ) = 0 , H ( a x ) = 0. (96) For a candidate value of c ( x ) c(x) c ( x ) a x a_x a x b x b_x b x
F π ∗ ( b x ∣ x ) = 1. (97) F_{\pi^*}(b_x \mid x) = 1. \tag*{(97)} F π ∗ ( b x ∣ x ) = 1. (97) This gives a one-dimensional shooting problem. The terminal value F π ∗ ( b x ∣ x ) F_{\pi^*}(b_x \mid x) F π ∗ ( b x ∣ x ) c ( x ) c(x) c ( x )
Discrete Reconstruction
In practice, the reference reward distribution is estimated from a finite bank of reference-model completions and is therefore discrete. Let
r 1 < r 2 < ⋯ < r m (98) r_1 < r_2 < \cdots< r_m \tag*{(98)} r 1 < r 2 < ⋯ < r m (98) be its distinct observed reward values, and define
q j = P π r e f [ R ( y , x ) = r j ] , j = 1 , … , m . (99) q_j = \mathbb{P}_{\pi_{\mathrm{ref}}}\left[R(y,x)=r_j\right], \qquad j=1,\ldots,m. \tag*{(99)} q j = P π ref [ R ( y , x ) = r j ] , j = 1 , … , m . (99) Let
p j = P π ∗ [ R ( y , x ) = r j ] , P j = ∑ ℓ = 1 j p ℓ (100) p_j = \mathbb{P}_{\pi^*}[R(y,x)=r_j], \qquad P_j = \sum_{\ell=1}^{j} p_\ell \tag*{(100)} p j = P π ∗ [ R ( y , x ) = r j ] , P j = ℓ = 1 ∑ j p ℓ (100) denote the optimal probability mass at reward r j r_j r j
Define
H j = H ( r j ) = k ∫ − ∞ r j F π ∗ ( u ∣ x ) k − 1 d u . (101) H_j=H(r_j)=k\int_{-\infty}^{r_j}F_{\pi^*}(u\mid x)^{k-1}\,du. \tag*{(101)} H j = H ( r j ) = k ∫ − ∞ r j F π ∗ ( u ∣ x ) k − 1 d u . (101) Since the optimal CDF is constant between consecutive reward values,
H 1 = 0 , H j = H j − 1 + k ( r j − r j − 1 ) P j − 1 k − 1 , j = 2 , … , m . (102) H_1=0,\qquad H_j=H_{j-1}+k(r_j-r_{j-1})P_{j-1}^{k-1},\quad j=2,\ldots,m. \tag*{(102)} H 1 = 0 , H j = H j − 1 + k ( r j − r j − 1 ) P j − 1 k − 1 , j = 2 , … , m . (102) Grouping Eq. (88) by reward value gives
p j = c ( x ) q j exp ( H j β ) . (103) p_j=c(x)q_j\exp\left(\frac{H_j}{\beta}\right). \tag*{(103)} p j = c ( x ) q j exp ( β H j ) . (103) Equations 102 and 103 yield a forward recursion for any candidate c = c ( x ) c=c(x) c = c ( x )
H 1 ( c ) = 0 , p 1 ( c ) = c q 1 , P 1 ( c ) = p 1 ( c ) , H j ( c ) = H j − 1 ( c ) + k ( r j − r j − 1 ) P j − 1 ( c ) k − 1 , p j ( c ) = c q j exp ( H j ( c ) β ) , P j ( c ) = P j − 1 ( c ) + p j ( c ) , j = 2 , … , m . (104) \begin{aligned}
H_1(c)&=0,\\
p_1(c)&=cq_1,\\
P_1(c)&=p_1(c),\\
H_j(c)&=H_{j-1}(c)+k(r_j-r_{j-1})P_{j-1}(c)^{k-1},\\
p_j(c)&=cq_j\exp\left(\frac{H_j(c)}{\beta}\right),\\
P_j(c)&=P_{j-1}(c)+p_j(c),\qquad j=2,\ldots,m. \tag*{(104)}
\end{aligned} H 1 ( c ) p 1 ( c ) P 1 ( c ) H j ( c ) p j ( c ) P j ( c ) = 0 , = c q 1 , = p 1 ( c ) , = H j − 1 ( c ) + k ( r j − r j − 1 ) P j − 1 ( c ) k − 1 , = c q j exp ( β H j ( c ) ) , = P j − 1 ( c ) + p j ( c ) , j = 2 , … , m . (104) The required normalization is
P m ( c ) = 1. (105) P_m(c)=1. \tag*{(105)} P m ( c ) = 1. (105) Each quantity in the recursion is continuous and strictly increasing in c c c
P m ( 0 ) = 0 , (106) P_m(0)=0, \tag*{(106)} P m ( 0 ) = 0 , (106) while, for q 1 > 0 q_1>0 q 1 > 0
P m ( 1 q 1 ) ≥ p 1 ( 1 q 1 ) = 1 , (107) P_m\left(\frac{1}{q_1}\right)\geq p_1\left(\frac{1}{q_1}\right)=1, \tag*{(107)} P m ( q 1 1 ) ≥ p 1 ( q 1 1 ) = 1 , (107) with strict inequality whenever m > 1 m>1 m > 1 (105) has a unique solution in
0 < c ≤ 1 q 1 , (108) 0<c\leq\frac{1}{q_1}, \tag*{(108)} 0 < c ≤ q 1 1 , (108) which can be found efficiently by bisection. Equivalently, one may treat the optimal probability of the lowest reward, p 1 = c q 1 p_1=cq_1 p 1 = c q 1
Once c c c p j p_j p j P j P_j P j H j H_j H j r ∈ [ r j , r j + 1 ) r\in[r_j,r_{j+1}) r ∈ [ r j , r j + 1 )
F π ∗ ( r ∣ x ) = P j , H ( r ) = H j + k ( r − r j ) P j k − 1 . (109) F_{\pi^*}(r\mid x)=P_j,\qquad H(r)=H_j+k(r-r_j)P_j^{k-1}. \tag*{(109)} F π ∗ ( r ∣ x ) = P j , H ( r ) = H j + k ( r − r j ) P j k − 1 . (109) Thus, the score is linearly interpolated between adjacent reward levels, with slope determined by the reconstructed optimal CDF.
Empirical Approximation
Given M M M
q ^ j = 1 M ∑ i = 1 M I [ R ( y i r e f , x ) = r j ] . (110) \widehat{q}_j=\frac{1}{M}\sum_{i=1}^{M}\mathbb{I}\left[R(y_i^{\mathrm{ref}},x)=r_j\right]. \tag*{(110)} q j = M 1 i = 1 ∑ M I [ R ( y i ref , x ) = r j ] . (110) Substituting these empirical masses into Eq. (104) produces an approximation of F π ∗ F_{\pi^*} F π ∗ H H H ( k , β ) (k,\beta) ( k , β )
Unlike the binary-reward case, where the reference distribution is summarized by the single scalar p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x )
Reward-Scale Assumptions for Continuous Rewards
The derivation in Appendix C requires only prompt-wise bounded rewards, or suitable integrability conditions in the unbounded case. The stronger shared-range assumption used in the main text is needed only for interpreting the resulting adaptation across prompts.
Shift invariance. Both objectives studied in this work are invariant, at the level of their optimizer, to prompt-dependent constant reward shifts. For
R c ( y , x ) = R ( y , x ) + c ( x ) , (111) R_c(y,x) = R(y,x) + c(x), \tag*{(111)} R c ( y , x ) = R ( y , x ) + c ( x ) , (111) we have
E y ∼ π ( ⋅ ∣ x ) [ R c ( y , x ) ] = E y ∼ π ( ⋅ ∣ x ) [ R ( y , x ) ] + c ( x ) (112) \mathbb{E}_{y\sim\pi(\cdot\mid x)}[R_c(y,x)] = \mathbb{E}_{y\sim\pi(\cdot\mid x)}[R(y,x)] + c(x) \tag*{(112)} E y ∼ π ( ⋅ ∣ x ) [ R c ( y , x )] = E y ∼ π ( ⋅ ∣ x ) [ R ( y , x )] + c ( x ) (112) and
E [ max i = 1 , … , k R c ( y i , x ) ] = E [ max i = 1 , … , k R ( y i , x ) ] + c ( x ) . (113) \mathbb{E}\left[\max_{i=1,\ldots,k} R_c(y_i,x)\right] = \mathbb{E}\left[\max_{i=1,\ldots,k} R(y_i,x)\right] + c(x). \tag*{(113)} E [ i = 1 , … , k max R c ( y i , x ) ] = E [ i = 1 , … , k max R ( y i , x ) ] + c ( x ) . (113) The shift therefore adds a policy-independent constant to Eqs. (1) and (3) , leaving their optimal policies unchanged. Consequently, if rewards are bounded below by r min r_{\min} r m i n R R R R − r min R-r_{\min} R − r m i n
Nonnegative rewards and effective regularization. Suppose that the reward support for prompt x x x a x a_x a x F π ∗ ( u ∣ x ) = 0 F_{\pi^*}(u\mid x)=0 F π ∗ ( u ∣ x ) = 0 a x a_x a x (8) can be written, for r > a x r>a_x r > a x
k β ∫ a x r F π ∗ ( u ∣ x ) k − 1 d u = r − a x β e f f ( r , x ) , (114) \frac{k}{\beta}\int_{a_x}^{r} F_{\pi^*}(u\mid x)^{k-1}\,\mathrm{d}u = \frac{r-a_x}{\beta_{\mathrm{eff}}(r,x)}, \tag*{(114)} β k ∫ a x r F π ∗ ( u ∣ x ) k − 1 d u = β eff ( r , x ) r − a x , (114) where
β e f f ( r , x ) = β ( r − a x ) k ∫ a x r F π ∗ ( u ∣ x ) k − 1 d u . (115) \beta_{\mathrm{eff}}(r,x) = \frac{\beta(r-a_x)}{k\int_{a_x}^{r} F_{\pi^*}(u\mid x)^{k-1}\,\mathrm{d}u}. \tag*{(115)} β eff ( r , x ) = k ∫ a x r F π ∗ ( u ∣ x ) k − 1 d u β ( r − a x ) . (115) Shifting the lower endpoint to zero gives the simpler expression in Eq. (9) . The nonnegative convention therefore provides a natural origin relative to which reward magnitude and effective regularization can be compared. Without such an origin, a definition based directly on r / β e f f r/\beta_{\mathrm{eff}} r / β eff
This convention is not required for the transformed-reward interpretation. More generally, for any anchor a a a
S x , a ( r ) = k ∫ a r F π ∗ ( u ∣ x ) k − 1 d u . (116) S_{x,a}(r) = k\int_{a}^{r} F_{\pi^*}(u\mid x)^{k-1}\,\mathrm{d}u. \tag*{(116)} S x , a ( r ) = k ∫ a r F π ∗ ( u ∣ x ) k − 1 d u . (116) Changing a a a S x , a ( r ) S_{x,a}(r) S x , a ( r )
Common reward range. Our cross-prompt interpretation treats the location of the reference reward distribution — and, informally, its mean
μ r e f ( x ) = E y ∼ π r e f ( ⋅ ∣ x ) [ R ( y , x ) ] (117) \mu_{\mathrm{ref}}(x) = \mathbb{E}_{y\sim\pi_{\mathrm{ref}}(\cdot\mid x)}[R(y,x)] \tag*{(117)} μ ref ( x ) = E y ∼ π ref ( ⋅ ∣ x ) [ R ( y , x )] (117) — as an indicator of task difficulty: prompts receiving lower rewards under the reference policy are regarded as more difficult. This interpretation is meaningful only when rewards can be compared across prompts. In particular, a reward value r 1 r_1 r 1 r 2 < r 1 r_2 < r_1 r 2 < r 1
We therefore assume a shared, semantically meaningful fixed range,
R ( y , x ) ∈ [ r min , r max ] for every x . (118) R(y,x) \in[r_{\min},r_{\max}] \quad\text{for every } x. \tag*{(118)} R ( y , x ) ∈ [ r m i n , r m a x ] for every x . (118) Without common anchors, the average reward cannot define prompt difficulty consistently. Indeed, because the objectives are invariant to prompt-dependent reward shifts, one could change the mean reward of any prompt arbitrarily without changing either its optimization problem or its optimal policy. A difference between two means would then reflect only a difference in reward coordinates rather than a difference in task difficulty. Prompt-dependent rescalings create the same comparability problem and additionally alter the balance between reward and KL regularization unless β \beta β
The mean provides an intuitive summary of difficulty but does not, by itself, guarantee a pointwise ordering of the adaptive transformations. A sufficient condition for the ordering illustrated in Figure 1 b and Figure 1 c can be stated in terms of the optimal reward CDFs. Consider an easy prompt x e x_e x e x h x_h x h [ 0 , r max ] [0,r_{\max}] [ 0 , r m a x ]
F π ∗ ( r ∣ x e ) ≤ F π ∗ ( r ∣ x h ) for all r ∈ [ 0 , r max ] . (119) F_{\pi^*}(r \mid x_e) \le F_{\pi^*}(r \mid x_h) \quad\text{for all } r \in[0,r_{\max}]. \tag*{(119)} F π ∗ ( r ∣ x e ) ≤ F π ∗ ( r ∣ x h ) for all r ∈ [ 0 , r m a x ] . (119) Then, for every common raw reward value,
S x e ( r ) ≤ S x h ( r ) , β e f f ( r , x e ) ≥ β e f f ( r , x h ) . (120) S_{x_e}(r) \le S_{x_h}(r), \qquad\beta_{\mathrm{eff}}(r,x_e) \ge\beta_{\mathrm{eff}}(r,x_h). \tag*{(120)} S x e ( r ) ≤ S x h ( r ) , β eff ( r , x e ) ≥ β eff ( r , x h ) . (120) Thus, the easy prompt receives a smaller transformed reward and stronger effective regularization. This condition formalizes the representative easy–difficult comparison used in the main text; it is not a universal definition of task difficulty. If the reward CDFs cross, the corresponding transformation and effective-regularization curves may also cross, and no global ordering is implied.
Explicit max@k Solution for Win-Rate Rewards
We now derive the explicit solution in Eqs. (11) –(12) . Fix a prompt x x x
F r e f ( r ∣ x ) = P y ′ ∼ π r e f ( ⋅ ∣ x ) [ R ( y ′ , x ) ≤ r ] (121) F_{\mathrm{ref}}(r \mid x) = \mathbb{P}_{y' \sim\pi_{\mathrm{ref}}(\cdot\mid x)}[R(y',x) \le r] \tag*{(121)} F ref ( r ∣ x ) = P y ′ ∼ π ref ( ⋅ ∣ x ) [ R ( y ′ , x ) ≤ r ] (121) denote the reference reward CDF, and define the reference quantile, or win rate, of a completion y y y
Q ( y ) = F r e f ( R ( y , x ) ∣ x ) . (122) Q(y) = F_{\mathrm{ref}}(R(y,x) \mid x). \tag*{(122)} Q ( y ) = F ref ( R ( y , x ) ∣ x ) . (122) Proposition 4.7 (restated). Suppose that the reference reward distribution is continuous. For k ≥ 1 k \ge1 k ≥ 1 β > 0 \beta> 0 β > 0 Q ( y ) Q(y) Q ( y )
π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) exp ( S x q ( R ( y , x ) ) β ) , \pi_q^*(y \mid x) = \pi_{\mathrm{ref}}(y \mid x) \exp\left(\frac{S_x^q(R(y,x))}{\beta}\right), π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) exp ( β S x q ( R ( y , x )) ) , where
S x q ( r ) = β log ( C + [ Ψ − 1 ( F r e f ( r ∣ x ) ) ] k β ) , Ψ ( z ) = ∫ 0 z d u C + u k / β , S_x^q(r) = \beta\log\left(C + \frac{[\Psi^{-1}(F_{\mathrm{ref}}(r \mid x))]^k}{\beta}\right), \qquad\Psi(z) = \int_0^z \frac{du}{C+u^k/\beta}, S x q ( r ) = β log ( C + β [ Ψ − 1 ( F ref ( r ∣ x )) ] k ) , Ψ ( z ) = ∫ 0 z C + u k / β d u , and C > 0 C>0 C > 0 Ψ ( 1 ) = 1 \Psi(1)=1 Ψ ( 1 ) = 1
Proof. Step 1: Deriving the optimal win-rate density. Since the reference reward distribution is atomless, the probability integral transform gives
Q ( y ) ∼ Unif [ 0 , 1 ] , y ∼ π r e f ( ⋅ ∣ x ) . (123) Q(y) \sim\operatorname{Unif}[0,1], \qquad y \sim\pi_{\mathrm{ref}}(\cdot\mid x). \tag*{(123)} Q ( y ) ∼ Unif [ 0 , 1 ] , y ∼ π ref ( ⋅ ∣ x ) . (123) Let
G ( q ) = P y ∼ π q ∗ ( ⋅ ∣ x ) [ Q ( y ) ≤ q ] (124) G(q) = \mathbb{P}_{y \sim\pi_q^*(\cdot\mid x)}[Q(y) \le q] \tag*{(124)} G ( q ) = P y ∼ π q ∗ ( ⋅ ∣ x ) [ Q ( y ) ≤ q ] (124) be the win-rate CDF under the optimal policy, and let g ( q ) = G ′ ( q ) g(q) = G'(q) g ( q ) = G ′ ( q ) § 14 , applied with Q ( y ) Q(y) Q ( y )
π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) w ( Q ( y ) ) , w ( q ) = 1 Z ( x ) exp ( k β ∫ 0 q G ( u ) k − 1 d u ) . \pi_q^*(y \mid x) = \pi_{\mathrm{ref}}(y \mid x) w(Q(y)), \qquad w(q) = \frac{1}{Z(x)} \exp\left(\frac{k}{\beta} \int_0^q G(u)^{k-1}\,\mathrm{d}u\right). π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) w ( Q ( y )) , w ( q ) = Z ( x ) 1 exp ( β k ∫ 0 q G ( u ) k − 1 d u ) . For k = 1 k = 1 k = 1 (2) .
The multiplier w ( Q ( y ) ) w(Q(y)) w ( Q ( y )) Q Q Q
G ( q ) = E y ∼ π r e f ( ⋅ ∣ x ) [ w ( Q ( y ) ) I { Q ( y ) ≤ q } ] = ∫ 0 q w ( u ) d u . G(q) = \mathbb{E}_{y\sim\pi_{\mathrm{ref}}(\cdot\mid x)}\left[w(Q(y)) \mathbb{I}\{Q(y) \le q\}\right] = \int_0^q w(u)\,\mathrm{d}u. G ( q ) = E y ∼ π ref ( ⋅ ∣ x ) [ w ( Q ( y )) I { Q ( y ) ≤ q } ] = ∫ 0 q w ( u ) d u . Differentiating gives g ( q ) = w ( q ) g(q) = w(q) g ( q ) = w ( q ) π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) g ( Q ( y ) ) \pi_q^*(y \mid x) = \pi_{\mathrm{ref}}(y \mid x) g(Q(y)) π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) g ( Q ( y ))
g ( q ) = 1 Z ( x ) exp ( k β ∫ 0 q G ( u ) k − 1 d u ) . (125) g(q) = \frac{1}{Z(x)} \exp\left(\frac{k}{\beta} \int_0^q G(u)^{k-1}\,\mathrm{d}u\right). \tag*{(125)} g ( q ) = Z ( x ) 1 exp ( β k ∫ 0 q G ( u ) k − 1 d u ) . (125) Differentiating Eq. (125) yields
g ′ ( q ) g ( q ) = k β G ( q ) k − 1 . (126) \frac{g'(q)}{g(q)} = \frac{k}{\beta}G(q)^{k-1}. \tag*{(126)} g ( q ) g ′ ( q ) = β k G ( q ) k − 1 . (126) Using G ′ ( q ) = g ( q ) > 0 G'(q) = g(q) > 0 G ′ ( q ) = g ( q ) > 0 g g g G G G
d g d G = g ′ ( q ) G ′ ( q ) = k β G k − 1 . (127) \frac{\mathrm{d}g}{\mathrm{d}G} = \frac{g'(q)}{G'(q)} = \frac{k}{\beta}G^{k-1}. \tag*{(127)} d G d g = G ′ ( q ) g ′ ( q ) = β k G k − 1 . (127) Integration therefore gives
g ( q ) = C + 1 β G ( q ) k , (128) g(q) = C + \frac{1}{\beta}G(q)^k, \tag*{(128)} g ( q ) = C + β 1 G ( q ) k , (128) where C = g ( 0 ) > 0 C = g(0) > 0 C = g ( 0 ) > 0
G ′ ( q ) = C + 1 β G ( q ) k , G ( 0 ) = 0 , G ( 1 ) = 1. (129) G'(q) = C + \frac{1}{\beta}G(q)^k, \qquad G(0) = 0, \qquad G(1) = 1. \tag*{(129)} G ′ ( q ) = C + β 1 G ( q ) k , G ( 0 ) = 0 , G ( 1 ) = 1. (129) Step 2: Determining the constant and the CDF. Separating variables gives
q = ∫ 0 G ( q ) d u C + u k / β = Ψ ( G ( q ) ) . (130) q = \int_0^{G(q)} \frac{\mathrm{d}u}{C + u^k/\beta} = \Psi(G(q)). \tag*{(130)} q = ∫ 0 G ( q ) C + u k / β d u = Ψ ( G ( q )) . (130) The terminal condition G ( 1 ) = 1 G(1) = 1 G ( 1 ) = 1
1 = ∫ 0 1 d u C + u k / β . (131) 1 = \int_0^1 \frac{\mathrm{d}u}{C + u^k/\beta}. \tag*{(131)} 1 = ∫ 0 1 C + u k / β d u . (131) The right-hand side is continuous and strictly decreasing in C > 0 C > 0 C > 0 C ↓ 0 C \downarrow0 C ↓ 0 k ≥ 1 k \ge1 k ≥ 1 C → ∞ C \to\infty C → ∞ (131) has a unique solution C > 0 C > 0 C > 0 C C C Ψ \Psi Ψ [ 0 , 1 ] [0,1] [ 0 , 1 ] [ 0 , 1 ] [0,1] [ 0 , 1 ]
G = Ψ − 1 . (132) G = \Psi^{-1}. \tag*{(132)} G = Ψ − 1 . (132) Step 3: Recovering and verifying the optimal policy. As established in Step 1, the optimal completion policy satisfies π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) g ( Q ( y ) ) \pi_q^*(y \mid x) = \pi_{\mathrm{ref}}(y \mid x) g(Q(y)) π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) g ( Q ( y )) (128) and Q ( y ) = F r e f ( R ( y , x ) ∣ x ) Q(y) = F_{\mathrm{ref}}(R(y,x) \mid x) Q ( y ) = F ref ( R ( y , x ) ∣ x )
π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) ( C + [ G ( F r e f ( R ( y , x ) ∣ x ) ) ] k β ) . (133) \pi_q^*(y \mid x) = \pi_{\mathrm{ref}}(y \mid x)\left(C + \frac{\left[G(F_{\mathrm{ref}}(R(y,x) \mid x))\right]^k}{\beta}\right). \tag*{(133)} π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) ( C + β [ G ( F ref ( R ( y , x ) ∣ x )) ] k ) . (133) Equivalently, defining
S x q ( r ) = β log ( C + [ G ( F r e f ( r ∣ x ) ) ] k β ) (134) S_x^q(r) = \beta\log\left(C + \frac{\left[G(F_{\mathrm{ref}}(r \mid x))\right]^k}{\beta}\right) \tag*{(134)} S x q ( r ) = β log ( C + β [ G ( F ref ( r ∣ x )) ] k ) (134) yields
π q ∗ ( y ∣ x ) = π r e f ( y ∣ x ) exp ( S x q ( R ( y , x ) ) β ) . (135) \pi_q^*(y \mid x) = \pi_{\mathrm{ref}}(y \mid x) \exp\left(\frac{S_x^q(R(y,x))}{\beta}\right). \tag*{(135)} π q ∗ ( y ∣ x ) = π ref ( y ∣ x ) exp ( β S x q ( R ( y , x )) ) . (135) Normalization is already incorporated because
E y ∼ π r e f ( ⋅ ∣ x ) [ exp ( S x q ( R ( y , x ) ) β ) ] = ∫ 0 1 g ( q ) d q = G ( 1 ) − G ( 0 ) = 1. (136) \mathbb{E}_{y\sim\pi_{\mathrm{ref}}(\cdot\mid x)}\left[\exp\left(\frac{S_x^q(R(y,x))}{\beta}\right)\right] = \int_0^1 g(q)\,dq = G(1)-G(0)=1. \tag*{(136)} E y ∼ π ref ( ⋅ ∣ x ) [ exp ( β S x q ( R ( y , x )) ) ] = ∫ 0 1 g ( q ) d q = G ( 1 ) − G ( 0 ) = 1. (136) To verify sufficiency, observe that the constructed policy induces win-rate density g g g G G G g = C + G k / β g=C+G^k/\beta g = C + G k / β G ′ = g G'=g G ′ = g g ′ / g = ( k / β ) G k − 1 g'/g=(k/\beta)G^{k-1} g ′ / g = ( k / β ) G k − 1 g ( 0 ) = C g(0)=C g ( 0 ) = C (125) , with normalizing constant 1 / C 1/C 1/ C § 14 . Strict concavity of the KL-regularized objective establishes uniqueness.
The quantities C C C Ψ \Psi Ψ G G G k k k β \beta β F r e f F_{\mathrm{ref}} F ref
For k = 1 k=1 k = 1
g ( q ) = exp ( q / β ) ∫ 0 1 exp ( u / β ) d u . (137) g(q)=\frac{\exp(q/\beta)}{\int_0^1 \exp(u/\beta)\,du}. \tag*{(137)} g ( q ) = ∫ 0 1 exp ( u / β ) d u exp ( q / β ) . (137) Ties and discrete rewards. In practical settings such as test-case pass-rate rewards, the reference reward distribution is discrete and contains ties. The deterministic win rate Q ( y ) = F r e f ( R ( y , x ) ∣ x ) Q(y)=F_{\mathrm{ref}}(R(y,x)\mid x) Q ( y ) = F ref ( R ( y , x ) ∣ x ) Z ( x ) = 1 Z(x)=1 Z ( x ) = 1 Z ( x ) = 1 Z(x)=1 Z ( x ) = 1
KAPPO Adaptation of GRPO
GRPO constructs its learning signal by centering and standardizing rewards within each prompt’s completion group. This section explains how that normalization interacts with the adaptive reward transformations derived in § 4 . For a group of rewards R 1 , … , R n R_1,\ldots,R_n R 1 , … , R n R ˉ = 1 n ∑ i R i \bar{R}=\frac{1}{n}\sum_i R_i R ˉ = n 1 ∑ i R i s R = std i ( R i ) s_R=\operatorname{std}_i(R_i) s R = std i ( R i )
Cancellation of binary reward scaling. For binary rewards, Corollary 4.4 gives the transformed reward R ~ i = a k ( x ) R i \widetilde{R}_i=a_k(x)R_i R i = a k ( x ) R i
a k ( x ) = β β e f f ( x ) = k ( 1 − p ∗ ( x ) ) k − 1 . (138) a_k(x)=\frac{\beta}{\beta_{\mathrm{eff}}(x)}=k\left(1-p^*(x)\right)^{k-1}. \tag*{(138)} a k ( x ) = β eff ( x ) β = k ( 1 − p ∗ ( x ) ) k − 1 . (138) Since this multiplier is shared by all completions of the same prompt, it scales both the centered rewards and their standard deviation. For a k ( x ) > 0 a_k(x)>0 a k ( x ) > 0
R ~ i − R ~ std j ( R ~ j ) = a k ( x ) ( R i − R ˉ ) a k ( x ) s R = R i − R ˉ s R . (139) \frac{\widetilde{R}_i-\widetilde{R}}{\operatorname{std}_j(\widetilde{R}_j)}
=\frac{a_k(x)(R_i-\bar{R})}{a_k(x)s_R}
=\frac{R_i-\bar{R}}{s_R}. \tag*{(139)} std j ( R j ) R i − R = a k ( x ) s R a k ( x ) ( R i − R ˉ ) = s R R i − R ˉ . (139) Thus, substituting the transformed binary rewards before group standardization cancels the intended prompt-dependent scaling.
Preserving the binary adaptation. The core issue here is that standardized GRPO does not generally estimate the policy gradient of the underlying expected-reward objective. Dividing each group by its empirical reward standard deviation reweights each prompt’s update by a factor that depends on the sampled completions and the current policy; consequently, the expected GRPO update generally differs from the policy gradient of the original objective. For the same reason, KAPPO’s adaptive reward scaling is discarded, as shown in Eq. (139) . Therefore, the failure of ordinary reward substitution in this case arises precisely because standardized GRPO does not follow the objective to which that connection applies.
Accordingly, we treat the KAPPO adaptation of standardized GRPO as an empirical extension of our theoretical result. To retain the derived redistribution of pressure while preserving GRPO’s normalization, we apply the adaptive multiplier after standardization:
A ^ i k = a k ( x ) R i − R ‾ s R . (140) \widehat{A}_i^k = a_k(x)\frac{R_i-\overline{R}}{s_R}. \tag*{(140)} A i k = a k ( x ) s R R i − R . (140) Since the clipped GRPO surrogate is positively homogeneous in the advantage, this construction multiplies each prompt’s contribution to the GRPO loss and gradient by exactly a k ( x ) a_k(x) a k ( x )
Two alternative implementations recover the theoretically grounded interfaces more directly. Removing standard-deviation normalization while retaining group centering allows the ordinary KAPPO reward scaling to survive. Alternatively, for KL-regularized GRPO, the regularization-side interface can be applied through the adaptive coefficient β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x )
Continuous reward transformations. For continuous rewards, group standardization does not cause the same cancellation. The transformation S x q S_x^q S x q k > 1 k>1 k > 1 Q k \mathcal{Q}_k Q k
For the win-rate-based transformation used in our experiments, the interaction is particularly clear. At the beginning of training, when the current policy is close to the reference policy, the population mean and standard deviation of the transformed rewards are approximately the same for every prompt. GRPO normalization is therefore trivial with respect to prompt adaptivity: it applies the same affine transformation across prompts and does not alter the redistribution of optimization pressure induced by KAPPO.
Indeed, write
Q x ( Y ) = F r e f ( R ( Y , x ) ∣ x ) , S x q ( R ( Y , x ) ) = T k , β ( Q x ( Y ) ) , (141) Q_x(Y)=F_{\mathrm{ref}}(R(Y,x)\mid x),\qquad S_x^q(R(Y,x))=T_{k,\beta}(Q_x(Y)), \tag*{(141)} Q x ( Y ) = F ref ( R ( Y , x ) ∣ x ) , S x q ( R ( Y , x )) = T k , β ( Q x ( Y )) , (141) where T k , β T_{k,\beta} T k , β § 17 , at initialization
Q x ( Y ) ∼ Unif [ 0 , 1 ] , Y ∼ π r e f ( ⋅ ∣ x ) . (142) Q_x(Y)\sim\operatorname{Unif}[0,1],\qquad Y\sim\pi_{\mathrm{ref}}(\cdot\mid x). \tag*{(142)} Q x ( Y ) ∼ Unif [ 0 , 1 ] , Y ∼ π ref ( ⋅ ∣ x ) . (142) Hence,
μ k , β = E U ∼ Unif [ 0 , 1 ] [ T k , β ( U ) ] , σ k , β 2 = Var U ∼ Unif [ 0 , 1 ] [ T k , β ( U ) ] (143) \mu_{k,\beta}=\mathbb{E}_{U\sim\operatorname{Unif}[0,1]}[T_{k,\beta}(U)],\qquad\sigma_{k,\beta}^{2}=\operatorname{Var}_{U\sim\operatorname{Unif}[0,1]}[T_{k,\beta}(U)] \tag*{(143)} μ k , β = E U ∼ Unif [ 0 , 1 ] [ T k , β ( U )] , σ k , β 2 = Var U ∼ Unif [ 0 , 1 ] [ T k , β ( U )] (143) do not depend on the prompt, and GRPO applies the common normalization
S x q ( r ) ⟼ S x q ( r ) − μ k , β σ k , β . (144) S_x^q(r)\longmapsto\frac{S_x^q(r)-\mu_{k,\beta}}{\sigma_{k,\beta}}. \tag*{(144)} S x q ( r ) ⟼ σ k , β S x q ( r ) − μ k , β . (144) Thus, the prompt-adaptive structure of S x q S_x^q S x q
As training proceeds, the current policy’s win-rate distribution can evolve differently across prompts, so the corresponding means and standard deviations become prompt dependent. GRPO normalization then becomes nontrivial and interacts with the KAPPO transformation, introducing GRPO’s usual prompt-dependent reweighting on top of the adaptive pressure supplied by KAPPO. Since S x q S_x^q S x q Q k \mathcal{Q}_k Q k
Experimental Details
General setup. We perform full-parameter training in bfloat16 using AdamW with zero weight decay. We use a cosine learning-rate schedule with a 3% warmup, a maximum prompt and response length of 4096 tokens, and sample training completions with temperature 1 and top-p = 1 p = 1 p = 1 M = 30 M = 30 M = 30
Mathematical reasoning. We train for at most 1000 steps. For @k-A⋆ \star ⋆ { 5 × 10 − 6 , 10 − 5 , 2 × 10 − 5 } \{5 \times10^{-6}, 10^{-5}, 2 \times10^{-5}\} { 5 × 1 0 − 6 , 1 0 − 5 , 2 × 1 0 − 5 } β \beta β { 2.5 × 10 − 4 , 5 × 10 − 4 , 10 − 3 } \{2.5 \times10^{-4}, 5 \times10^{-4}, 10^{-3}\} { 2.5 × 1 0 − 4 , 5 × 1 0 − 4 , 1 0 − 3 } { 10 − 6 , 5 × 10 − 6 , 10 − 5 } \{10^{-6}, 5 \times10^{-6}, 10^{-5}\} { 1 0 − 6 , 5 × 1 0 − 6 , 1 0 − 5 } n = k n = k n = k n = 8 n = 8 n = 8 { 5 × 10 − 6 , 10 − 5 , 2 × 10 − 5 } \{5 \times10^{-6}, 10^{-5}, 2 \times10^{-5}\} { 5 × 1 0 − 6 , 1 0 − 5 , 2 × 1 0 − 5 }
For standard DPO, we search learning rates in { 5 × 10 − 6 , 10 − 5 , 2 × 10 − 5 } \{5 \times10^{-6}, 10^{-5}, 2 \times10^{-5}\} { 5 × 1 0 − 6 , 1 0 − 5 , 2 × 1 0 − 5 } β \beta β { 0.01 , 0.02 , 0.05 } \{0.01, 0.02, 0.05\} { 0.01 , 0.02 , 0.05 } β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) { 10 − 4 , 2 × 10 − 4 , 5 × 10 − 4 } \{10^{-4}, 2 \times10^{-4}, 5 \times10^{-4}\} { 1 0 − 4 , 2 × 1 0 − 4 , 5 × 1 0 − 4 } α = 0.5 \alpha= 0.5 α = 0.5 β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) § 21 .§ 2 . The other mathematical-reasoning experiments use 128 prompts per training step.
Code generation. We train Qwen3-4B for at most 500 steps using 128 prompts per step. For win-rate @k-A⋆ \star ⋆ { 10 − 5 , 2 × 10 − 5 } \{10^{-5}, 2 \times10^{-5}\} { 1 0 − 5 , 2 × 1 0 − 5 } β \beta β { 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } \{5 \times10^{-4}, 10^{-3}, 5 \times10^{-3}\} { 5 × 1 0 − 4 , 1 0 − 3 , 5 × 1 0 − 3 } { 10 − 6 , 5 × 10 − 6 , 10 − 5 } \{10^{-6}, 5 \times10^{-6}, 10^{-5}\} { 1 0 − 6 , 5 × 1 0 − 6 , 1 0 − 5 } n = k n = k n = k n = 8 n = 8 n = 8
For win-rate @k-GRPO, we search learning rates in { 5 × 10 − 6 , 10 − 5 } \{5 \times10^{-6}, 10^{-5}\} { 5 × 1 0 − 6 , 1 0 − 5 } β \beta β { 5 × 10 − 4 , 10 − 3 , 5 × 10 − 3 } \{5 \times10^{-4}, 10^{-3}, 5 \times10^{-3}\} { 5 × 1 0 − 4 , 1 0 − 3 , 5 × 1 0 − 3 } n = 8 n = 8 n = 8
Compute Resources
All training runs used 32 NVIDIA GH200 GPUs. Both mathematical reasoning and code-generation runs took approximately 10–12 hours, corresponding to 320–384 GPU-hours per run. The overall computational cost of the study, including hyperparameter sweeps, exploratory runs, and evaluation exceeded 100,000 GPU-hours.
Reference Calibration
KAPPO requires prompt-level statistics of the reference policy. We estimate them once before training from a fixed bank of M M M
For binary rewards, calibration estimates
p ^ r e f ( x ) = 1 M ∑ j = 1 M I [ R ( y j r e f , x ) = 1 ] . (145) \hat{p}_{\mathrm{ref}}(x) = \frac{1}{M} \sum_{j=1}^{M} \mathbb{I}[R(y^{\mathrm{ref}}_j, x) = 1]. \tag*{(145)} p ^ ref ( x ) = M 1 j = 1 ∑ M I [ R ( y j ref , x ) = 1 ] . (145) For each choice of k k k β \beta β (6) for p ∗ ( x ) p^*(x) p ∗ ( x ) β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) Z ( x ) Z(x) Z ( x ) (5) ). The @k-DPO instantiation uses β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) ⋆ \star ⋆ Z ( x ) Z(x) Z ( x ) β / β e f f ( x ) \beta/ \beta_{\mathrm{eff}}(x) β / β eff ( x )
For continuous rewards, the same bank defines the empirical reference CDF
F ^ r e f ( r ∣ x ) = 1 M ∑ j = 1 M I [ R ( y x , j r e f , x ) ≤ r ] . (146) \widehat{F}_{\mathrm{ref}}(r \mid x) = \frac{1}{M}\sum_{j=1}^{M} \mathbb{I}[R(y^{\mathrm{ref}}_{x,j},x) \le r]. \tag*{(146)} F ref ( r ∣ x ) = M 1 j = 1 ∑ M I [ R ( y x , j ref , x ) ≤ r ] . (146) It is substituted into the explicit transformation S x q S^{q}_{x} S x q (12) . The remaining quantities C C C Ψ \Psi Ψ k k k β \beta β k k k β \beta β
We use M = 30 M = 30 M = 30 30 × 7,500 = 225,000 30 \times7{,}500 = 225{,}000 30 × 7 , 500 = 225 , 000
0.5 10 – 12 ≈ 4.2 % – 5.0 % (147) \frac{0.5}{10\text{--}12} \approx4.2\%\text{--}5.0\% \tag*{(147)} 10 – 12 0.5 ≈ 4.2% – 5.0% (147) of its wall-clock time. More importantly, this is a one-time cost shared by every run using the same model and dataset. Amortized over N N N ( 4.2 % – 5.0 % ) / N (4.2\%\text{--}5.0\%)/N ( 4.2% – 5.0% ) / N
Calibration Algorithms
Algorithms 1 and 2 summarize reference calibration for binary and continuous win-rate rewards, respectively. Both start from the same reusable reference bank: for each training prompt x ∈ X x \in\mathcal{X} x ∈ X y x , 1 r e f , … , y x , M r e f ∼ i . i . d . π r e f ( ⋅ ∣ x ) y^{\mathrm{ref}}_{x,1},\ldots,y^{\mathrm{ref}}_{x,M} \mathrel{\stackrel{\mathrm{i.i.d.}}{\sim}} \pi_{\mathrm{ref}}(\cdot\mid x) y x , 1 ref , … , y x , M ref ∼ i.i.d. π ref ( ⋅ ∣ x ) r x , j r e f = R ( y x , j r e f , x ) r^{\mathrm{ref}}_{x,j} = R(y^{\mathrm{ref}}_{x,j},x) r x , j ref = R ( y x , j ref , x ) k k k β \beta β
Binary rewards. The calibration depends on each prompt only through its reference success count N ( x ) N(x) N ( x ) M + 1 M+1 M + 1 A ⋆ A^\star A ⋆ ℓ ( x ) = β log Z ^ ( x ) \ell(x) = \beta\log\widehat{Z}(x) ℓ ( x ) = β log Z ( x ) Z ^ ( x ) \widehat{Z}(x) Z ( x )
During training, binary @k-A⋆PO uses the regression target a ^ k ( x ) R ( y , x ) − ℓ ( x ) \widehat{a}_k(x)R(y,x) - \ell(x) a k ( x ) R ( y , x ) − ℓ ( x ) β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) § 21 .§ 2 .
Continuous win-rate rewards. An empirical CDF constructed from M M M { 0 , 1 / M , … , 1 } \{0,1/M,\ldots,1\} { 0 , 1/ M , … , 1 }
For an online completion with reward r r r S ^ x q ( r ) = T M F ^ r e f ( r ∣ x ) \widehat{S}^{q}_{x}(r) = T_{M}\widehat{F}_{\mathrm{ref}}(r\mid x) S x q ( r ) = T M F ref ( r ∣ x )
Continuous @k-A⋆PO uses this value as its regression target, retaining Z ( x ) = 1 Z(x) = 1 Z ( x ) = 1 § 17 . @k-GRPO uses the same transformed rewards before its usual group normalization. The parameter β \beta β
Boundary Handling
A finite reference bank can contain no successes or only successes even when the true success probability lies strictly between zero and one. The implementation therefore needs explicit conventions
Algorithm 1 Reference calibration for binary rewards
Require: Reference rewards $\{r^{\mathrm{ref}}_{x,j}\}$, bank size $M$, training parameter $k \ge 1$, $\beta > 0$, method $m \in \{A^\star\mathrm{PO},\mathrm{DPO}\}$
1: $\alpha \leftarrow 0$ for $A^\star\mathrm{PO}$, or $\alpha \leftarrow 0.5$ for DPO
2: for $j = 0, \ldots, M$ do
3: $q_j \leftarrow j/(M + \alpha)$
4: if $q_j = 0$ then
5: $p_j \leftarrow 0$; $a_j \leftarrow k$; $\ell_j \leftarrow 0$
6: else if $q_j = 1$ then
7: $p_j \leftarrow 1$
8: $a_j \leftarrow 1$ if $k = 1$, otherwise $a_j \leftarrow 0$
9: $\ell_j \leftarrow a_j$
10: else
11: Solve {xref}`Eq. 6` with $p_{\mathrm{ref}} = q_j$ for $p_j$
12: $a_j \leftarrow k(1-p_j)^{k-1}$
13: $\ell_j \leftarrow \beta \log(1-q_j+q_j\exp(a_j/\beta))$
14: end if
15: if $m = \mathrm{DPO}$ then
16: $b_j \leftarrow \min\{\beta/a_j,2\}$
17: end if
18: end for
19: for each prompt $x \in \mathcal{X}$ do
20: $N(x) \leftarrow \sum_{j=1}^{M}\mathbb{I}[r^{\mathrm{ref}}_{x,j}=1]$
21: $\widehat{p}^{\ast}(x) \leftarrow p_{N(x)}$
22: if $m = A^\star\mathrm{PO}$ then
23: $\widehat{a}_k(x) \leftarrow a_{N(x)}$; $\ell(x) \leftarrow \ell_{N(x)}$
24: else
25: $\widehat{\beta}_{\mathrm{eff}}(x) \leftarrow b_{N(x)}$
26: end if
27: end for
28: return $\{\widehat{a}_k(x),\ell(x)\}_{x\in\mathcal{X}}$ for $A^\star\mathrm{PO}$, or $\{\widehat{\beta}_{\mathrm{eff}}(x)\}_{x\in\mathcal{X}}$ for DPOfor these empirical endpoints, distinct from the population boundary cases in Appendix § 13 . Let N ( x ) N(x) N ( x ) M M M
Binary @k-A ⋆ A^\star A ⋆ We use the unsmoothed estimate p ^ r e f ( x ) = N ( x ) / M \widehat{p}_{\mathrm{ref}}(x)=N(x)/M p ref ( x ) = N ( x ) / M a k ( x ) = β / β e f f ( x ) a_k(x)=\beta/\beta_{\mathrm{eff}}(x) a k ( x ) = β / β eff ( x ) k > 1 k>1 k > 1
( p ^ ∗ ( x ) , a ^ k ( x ) , β log Z ^ ( x ) ) = { ( 0 , k , 0 ) , N ( x ) = 0 , ( 1 , 0 , 0 ) , N ( x ) = M . (148) \left(\widehat{p}^{\ast}(x),\widehat{a}_k(x),\beta\log\widehat{Z}(x)\right)=
\begin{cases}
(0,k,0), & N(x)=0,\\
(1,0,0), & N(x)=M.
\end{cases} \tag*{(148)} ( p ∗ ( x ) , a k ( x ) , β log Z ( x ) ) = { ( 0 , k , 0 ) , ( 1 , 0 , 0 ) , N ( x ) = 0 , N ( x ) = M . (148) We compute the regression target a ^ k ( x ) R ( y , x ) − β log Z ^ ( x ) \widehat{a}_k(x)R(y,x)-\beta\log\widehat{Z}(x) a k ( x ) R ( y , x ) − β log Z ( x ) β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) N ( x ) = M N(x)=M N ( x ) = M k R ( y , x ) kR(y,x) k R ( y , x ) k > 1 k>1 k > 1 k = 1 k=1 k = 1 N ( x ) = 0 N(x)=0 N ( x ) = 0 N ( x ) = M N(x)=M N ( x ) = M R ( y , x ) R(y,x) R ( y , x ) R ( y , x ) − 1 R(y,x)-1 R ( y , x ) − 1
Binary @k-DPO. Unlike the regression target, the DPO loss requires a finite coefficient. We therefore add a failure pseudo-count and solve the fixed-point equation using p ^ r e f ( x ) = N ( x ) / ( M + α ) \widehat{p}_{\mathrm{ref}}(x)=N(x)/(M+\alpha) p ref ( x ) = N ( x ) / ( M + α ) α = 0.5 \alpha=0.5 α = 0.5
β ^ e f f ( x ) = min { β k ( 1 − p ^ ∗ ( x ) ) k − 1 , 2 } . (149) \widehat{\beta}_{\mathrm{eff}}(x)=\min\left\{\frac{\beta}{k(1-\widehat{p}^{\ast}(x))^{k-1}},2\right\}. \tag*{(149)} β eff ( x ) = min { k ( 1 − p ∗ ( x ) ) k − 1 β , 2 } . (149) Smoothing prevents the infinite coefficient at an empirical success probability of one, while the cap limits very large coefficients near this boundary.
Algorithm 2 Reference calibration for continuous win-rate rewards
Require: Reference rewards { r x , j r e f } \{r^{\mathrm{ref}}_{x,j}\} { r x , j ref } M M M k ≥ 1 k \ge1 k ≥ 1 β > 0 \beta> 0 β > 0
if k = 1 k = 1 k = 1 then
log C ← − log β − β − 1 − log ( 1 − e − 1 / β ) \log C \leftarrow-\log\beta- \beta^{-1} - \log(1 - e^{-1/\beta}) log C ← − log β − β − 1 − log ( 1 − e − 1/ β )
for j = 0 , … , M j = 0, \ldots, M j = 0 , … , M do
T j ← j / M + β log C T_j \leftarrow j/M + \beta\log C T j ← j / M + β log C
end for
else
Solve ∫ 0 1 d u C + u k / β = 1 \int_0^1 \frac{du}{C + u^{k/\beta}} = 1 ∫ 0 1 C + u k / β d u = 1 C > 0 C > 0 C > 0
Define Ψ ( z ) = ∫ 0 z d u C + u k / β \Psi(z) = \int_0^z \frac{du}{C + u^{k/\beta}} Ψ ( z ) = ∫ 0 z C + u k / β d u
T 0 ← β log C T_0 \leftarrow\beta\log C T 0 ← β log C T M ← β log ( C + 1 / β ) T_M \leftarrow\beta\log(C + 1/\beta) T M ← β log ( C + 1/ β )
for j = 1 , … , M − 1 j = 1, \ldots, M - 1 j = 1 , … , M − 1 do
Solve Ψ ( z j ) = j / M \Psi(z_j) = j/M Ψ ( z j ) = j / M z j ∈ ( 0 , 1 ) z_j \in(0,1) z j ∈ ( 0 , 1 )
T j ← β log ( C + z j k / β ) T_j \leftarrow\beta\log(C + z_j^{k}/\beta) T j ← β log ( C + z j k / β )
end for
end if
return Stored reference rewards and the shared table { T j } j = 0 M \{T_j\}_{j=0}^{M} { T j } j = 0 M
function TRANSFORMREWARD ( x , r ) (x,r) ( x , r )
j ← ∑ i = 1 M I [ r x , i r e f ≤ r ] j \leftarrow\sum_{i=1}^{M} \mathbb{I}[r^{\mathrm{ref}}_{x,i} \le r] j ← ∑ i = 1 M I [ r x , i ref ≤ r ]
return T j T_j T j
end function
Win-rate endpoints and ties. For continuous @k-A⋆PO and @k-GRPO, we use the inclusive empirical CDF without smoothing or clipping its values away from zero and one. The transformation remains finite at both endpoints:
S x q ( r ) = { β log C , F ^ r e f ( r ∣ x ) = 0 , β log ( C + 1 / β ) , F ^ r e f ( r ∣ x ) = 1. (150) S_x^q(r) =
\begin{cases}
\beta\log C, & \widehat{F}_{\mathrm{ref}}(r \mid x) = 0,\\
\beta\log(C + 1/\beta), & \widehat{F}_{\mathrm{ref}}(r \mid x) = 1.
\end{cases} \tag*{(150)} S x q ( r ) = { β log C , β log ( C + 1/ β ) , F ref ( r ∣ x ) = 0 , F ref ( r ∣ x ) = 1. (150) Rewards below the smallest reference reward receive the first value; rewards at or above the largest receive the second. Ties retain the inclusive-CDF convention, including for a constant-reward reference bank. For A⋆PO, we retain Z ( x ) = 1 Z(x) = 1 Z ( x ) = 1 § 17 .
Numerical stability. For interior binary estimates, we solve the scalar equation in log-odds coordinates and compute a k a_k a k 1 − p ∗ 1 - p^* 1 − p ∗ k = 1 k = 1 k = 1 β \beta β
Degenerate online groups. In @k-GRPO, we divide centered transformed rewards by their group standard deviation plus 10 − 6 10^{-6} 1 0 − 6
Sensitivity to Finite Reference Calibration
Reference calibration computes the quantities and mappings required by KAPPO before training. For binary rewards, we first estimate p r e f ( x ) p_{\mathrm{ref}}(x) p ref ( x ) (6) for p ∗ ( x ) p^*(x) p ∗ ( x ) a k ( x ) a_k(x) a k ( x ) β e f f ( x ) \beta_{\mathrm{eff}}(x) β eff ( x ) F r e f ( r ∣ x ) F_{\mathrm{ref}}(r \mid x) F ref ( r ∣ x ) S x q ( r ) S_x^q(r) S x q ( r ) p r e f p_{\mathrm{ref}} p ref F r e f F_{\mathrm{ref}} F ref M = 30 M = 30 M = 30 β = 0.1 \beta= 0.1 β = 0.1
Exact finite-sample calculation. Fix a prompt and let q q q q = p r e f ( x ) q = p_{\mathrm{ref}}(x) q = p ref ( x ) q = F r e f ( r ∣ x ) q = F_{\mathrm{ref}}(r \mid x) q = F ref ( r ∣ x ) r r r R ( y , x ) ≤ r R(y,x) \le r R ( y , x ) ≤ r N N N M M M N ∼ Binomial ( M , q ) N \sim\operatorname{Binomial}(M,q) N ∼ Binomial ( M , q ) q ^ = N / M \hat{q} = N/M q ^ = N / M
Let θ ( q ) \theta(q) θ ( q ) q q q p ∗ p^{*} p ∗ a k a_k a k β e f f \beta_{\mathrm{eff}} β eff S x q ( r ) S_x^{q}(r) S x q ( r ) θ ^ j \hat{\theta}_j θ ^ j N = j N = j N = j Pr ( N = j ) = w j ( q ) : = ( M j ) q j ( 1 − q ) M − j \Pr(N = j) = w_j(q) := \binom{M}{j}q^j(1-q)^{M-j} Pr ( N = j ) = w j ( q ) := ( j M ) q j ( 1 − q ) M − j
μ θ ( q ) = ∑ j = 0 M w j ( q ) θ ^ j , σ θ ( q ) = ( ∑ j = 0 M w j ( q ) [ θ ^ j − μ θ ( q ) ] 2 ) 1 / 2 , b θ ( q ) = μ θ ( q ) − θ ( q ) . (151) \mu_{\theta}(q) = \sum_{j=0}^{M} w_j(q)\hat{\theta}_j,\qquad
\sigma_{\theta}(q) = \left(\sum_{j=0}^{M} w_j(q)[\hat{\theta}_j-\mu_{\theta}(q)]^2\right)^{1/2},\qquad
b_{\theta}(q) = \mu_{\theta}(q)-\theta(q). \tag*{(151)} μ θ ( q ) = j = 0 ∑ M w j ( q ) θ ^ j , σ θ ( q ) = ( j = 0 ∑ M w j ( q ) [ θ ^ j − μ θ ( q ) ] 2 ) 1/2 , b θ ( q ) = μ θ ( q ) − θ ( q ) . (151) We report absolute noise and bias for p ∗ p^{*} p ∗ a k a_k a k β e f f \beta_{\mathrm{eff}} β eff σ θ ( q ) / θ ( q ) \sigma_{\theta}(q)/\theta(q) σ θ ( q ) / θ ( q ) b θ ( q ) / θ ( q ) b_{\theta}(q)/\theta(q) b θ ( q ) / θ ( q ) (151) sums over all M + 1 M+1 M + 1
Binary Rewards
Suppressing the prompt argument, set q = p r e f ( x ) q = p_{\mathrm{ref}}(x) q = p ref ( x ) g ( q ) g(q) g ( q ) (6) . The exact maps applied during reference calibration are
p ∗ = g ( q ) , a k ( q ) = k [ 1 − g ( q ) ] k − 1 , β e f f ( q ) = β a k ( q ) , (152) p^{*}=g(q),\qquad a_k(q)=k[1-g(q)]^{k-1},\qquad\beta_{\mathrm{eff}}(q)=\frac{\beta}{a_k(q)}, \tag*{(152)} p ∗ = g ( q ) , a k ( q ) = k [ 1 − g ( q ) ] k − 1 , β eff ( q ) = a k ( q ) β , (152) where a k a_k a k k > 1 k > 1 k > 1 g ( 0 ) = 0 g(0)=0 g ( 0 ) = 0 g ( 1 ) = 1 g(1)=1 g ( 1 ) = 1 a k ( 0 ) = k a_k(0)=k a k ( 0 ) = k a k ( 1 ) = 0 a_k(1)=0 a k ( 1 ) = 0 k = 1 k = 1 k = 1
Sensitive regions. Figure 5 shows that g g g p ∗ p^{*} p ∗ k = 4 k = 4 k = 4 g ( 0 ) = 0 g(0) = 0 g ( 0 ) = 0 g ( 1 / 30 ) ≈ 0.553 g(1/30) \approx0.553 g ( 1/30 ) ≈ 0.553 p ^ \hat{p} p ^ ( 1 − q ) M (1-q)^M ( 1 − q ) M M = 30 M = 30 M = 30 0.740 0.740 0.740 q = 0.01 q = 0.01 q = 0.01 0.042 0.042 0.042 q = 0.1 q = 0.1 q = 0.1
Figure 5. Exact optimal success probability as a function of the true reference success probability, for β = 0.1 \beta= 0.1 β = 0.1 j / M j/M j / M M = 30 M = 30 M = 30 1 / M 1/M 1/ M
Estimators used in the analysis. For p ∗ p^* p ∗ a k a_k a k q ^ = N / M \hat{q}=N/M q ^ = N / M ⋆ ^\star ⋆ p ^ ∗ = g ( q ^ ) \hat{p}^*=g(\hat{q}) p ^ ∗ = g ( q ^ ) a ^ k = a k ( q ^ ) \hat{a}_k=a_k(\hat{q}) a ^ k = a k ( q ^ ) Figure 6 fixes the true p r e f p_{\mathrm{ref}} p ref β e f f \beta_{\mathrm{eff}} β eff
Figure 6. Exact finite-bank noise (left) and signed bias (right) for M = 30 M = 30 M = 30 β = 0.1 \beta= 0.1 β = 0.1 q ^ = N / M \hat{q} = N/M q ^ = N / M (153) . Coefficient errors are divided by the exact β e f f ( q ) \beta_{\mathrm{eff}}(q) β eff ( q ) q = 1 q = 1 q = 1 k > 1 k > 1 k > 1 [ − 10 − 6 , 10 − 6 ] [-10^{-6}, 10^{-6}] [ − 1 0 − 6 , 1 0 − 6 ]
q ^ D = N M + 0.5 , β ^ e f f = min { β a k ( q ^ D ) , 2 } . (153) \hat{q}_{\mathrm{D}}=\frac{N}{M+0.5},\qquad\hat{\beta}_{\mathrm{eff}}=\min\left\{\frac{\beta}{a_k(\hat{q}_{\mathrm{D}})},2\right\}. \tag*{(153)} q ^ D = M + 0.5 N , β ^ eff = min { a k ( q ^ D ) β , 2 } . (153) The smoothing keeps q ^ D < 1 \hat{q}_{\mathrm{D}}<1 q ^ D < 1 q ^ D = 0 \hat{q}_{\mathrm{D}}=0 q ^ D = 0 N = M N=M N = M k > 1 k>1 k > 1 β e f f ( q ) \beta_{\mathrm{eff}}(q) β eff ( q )
Consistency and convergence rates.
Lemma K.1 (Consistency of binary calibration). Fix an integer k ≥ 1 k\ge1 k ≥ 1 β > 0 \beta>0 β > 0 q ∈ [ 0 , 1 ] q\in[0,1] q ∈ [ 0 , 1 ] p ^ ∗ = g ( q ^ ) \hat{p}^*=g(\hat{q}) p ^ ∗ = g ( q ^ ) a ^ k = a k ( q ^ ) \hat{a}_k=a_k(\hat{q}) a ^ k = a k ( q ^ ) q ∈ ( 0 , 1 ) q\in(0,1) q ∈ ( 0 , 1 ) θ ∈ { g , a k } \theta\in\{g,a_k\} θ ∈ { g , a k }
σ θ ( q ) = ∣ θ ′ ( q ) ∣ q ( 1 − q ) M + o ( M − 1 / 2 ) , b θ ( q ) = q ( 1 − q ) 2 M θ ′ ′ ( q ) + o ( M − 1 ) . (154) \sigma_\theta(q)=\lvert\theta'(q)\rvert\sqrt{\frac{q(1-q)}{M}}+o(M^{-1/2}),\qquad b_\theta(q)=\frac{q(1-q)}{2M}\theta''(q)+o(M^{-1}). \tag*{(154)} σ θ ( q ) = ∣ θ ′ ( q )∣ M q ( 1 − q ) + o ( M − 1/2 ) , b θ ( q ) = 2 M q ( 1 − q ) θ ′′ ( q ) + o ( M − 1 ) . (154) At q = 0 q=0 q = 0 q = 1 q=1 q = 1 M M M
Proof. By the strong law of large numbers, q ^ → q \hat{q}\to q q ^ → q (6) implies that g g g [ 0 , 1 ] [0,1] [ 0 , 1 ] g ( q ) → 0 g(q)\to0 g ( q ) → 0 q ↓ 0 q\downarrow0 q ↓ 0 g ( q ) → 1 g(q)\to1 g ( q ) → 1 q ↑ 1 q\uparrow1 q ↑ 1 a k a_k a k q ^ \hat{q} q ^
For interior q q q g g g (6) , with respect to the candidate success probability p p p
β p ( 1 − p ) + k ( k − 1 ) ( 1 − p ) k − 2 > 0. \frac{\beta}{p(1-p)}+k(k-1)(1-p)^{k-2}>0. p ( 1 − p ) β + k ( k − 1 ) ( 1 − p ) k − 2 > 0. Consequently, both g g g a k a_k a k ( 0 , 1 ) (0,1) ( 0 , 1 )
For either map θ \theta θ q q q
E [ q ^ − q ] = 0 , E [ ( q ^ − q ) 2 ] = q ( 1 − q ) M , \mathbb{E}[\hat{q}-q]=0,\qquad\mathbb{E}[(\hat{q}-q)^2]=\frac{q(1-q)}{M}, E [ q ^ − q ] = 0 , E [( q ^ − q ) 2 ] = M q ( 1 − q ) , gives
E [ θ ( q ^ ) ] − θ ( q ) = θ ′ ′ ( q ) q ( 1 − q ) 2 M + o ( M − 1 ) , Var ( θ ( q ^ ) ) = θ ′ ( q ) 2 q ( 1 − q ) M + o ( M − 1 ) . \mathbb{E}[\theta(\hat{q})]-\theta(q)=\frac{\theta''(q)q(1-q)}{2M}+o(M^{-1}),\qquad\operatorname{Var}(\theta(\hat{q}))=\frac{\theta'(q)^2q(1-q)}{M}+o(M^{-1}). E [ θ ( q ^ )] − θ ( q ) = 2 M θ ′′ ( q ) q ( 1 − q ) + o ( M − 1 ) , Var ( θ ( q ^ )) = M θ ′ ( q ) 2 q ( 1 − q ) + o ( M − 1 ) . These expansions remain valid despite possible boundary behavior of the derivatives: outside a fixed interior neighborhood of q q q g g g a k a_k a k (154) . Finally, when q ∈ { 0 , 1 } q\in\{0,1\} q ∈ { 0 , 1 } q ^ = q \hat{q}=q q ^ = q M M M
Thus, sampling noise is O ( M − 1 / 2 ) O(M^{-1/2}) O ( M − 1/2 ) O ( M − 1 ) O(M^{-1}) O ( M − 1 ) M = 30 M=30 M = 30 (151) .
Effect of coefficient stabilization.
Lemma K.2 (Consistency under smoothing and clipping). Fix an integer k ≥ 1 k\ge1 k ≥ 1 β > 0 \beta>0 β > 0 q ∈ ( 0 , 1 ) q\in(0,1) q ∈ ( 0 , 1 ) (153) satisfies
β ^ e f f → M → ∞ a . s . min { β e f f ( q ) , 2 } . \hat{\beta}_{\mathrm{eff}}\xrightarrow[M\to\infty]{\mathrm{a.s.}}\min\{\beta_{\mathrm{eff}}(q),2\}. β ^ eff a.s. M → ∞ min { β eff ( q ) , 2 } . If β e f f ( q ) < 2 \beta_{\mathrm{eff}}(q) < 2 β eff ( q ) < 2 O ( M − 1 / 2 ) O(M^{-1/2}) O ( M − 1/2 ) O ( M − 1 ) O(M^{-1}) O ( M − 1 ) β e f f ( q ) > 2 \beta_{\mathrm{eff}}(q) > 2 β eff ( q ) > 2 2 − β e f f ( q ) 2-\beta_{\mathrm{eff}}(q) 2 − β eff ( q ) M M M β e f f ( q ) = 2 \beta_{\mathrm{eff}}(q) = 2 β eff ( q ) = 2 k > 1 k > 1 k > 1 M − 1 / 2 M^{-1/2} M − 1/2
Proof. The smoothed estimate can be written as
q ^ D = M M + 0.5 q ^ . \widehat{q}_{D} = \frac{M}{M+0.5}\widehat{q}. q D = M + 0.5 M q . It therefore converges almost surely to q q q
E [ q ^ D ] − q = − q 2 M + 1 = O ( M − 1 ) . \mathbb{E}[\widehat{q}_{D}] - q = -\frac{q}{2M+1} = O(M^{-1}). E [ q D ] − q = − 2 M + 1 q = O ( M − 1 ) . Continuity of the capped coefficient at every interior q q q
When β e f f ( q ) < 2 \beta_{\mathrm{eff}}(q) < 2 β eff ( q ) < 2 q q q q ^ D \widehat{q}_{D} q D O ( M − 1 / 2 ) O(M^{-1/2}) O ( M − 1/2 ) O ( M − 1 ) O(M^{-1}) O ( M − 1 )
When β e f f ( q ) > 2 \beta_{\mathrm{eff}}(q) > 2 β eff ( q ) > 2 q q q q ^ D \widehat{q}_{D} q D q q q
Finally, suppose that k > 1 k > 1 k > 1 β e f f ( q ) = 2 \beta_{\mathrm{eff}}(q) = 2 β eff ( q ) = 2
M ( β ^ e f f − 2 ) → d β e f f ′ ( q ) min { Z , 0 } , Z ∼ N ( 0 , q ( 1 − q ) ) . \sqrt{M}(\widehat{\beta}_{\mathrm{eff}}-2) \xrightarrow{d} \beta_{\mathrm{eff}}'(q)\min\{Z,0\}, \qquad Z \sim N(0,q(1-q)). M ( β eff − 2 ) d β eff ′ ( q ) min { Z , 0 } , Z ∼ N ( 0 , q ( 1 − q )) . Binomial moment bounds and local smoothness also justify convergence of the first moments. Hence,
E [ β ^ e f f ] − 2 = − β e f f ′ ( q ) q ( 1 − q ) 2 π M + o ( M − 1 / 2 ) , \mathbb{E}[\widehat{\beta}_{\mathrm{eff}}]-2 = -\beta_{\mathrm{eff}}'(q)\sqrt{\frac{q(1-q)}{2\pi M}} + o(M^{-1/2}), E [ β eff ] − 2 = − β eff ′ ( q ) 2 π M q ( 1 − q ) + o ( M − 1/2 ) , establishing the downward bias of order M − 1 / 2 M^{-1/2} M − 1/2
The failure pseudo-count thus vanishes asymptotically, whereas the fixed coefficient cap does not. Dividing the noise and bias by the fixed, finite β e f f ( q ) \beta_{\mathrm{eff}}(q) β eff ( q )
Implications. For p ∗ p^* p ∗ a k a_k a k k = 4 k = 4 k = 4 p r e f = 0.1 p_{\mathrm{ref}} = 0.1 p ref = 0.1 p ^ ∗ \widehat{p}^* p ∗ − 0.028 -0.028 − 0.028 p r e f = 0.5 p_{\mathrm{ref}} = 0.5 p ref = 0.5 p ^ ∗ = 0 \widehat{p}^* = 0 p ∗ = 0 a ^ k = k \widehat{a}_k = k a k = k M → ∞ M \to\infty M → ∞ ( 1 − p r e f ) M (1-p_{\mathrm{ref}})^M ( 1 − p ref ) M M M M M p r e f Mp_{\mathrm{ref}} M p ref p r e f = 0 p_{\mathrm{ref}} = 0 p ref = 0
For k > 1 k > 1 k > 1 p r e f = 1 p_{\mathrm{ref}} = 1 p ref = 1 β e f f \beta_{\mathrm{eff}} β eff M M M
Figure 7. Win-rate transformation and exact finite-bank noise and bias for M = 30 M = 30 M = 30 β = 0.1 \beta= 0.1 β = 0.1 k k k Δ T = T ( 1 ) − T ( 0 ) \Delta_T = T(1) - T(0) Δ T = T ( 1 ) − T ( 0 ) T ( 0 ) T(0) T ( 0 ) Figure 6 .
Continuous Rewards
For continuous rewards, reference calibration estimates the CDF value at each reward threshold and substitutes it into the win-rate-based transformation. At a fixed prompt and raw reward r r r (12) as
S q ( r ) = T ( q ) , q = F r e f ( r ∣ x ) , T ( q ) = β log ( C + G ( q ) k β ) , G = Ψ − 1 . (155) S^q(r) = T(q), \qquad q = F_{\mathrm{ref}}(r \mid x), \qquad T(q) = \beta\log\left(C + \frac{G(q)^k}{\beta}\right), \qquad G = \Psi^{-1}. \tag*{(155)} S q ( r ) = T ( q ) , q = F ref ( r ∣ x ) , T ( q ) = β log ( C + β G ( q ) k ) , G = Ψ − 1 . (155) Here, N N N r r r N ∼ Binomial ( M , q ) N \sim\operatorname{Binomial}(M,q) N ∼ Binomial ( M , q ) q ^ = N / M \hat{q} = N/M q ^ = N / M T ( q ^ ) T(\hat{q}) T ( q ^ ) (151) by setting θ ( q ) = T ( q ) \theta(q) = T(q) θ ( q ) = T ( q ) θ ^ j = T ( j / M ) \hat{\theta}_j = T(j/M) θ ^ j = T ( j / M ) k k k Δ T = T ( 1 ) − T ( 0 ) \Delta_T = T(1) - T(0) Δ T = T ( 1 ) − T ( 0 )
Sensitivity and convergence guarantees.
Lemma K.3 (Consistency and error bounds for win-rate calibration). Fix an integer k ≥ 1 k \ge1 k ≥ 1 β > 0 \beta> 0 β > 0 T T T (155) is convex and k k k q ∈ [ 0 , 1 ] q \in[0,1] q ∈ [ 0 , 1 ] M M M T ( q ^ ) T(\hat{q}) T ( q ^ )
σ T ( q ) ≤ k q ( 1 − q ) M , 0 ≤ b T ( q ) ≤ k ( k − 1 ) ( C + β − 1 ) q ( 1 − q ) 2 M . (156) \sigma_T(q) \le k\sqrt{\frac{q(1-q)}{M}}, \qquad0 \le b_T(q) \le\frac{k(k-1)(C+\beta^{-1})q(1-q)}{2M}. \tag*{(156)} σ T ( q ) ≤ k M q ( 1 − q ) , 0 ≤ b T ( q ) ≤ 2 M k ( k − 1 ) ( C + β − 1 ) q ( 1 − q ) . (156) At fixed interior q q q (154) apply with θ = T \theta= T θ = T q = 0 q = 0 q = 0 q = 1 q = 1 q = 1
For each fixed prompt, the estimated mapping S ^ x q ( r ) : = T ( F ^ r e f ( r ∣ x ) ) \widehat{S}_{x}^{q}(r) := T(\widehat{F}_{\mathrm{ref}}(r \mid x)) S x q ( r ) := T ( F ref ( r ∣ x )) S x q ( r ) S_{x}^{q}(r) S x q ( r ) δ ∈ ( 0 , 1 ) \delta\in(0,1) δ ∈ ( 0 , 1 )
Pr ( sup r ∣ S ^ x q ( r ) − S x q ( r ) ∣ > k log ( 2 / δ ) 2 M ) ≤ δ . (157) \Pr\left(\sup_{r}\left|\widehat{S}_{x}^{q}(r)-S_{x}^{q}(r)\right|>k\sqrt{\frac{\log(2/\delta)}{2M}}\right)\leq\delta. \tag*{(157)} Pr ( r sup S x q ( r ) − S x q ( r ) > k 2 M log ( 2/ δ ) ) ≤ δ . (157) These estimation guarantees hold for both continuous and discrete reference reward distributions.
Proof. Shape of the transformation. Using G ′ ( q ) = C + G ( q ) k / β G'(q)=C+G(q)^k/\beta G ′ ( q ) = C + G ( q ) k / β k > 1 k>1 k > 1
T ′ ( q ) = k G ( q ) k − 1 , T ′ ′ ( q ) = k ( k − 1 ) G ( q ) k − 2 ( C + G ( q ) k β ) . (158) T'(q)=kG(q)^{k-1},\qquad T''(q)=k(k-1)G(q)^{k-2}\left(C+\frac{G(q)^k}{\beta}\right). \tag*{(158)} T ′ ( q ) = k G ( q ) k − 1 , T ′′ ( q ) = k ( k − 1 ) G ( q ) k − 2 ( C + β G ( q ) k ) . (158) Since G ( q ) ∈ [ 0 , 1 ] G(q)\in[0,1] G ( q ) ∈ [ 0 , 1 ]
0 ≤ T ′ ( q ) ≤ k , 0 ≤ T ′ ′ ( q ) ≤ k ( k − 1 ) ( C + β − 1 ) . 0\leq T'(q)\leq k,\qquad0\leq T''(q)\leq k(k-1)(C+\beta^{-1}). 0 ≤ T ′ ( q ) ≤ k , 0 ≤ T ′′ ( q ) ≤ k ( k − 1 ) ( C + β − 1 ) . Thus, T T T k k k k = 1 k=1 k = 1 T ′ ( q ) = 1 T'(q)=1 T ′ ( q ) = 1 T ′ ′ ( q ) = 0 T''(q)=0 T ′′ ( q ) = 0
Pointwise noise and bias. The Lipschitz property gives
Var ( T ( q ^ ) ) ≤ E [ ( T ( q ^ ) − T ( q ) ) 2 ] ≤ k 2 E [ ( q ^ − q ) 2 ] = k 2 q ( 1 − q ) M . \operatorname{Var}(T(\widehat{q}))\leq\mathbb{E}\left[(T(\widehat{q})-T(q))^2\right]\leq k^2\mathbb{E}\left[(\widehat{q}-q)^2\right]=\frac{k^2q(1-q)}{M}. Var ( T ( q )) ≤ E [ ( T ( q ) − T ( q ) ) 2 ] ≤ k 2 E [ ( q − q ) 2 ] = M k 2 q ( 1 − q ) . Taking square roots proves the noise bound. Convexity and E [ q ^ ] = q \mathbb{E}[\widehat{q}]=q E [ q ] = q b T ( q ) ≥ 0 b_T(q)\geq0 b T ( q ) ≥ 0 T ′ ′ T'' T ′′
T ( q ^ ) ≤ T ( q ) + T ′ ( q ) ( q ^ − q ) + k ( k − 1 ) ( C + β − 1 ) 2 ( q ^ − q ) 2 . T(\widehat{q})\leq T(q)+T'(q)(\widehat{q}-q)+\frac{k(k-1)(C+\beta^{-1})}{2}(\widehat{q}-q)^2. T ( q ) ≤ T ( q ) + T ′ ( q ) ( q − q ) + 2 k ( k − 1 ) ( C + β − 1 ) ( q − q ) 2 . Taking expectations proves the upper bound on bias. At fixed interior q q q (154) . At either endpoint, q ^ = q \widehat{q}=q q = q
Uniform consistency. The Lipschitz property also implies
sup r ∣ S ^ x q ( r ) − S x q ( r ) ∣ ≤ k sup r ∣ F ^ r e f ( r ∣ x ) − F r e f ( r ∣ x ) ∣ . \sup_r\left|\widehat{S}_{x}^{q}(r)-S_{x}^{q}(r)\right|\leq k\sup_r\left|\widehat{F}_{\mathrm{ref}}(r\mid x)-F_{\mathrm{ref}}(r\mid x)\right|. r sup S x q ( r ) − S x q ( r ) ≤ k r sup F ref ( r ∣ x ) − F ref ( r ∣ x ) . Applying the Dvoretzky–Kiefer–Wolfowitz–Massart inequality ([15 ] ) proves Eq. (157) . More generally, for every ε > 0 \varepsilon>0 ε > 0
Pr ( sup r ∣ S ^ x q ( r ) − S x q ( r ) ∣ > ε ) ≤ 2 exp ( − 2 M ε 2 k 2 ) . \Pr\left(\sup_r\left|\widehat{S}_{x}^{q}(r)-S_{x}^{q}(r)\right|>\varepsilon\right)\leq2\exp\left(-\frac{2M\varepsilon^2}{k^2}\right). Pr ( r sup S x q ( r ) − S x q ( r ) > ε ) ≤ 2 exp ( − k 2 2 M ε 2 ) . The right-hand side is summable over M M M
Where quantile error matters. Equation (158) shows that sensitivity increases toward the upper reference quantiles. Convexity also explains the upward bias: although the empirical CDF is unbiased at each fixed threshold, the transformed reward generally is not. For k = 1 k=1 k = 1 T T T q ( 1 − q ) / M \sqrt{q(1-q)/M} q ( 1 − q ) / M
The finite-sample bounds establish O ( M − 1 / 2 ) O(M^{-1/2}) O ( M − 1/2 ) O ( M − 1 ) O(M^{-1}) O ( M − 1 ) k k k β \beta β (154) gives the leading constants at fixed interior quantiles. These guarantees concern estimation of the mapping T ( F r e f ( r ∣ x ) ) T(F_{\mathrm{ref}}(r\mid x)) T ( F ref ( r ∣ x ))
Implications. The win-rate transformation has the opposite sensitive region: both noise and bias are small at low and intermediate quantiles but increase toward the upper tail, where T T T q = 1 q = 1 q = 1 k = 4 k = 4 k = 4 q = 0.5 q = 0.5 q = 0.5 q = 0.9 q = 0.9 q = 0.9 q = 0.99 q = 0.99 q = 0.99 Pr ( q ^ = 1 ) = q 30 ≈ 0.740 \Pr(\widehat{q} = 1) = q^{30} \approx0.740 Pr ( q = 1 ) = q 30 ≈ 0.740
Both effects vanish as M → ∞ M \to\infty M → ∞ q ^ \widehat{q} q q q q q M q^M q M M ( 1 − q ) M(1 - q) M ( 1 − q ) M = 30 M = 30 M = 30 § 17 .
Comparison with Direct Multi-Completion Optimization
Checkpoint selection. For each model, method, and training value k t r a i n ∈ { 2 , 4 } k_{\mathrm{train}} \in\{2, 4\} k train ∈ { 2 , 4 } k t r a i n k_{\mathrm{train}} k train Figure 3 reports the complete validation curves of the selected checkpoints, and Table 1 reports the corresponding numerical values.
Table 1. MATH Level 5 validation performance (in %) of the starting checkpoints and selected @k-A ⋆ A^{\star} A ⋆ Figure 3 .
At the evaluation budget matching the training objective, @k-A⋆PO is within 0.5 percentage points of PKPO in every comparison. The two methods also produce similar performance across the remaining evaluation budgets. This supports the central prediction of our analysis: adaptive pressure applied through an ordinary single-completion optimizer can reproduce direct multi-completion pass@k optimization.
Figure 8 extends this comparison to AIME 2026. Across both model families and training values of k k k
Figure 8. AIME 2026 test performance of the starting checkpoints and the selected @k-A⋆PO and PKPO checkpoints across evaluation sampling budgets. Each subfigure reports Qwen3-4B (left) and OLMo-3-7B (right).
Rollout-compute comparison. With a prompt batch of 128 and 1,000 training steps, @k-A⋆PO generates 128,000 on-policy completions per run. PKPO generates n n n n n n n = k n = k n = k k = 2 k = 2 k = 2 k = 4 k = 4 k = 4
@k-A ⋆ A^{\star} A ⋆ § 21 . In our setup, generating the complete reference bank took approximately half an hour, compared with 10–12 hours for one training run, and the same bank is reused across all configurations of a given model and dataset. Even when amortized over only six runs, calibration therefore adds approximately five minutes per run, or less than 1% of training time. Its contribution is consequently small relative to the savings from removing PKPO’s k k k
