SD-E$^2$: Semantic Exploration for Reasoning Under Token Budgets

TL;DR

SD-E$^2$ tackles the challenge of enabling accurate reasoning in small language models under limited token budgets by introducing a semantics-aware reinforcement learning framework. It imposes a semantic diversity gate, computed with a frozen sentence encoder, to reward meaningfully distinct reasoning trajectories and switches to exploitation once a correct strategy is found, thereby balancing exploration and efficiency. The approach combines Div$(H)$ and Uniq$(H;)$ into a multi-objective GRPO objective with KL regularization, achieving substantial gains on GSM8K, AIME, and MedMCQA, and demonstrating transferability across backbones. These results suggest that explicitly rewarding semantic novelty yields a compute-efficient signal for training reasoning-capable SLMs and that cognitive adaptation—altering the high-level reasoning process rather than per-token computation—offers a complementary path to efficiency in resource-constrained models.

Abstract

Small language models (SLMs) struggle with complex reasoning because exploration is expensive under tight compute budgets. We introduce Semantic Diversity-Exploration-Exploitation (SD-E$^2$), a reinforcement learning framework that makes exploration explicit by optimizing semantic diversity in generated reasoning trajectories. Using a frozen sentence-embedding model, SD-E$^2$ assigns a diversity reward that captures (i) the coverage of semantically distinct solution strategies and (ii) their average pairwise dissimilarity in embedding space, rather than surface-form novelty. This diversity reward is combined with outcome correctness and solution efficiency in a z-score-normalized multi-objective objective that stabilizes training. On GSM8K, SD-E$^2$ surpasses the base Qwen2.5-3B-Instruct and strong GRPO baselines (GRPO-CFL and GRPO-CFEE) by +27.4, +5.2, and +1.5 percentage points, respectively, while discovering on average 9.8 semantically distinct strategies per question. We further improve MedMCQA to 49.64% versus 38.37% for the base model and show gains on the harder AIME benchmark (1983-2025), reaching 13.28% versus 6.74% for the base. These results indicate that rewarding semantic novelty yields a more compute-efficient exploration-exploitation signal for training reasoning-capable SLMs. By introducing cognitive adaptation-adjusting the reasoning process structure rather than per-token computation-SD-E$^2$ offers a complementary path to efficiency gains in resource-constrained models.

Figures (1)

• Figure 1: Problem and approach overview on a GSM8K example.Left: Outcome-driven baselines (e.g., GRPO-CFL) and the non-semantic GRPO-CFEE can generate multiple, near-duplicate strategies, leading to redundant exploration (Tokens ↑, Diversity ↓). Right: SD-E$^{2}$ encodes each <reasoning> with a frozen sentence encoder to compute (i) $\mathrm{Div}(H)=1-\text{avg cosine}$ and (ii) $\mathrm{Uniq}(H;\delta)$, rewarding exploration only when strategies are semantically distinct; upon any correct strategy, it switches to an exploitation bonus. Normalized components $R_{\mathrm{oc}},R_{\mathrm{re}},R_{\mathrm{fa}},R_{\mathrm{sd}}$ are combined under a GRPO objective with KL regularization, yielding higher ACC under the same max-token decoding budget.