SAGE: Mitigating Long-Horizon Reasoning Biases via Topological Guidance
Xinyue Zeng, Jiawei Zhang, Yujun Yan, Dawei Zhou
NeurIPS 2026
cs.AI
2026-09-25
SAGE injects algebraic and hyperbolic priors into GRPO-style post-training, hitting 64.86% math avg at 35B and nearly 8x Lean proofs vs base on Andrews-Curtis.
Outcome-only post-training, GRPO and its cousins, collapses on long-horizon tasks with sparse rewards. The trajectory is scored at the end. Intermediate steps get almost no usable signal.
The paper splits the collapse into two coupled biases. Exploration bias is a volume problem: as depth grows, locally legal branches that never reach success soak up the sampling budget, so the policy learns trajectories that are easy to draw, not trajectories that stay extendable. Compounding bias is a delay problem: an early, locally plausible detour only fails many steps later, and KL regularization pins updates to the reference policy. Theorem 3.5 makes this quantitative. If success has probability p under the reference, the full-support KL optimizer stays within total variation (e^{Rmax/λ}-1)p of that reference. When p is tiny, the policy barely moves.
The running example is Andrews-Curtis trivialization: take a balanced presentation of the trivial group and reduce it with legal moves. Many moves are legal. Few actually simplify the relators.
Symbolic Closure Analysis (SCA) defines feasibility as local admissibility. A prefix enters the feasible set F only if every step is well-formed under the task interface. F is prefix-closed. It does not hand out gold steps. It is a structural proxy.
SAGE (Structural Admissibility-Guided Exploration) turns the two SCA requirements into training-time potentials and injects them into group-relative policy updates. At inference the trained policy runs alone. No scoring, no search.
Algebraic sparsification attacks exploration bias. The unresolved residual rt is projected onto the subspace Sj of a candidate operator. The score is the fraction of residual energy that projection explains. A high score means the operator is working on structure that is still open. On Andrews-Curtis the residual comes from generator counts, relator lengths, and leftover components. On math it comes from open variables, equation structure, and answer-schema status. On free-form reasoning it comes from prompt requirements the prefix has not covered. Operator subspaces are the top principal directions of residuals from training rollouts of that operator type.
Hyperbolic structural guidance attacks compounding bias. States and a target anchor are mapped into the Poincaré ball. Negative curvature supplies a dense, depth-aware signal: tree-like drift shows up in distance immediately. For Andrews-Curtis the anchor is the trivial presentation. For math and free-form tasks the anchor is prompt-conditioned and fitted on training rollouts, never on test labels or gold traces.
Sampling reweights a finite set of candidate steps proposed by the old policy, with probability proportional to length-normalized log-probability plus λ times the SAGE potential. The reward adds the mean potential along the trajectory to the terminal outcome, then uses group-relative advantages. If every terminal reward in a group is zero, a non-zero within-group variance in the potential still yields an update (Proposition 4.1). The structural preference is absorbed into the policy, so inference pays nothing extra.
On seven math benchmarks (MATH, Minerva, Olympiad, AIME24, AMC23, GSM8K, Putnam) with Qwen3.5:
| Scale | Base avg | Strongest baseline | SAGE |
| 2B | 27.24 | EMPO 40.28 | 42.11 |
| 9B | 29.01 | SFT 44.93 | 47.95 |
| 35B | 45.21 | EMPO 62.37 | 64.86 |
Margins over the strongest baseline are +1.83, +3.02, and +2.49. The 35B SAGE run at 64.86% beats Llama-3.3-70B-Instruct at 38.48%. Putnam moves from 11.36 to 22.62. On AIME24, EMPO is 62.31 and SAGE is 62.04; SAGE loses that column.
Free-form at 35B: MMLU-Pro 59.33 (GRPO 57.66), BBH-H 69.07 (EMPO 65.02), ARC-C 66.41 (EMPO 62.97). At 9B those three go from 37.91 / 44.04 / 39.73 to 39.99 / 45.31 / 42.04.
On Andrews-Curtis, AC Validity rises +19.2 to +26.0 points over the base models. Lean-verified proofs rise by more than 13 points on every backbone; Qwen3 is nearly 8× the base. On Qwen3.5-35B, SAGE Lean is 23.69 against GRPO 14.64 and GRPO-PRM 17.36. Dropping either potential hurts Validity, Path, and Lean together. On Olympiad, Pass@64 is about 87% for SAGE and about 84% for GRPO.
If the goal is long-horizon sparse-reward post-training without a step-labeled process reward model, SAGE is a training-time structural prior with zero inference overhead. The multiplicative gap on Andrews-Curtis, an open symbolic task, is the evidence that geometry-aware priors beat generic dense shaping. Gains on math and QA are two to three points. That is incremental. The cost sits in training: candidate-step scoring, residual probes, and hyperbolic distances all need a task interface or a learned proxy.
The authors are explicit. Local admissibility is exact in symbolic domains; residuals, anchors, and operator subspaces on math and free-form tasks are learned proxies. The concentration theorems apply directly only in the symbolic case. Training is more expensive. The abstract says 12 benchmarks and 7 model families; the contributions list 13 and 8. SAGE does not win AIME24 against EMPO. How structural the free-form prior is depends on probe pseudo-labels, and the paper never reports alignment between those probes and true admissibility. The near-8× Lean figure is relative to the un-post-trained base and lives in Figure 3; the main text does not print the base Lean percentage.