Decoding in Geometry: Alleviating Embedding-Space Crowding for Complex Reasoning

TL;DR

The paper identifies embedding-space crowding as a decoding-time phenomenon where next-token probability mass concentrates among geometrically similar embeddings, and demonstrates its negative association with mathematical reasoning success. It then introduces CraEG, a training-free, geometry-guided reweighting method that downweights high-probability, highly crowded tokens in a step-adaptive manner, improving robustness and diversity without extra forward passes. Empirical results across multiple models and benchmarks (AIME, HMMT) show CraEG consistently enhances Avg@32 and Pass@8 while boosting diversity metrics, with ablations clarifying the roles of nonlinear weighting and correction strength. Overall, the work provides a practical framework for geometry-informed decoding that can be integrated with standard sampling strategies to improve complex reasoning in LLMs.

Abstract

Sampling-based decoding underlies complex reasoning in large language models (LLMs), where decoding strategies critically shape model behavior. Temperature- and truncation-based methods reshape the next-token distribution through global probability reweighting or thresholding to balance the quality-diversity tradeoff. However, they operate solely on token probabilities, ignoring fine-grained relationships among tokens in the embedding space. We uncover a novel phenomenon, embedding-space crowding, where the next-token distribution concentrates its probability mass on geometrically close tokens in the embedding space. We quantify crowding at multiple granularities and find a statistical association with reasoning success in mathematical problem solving. Motivated by this finding, we propose CraEG, a plug-and-play sampling method that mitigates crowding through geometry-guided reweighting. CraEG is training-free, single-pass, and compatible with standard sampling strategies. Experiments on multiple models and benchmarks demonstrate improved generation performance, with gains in robustness and diversity metrics.

Figures (6)

• Figure 1: Examples of embedding-space crowding in next-token distributions. We compare two decoding states from the same AIME25 solution trajectory with similar entropy but different crowding levels (higher: Step 447; lower: Step 2003). Left panels show the top-10 probabilities and similarity to the top candidate; right panels show a UMAP projection of the top-200 tokens colored by probability. Higher crowding concentrates probability mass on geometrically similar tokens; lower crowding is more dispersed.
• Figure 2: Sequences are binned into low/mid/high groups by quantiles of sequence-level crowding. Bars report accuracy per bin, which decreases monotonically with crowding.
• Figure 3: Step-level example of probability reallocation. Top-10 next-token probabilities and cosine similarity to the top-1 candidate (baseline vs. CraEG).
• Figure 4: Trajectory-level token crowding analysis (top-30 per step). We report $\mathbb{E}[p \mid \text{crowding}]$ and the mean crowding aggregated over decoding steps. CraEG reduces the mean from 0.1934 to 0.1864.
• Figure 5: ECDFs of step-level embedding-space crowding. Empirical cumulative distribution functions of step-level crowding for decoding steps from correct and incorrect samples. Steps from incorrect samples are generally right-shifted toward higher crowding, consistent with elevated crowding being broadly present across steps rather than confined to a few extremes.

Theorems & Definitions (3)

• Definition 3.1: Token-Level Crowding Score
• Definition 3.2: Step-Level Crowding Score
• Definition 3.3: Sequence-Level Crowding Score