Soft Reasoning: Navigating Solution Spaces in Large Language Models through Controlled Embedding Exploration

1.Why Embedding Search Outperforms Discrete Sampling

A theoretical provement of "searching for a start token in the continuous embedding space" being better than "sampling every token in the discrete token space" can further improve the soundness of the paper.

Thank you for the suggestion. Searching for a start token in the continuous embedding space is preferable to sampling every token in the discrete token space from the theoretical, interpretability and the cost perspectives:

• Optimal solution to the objective. The set of the embeddings of all tokens forms a discrete subset of our continuous space. Importantly, this means that the optimal solution (maximiser) to our objective function may not lie within this discrete subset.
• Role of the soft token. As the soft token in our framework is not intended to approximate a specific discrete token, but rather to serve as a functional control token, influencing the overall response. Optimising in the continuous space allows exploration of representations that may lie between or beyond discrete tokens, enabling smoother control and better performance than fixed-token selection.
• High computational cost. With the vocabulary size |V|>30k, sampling every token in the discrete token space requires generating the output for each token x and computing the corresponding f(x), making it computationally expensive at O(|V|).

2. Computational efficiency

The authors only provide a comparison with RAP in terms of token usage, but I believe comparisons with other methods in Table 1 and comparison on inference time cost can be added to further prove the efficiency of this method.

As the reviewer suggested, we additionally report inference time(min) cost across all baselines:

Our method consistently achieves the lowest inference time across all tasks, further demonstrating its efficiency beyond token-level savings.

3. Neuron activation

It is intriguing that the authors have found that applying this approach can increase the activation rate of neurons by 3–4%, how do you understand this phenomenon, and have you tried to analyze why this occurs?

Thank you for the insightful question. We believe this phenomenon occurs because our perturbation strategy encourages broader exploration within the neuron space, thereby increasing the likelihood of activating previously dormant or marginal neurons across the architecture. In our setting, we define a neuron as activated if its activation value is greater than zero. As our perturbations produce more diverse token embeddings, the resulting activations tend to be more widely distributed across different neurons as well, thus increasing the overall number of activated neurons.

However, we would like to stress that the overall increase in activation rate does not necessarily translate to improved final answer accuracy. Rather, the key contributing factor is the ability to activate more neurons identified as critical for the current problem instance [1]. The increased activation rate simply raises the likelihood of activating these critical neurons. The BO process then further optimises the soft embedding to consistently target them, as illustrated in Figure 4 of our paper.

To further illustrate this point and to analyse the functional importance of these neurons, we have conducted an ablation study and found that masking critical neurons reduced accuracy from 62.14% to 13.27%, whereas random masking only dropped it to 41.89%. This confirms that the activated neurons are non-random and impactful, consistent with prior findings on causal neurons [2].

[1] Andy Zou et al., 2025. Representation Engineering: A Top-Down Approach to AI Transparency.

[2] Kevin Meng et al., 2022. Locating and Editing Factual Associations in GPT. NeurIPS

4.Interpretation of the Soft Token

Have you tried to interpret the sampled first token in your framework?

We appreciate the reviewer’s suggestion. We conducted nearest-neighbor analysis to interpret the optimised soft token, but found it to be a clear outlier in the embedding space. Its Euclidean distance from all vocabulary tokens are several orders of magnitude beyond normal variation (e.g., z-score > 3000), suggesting that it does not correspond meaningfully to any real token. This is expected, as the soft token in our framework is not intended to approximate a specific discrete token, but rather to serve as a functional control token, influencing the overall response.

1. Verifier Setup

Regarding the questions about details of the verifier, NO separate or stronger LLM is used for verification (Line 31-33, right column); the same model as the generator is employed (i.e. LLaMA3-8B-Ins, Qwen2-7B-Ins, and Mistral-7B-Ins).

• The verifier score $r_{verifier}(y) = \mathbb{1}_{y_v = y}$ is a binary indicator, where $y_v$ is the model's regenerated answer based on the current question and previous responses. Specifically, $y_v$ is obtained by prompting the same LLM with a verification query:

  Based on the given question and the previous answers, please provide your analysis and final answer.

  The exact prompts used are listed in Appendix B.3.

• We also conducted ablation studies (see Verifier Comparison: Judgment vs. Generation, Line 427, left column) to compare different verifier strategies and assess their impact on performance.

This setup ensures a fair comparison, as no external or more capable model is used. It also highlights one of the main contributions of this work: our BO framework's ability to enhance reasoning capabilities using a single unified LLM without additional verifiers. We will add these details on the setup of the verifier to our paper in the next revision.

2. Related Works

Lack of discussions and comparisons to other controlled generation works [1,2]

Thanks, we will add the comparison in the revision. Our goal is to propose an efficient reasoning framework that does not require additional or stronger verifiers, nor task-specific model fine-tuning. Instead, we leverage verification from the model itself to enhance the reasoning ability of LLMs.

