🤖 machine learning

Skip-Connected Policy Optimization for Implicit Advantage

This paper introduces Skip-Connected Policy Optimization (SKPO), a novel reinforcement learning framework that decomposes reasoning into upstream and downstream phases with a skip connection to resolve the high-variance advantage estimation issues in dense reward settings, thereby achieving significant performance gains over Group Relative Policy Optimization (GRPO) on mathematical and general reasoning benchmarks.

Original authors: Fengwei Teng, Jinyi Bai, Xinhao Yao, Demi Ruohan Wang, Jiahao Zhao, Zhijiang Guo

Published 2026-04-13

📖 5 min read🧠 Deep dive

CC BY 4.0

Original authors: Fengwei Teng, Jinyi Bai, Xinhao Yao, Demi Ruohan Wang, Jiahao Zhao, Zhijiang Guo

Original paper licensed under CC BY 4.0 (http://creativecommons.org/licenses/by/4.0/). ✨ This is an AI-generated explanation of the paper below. It is not written or endorsed by the authors. For technical accuracy, refer to the original paper. Read full disclaimer

The Big Problem: The "Guessing Game" of AI Reasoning

Imagine you are teaching a student (an AI) to solve a complex math problem.

The Old Way (GRPO): You let the student write the whole solution from start to finish. If they get the final answer right, you give them a gold star. If they get it wrong, you give them a red X.
- The Flaw: The student doesn't know which step was good or bad. Did they mess up the very first sentence? Or just the last calculation? They have to guess.
The "Fine-Grained" Dream: Ideally, you'd want to grade every single sentence. "Great job on step 3! But step 5 was wrong."
- The Reality Check: To do this, you'd have to ask the student to write the solution 100 times, stop at step 5, and see how often they get the right answer from there. This takes forever and costs a fortune in computer power. If you don't do it enough times, your grading is just random noise (high variance), and the student gets confused, learning worse than if you just graded the final answer.

The Paper's Discovery: Trying to grade every single step with limited computer power is actually worse than just grading the final answer because the "grades" are too noisy and often wrong.

The Solution: SKPO (The "Smart Coach" Strategy)

The authors propose a new method called SKPO (Skip-Connected Policy Optimization). Think of it as a two-phase coaching strategy that gets the best of both worlds without the massive cost.

Phase 1: The "Sketch" (Upstream)

Instead of asking the student to write the whole essay, you ask them to write just the first half (the outline or the setup).

The Trick: You don't grade this half immediately. Instead, you take that half, and you ask the student to finish the story 8 different ways based on that sketch.
The Reward: If 6 out of those 8 endings are correct, you know the sketch was good! You give the sketch a "good job" reward. If only 1 is correct, the sketch was shaky.
Why it works: You get a "dense" reward (a grade for the middle of the story) without having to stop and grade every single word of the story.

Phase 2: The "Finish" (Downstream)

Now, the student takes that same sketch and writes the full solution.

The "Skip Connection" (The Magic Ingredient): This is the most important part. When the student starts writing the second half, they are allowed to skip the sketch they just wrote and go straight back to the original problem if they want.
- Analogy: Imagine the student is writing a story. They wrote a great opening paragraph (the sketch). But then they realize, "Wait, my opening paragraph is leading me to a boring ending."
- In a normal system, they are stuck with that opening.
- In SKPO, the student has a "magic eraser." They can look at their opening, decide it's bad, and say, "Actually, let's ignore that and start fresh from the original prompt."
- Why this matters: This prevents the AI from getting "stuck" in a bad thought pattern. It keeps the freedom to explore new ideas while still learning from the good parts of the sketch.

The Engineering Magic: Doing it Fast

Usually, doing this two-step process would take twice as long (write half, stop, write full). The authors built a clever engineering trick (using something called "KV Cache rewriting") that lets the computer do both steps in one single pass. It's like a chef who can chop vegetables and cook the soup simultaneously without slowing down the kitchen.

Why is this a Big Deal? (The "Implicit Advantage")

The paper found a surprising side effect. Even though the AI wasn't explicitly told to write "better middle steps," the AI trained with SKPO started writing higher quality middle steps automatically.

The Analogy: Imagine two runners.
- Runner A (Old Method) runs fast but stumbles a lot in the middle of the track because they aren't sure where to step.
- Runner B (SKPO) learns to take smoother, more confident steps in the middle, even though the finish line is the only thing that matters.
The Result: SKPO beats the current state-of-the-art methods on math tests, coding tasks, and general reasoning. It's not just getting the right answer; it's getting there via a smarter, more logical path.

Summary in One Sentence

SKPO teaches AI to write better reasoning steps by giving it a "draft" phase that gets graded based on how well it leads to a solution, while giving the AI the freedom to ignore that draft if it turns out to be a dead end—all without slowing down the training process.

1. Problem Statement

The paper addresses a critical bottleneck in Reinforcement Learning with Verifiable Rewards (RLVR) for Large Language Models (LLMs), specifically within the context of Group Relative Policy Optimization (GRPO).

The Dilemma: While fine-grained, dense rewards (assigning specific values to individual reasoning steps) are theoretically superior to coarse-grained, outcome-only rewards, they are practically difficult to implement.
The Failure of Monte Carlo (MC) Estimation: To obtain dense rewards, researchers often use MC sampling (generating multiple continuations from a specific token to estimate its value). However, under practical sampling budgets (e.g., $K=8$ ), MC estimation yields high-variance advantages with frequent sign errors for early reasoning tokens.
The Paradox: Due to these estimation errors, methods attempting to use fine-grained MC rewards often perform worse than standard outcome-only GRPO. Existing methods either sacrifice fine-grained information (using only outcome rewards) or incur prohibitive computational costs to get reliable estimates.

2. Methodology: Skip-Connected Policy Optimization (SKPO)

The authors propose SKPO, a novel architecture that decomposes the reasoning process into upstream and downstream phases to balance fine-grained learning with computational efficiency.

Core Architecture

Upstream Phase (Single-Stream):
- The model generates an early-stopped reasoning segment ( $s$ ) from the original problem ( $q$ ).
- Optimization: Uses Single-Stream Policy Optimization (SPO). Since only one segment is generated per prompt in this phase, group-relative estimation is impossible. Instead, it uses a temporal baseline (a per-prompt value tracker updated via Exponential Moving Average) to stabilize training.
- Reward: The upstream segment receives a dense reward derived from the aggregated outcomes of the downstream phase.
Downstream Phase (Group-Relative):
- The model generates $G$ complete responses starting from a Skip Connection input: the concatenation of the upstream segment and the original problem, denoted as $[s, q]$ .
- Optimization: Uses standard Group Relative Policy Optimization (GRPO) with $G$ samples.
- Skip Connection Logic: The input $[s, q]$ $[s, q]$ serves two purposes:
  - It allows the model to leverage the upstream reasoning ( $s$ ).
  - It preserves exploration freedom by keeping the original problem ( $q$ ) accessible, allowing the model to implicitly bypass or correct flawed upstream reasoning if necessary.
Engineering Optimization (Single-Pass Rollout):
- A naive implementation would require two separate GPU batches (one for upstream, one for downstream), doubling latency.
- SKPO achieves computational parity with standard GRPO by modifying the inference engine (vLLM) to perform mid-generation KV cache rewriting. It generates $G$ parallel responses, pauses at the split point, selects the best upstream segment (based on median negative log-likelihood), redirects the KV cache pointers to this single segment, and continues generation for all $G$ paths in a single batch.

Asymmetric Optimization Strategy

Upstream: Optimized with SPO (temporal baseline) to handle the lack of concurrent samples.
Downstream: Optimized with GRPO (group relative advantage) to leverage the $G$ samples for stable estimation.

3. Key Contributions

Diagnosis of MC Failure: The paper rigorously demonstrates that under practical sampling budgets, MC estimation for early tokens suffers from high variance and sign inconsistency, explaining why previous attempts at dense rewards often fail.
SKPO Architecture: Introduction of a two-phase framework that decouples the sampling requirements: using single-stream optimization for the data-scarce upstream phase and group-relative optimization for the data-rich downstream phase.
Implicit Advantage Discovery: The authors reveal an "implicit advantage" phenomenon: even when final correctness is matched, SKPO-trained models generate trajectories with significantly higher intermediate-step quality. This is verified via off-policy analysis using an external evaluator (GPT-5-nano).
Efficiency: The method achieves these gains without increasing inference costs compared to standard GRPO, thanks to the single-pass KV cache optimization.

4. Experimental Results

The authors evaluated SKPO on Qwen2.5-Math-7B and Llama-3.2-3B-Instruct across mathematical benchmarks (AIME, AMC, MATH, Olympiad) and out-of-domain tasks (MMLU-Pro, LiveCodeBench).

Performance Gains:
- Qwen2.5-Math-7B: SKPO achieved an average score of 42.5%, outperforming the second-best baseline (CISPO at 40.9%) by 1.6 points and vanilla GRPO (30.2%) by 12.3 points.
- Llama-3.2-3B: SKPO achieved 24.1%, surpassing the second-best (DAPO at 22.7%) by 1.4 points and GRPO (18.9%) by 5.2 points.
Generalization: The improvements transferred effectively to out-of-domain tasks (general reasoning and code generation), indicating the learned reasoning structure is robust.
Ablation Studies:
- Removing the Skip Connection (forcing downstream to rely solely on the upstream prefix) caused a significant drop in performance due to constrained exploration.
- Removing MC Estimation (giving upstream only its own outcome reward) prevented the model from learning patterns that facilitate downstream solving.
- Adaptive Splitting (sampling split points dynamically) outperformed fixed split positions.

5. Significance

Bridging the Gap: SKPO successfully bridges the gap between the theoretical benefits of fine-grained rewards and the practical constraints of sampling budgets.
Structural Innovation: By structurally partitioning the problem rather than just tuning hyperparameters, SKPO avoids the "sign noise" that plagues token-level MC estimation.
Quality over Quantity: The discovery of implicit advantage suggests that SKPO teaches models to generate higher-quality intermediate reasoning steps, not just correct final answers. This is crucial for complex reasoning tasks where the path matters as much as the destination.
Scalability: The engineering optimization ensures that these advanced training dynamics do not come at the cost of increased training time, making the method viable for large-scale deployment.

In conclusion, SKPO represents a significant step forward in RLVR by providing a stable, efficient, and high-performing framework for integrating dense, step-level rewards into LLM training without the prohibitive costs previously associated with Monte Carlo estimation.