Fibration Policy Optimization

Chang Li; Tshihao Tsu; Yaren Zhang; Chao Xue; Xiaodong He

cs.LGMar 9, 2026

Fibration Policy Optimization

Authors: Chang Li, Tshihao Tsu, Yaren Zhang, Chao Xue, Xiaodong He

Organizations: JD Explore Academy · Carleton University

Abstract

Large language models are increasingly trained as heterogeneous systems spanning multiple domains, expert partitions, and agentic pipelines, yet prevalent proximal objectives operate at a single scale and lack a principled mechanism for coupling token-level, trajectory-level, and higher-level hierarchical stability control. To bridge this gap, we derive the Aggregational Policy Censoring Objective (APC-Obj), the first exact unconstrained reformulation of sample-based TV-TRPO, establishing that clipping-based surrogate design and trust-region optimization are dual formulations of the same problem. Building on this foundation, we develop Fiber Bundle Gating (FBG), an algebraic framework that organizes sampled RL data as a fiber bundle and decomposes ratio gating into a base-level gate on trajectory aggregates and a fiber-level gate on per-token residuals, with provable first-order agreement with the true RL objective near on-policy. From APC-Obj and FBG we derive Fibration Policy Optimization (or simply, FiberPO), a concrete objective whose Jacobian is block-diagonal over trajectories, reduces to identity at on-policy, and provides better update direction thus improving token efficiency. The compositional nature of the framework extends beyond the trajectory-token case: fibrations compose algebraically into a Fibration Gating Hierarchy (FGH) that scales the same gating mechanism to arbitrary hierarchical depth without new primitives, as demonstrated by FiberPO-Domain, a four-level instantiation with independent trust-region budgets at the domain, prompt group, trajectory, and token levels. Together, these results connect the trust-region theory, a compositional algebraic structure, and practical multi-scale stability control into a unified framework for LLM policy optimization.

Figures & tables

Figure 2: Base weight visualization. Panel (a) shows the aggregate gate $g^{\rm agg}$ (Eq. 10 ) with three regimes: pass-through ( $|x|\leq C$ , slope $1$ ), rollback ( $C<|x|<C^{*}$ , slope $-k$ ), and zeroed ( $|x|\geq C^{*}$ , output $0$ ). As $k=T_{\tau}$ increases, the rollback regime narrows (width $C/T_{\tau}$ ) and $g^{\rm agg}$ approaches a hard clip at $\pm C$ . Panel (b) shows the base weight $\log w_{\tau}^{\rm base}$ (Eq. 8 ) in $(\log s^{+},\log s^{-})$ -space with asymmetric thresholds $C^{+}=0.20$ , $C^{-}=0.1$ ( $k=2$ ). Dashed lines mark the budget boundaries $C^{\pm}$ (onset of rollback); dotted lines mark the full-gating thresholds $C^{*\pm}$ (onset of zeroing). The five global regimes (G-I through G-III) follow a non-monotonic pattern: $|\log w|$ rises through the rollback onset (G-II,r), peaks when one channel is fully gated (G-II), declines under mutual rollback (G-III,r), and collapses to zero when both channels are fully gated (G-III, $w^{\rm base}_{\tau}=1$ ).

Figure 3: Regime map on the probability simplex given by all possible configurations of a trajectory’s policy ratios (subject to probabilistic constraint ( $\frac{1}{T_{\rm trajectory}}\sum_{t=0}^{T_{\rm trajectory}-1}r_{t}=1$ ) where $T{=}3$ , $\varepsilon{=}0.025$ , $C^{+}{=}0.15$ , $C^{-}{=}0.09$ , $K{=}3$ , $C^{*\pm}{=}(1{+}1/K)\,C^{\pm}$ ). (a) Local branch (zoomed $\pm 0.18$ around centroid): L-I (no clip, central polytope inside teal boundary), L-II (partial clip, between teal solid and purple solid boundaries), and L-III (full clip, outside purple boundary). Red and blue dashed contours show the rollback thresholds $N{=}C^{-}$ and $P{=}C^{+}$ . (b) Global branch (full simplex): five concentric regimes—G-I (inside red solid $N{=}C^{-}$ ), G-II,r (rollback transition, between red solid $N{=}C^{-}$ and red dotted $N{=}C^{*-}$ ), G-II (between red dotted $N{=}C^{*-}$ and blue solid $P{=}C^{+}$ ), G-III,r (rollback transition, between blue solid $P{=}C^{+}$ and blue dotted $P{=}C^{*+}$ ), and G-III (outside blue dotted $P{=}C^{*+}$ ).

Figure 4: Per-token and total FiberPO objective under parameterized drift $r_{i}=1+t\hat{\Delta}_{i}$ ( $T{=}10$ , $k{=}32$ , $\epsilon{=}0.04$ , $C^{+}{=}0.12$ , $C^{-}{=}0.05$ ), with $\operatorname{sign}(\hat{\Delta}_{i})=\operatorname{sign}(\hat{A}_{i})$ , reflecting the typical policy update in which tokens with positive advantage increase their probability ( $\hat{\Delta}_{i}>0$ ) and tokens with negative advantage decrease their probability ( $\hat{\Delta}_{i}<0$ ). GRPO in the graph uses the same $\epsilon$ -clip as FiberPO. (a) Trajectory in $(\log s^{+},\log s^{-})$ -space overlaid on the $\log w$ heatmap; as drift increases, the trajectory moves through the global regimes of Figure 2 (b). (b) Per-token fiber correction at $t{=}1$ : fiber residual (red) and base-weight contribution (orange) stack to form the full gated log-ratio. (c) Per-token objective for two representative tokens. (d) Total objective: FiberPO’s $g^{\rm agg}$ and $\operatorname{logclip}$ saturate, causing the objective to return to baseline. Vertical markers: $\epsilon$ -clip transitions (gray, L-I $\to$ L-II), $C^{-}$ rollback (red, G-II,r onset), $C^{+}$ rollback (blue, G-II,r onset).

