Right Makes Might: Aligning Verified Hidden States Empowers RL Reasoning

2026-06-02

AI Summary

This work addresses a key limitation in existing reinforcement learning approaches for mathematical reasoning, which rely solely on binary rewards from correct solution trajectories and ignore the geometric structure among hidden states. The authors propose Hidden-Align, an auxiliary loss that aligns the last-layer hidden states of correct trajectories at an anchor position just before the answer token, leveraging their inherent convergence property to construct a unified representation of “correct decisions.” Built upon the RLVR framework, Hidden-Align integrates cosine-similarity-guided alignment, adaptive anchor positioning, and multi-scale training without introducing additional computational overhead. Evaluated across eight mathematical reasoning benchmarks, the method consistently improves performance: Qwen3 models (1.7B, 4B, and 14B) achieve average pass@1 gains of 3.8, 6.2, and 5.4 percentage points, respectively, with pass@k results uniformly surpassing baseline methods.

Abstract

Reinforcement Learning from Verifiable Rewards (RLVR) has become the dominant approach for improving mathematical reasoning in large language models, yet current methods reduce each correct rollout to a single reward bit, ignoring the geometric structure shared among their hidden states. Investigating this structure, we find that at the anchor token (the position immediately before the answer marker), correct rollouts converge naturally because they must produce the same answer (cosine similarity ~0.84), yet each retains residual variance from its unique reasoning path. Encouraging full alignment at this point pushes the model to extract a unified "correct decision" representation, reducing sensitivity to which reasoning path was taken. Based on this observation, we propose Hidden-Align, an auxiliary loss function that aligns the last-layer hidden states of correct rollouts at the anchor token during RL training, with zero overhead in both training and inference. On eight mathematical reasoning benchmarks, Hidden-Align improves average pass@1 over the DAPO baseline by 3.8, 6.2, and 5.4 percentage points on Qwen3-1.7B, 4B, and 14B respectively, with consistent pass@k gains across all three scales, supported by ablations on loss type, anchor position, layer depth, and loss weight.

Problem

• Reinforcement Learning
• Mathematical Reasoning
• Hidden States
• Reward Signal
• Representation Alignment

Innovation

• Hidden-Align
• reinforcement learning
• hidden state alignment
• mathematical reasoning
• RLVR