Supporting FIPO (Future-KL Influenced Policy Optimization)

[SeungyounShin]

Summary

Add FIPO (Future-KL Influenced Policy Optimization) as a built-in loss type, enabling dense token-level credit assignment for RL training without a value network.

Ma et al., "FIPO: Eliciting Deep Reasoning with Future-KL Influenced Policy Optimization" (arXiv:2603.19835)

Method

FIPO extends standard clipped PPO by re-weighting per-token advantages with Future-KL influence weights. While GRPO broadcasts a uniform sequence-level advantage to all tokens, FIPO modulates each token's contribution based on how the future trajectory evolves:

$$\text{FutureKL}_t = \sum_{k=t}^{T} M_k \cdot \gamma^{k-t} \cdot \Delta \log p_k$$

where $\Delta \log p_k = \log \pi_\theta(y_k \mid y_{<k}) - \log \pi_{\theta_{\text{old}}}(y_k \mid y_{<k})$ and $\gamma = 2^{-1/\tau}$ is an exponential decay factor.

The influence weight is computed as:

$$f_t = \text{clip}\left(\exp(\text{FutureKL}_t),; 1 - \epsilon_{f},; 1 + \epsilon_{f}\right)$$

The final FIPO loss modifies the clipped PPO objective with re-weighted advantages:

$$\mathcal{L}_t^{\text{FIPO}} = \min\left(r_t \cdot \hat{A}_t \cdot f_t,;; \text{clip}(r_t, 1{-}\epsilon, 1{+}\epsilon) \cdot \hat{A}_t \cdot f_t\right)$$

Key design choices

• influence_weights are detached from the computation graph — FIPO does not backpropagate through Future-KL
• Requires multi-step training per rollout (global_batch_size < rollout_batch_size × n_samples_per_prompt) so that $\theta$ drifts from $\theta_{\text{old}}$, making Future-KL non-trivial
• Uses chunked matrix multiplication for memory-efficient $O(B \cdot L^2)$ Future-KL computation

Changes

File	Description
slime/backends/megatron_utils/loss.py	_compute_future_kl() and fipo_loss_function() with dual-clip PPO, sequence filtering, and FIPO metrics
slime/utils/arguments.py	--loss-type fipo_loss and 6 FIPO-specific args (--fipo-decay-rate, --fipo-chunk-size, --fipo-clip-ratio, --fipo-clip-high-only, --fipo-safety-thresh, --fipo-dual-clip-c)
scripts/models/qwen3.5-2B.sh	Qwen3.5-2B model config
examples/fipo/fipo_qwen3.5_2b.sh	Training example script

Usage

# Use with GRPO advantage estimator
--loss-type fipo_loss \
--advantage-estimator grpo \
--fipo-decay-rate 32.0 \
--fipo-chunk-size 128 \
--fipo-clip-ratio 0.2 \
--fipo-safety-thresh 3.0 \
--fipo-dual-clip-c 10.0 \
--global-batch-size 64   # must be < rollout_batch_size × n_samples_per_prompt

FIPO hyperparameter guide

Parameter	7B and below	32B and above	Description
--fipo-decay-rate	32	32	Half-life $\tau$ for exponential decay
--fipo-clip-ratio	0.2	0.2	Influence weight clip range
--fipo-clip-high-only	✗	✓	Clip only upper bound $[1.0, 1.2]$ for larger models
--fipo-safety-thresh	3.0	10.0	Cap high-IS negative samples

Results

FIPO consistently increases raw reward (Qwen3.5-4B-Base)

A token-level importance-ratio heatmap under FIPO. The strong emphasis on “Alternative” suggests that FIPO reinforces branching points that trigger self-reflection and exploration of alternative solution paths, rather than treating all tokens uniformly.

cc. @zhuzilin