Skip-Connected Policy Optimization for Implicit Advantage — part 2/4
Cleaned-text code points [13842, 29395). Bibliography omitted here; source.md includes it.
To provide fine-grained rewards for reasoning step quality, we consider token-level Monte Carlo reward estimation through trajectory sampling. For any token position t t in a reasoning trajectory o 1 : T o_{1:T} , we estimate dense token-level rewards by sampling K K continuations from the prefix o 1 : t o_{1:t} and averaging their outcome rewards: ℛ MC ( o t , a ) \displaystyle\mathcal{R}_{\mathrm{MC}}(o_{t},a) = 1 K ∑ k = 1 K ℛ score ( o ( k ) , a ) , \displaystyle=\frac{1}{K}\sum_{k=1}^{K}\mathcal{R}_{\mathrm{score}}(o^{(k)},a), (1) where o ( k ) ∼ π 𝜽 ( ⋅ | o 1 : t ) , \displaystyle\text{where }o^{(k)}\sim\pi_{\boldsymbol{\theta}}(\cdot|o_{1:t}), and ℛ score ( o , a ) \mathcal{R}_{\mathrm{score}}(o,a) is a simple scoring function based on outcome correctness. In this idealized formulation, each token receives a distinct reward signal reflecting how well its prefix facilitates correct completions. Notably, as t t approaches the trajectory end T T , the Monte Carlo estimates naturally converge toward the outcome reward ℛ score ( o , a ) \mathcal{R}_{\mathrm{score}}(o,a) , since diminishing remaining uncertainty leaves little room for divergent continuations. This convergence property places outcome-based methods such as GRPO ( Shao et al., 2024 ) , which directly broadcast the outcome reward to all tokens without any intermediate sampling, as the limiting behavior that Monte Carlo estimation approaches for late-trajectory tokens. While theoretically appealing for offering dense rewards at every reasoning step, this ideal token-level Monte Carlo approach faces severe practical challenges when combined with group-relative advantage computation. Computing group-relative advantages at intermediate positions requires forming stable comparison groups through extensive continuation sampling at each token location. For early reasoning tokens, accurately estimating advantage signs demands prohibitively large K K values, and insufficient sampling causes token-level rewards to paradoxically underperform outcome-level baselines. Conversely, GRPO’s broadcasting of outcome rewards, while sacrificing fine-grained discrimination, provides a conservative estimate that avoids the catastrophic sign noise of undersampled Monte Carlo. This perspective also explains why curriculum learning ( Xu and Ding, 2025 ) and resampling ( Yu et al., 2025b ) improve stability: prompts with intermediate accuracy exhibit smaller reward and advantage estimation errors (see Section A.3 ).
#### Upstream-Downstream Instantiation: A Cost-Effective Step Beyond.
To balance the computational feasibility with fine-grained rewards, we instantiate Monte Carlo trajectory sampling at a single intermediate position through an upstream-downstream architecture. The policy model π 𝜽 \pi_{\boldsymbol{\theta}} first generates an early-stopped reasoning segment s ∈ 𝒮 s\in\mathcal{S} in the upstream phase, with the segmentation position determined adaptively based on cumulative token-level statistics rather than fixed formatting rules. Specifically, for each prompt q q , we uniformly sample a termination position t q ∼ 𝒰 [ 1 6 L q , 1 2 L q ] t_{q}\sim\mathcal{U}[\frac{1}{6}L_{q},\frac{1}{2}L_{q}] , where L q L_{q} is the per-prompt average response length tracked via a similar KL-adaptive mechanism as the upstream value tracker (see Section 3.2 ), without warm-up to ensure cost parity with baselines. For responses that terminate before t q t_{q} , we use half of the actual response length as the split position. We then sample G G continuations from the augmented input [ s , q ] [s,q] , the concatenation of the segment and original problem, to obtain complete responses { o i } i = 1 G \{o_{i}\}_{i=1}^{G} in the downstream phase. This skip connection [ s , q ] [s,q] serves dual functions: ❶ it preserves the upstream segment’s influence on downstream reasoning, enabling valid Monte Carlo estimation ; ❷ it maintains exploration freedom by providing a shortcut to the original problem , allowing the model to implicitly correct or bypass unhelpful upstream reasoning. Formally: I : \displaystyle\mathrm{I}: s ∼ π 𝜽 ( ⋅ | q ) , \displaystyle s\sim\pi_{\boldsymbol{\theta}}(\cdot|q), (2) II : \displaystyle\mathrm{II}: { o i } i = 1 G ∼ π 𝜽 ( ⋅ | [ s , q ] ) . \displaystyle\{o_{i}\}_{i=1}^{G}\sim\pi_{\boldsymbol{\theta}}(\cdot|[s,q]). where q ∈ 𝒬 q\in\mathcal{Q} denotes the problem, s ∈ 𝒮 s\in\mathcal{S} the reasoning segment, and o ∈ 𝒪 o\in\mathcal{O} the final output. We define reward signals accordingly: downstream reward ℛ II ( o , a ) = 2 ⋅ 𝟙 [ o = a ] − 1 \mathcal{R}_{\mathrm{II}}(o,a)=2\cdot\mathbbm{1}[o=a]-1 measures final correctness, while upstream reward ℛ I ( s , a ) \mathcal{R}_{\mathrm{I}}(s,a) implements the Monte Carlo estimation: ℛ I ( s , a ) = 1 G ∑ i = 1 G ℛ II ( o i , a ) , where o i ∼ π 𝜽 ( ⋅ ∣ [ s , q ] ) . \mathcal{R}_{\mathrm{I}}(s,a)=\frac{1}{G}\sum_{i=1}^{G}\mathcal{R}_{\mathrm{II}}(o_{i},a),\quad\text{where }o_{i}\sim\pi_{\boldsymbol{\theta}}(\cdot\mid[s,q]). (3) This design achieves a favorable tradeoff: whereas GRPO requires G G complete responses per problem with no intermediate feedback, and full token-level Monte Carlo would require K K continuations for each token across all G G responses, our approach requires only G G samples total while providing one critical intermediate reward signal.
Engineering Optimization. A naive implementation would first generate 1 1 upstream segment, then spawn G G downstream continuations, requiring two separate GPU batches. Although this saves the complexity of ( G − 1 ) (G-1) upstream segments compared to generating G G complete responses, the overhead of two separate batch dispatches actually increases overall training time significantly. Figure 11 in Appendix B illustrates the typical difference between two-batch approaches and our optimized single-pass implementation. Notably, Critique-GRPO, which first rolls out 7 direct responses then generates 1 self-critique response in a separate batch, exhibits substantially higher per-step training time due to this sequential dispatch overhead, with time complexity similar to our naive two-batch design. To achieve true computational parity with mainstream GRPO variants, we modify the vLLM ( Kwon et al., 2023 ) inference engine to perform mid-generation KV cache rewriting within a single rollout pass. Concretely, all G G responses are generated in parallel from q q , paused at the shared segmentation position t q t_{q} . At this point, each sequence i i maintains its own KV cache for segment s i s_{i} , while sharing the KV cache for prompt q q via pointer references. We then select the segment closest to the median deviation as the shared upstream prefix s = s k ∗ s=s_{k^{*}} , where k ∗ = arg min k | d ( s k ) − median ( { d ( s i ) } ) | k^{*}=\arg\min_{k}|d(s_{k})-\mathrm{median}(\{d(s_{i})\})| and d ( s ) = 1 | s | ∑ t − log π 𝜽 ( s t | q , s < t ) d(s)=\frac{1}{|s|}\sum_{t}-\log\pi_{\boldsymbol{\theta}}(s_{t}|q,s_{<t}) is the mean negative log-probability. The key operation is KV cache pointer redirection : we redirect all G G sequences’ segment KV cache pointers to the single selected KV ( s ) \text{KV}(s) , then recompute the KV cache for the reordered prefix [ s , q ] [s,q] . This enables all subsequent outcome tokens o i o_{i} to be generated conditional on the identical prefix, ensuring valid Monte Carlo estimation while requiring only one shared copy of KV ( s ) \text{KV}(s) in GPU memory.
### 3.2 Asymmetric Optimization Strategy
Given a dataset 𝒟 \mathcal{D} with ground-truth answers a a , the skip-connected optimization objective maximizes the expected advantages at both stages: max 𝜽 𝔼 ( q , a ) ∼ 𝒟 [ 𝒥 I ( 𝜽 , q , a ) + 𝒥 II ( 𝜽 , q , s , a ) ] , \underset{\boldsymbol{\theta}}{\max}\;\mathbb{E}_{(q,a)\sim\mathcal{D}}\Big[\mathcal{J}_{\mathrm{I}}(\boldsymbol{\theta};q,a)+\mathcal{J}_{\mathrm{II}}(\boldsymbol{\theta};q,s,a)\Big], (4) where 𝒥 I \mathcal{J}_{\mathrm{I}} and 𝒥 II \mathcal{J}_{\mathrm{II}} are the upstream and downstream policy gradient objectives defined using advantages computed from rewards ℛ I \mathcal{R}_{\mathrm{I}} and ℛ II \mathcal{R}_{\mathrm{II}} respectively.
Given that the upstream response density is multiplicatively lower than that of the downstream, we employ an asymmetric optimization strategy using different algorithms for each phase, as detailed below.
#### Downstream Optimization: Group Relative Policy Optimization.
For the downstream phase, we employ group-relative policy optimization with clipped importance sampling. For each augmented input q ~ i = [ s i ; q i ] \tilde{q}_{i}=[s_{i};q_{i}] (where s i s_{i} denotes the selected segment for problem q i q_{i} , distinct from the G − 1 G-1 discarded segments during parallel generation), we sample a group of G G responses and update the policy according to: 𝒥 II ( 𝜽 ) \displaystyle\mathcal{J}_{\mathrm{II}}(\boldsymbol{\theta}) = 𝔼 ( q , a ) ∼ 𝒟 , { o i } i = 1 G ∼ π 𝜽 old ( ⋅ | q ~ i ) \displaystyle=\mathbb{E}_{(q,a)\sim\mathcal{D},\{o_{i}\}_{i=1}^{G}\sim\pi_{\boldsymbol{\theta}_{\text{old}}}(\cdot|\tilde{q}_{i})} (5) [ 1 ∑ i | o i | ∑ i = 1 G ∑ t = 1 | o i | sg ( r ^ i , t ) A ^ i log π 𝜽 ( o i , t ∣ q ~ i , o i , < t ) ] , \displaystyle\Bigg[\frac{1}{\sum_{i}|o_{i}|}\sum_{i=1}^{G}\sum_{t=1}^{|o_{i}|}\texttt{{sg}}(\hat{r}_{i,t})\,\hat{A}_{i}\log\pi_{\boldsymbol{\theta}}(o_{i,t}\mid\tilde{q}_{i},o_{i,<t})\Bigg], where sg ( ⋅ ) \texttt{{sg}}(\cdot) denotes the stop-gradient operation, r ^ i , t = clip ( r i , t , 1 − ε , 1 + ε ) \hat{r}_{i,t}=\text{clip}(r_{i,t},1-\varepsilon,1+\varepsilon) with r i , t = π 𝜽 ( o i , t ∣ q ~ i , o i , < t ) π 𝜽 old ( o i , t ∣ q ~ i , o i , < t ) r_{i,t}=\frac{\pi_{\boldsymbol{\theta}}(o_{i,t}\mid\tilde{q}_{i},o_{i,<t})}{\pi_{\boldsymbol{\theta}_{\text{old}}}(o_{i,t}\mid\tilde{q}_{i},o_{i,<t})} , and A ^ i = ( R i − R ¯ ) / σ R \hat{A}_{i}=(R_{i}-\bar{R})/\sigma_{R} .
#### Upstream Optimization: Single-stream Policy Optimization.
Since only one selected segment per problem participates in upstream optimization (with the remaining G − 1 G-1 segments from parallel generation discarded after contributing to median-based selection), group-relative advantage estimation is infeasible. We adopt SPO ( Xu and Ding, 2025 ) , which replaces the spatial baseline, i.e., group mean from concurrent rollouts, with a temporal baseline: a per-prompt tabular tracker V ^ ( q ) \hat{V}(q) updated via exponential moving average V ^ ( q ) ← V ^ − 1 ( q ) + η ⋅ ( ℛ I ( s , a ) − V ^ − 1 ( q ) ) \hat{V}(q)\leftarrow\hat{V}_{-1}(q)+\eta\cdot(\mathcal{R}_{\mathrm{I}}(s,a)-\hat{V}_{-1}(q)) , where V ^ − 1 ( q ) \hat{V}_{-1}(q) denotes the pre-update baseline (initialized to a global default, e.g., 0 0 , and immediately overwritten upon the first observation for each prompt), and the adaptive learning rate η = ( ρ N + 1 ) − 1 \eta=(\rho N+1)^{-1} depends on a KL-adaptive forgetting factor ρ = 2 − D KL / τ half \rho=2^{-D_{\mathrm{KL}}/\tau_{\mathrm{half}}} and the discounted cumulative observation count N N . 1 1 1 While the original SPO formulation assumes rewards in [ 0 , 1 ] [0,1] , our rewards ℛ I ∈ [ − 1 , + 1 ] \mathcal{R}_{\mathrm{I}}\in[-1,+1] are affinely mapped to [ 0 , 1 ] [0,1] via ( r + 1 ) / 2 (r+1)/2 for tracker updates, then mapped back to [ − 1 , + 1 ] [-1,+1] when retrieving baselines. This preserves the EMA equivalence while maintaining numerical stability. The advantage A i = ℛ I ( s i , a i ) − V ^ − 1 ( q i ) A_{i}=\mathcal{R}_{\mathrm{I}}(s_{i},a_{i})-\hat{V}_{-1}(q_{i}) uses the pre-update baseline to ensure unbiasedness, then normalized across the batch: A ~ i = ( A i − μ ℬ ) / σ ℬ \tilde{A}_{i}=(A_{i}-\mu_{\mathcal{B}})/\sigma_{\mathcal{B}} . The objective is: 𝒥 I ( 𝜽 ) \displaystyle\mathcal{J}_{\mathrm{I}}(\boldsymbol{\theta}) = 𝔼 ( q , a ) ∼ 𝒟 , s i ∼ π 𝜽 old ( ⋅ | q ) \displaystyle=\mathbb{E}_{(q,a)\sim\mathcal{D},s_{i}\sim\pi_{\boldsymbol{\theta}_{\text{old}}}(\cdot|q)} (6) [ 1 | s i | ∑ t = 1 | s i | min ( r i , t ( 𝜽 ) A ~ i , r ^ i , t A ~ i ) ] , \displaystyle\Bigg[\frac{1}{|s_{i}|}\sum_{t=1}^{|s_{i}|}\min\big(r_{i,t}(\boldsymbol{\theta})\,\tilde{A}_{i},\,\hat{r}_{i,t}\,\tilde{A}_{i}\big)\Bigg], where r i , t ( 𝜽 ) = π 𝜽 ( s i , t ∣ q i , s i , < t ) / π 𝜽 old ( s i , t ∣ q i , s i , < t ) r_{i,t}(\boldsymbol{\theta})=\pi_{\boldsymbol{\theta}}(s_{i,t}\mid q_{i},s_{i,<t})/\pi_{\boldsymbol{\theta}_{\text{old}}}(s_{i,t}\mid q_{i},s_{i,<t}) and r ^ i , t = clip ( r i , t , 1 − ε , 1 + ε ) \hat{r}_{i,t}=\text{clip}(r_{i,t},1-\varepsilon,1+\varepsilon) .
## 4 Experiments
#### Training Setup.
We train SKPO on the dapo-math-17k dataset ( Yu et al., 2025b ) , a curated mathematical reasoning dataset containing 17k problems spanning various difficulty levels. All models are trained for 500 optimization steps with a batch size of 128 prompts per step. We set the maximum prompt length to 1024 tokens and the maximum completion length to 3072 tokens for both training and evaluation. For fair comparison, all methods using group-relative advantage estimation adopt G = 8 G=8 rollouts per prompt; with our engineering optimization, SKPO generates G = 8 G=8 upstream segments in parallel, selects one via median-based selection, and continues with G = 8 G=8 downstream responses within a single GPU batch. For SPO which does not use group-relative estimation, we scale the batch size by 8 to ensure equivalent computational cost; we also omit the warm-up phase in our SPO reproduction to avoid introducing a cost-misaligned initialization stage. We use Qwen2.5-Math-7B and Llama-3.2-3B-Instruct as base models. Complete hyperparameters and implementation details are provided in Appendix C .
### 4.1 Main Results
We evaluate SKPO on diverse benchmarks covering In-Domain mathematical reasoning and Out-of-Domain general and code reasoning. For In-Domain evaluation, we conduct extensive experiments on mathematical benchmarks. For Out-of-Domain evaluation, we employ MMLU-Pro ( Wang et al., 2024 ) to assess general reasoning capabilities and LiveCodeBench ( Jain et al., 2025a ) to evaluate code generation skills. To reduce variance from stochastic sampling, we repeat evaluations 32 times for AIME24 and AIME25 due to their small dataset sizes, and 3 times for all other benchmarks. The main experimental results are summarized in Table 1 , where we compare SKPO against state-of-the-art baselines including GRPO ( Shao et al., 2024 ) , GSPO ( Zheng et al., 2025 ) , SPO ( Xu and Ding, 2025 ) , Critique-GRPO ( Zhang et al., 2025b ) , PRIME ( Cui et al., 2025 ) , DAPO ( Yu et al., 2025b ) , SAPO ( Gao et al., 2025 ) , and CISPO ( Chen et al., 2025a ) . These methods all rely solely on the policy model’s own reasoning without external information, ensuring comparable inference costs, and represent a collection of recent influential approaches that have been widely adopted. Detailed experimental settings and hyperparameters are provided in Appendix C .
† The original Critique-GRPO uses external strong models (e.g., GPT-4o) to generate critiques. To align computational costs across all methods, we implement a self-critique variant where the policy model generates critiques itself, rather than the external-critique approach advocated in the original paper. We use 7 direct rollouts plus 1 self-reflection rollout (8 total, matching other methods). We failed to obtain stable training curves on Llama-3.2-3B-Instruct.