The methods in [1,2] serve different purposes: [1] uses trainable prefix scorers to guide decoding, while [2] proposes a fine-tuning objective aimed at improving robustness against adversarial prompts. Both approaches require additional training, making direct comparison with our training-free method less straightforward.

We implemented a baseline based on [1] using a blockwise best-of-8 decoding (block size=16), where the model scores its own without any fine-tuning. To improve the scorer's quality, we applied self-consistency. For [2], we performed token-wise constrained fine-tuning on LLaMA3-8B-Ins model using LoRA (5 epochs, learning rate 2e-5, rank 16).

As shown, while [2] achieves the best performance on the StrategyQA task with additional training, our method—without requiring any extra training—achieves superior results on the other three tasks, particularly in few-shot. In a fairer comparison (i.e. without training), our method consistently outperforms the prefix scorer [1] across all tasks.

[1] Controlled Decoding from Language Models. ICML

[2] Safety alignment should be made more than just a few tokens deep.

3. Why choose EI?

There exist multiple acquisition functions for BO, such as UCB score, etc., is there any theoratical/empirical results that support the choice of EI?

Indeed, PI and GP-UCB can also be used as the acquisition function. The cumulative regret for GP-UCB has the same rate (Theorem A.1) as EI [3]. Unlike EI or GP-UCB, PI only considers the probability of improvement, without accounting for its magnitude. It is thus less theoretically grounded and more prone to premature exploitation [4].

The experiments show that PI consistently underperforms compared to EI, as expected from the theoretical discussion above. For GP-UCB, its performance is sensitive to the choice of the exploration parameter and is, in most settings, worse than EI. We would also like to mention that the optimal parameter choice for GP-UCB varies across different tasks, making it difficult to guarantee good performance in unseen settings. In contrast, EI performs robustly without requiring task-specific tuning.

[3] Taking the human out of the loop: A review of Bayesian optimization. Proc. IEEE

[4] Gaussian process optimization in the bandit setting: No regret and experimental design. ICML

In my opinion this work could get even more impact and recognition if authors might show the effectiveness of this approach in terms of sample efficiency when its used during training with preference optimization or reward-based training.

Thank you for this insightful and forward-looking suggestion. We agree that extending our embedding-based exploration method to training—particularly within preference optimisation frameworks—to enhance sample efficiency could be a promising direction.

Specifically, our Bayesian optimisation approach to embedding perturbations can be integrated into preference optimisation as an efficient and targeted data generation mechanism. A key bottleneck in current reward-based training frameworks lies in the use of standard generation methods (e.g., temperature sampling), which often produce outputs that, while differing in perplexity, appear similarly plausible—making it difficult for reward models or human annotators to assign differentiated scores. Our proposed method could indeed offer an alternative approach. Instead of relying on random sampling or temperature-based decoding to generate candidate answers, our method explores promising regions in the embedding space, producing diverse and high-quality outputs with fewer trials. During training, we hope that these outputs can be evaluated using human preferences or reward models to construct preference pairs or ordered sequence with more significant distinct scores, which is able to provide more efficient supervision signals for optimising the model.

If it is possible, by having those more distinctive generated samples and preference-aligned candidates in each iteration, the training process can access richer supervision without increasing the number of queries or annotations. This could lead to better utilisation of limited feedback and help the model learn more effectively from each batch of generated candidates. Additionally, one of the main properties of our proposed framework is that it is model-agnostic and lightweight. It can be seamlessly integrated into standard preference optimisation pipelines without modifying model architectures or requiring access to gradients.

We appreciate the suggestion and view this integration as a promising approach for generating more preference-aligned data during training. We plan to explore its impact on training efficiency and reasoning performance in future work.

1. Verifier Reliability

Thanks for the comment. You are right in saying that incorporating more reliable, domain-specific verifiers could further improve performance. Our framework is flexible enough to integrate such tools, and we agree that doing so would likely yield gains in tasks where those verifiers are applicable [1]. That said, our current focus is on improving reasoning in scenarios where such external tools are not available or applicable — for instance, in common-sense or social reasoning, where exact computation isn’t useful. By relying on internal LLM verifier signals, we aim to explore the model's capacity to self-reflect and reason without additional guidance (Line 31-33, right column).

[1] Large Language Models for Mathematical Reasoning: Progresses and Challenges. ACL

2. Scaling to Larger Models

Thanks for the suggestion. Note that our method involves modifications within the model’s parameter space, and due to the closed-source nature of GPT-4, it is not possible to test on it. Nonetheless, the point is well-taken and we have conducted additional experiments using Qwen2-72B-Instruct, a significantly larger model (72B) than those used in the main paper (7B and 8B), across both zero-shot and few-shot settings:

While we still see some improvements, the gains are indeed smaller on easier benchmarks like GSM8K and SVAMP, as large models can already achieve very high accuracy. To further stress-test our method, we evaluated it on the more challenging AIME-2024 dataset. Here, we observe that our approach continues to yield significant gains even at the 72B scale, showing its robustness and scalability to harder tasks and more capable models. We will add these results to clarify the method’s applicability beyond smaller-scale settings.

