🤖 AI Summary
This work addresses the challenge in reinforcement learning where correct reasoning trajectories, due to their extremely low sampling probability, fail to provide effective training signals—particularly in zero-shot settings with no successful demonstrations. The paper introduces the first first-order gradient-based latent-space correction method tailored for this scenario: leveraging failed reasoning attempts and ground-truth answers as anchors, it optimizes the input embeddings of reasoning prefixes under convex hull constraints over token embeddings. This approach guides the model toward generating longer, self-reflective, and accurate reasoning paths. By integrating gradient-driven latent-space updates, reuse of failed trajectories, and joint training that combines supervised fine-tuning (SFT) with reinforcement learning from verifiable rewards (RLVR), the method achieves substantial performance gains over baselines on mathematical reasoning benchmarks and significantly enhances both SFT and RLVR effectiveness.
📝 Abstract
Reinforcement learning with verifiable rewards (RLVR) is bottlenecked by hard prompts on which correct trajectories have low probability, so sampling misses them within a practical budget and leaves the policy update with little useful signal. We frame such zero-hit prompts as RLVR's sampling frontier, where new reasoning behavior is most valuable yet least likely to be sampled. Importantly, failed rollouts can be informative: they expose where the model's reasoning went wrong. We introduce LatentRevise, a first-order latent revision method that recovers training signal for this zero-hit regime. Given a failed rollout and the gold answer as an anchor, LatentRevise optimizes the input embeddings of its reasoning prefix under two complementary gradients, moving the prefix away from the failed continuation and toward the gold answer. The optimization is constrained to the convex hull of the model's vocabulary embeddings, so each update moves the latent toward a real token embedding rather than an arbitrary feature direction. We find that continuations from the revised prefix lengthen, exhibit self-reflection, and reach correct answers missed by the original rollouts. Used as training data, these trajectories improve SFT and RLVR on math benchmarks over standard baselines.