Policy Optimization

Policy optimization develops algorithms for improving decision rules in reinforcement learning, language models, and sequential decision systems.

See all 224 papers

Dec 28, 2025

Trust Region Masking for Long-Horizon LLM Reinforcement Learning

Policy gradient methods for Large Language Models optimize a policy

π_θ

via a surrogate objective computed from samples of a rollout policy

π_{\text{roll}}

. However, modern LLM-RL pipelines suffer from unavoidable implementation divergences -- backend discrepancies, Mixture-of-Experts routing discontinuities, and distributed training staleness -- causing off-policy mismatch (

π_{\text{roll}} \neq π_θ

) and approximation errors between the surrogate and the true objective. We demonstrate that classical trust region bounds on this error scale as

O(T^2)

with sequence length

T

, rendering them vacuous for long-horizon tasks. To address this, we derive a family of bounds -- both KL-based and TV-based -- including a Pinsker-Marginal bound (

O(T^{3/2})

), a Mixed bound (

O(T)

), and an Adaptive bound that strictly generalizes the Pinsker-Marginal bound via per-position importance-ratio decomposition. Taking the minimum over all bounds yields the tightest known guarantee across all divergence regimes. Crucially, all bounds depend on the maximum token-level divergence

D_{\mathrm{KL}}^{\mathrm{tok,max}}

(or

D_{\mathrm{TV}}^{\mathrm{tok,max}}

), a sequence-level quantity that cannot be controlled by token-independent methods like PPO clipping. We propose Trust Region Masking (TRM), which masks entire sequences violating the trust region, enabling the first non-vacuous monotonic improvement guarantees for long-horizon LLM-RL.

Yingru Li, Jiacai Liu, Jiawei Xu +4

Fudan University · The Chinese University of Hong Kong, Shenzhen

RL for Language Models Reinforcement Learning Long-Horizon RL Policy Gradient Methods Trust Region Policy Optimization Policy Learning+1

Ask PDF

Feb 6, 2026

The Optimal Token Baseline: Variance Reduction for Long-Horizon LLM-RL

Reinforcement Learning (RL) for Large Language Models (LLMs) often suffers from training collapse in long-horizon tasks due to exploding gradient variance. To mitigate this, a baseline is commonly introduced for advantage computation; however, traditional value models remain difficult to optimize, and standard group-based baselines overlook sequence heterogeneity. Although classic optimal baseline theory can achieve global variance reduction, it neglects token heterogeneity and requires prohibitive gradient-based computation. In this work, we derive the Optimal Token Baseline (OTB) from first principles, proving that gradient updates should be weighted inversely to their cumulative gradient norm. To ensure efficiency, we propose the Logit-Gradient Proxy that approximates the gradient norm using only forward-pass probabilities. Our method achieves training stability and matches the performance of large group sizes (

N=32

) with only

N=4

, reducing token consumption by over 65% across single-turn and tool-integrated reasoning tasks.

Yingru Li, Jiawei Xu, Ziniu Li +10

The Chinese University of Hong Kong, Shenzhen · Fudan University · Hong Kong University of Science and Technology +4

RL for Language Model Reasoning Policy Gradient Methods RL Post-Training

Ask PDF

Feb 17, 2026

STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens

Reinforcement Learning (RL) has significantly improved large language model reasoning, but existing RL fine-tuning methods rely heavily on heuristic techniques such as entropy regularization and reweighting to maintain stability. In practice, they often suffer from late-stage performance collapse, leading to degraded reasoning quality and unstable training. We identify a key factor behind this instability: a small fraction of tokens, termed spurious tokens (around 0.01%), which contribute little to the reasoning outcome but receive disproportionately amplified gradient updates due to inheriting the full sequence-level reward. We present a unified framework for evaluating token-level optimization impacts across spurious risk, gradient norms, and entropy changes. Building on the analysis of token characteristics that severely disrupt optimization, we propose the Silencing Spurious Tokens (S2T) mechanism to efficiently suppress their gradient perturbations. Incorporating this mechanism into a group-based objective, we propose Spurious-Token-Aware Policy Optimization (STAPO), which promotes stable and effective large-scale model refinement. Across six mathematical reasoning benchmarks using Qwen 1.7B, 8B, and 14B base models, STAPO consistently demonstrates superior entropy stability and achieves an average performance improvement of 11.49% (

ρ_{\mathrm{T}}

=1.0, top-p=1.0) and 3.73% (

ρ_{\mathrm{T}}

=0.7, top-p=0.9) over GRPO, 20-Entropy, and JustRL.

Shiqi Liu, Zeyu He, Guojian Zhan +10

School of Vehicle and Mobility & College of AI, Tsinghua University · Didi Voyager Labs, DiDi Autonomous Driving

Reinforcement Learning RL for Language Model Reasoning RL Fine-Tuning Token-Level Credit Assignment

Ask PDF

Abstract

Policy Optimization

Explore similar work

Trust Region Masking for Long-Horizon LLM Reinforcement Learning

The Optimal Token Baseline: Variance Reduction for Long-Horizon LLM-RL

STAPO: Stabilizing Reinforcement Learning for LLMs by Silencing Rare Spurious Tokens