3. Optimisation Scope

To investigate this, we conducted additional experiments where we optimised embeddings for the first k tokens (instead of just the first):

The performance generally degrades as k increases, especially beyond 5 tokens. This suggests that naively extending to multiple tokens can introduce instability or overfitting. We also compare with RAP (one of our baselines), a tree-search-based method that operates at the sequence level rather than token-by-token—though it shares similar ideas with token-by-token search. While RAP achieves strong performance, it incurs substantially higher cost and still underperforms our approach. Developing a multi-token optimisation strategy that can achieve both high accuracy and cost-effectiveness would require deeper investigation and extensive experimentation. It remains an interesting direction for future work and would warrant a separate paper.

4. Coherence Evaluation

We appreciate the reviewer’s observation. To address this, we evaluated coherence using two metrics: perplexity and a coherence score rated by DeepSeek-R1-Distill-Llama-70B, prompted to rate responses from 1 (poor) to 5 (excellent). We tested 800 samples (200 from each of four datasets) using LLaMA3-8B-Ins as the base model:

We would like to mention that coherence is not the main objective of our method, but serves as a filtering criterion to discard low-quality completions and support answer correctness.

5. Related Works

We thank the reviewer for the suggestions. We will discuss them in the final version.

1. Verifier Setup

Thanks for the comments. Regarding the questions about details of the verifier, NO separate or stronger LLM is used for verification (Line 31-33, right column); the same model as the generator is employed (i.e. LLaMA3-8B-Ins, Qwen2-7B-Ins, and Mistral-7B-Ins).

• The verifier score $r_{verifier}(y) = \mathbb{1}_{y_v = y}$ is a binary indicator, where $y_v$ is the model's regenerated answer based on the current question and previous responses. Specifically, $y_v$ is obtained by prompting the same LLM with a verification query:

  "Based on the given question and the previous answers, please provide your analysis and final answer."

  The exact prompts used are listed in Appendix B.3.

• We also conducted ablation studies (see Verifier Comparison: Judgment vs. Generation, Line 427, left column) to compare different verifier strategies and assess their impact on performance.

This setup ensures a fair comparison, as no external or more capable model is used. It also highlights one of the main contributions of this work: our BO framework's ability to enhance reasoning capabilities using a single unified LLM without additional verifiers. Details on the verifier will be added to our paper.

2. Dimension Reduction

While dimensionality reduction alleviates the curse of dimensionality, random projection may lose critical information. Further validation is needed to confirm that the negative impacts of dimensionality reduction are controllable.

We agree that dimensionality reduction may lead to some information loss. However, it offers a practical tradeoff: retaining more information in high-dimensional spaces often results in poor sample efficiency for Bayesian optimisation, making it harder to find good configurations under a limited evaluation budget. As shown in Figure 6, we experimented with several projection dimensions and found that our current setting (d=50) provides a good balance between optimisation performance and computational cost.

To assess the stability of random projection, we ran simulations using 50 different random projection matrices. The box chart below shows the distribution of results:

0.850 ┤   ┐      ◀ Max
      │   │
0.846 ┤ ┌─┘─┐    ◀ Q3
      │ │   │
0.842 ┤ │ ─ │    ◀ Median
      │ │   │
0.838 ┤ └─┐─┘    ◀ Q1
      │   │
0.835 ┤   ┘      ◀ Min

The performance remains stable across multiple runs, indicating that random projection does not introduce significant variance or instability.

3. Critical neurons

The definition of critical neurons relies on statistical methods; citing related work would help justify this choice.

We appreciate the reviewer’s concern. Prior works have shown that it is possible to trace information flow within transformers and identify neurons with causal influence on model predictions by applying targeted interventions, such as activation replacement or ablation [1, 2].

Following this line of work, we mask the critical neurons identified for each input. This leads to a significant drop in accuracy to 13.27% (from 62.14% before masking). For comparison, we randomly masked the same number of neurons and repeated the experiment under identical settings, resulting in an average accuracy of 41.89%. This substantial gap (41.89%→13.27%) demonstrates that the identified neurons are indeed functionally important, beyond what would be expected by chance.

[1] Damai Dai et al., 2022. Knowledge Neurons in Pretrained Transformers. ACL

[2] Kevin Meng et al., 2022. Locating and Editing Factual Associations in GPT. NeurIPS

4. Computational efficiency

Regarding computational efficiency, readers may expect comparisons of time and space costs rather than solely token count statistics.

We thank the reviewer for the suggestion. We supplement the token count statistics with comparisons between our method and RAP on inference time and memory usage. To evaluate the latter, we focus on the two variable components: (1) KV cache, and (2) intermediate activations, since model weights remain constant across methods. Using vLLM’s block-based memory tracking, we report both average and peak usage, sampled at 1-second intervals.

The results show that our method's inference time is only 12.30% of RAP's while also consuming significantly less memory, further validating its computational efficiency advantage.