Title: Inference-time Alignment in Continuous Space
Venue: NeurIPS 2025
Accepted: yes
Keywords: Alignment
Abstract: Aligning large language models with human feedback at inference time has received increasing attention due to its flexibility. Existing methods rely on generating multiple responses from the base policy for search using a reward model, which can be considered as searching in a discrete response space. However, these methods struggle to explore informative candidates when the base policy is weak or the candidate set is small, resulting in limited effectiveness. In this paper, to address this problem, we propose Simple Energy Adaptation ($\textbf{SEA}$), a simple yet effective algorithm for inference-time alignment. In contrast to expensive search over the discrete space, SEA directly adapts original responses from the base policy toward the optimal one via gradient-based sampling in continuous latent space. Specifically, SEA formulates inference as an iterative optimization procedure on an energy function over actions in the continuous space defined by the optimal policy, enabling simple and effective alignment. For instance, despite its simplicity, SEA outperforms the second-best baseline with a relative improvement of up to $ \textbf{77.51\%}$ on AdvBench and $\textbf{16.36\%}$ on MATH. Code is publicly available at [this link](https://github.com/yuanyige/sea).


Peer reviews:
{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper introduces a very intriguing approach to inference-time alignment for large language models. It points out a potential limitation of traditional methods that search over a discrete space, such as best‑of‑N sampling, and proposes an elegant solution that operates directly in the continuous space of the model’s logits. \n\nBy framing the problem as an energy-based optimization, the authors show how to adapt an initial response towards one that better maximizes a reward model, using gradient-based Langevin dynamics. In doing so, the method allows the generation process to move beyond the constraints of the base model’s candidate set, yielding higher quality, better aligned responses. I find this paper very insightful: it is compelling to move from brute-force search towards a more principled, continuous optimization approach, and it demonstrates benefits across a range of alignment tasks.\n\nHowever, I still have some questions that I would like to clarify before I can fully recommend this paper for acceptance at NeurIPS.\n\nSTRENGTHS AND WEAKNESSES:\n## Strengths\n- This paper introduces a very interesting and compelling approach by shifting inference‑time alignment from a discrete candidate search to a continuous optimization of logits. To the best of my understanding, this is a very natural and promising idea. Although I am not deeply familiar with this area and cannot definitively say if it has been attempted before, the approach feels highly novel and valuable.\n- The paper is generally well‑written, making it relatively easy for readers to understand both the motivations and the proposed technique. The illustrations (such as Figure 1 and Figure 3) and the structured experimental results further aid comprehension. Although I found some parts of Figure 2 — for example, subfigure (b) — a bit hard to interpret clearly, overall the paper is very clearly presented and well‑structured.\n\n\n## Weaknesses\n- **Evaluation Scope**: The paper primarily showcases results on a relatively small set of reasoning tasks (GSM8K and MATH) and uses a single reward model for evaluation. Although the approach is novel and promising, it would be more convincing if it were compared against a more comprehensive range of state‑of‑the‑art methods for reasoning, especially recent advances that build upon best‑of‑N (I do not really care whether your proposed method is SOTA or not, but you should have included and discussed these baselines here). Currently, Table 3 doesn’t seem to include many strong baselines in this space. For instance, on GSM8K, to the best of my knowledge, the best results for LLaMA‑3.2‑1B‑Instruct are far beyond the 58% reported here. I’m also not sure about the role of the reward signal here. Reward hacking is a well‑known issue, and a high reward doesn’t necessarily imply a genuinely better or more aligned output. Moreover, for the safety tasks, I feel that an inference‑time optimization method is not primarily designed to solve safety concerns. It might be more convincing if the paper put less emphasis on safety evaluations and focused more on reasoning tasks, where the approach can truly demonstrate its value.\n- **Cost**: SEA introduces an iterative gradient‑based optimization at inference time. Compared to best‑of‑N sampling, this is a much more computationally expensive process. However, the paper provides no concrete discussion of inference cost, scalability. In fact, as model size grows, this cost is likely to scale more than linearly. How do you plan to address this?\n- **Learning Rate Sensitivity**: This is actually very important. It appears that the method’s success depends on a fixed learning rate. In practice, this can be challenging, and inappropriate lr settings could degrade performance significantly, especially given the high dimensionality and sparsity of the logits space.\n- **Greedy Decoding > meaningful outputs?**: The approach ultimately decodes the final optimized logits using a greedy or argmax strategy. This raises concerns that the final output may overfit the reward signal and result in degenerate or nonsensical text, especially when the reward landscape is highly irregular. The paper doesn’t clearly justify why or how SEA can reliably produce human‑understandable sequences under these conditions.\n\nQUESTIONS:\nIn token sampling methods, the discrete nature of sampling can act as a regularizer that limits extreme exploitation of the reward signal (*i.e.,* reward hacking). In contrast, your method directly optimizes in the continuous logits space, which seems like it could make the optimization more prone to reward hacking. Could you clarify why operating in the continuous space does not exacerbate reward hacking, or whether you observed any tendencies for the optimized outputs to become degenerate or misaligned with human preferences despite achieving higher rewards?\n\n> “the reward model should be based on a model of the same lineage as the policy model, or else downstream performance can degrade significantly. Therefore, simply taking the highest-scoring reward model on a benchmark will not ensure a good post-RLHF model.”\n> — Malik et al., *RewardBench 2*\n\nLIMITATIONS:\nSee above.", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 4, "CLARITY": 4, "IMPACT": 3, "ORIGINALITY": 4, "DATE": "2025-06-24"}

{"IS_META_REVIEW": false, "comments": "SUMMARY:\nThis paper proposes an inference-time alignment algorithm called Simple Energy Adaption (SEA), which leverages gradient-based searching to find optimal responses guided by some energy function. Instead of directly sampling tokens from the base model, SEA manipulates with logits in the continuous latent space. Through interative optimization, the final response achieves high score from the reward model.\n\nSTRENGTHS AND WEAKNESSES:\nStrengths:\n\nThe writing and layout of the paper is clearly structured and easy to follow, which is nice. The landscape illustration is effective, and the analysis of optimization dynamics (e.g., KL divergence and token position) provides useful insights.  The idea of optimizing in the continuous latent space is interesting and well-supported by energy-based model (EBM) theory.\n\nWeaknesses:\n\n- SEA extends token sampling into a continuous logits space, which significantly increases the search space.  I might missed some points, but as optimization relies on sparse rewards (evaluating full sequences), how does the method ensure fluency and coherence when directly modifying logits?\n\n- The same reward model is used for both optimization and evaluation, which risks reward hacking and may not provide a fair assessment. Using third-party reward models (or GPT-4 as a judge) would offer more reliable and unbiased comparisons.\n\n- The gradient-based sampling approach appears similar to COLD [1], with the main difference lying in the energy function design (SEA uses RLHF’s optimal policy, while COLD employs hard-coded constraints). A clearer discussion on the methodological distinctions and motivations would help position SEA’s contributions more effectively.\n\n[1] Qin, Lianhui, et al. \"Cold decoding: Energy-based constrained text generation with langevin dynamics.\" Advances in Neural Information Processing Systems 35 (2022): 9538-9551.\n\nQUESTIONS:\n1. what is the temperature during decoding for SEA and the other baselines?\n2. How many iterations are applied to SEA for the figures reported in Tables1-4?\n3. Do SEA-generated responses tend to be longer than baseline methods? If so, could this artificially inflate reward scores, since longer outputs may exploit length-related biases in the reward model?\n\nLIMITATIONS:\nyes", "TITLE": "Official Review", "RECOMMENDATION": 4, "REVIEWER_CONFIDENCE": 3, "CLARITY": 3, "IMPACT": 2, "ORIGINALITY": 2, "DATE": "2025-06-25"}

comments: SUMMARY:
This paper aims to improve test-time compute methods by addressing two key limitations: 1) the need for large number of samples and 2) having a weak sampling, reference policy. Inspired by the connection between the optimal KL regularized RLHF formulation and energy based models, the authors construct an MCMC procedure to generate samples at test-time to improve beyond existing test time search methods. Their experiments show great success in maximizing the reward of generations from small, weaker LLMs. Furthermore they conduct extensive experiments testing the effect of KL minimization, shallow alignment, and a computational comparison to pre-existing methods.

STRENGTHS AND WEAKNESSES:
This paper is very clearly written with an incredibly thorough experimentation section. Any weaknesses I present are minor nits:
1. The font on figures are quite small making the plots less readable
2. If possible, it would be great to check other model families such as a Qwen 0.7B 
3. How does quality/size of the reward model affect downstream performance as well as computation complexity? Some high level analysis on how the computational needs for best of N vs SEA on reward model size would also be very valuable. I would expect there to be worse memory scaling for SEA compared to BoN with Reward Model size increase but hopefully offset by needing fewer steps than best of N to achieve the same performance.
4. Rather than time comparisons, perhaps do flop computation

QUESTIONS:
I have listed by questions in the strengths and weaknesses section.

LIMITATIONS:
There was no specific discussion about limitations. Perhaps a brief discussion in the conclusion would be helpful.
TITLE: Official Review
CLARITY: 4
IMPACT: 4
ORIGINALITY: 4
DATE: 2025-07-02

comments: SUMMARY:
The paper focuses on a new method for inference-time alignment. Roughly, there are two general strategies for aligning language models. The first is fine-tuning-based, i.e., where a language model is fine-tuned given human-annotated data. The second is inference-based, where additional compute is spent at inference time to align the model. The classic example of this is best-of-N inference. 

As I understood the paper, the idea is to use the Bayesian characterization of reinforcement learning to get a posterior. Then, the novelty is to apply Langevin dynamics to sample from the model. Specifically, the paper makes use of the COLD algorithm by Yejin Choi's group. 

The experimental results show improvements.

STRENGTHS AND WEAKNESSES:
I liked the paper. It was easy for me to read it once and understand what was happening. This is the paper's biggest strength. I am not that involved in related work, so I can't speak to its novelty but, I trust that the authors did a literature review. 

My concern is that the paper does not acknowledge that COLD is not a valid sampler. Two other recent papers (https://arxiv.org/abs/2306.03061) and (https://arxiv.org/abs/2312.17710) develop correct samplers, i.e., MCMC samplers that sample int he the limit. I would really like to understand how much the bias of COLD affects the results. I also think this limitation should be acknowledge in the paper.

QUESTIONS:
* Why did you choose COLD over other alternatives (https://arxiv.org/abs/2205.12558) or the papers listed above?
* Can you determine how much burn-in you need in your MCMC chain to get samples? Does the biasedness of COLD matter?

LIMITATIONS:
None
TITLE: Official Review
CLARITY: 3
IMPACT: 3
ORIGINALITY: 3
DATE: 2025-07-03