PaperlayerView exact revision on arXiv

Source-linked AI summary

Seek in the Dark: Reasoning via Test-Time Instance-Level Policy Gradient in Latent Space

Hengli Li, Chenxi Li, Tong Wu, Xuekai Zhu, Yuxuan Wang, Zhaoxin Yu, Eric Hanchen Jiang, Song-Chun Zhu, Zixia Jia, Ying Nian Wu, Zilong Zheng

arXiv:2505.13308v3cs.LGcs.AIcs.CL

TL;DR

LLMs still struggle with structured, step-by-step reasoning, while parameter-updating approaches incur substantial computational costs. LatentSeek optimizes instance-specific latent representations at test time using policy gradients without changing model parameters, and it outperforms baselines across reasoning benchmarks.

  • Problem

    LLMs still struggle with structured, step-by-step reasoning, while parameter-updating approaches incur substantial computational costs and risks such as catastrophic forgetting.

    Source: §1 · pp. 1–2

  • Method

    LatentSeek uses policy gradients to optimize instance-specific latent representations at test time, steering reasoning without modifying model parameters.

    Source: §1 · p. 2

  • Results

    LATENTSEEK surpasses every baseline across all 30 settings and achieves the best average performance; its average GSM8K gain over BoN is 15.23 points.

    Source: §1 · p. 3, §3.2 · p. 7

  • Takeaways & Limitations

    The findings point to latent-space optimization as a possible avenue for test-time scaling, with rapid convergence on average-level problems.

    Source: §5 · p. 13

    The approach assumes independence among latent components, and the authors note that this assumption may introduce potential loss.

    Source: §2.2 · p. 4

Abstract

from arXiv · show

Reasoning ability, a core component of human intelligence, continues to pose a significant challenge for Large Language Models (LLMs) in the pursuit of AGI. Although model performance has improved under the training scaling law, significant challenges remain, particularly with respect to training algorithms, such as catastrophic forgetting, and the limited availability of novel training data. As an alternative, test-time scaling enhances reasoning performance by increasing test-time computation without parameter updating. Unlike prior methods in this paradigm focused on token space, we propose leveraging latent space for more effective reasoning and better adherence to the test-time scaling law. We introduce LatentSeek, a novel framework that enhances LLM reasoning through Test-Time Instance-level Adaptation (TTIA) within the model's latent space. Specifically, LatentSeek leverages policy gradient to iteratively update latent representations, guided by self-generated reward signals. LatentSeek is evaluated on a range of reasoning benchmarks, including GSM8K, MATH-500, and AIME2024, across multiple LLM architectures. Results show that LatentSeek consistently outperforms strong baselines, such as Chain-of-Thought prompting and fine-tuning-based methods. Furthermore, our analysis demonstrates that LatentSeek is highly efficient, typically converging within a few iterations for problems of average complexity, while also benefiting from additional iterations, thereby highlighting the potential of test-time scaling in the latent space. These findings position LatentSeek as a lightweight, scalable, and effective solution for enhancing the reasoning capabilities of LLMs.

1. Introduction

The introduction presents LATENTSEEK as a test-time method for improving instance-level reasoning without modifying model parameters. It optimizes latent representations using self-generated rewards and reports gains across reasoning benchmarks.

1. Introduction

  • LLMs still struggle with structured, step-by-step reasoning, while parameter-updating approaches can be costly, risk catastrophic forgetting, reduce exploration, or produce overly verbose responses.The introduction motivates test-time instance-level adaptation as an alternative that operates per instance without updating model parameters.Source: §1 · pp. 1–2
  • Earlier latent-space fine-tuning studies (Deng et al., 2022; Hao et al., 2024) support the semantic adequacy of hidden states, but their latent-space modifications leave performance below CoT.Prompt engineering, a prevalent test-time instance-level adaptation approach, is described as limited in expressiveness.Source: §1 · p. 2
  • LATENTSEEK uses policy gradient to update token-wise, instance-specific latent representations as a planning mechanism, decoding them for reward evaluation and stopping at a threshold or iteration limit.Its reward relies only on the model’s internal capabilities, without external information; the latent space is the transformer output space before the final language-model head.Source: §1 · pp. 2–3
  • LATENTSEEK gains 15.23 points over BoN on GSM8K, 4.72 points over Self-Reflection on MATH-500, and 6.67 points over CoT on AIME2024.These are reported average gains; the passage also cites Table 1.Source: §1 · p. 3
  • With LLaMA3.1-8B-Instruct, LATENTSEEK surpasses SimpleRL-Zoo (+18.1) and Genius (+12.7) on GSM8K and MATH-500, while more latent updates improve test-time performance.With an ideal ground-truth verifier, it averages a 19.12-point gain over CoT; case studies report correct answers despite unintelligible intermediate reasoning.Source: §1 · p. 3

2. Test-Time Instance-Level Policy Gradient in Latent Space

The section formulates instance-level reasoning as reward-guided search, then derives latent-space policy-gradient optimization and the LatentSeek procedure. Its method initializes from chain-of-thought latents, refines them using self-reward, and can stop early at a reward threshold.

2.1. Problem Formulation: Test-Time Instance-Level Reasoning

  • Reasoning is framed as finding an optimal token path for a problem context, scored by a reward function when ground truth is unavailable.Each token is associated with a latent representation produced by the Transformer backbone from the context and preceding tokens.Source: §2.1 · p. 3

2.2. Reasoning via Policy Gradient in Latent Space

  • The method optimizes a sequence of latent representations instead of searching directly over tokens, decoding them into tokens before continuing autoregressive generation.It uses a REINFORCE-based policy gradient to update the latents from sampled reasoning sequences and their rewards.Source: §2.2 · pp. 3–4
  • Assuming latent representations are independent expands exploration by decoupling latent components, while the base model’s follow-up generation and reward model’s evaluation can mitigate potential loss.The authors motivate this assumption by noting that autoregressive dependence can concentrate optimization on the first latent token and collapse the effective search space.Source: §2.2 · p. 4

2.3. LATENTSEEK Algorithm

  • LatentSeek initializes latents from chain-of-thought, optimizes only a fraction of the sequence, and iteratively refines latents using self-reward until exceeding a threshold or typically fewer than three iterations.The optimized fraction balances exploration capacity against update cost; modifying incoherent decoded outputs can make the reward less reliable.Source: §2.3 · p. 5

3. Empirical Results

The empirical results evaluate LATENTSEEK across mathematical-reasoning benchmarks, model backbones, and comparison methods. They report broad benchmark gains and examine sparse rewards, iteration scaling, reward guidance, and latent-space reasoning patterns.

3.1. Experimental Setup

  • Experiments assess GSM8K, MATH-500, and AIME2024 across Qwen and LLaMA backbones, using two answer formats and self-generated reward prompts.The comparison set includes prompting, explicit-search, reinforcement-learning, latent-chain-of-thought, and supervised-fine-tuning baselines.Source: §3.1 · pp. 5–6

3.2. State-of-the-art Test-time Reasoning Performance

  • LATENTSEEK surpasses every baseline across all 30 settings and achieves the best average; across backbones, it gains 15.23 points over BoN on GSM8K and 4.72 over Self-Reflection on MATH-500.On AIME2024, its average gain over CoT is 6.34 points. Across 7–8B models, it exceeds BoN by 7.67 points on Qwen2, Self-Reflection by 4.82 on Qwen2.5, and Self-Consistency by 9.65 on LLaMA3.1.Source: Table 1 · p. 6, Table 2 · p. 7, §3.2 · p. 7

3.3. Ideal Experiment: Perfect Sparse Reward Model

  • Under an all-or-nothing Perfect Sparse Reward Model, LATENTSEEK outperforms its self-reward variant by an average of 10.67 points across evaluated settings.Because the reward is 0 only for an exactly correct final answer and −1 otherwise, the authors interpret the gains as evidence that latent-space exploration can help despite nearly directionless feedback.Source: Table 3 · p. 8, §3.3 · p. 8

3.4. Test-Time Scaling: scaling up the iteration of LATENTSEEK

  • With PSRM, increasing LATENTSEEK iterations produces continued GSM8K improvement, whereas self-reward gains plateau before improving again.With extreme scaling, Qwen2.5-1.5B-Instruct gains 14% point performance over GPT-4o on AIME2024 and trails o1-preview by 2.7% point on MATH-500, using an average of 61.8 iterations on MATH-500.Source: §3.4 · p. 9, Table 4 · p. 9

3.5. Comparison with Unguided Search

  • LATENTSEEK outperforms single-step stochastic exploration on GSM8K by an average of 13.66 points across noise levels σ2 = 0.5, 0.75, and 1.0.The comparison tests whether iterative gradient guidance, rather than unguided latent perturbation, accounts for the gains.Source: Table 5 · p. 9, §3.5 · p. 9

3.6. Analysis of Self-Reward Mechanism

  • The self-reward mechanism is reliable in approximately 80% of cases, while replacing it with constant −1 rewards causes an average 14.21-point performance drop.For Qwen2.5-7B-Instruct, the reported example falls from 85.06% to 64.52%.Source: §3.6 · p. 10, Table 6 · p. 10, Table 7 · p. 10

3.7. Algorithmic Statistics

  • Table 8 reports the ratio of answer length to CoT length and the average number of update iterations.The supplied passage provides no values for these statistics.Source: Table 8 · p. 10

3.8. Qualitative Analysis

  • LATENTSEEK can reach correct answers despite linguistically anomalous reasoning, and its latent optimization trajectories include zig-zag, unidirectional, and V-shaped patterns.The analyzed first three tokens under Prompt 1 included “thecy,” “theella,” and “Sid.”Source: §3.8 · pp. 11–12

4. Related Work

Related work spans reasoning, reinforcement learning, controllable generation, and prompt tuning; LatentSeek is distinguished by instance-level test-time optimization of latent representations rather than static prompting, parameter updates, discrete search, or trained prompts.

Reasoning in Language Models

Reinforcement Learning for Language Models

Controllable Generation and Test-Time Optimization

Prompt Tuning and Soft Prompt

  • Unlike prompt and prefix tuning, LatentSeek requires neither labeled data nor model updates.The passage characterizes prompt and prefix tuning as requiring labeled data and full backpropagation, while LatentSeek uses latent-space manipulation for adaptation.Source: §4 · p. 12

5. Conclusion

LatentSeek uses policy-gradient updates to latent representations for test-time instance-level adaptation without updating model parameters. It outperforms existing baselines across reasoning benchmarks and converges rapidly on average-level problems.

Conclusion

  • LatentSeek applies policy gradient to optimize latent representations through test-time instance-level adaptation, avoiding parameter updates and substantial retraining or reinforcement learning.The conclusion presents this as an alternative to methods that require retraining or reinforcement learning.Source: §5 · p. 12
  • Across multiple reasoning benchmarks, LatentSeek consistently outperforms existing baselines, including Chain-of-Thought and reinforcement learning-based techniques.The conclusion describes the results as superior performance compared with these baselines.Source: §5 · pp. 12–13
  • LatentSeek is computationally efficient, with rapid convergence on average-level problems.The conclusion identifies rapid convergence as a computational-efficiency benefit.Source: §5 · p. 13
  • The framework demonstrates a possible avenue for test-time scaling in latent space.The conclusion frames latent-space test-time scaling as a direction enabled by this work.Source: §5 · p. 13

Appendices · A. Discussion and future works

The discussion identifies limitations in reward guidance, latent optimization, and model scale, and presents these as directions for future research.

A. Discussion and future works

  • Self-rewarding constrains optimization to the base model’s evaluation capabilities and potential biases, motivating objective outcome assessment and more reliable verifiers.The authors describe a principled Outcome Reward Model (ORM) as an ideal source of unambiguous final-answer guidance.Source: A. Discussion and future works · p. 19
  • Publicly available ORMs were not robust or generalizable enough for this task, and their noisy signals guided latent optimization less effectively than self-reward.Developing more powerful verifiers for complex reasoning is identified as future work.Source: A. Discussion and future works · p. 19
  • The implementation uses standard policy-gradient methods, leaving Proximal Policy Optimization and latent-space-specific optimization for future exploration.The authors name PPO as a more advanced reinforcement-learning algorithm to investigate.Source: A. Discussion and future works · p. 19
  • Experiments cover base models up to 14B parameters because of computational-resource constraints, leaving larger-model scaling for future investigation.The authors identify scaling to larger base models as an important avenue for future work.Source: A. Discussion and future works · p. 19

B. Methods of Test-Time Instance-Level Reasoning · C. Theoretical Analysis · C.1. Preliminaries: Multiple Prover Interactive Proofs and NEXP

The sections define two test-time reasoning baselines and set out a theoretical framework centered on latent-variable independence. They also introduce MIP and nondeterministic-time preliminaries used in the analysis.

B. Methods of Test-Time Instance-Level Reasoning

C. Theoretical Analysis

  • LatentSeek’s theoretical framework focuses on justifying the expressiveness of independence among latent variables.The section presents this analysis as a starting point for understanding and inspecting the algorithm.Source: C. Theoretical Analysis · p. 19

C.1. Preliminaries: Multiple Prover Interactive Proofs and NEXP

  • MIP provers may communicate before a proof, but during it they interact privately with the verifier and cannot communicate with one another.The preliminaries introduce MIP and NEXP for completeness and cite Vadhan and Zhang (2002) and Arora and Barak (2006) as expositions.Source: §C.1 · p. 19
  • A k-MIP verifier runs in polynomial time with polynomially many and sized messages, and the definition requires completeness at least 2/3 and soundness at most 1/3.MIP is defined as the union of k-MIP over k; the k-MIP definition is attributed to Vadhan and Zhang (2002).Source: §C.1 · p. 20
  • The preliminaries define NTIME(T(n)) by nondeterministic Turing-machine recognition within cT(n) time for a constant c.This definition is presented in introducing NEXP and cites Arora and Barak (2006).Source: §C.1 · p. 20

C.2. Theoretical Analysis: Independent Updating

The analysis argues that independent latent-token updates can retain theoretical expressive power through a bounded multi-prover formulation. It presents this as a heuristic rather than rigorous validation, and says the independence deficit can be mitigated under favorable model and reward conditions.

C.2. Theoretical Analysis: Independent Updating

  • LatentSeek models independently updated tokens as noncommunicating provers, autoregressive decoding as the verifier, and CoT initialization as pre-proof communication; each prover emits one bounded token.A polynomial proportion of tokens may be optimized, yielding a polynomial number of provers.Source: §C.2 · pp. 20–21
  • The corollary says a verifier can control dependence between provers, and a strong base model with a well-aligned reward model can mitigate the independence deficit.The verifier may, for example, request dependence between generated tokens.Source: §C.2 · pp. 21–22
  • The analysis is a heuristic guideline, not rigorous validation, because it omits gradient-update changes and cannot guarantee every token is reachable at each step.This limitation is especially relevant when provers rely only on gradients from the language-model head.Source: §C.2 · p. 22
  • MIP-Bounded = MIP = NEXP, and the corollary gives NP ⊂ NEXP = MIP-Bounded (Babai et al., 1990).The paper derives the equivalence from the classical MIP result and uses it to characterize the expressive power of its framework.Source: §C.2 · p. 23

D. Derivation of Policy Gradient … E.4. Greedy Decoding vs Sampling

The appendix derives the latent-space policy gradient and reports additional results showing LATENTSEEK’s benchmark performance, token efficiency, and decoding behavior. Across the reported comparisons, LATENTSEEK outperforms BoN, while greedy decoding is selected for the main experiments.

D. Derivation of Policy Gradient

  • The policy-gradient derivation differentiates the objective with respect to z, applies the log-policy gradient identity, and removes a term independent of z.The derivation targets Equation (7) and identifies the resulting term for each t.Source: D. Derivation of Policy Gradient · pp. 23–24

E.1. LATENTSEEK vs BoN (N=10)

  • LATENTSEEK remains superior to BoN (N=10) on all datasets and backbone models, despite BoN requiring at least five times more sequence-level computation.The comparison uses Prompt 2.Source: Table 10 · p. 25, §E.1 · p. 25

E.2. Qwen3 AIME Performance

  • Under Prompt 1 on AIME2024, LATENTSEEK reaches 73.3% accuracy, a 10-point uplift over CoT and BoN, which both score 63.3%.These results are for Qwen3-4B-Instruct-2507.Source: §E.2 · p. 25, Figure 4 · p. 26
  • Under Prompt 2 on AIME2024, LATENTSEEK scores 56.7% accuracy, surpassing CoT by 16.7 points and BoN by 3.4 points.These results are for Qwen3-4B-Instruct-2507.Source: §E.2 · p. 25, Figure 4 · p. 26

E.3. Token Efficiency Comparison with BoN Baseline

  • LATENTSEEK achieves higher token efficiency than BoN across model families and sizes while attaining superior reasoning accuracy, without additional token overhead.The token-efficiency comparison covers GSM8K and MATH-500 under Prompt 2.Source: §E.3 · p. 27, Figure 5 · p. 27

E.4. Greedy Decoding vs Sampling

  • Greedy decoding performs slightly better than categorical sampling and is used in the main experiments to reduce stochastic error.The paper hypothesizes that sampling can amplify approximation error when replacing an expected mean with an empirical one; Table 11 compares the methods at temperature 0.7 on MATH-500.Source: §E.4 · p. 27, Table 11 · p. 27

E.5. Limitations of Outcome Reward Model … G.3. Total FLOPs for LATENTSEEK

The appendix reports limits of outcome-based rewards, reward and fraction-ratio trends, and the advantage of optimizing initial-stage rather than middle-stage latent tokens. It also documents experimental prompts, models, baselines, dataset procedures, and FLOPs estimation methods.

E.5. Limitations of Outcome Reward Model

  • Table 12 reports an average score of 64.79 for the outcome reward model, 3.23 points above CoT but below self-rewarding at 77.18.The authors identify the gap as evidence of outcome-based reward model limitations when ground-truth data are unavailable.Source: §E.5 · pp. 27–28, Table 12 · p. 28

E.6. Trends of Reward

  • On GSM8K with a JSON-formatted prompt, LATENTSEEK’s reward trends move from a low-reward region toward a higher-reward region.The figure presents reward trends from the update dynamics under Prompt 2.Source: Figure 6 · p. 28, §E.6 · p. 28

E.7. Comparison of Optimizing Initial-Stage Latent Tokens and Optimizing Middle-Stage Latent Tokens

  • Middle-stage token optimization improves average performance by 5.16 points over CoT but trails initial-stage optimization by 7.58 points.The authors attribute the gap to prefix tokens being independent of the optimized middle-stage tokens, limiting use of improved latent representations.Source: §E.7 · pp. 28–29, Table 13 · p. 29

E.8. Performance vs Fraction Ratio

  • Qwen2.5-14B-Instruct remains above 91% across fraction ratios 0.1–0.8, while smaller models—especially Qwen2.5-1.5B-Instruct—degrade at higher ratios.Mistral-7B-Instruct and LLaMA3.1-8B-Instruct show milder downward trends; Qwen2-7B-Instruct declines across the range.Source: §E.8 · p. 29, Figure 7 · p. 30

F. Experimental Details; F.1. Prompt Designation; F.2. Backbone

  • The experimental appendix provides prompt configurations and reward prompts, and lists the model backbones used in the experiments.Prompt 1 and Prompt 2 are specified in Tables 15–16 and Table 14, respectively; reward prompts appear in Tables 17–20.Source: §F.1 · p. 30, §F.2 · p. 33

F.3. Baselines

  • The baseline set includes CoT, Few-Shot CoT, Best-of-N, Self-Consistency, Self-Rewarding, and supervised fine-tuning.Best-of-N and Self-Consistency use three candidates; the former uses LATENTSEEK’s self-reward mechanism, and Few-Shot CoT results are adopted from Lambert et al. (2024).Source: §F.3 · p. 34

F.4. GSM8K; F.5. MATH-500; F.6. AIME2024

  • Across all backbones and both prompts, inference uses greedy decoding; self-reward evaluates final-answer correctness, problem-comprehension accuracy, numerical-calculation correctness, and clarity, weighted 1:1:2:2 and normalized to [-1, 0].Source: §F.4 · p. 35, §F.5 · p. 37, §F.6 · p. 38

G. Detailed FLOPs Calculation; G.1. Forward Pass FLOPs Estimation; G.2. Total FLOPs for Genius; G.3. Total FLOPs for LATENTSEEK

  • The FLOPs appendix estimates forward-pass costs for LLaMA3.1-8B-Instruct and contrasts Genius training-and-testing computation with LATENTSEEK’s test-time iterations.Genius is described as training on 100,000 instances and testing on 1,319, with backward passes estimated at twice forward-pass cost; LATENTSEEK averages 1.27 GSM8K iterations, each using two forward passes and a partial backward pass through the LM head.Source: Table 27 · p. 39, §G.2 · p. 39, §G.3 · p. 39

G.4. Efficiency Threshold Analysis · H. Latent Space Trajectories

LATENTSEEK remains more efficient than Genius up to about 1.94 × 10^5 inferences, beyond Genius’s 1.00 × 10^5-example training set. Visualized latent trajectories reveal several distinct token-specific search patterns.

G.4. Efficiency Threshold Analysis

  • The efficiency threshold is calculated by equating the total computational costs of Genius and LATENTSEEK across x inference instances.Source: §G.4 · p. 40
  • LATENTSEEK remains more efficient than Genius through approximately 1.94 × 10^5 inference instances, exceeding Genius’s 1.00 × 10^5-example training set.Source: §G.4 · p. 40

H. Latent Space Trajectories

I. Qualitative Analysis and Case Studies · J. Computational Resources · K. The Use of Large Language Models (LLMs)

The qualitative cases suggest that machine reasoning can diverge from human-like, coherent text while still reaching correct answers through small latent-space adjustments. The paper also reports its experimental hardware and limits LLM use to language polishing and figure generation.

I. Qualitative Analysis and Case Studies

  • Generated token patterns and obscure tokens suggest that language-model reasoning can diverge from human reasoning and may be better modeled in latent space.In the analyzed GSM8K outputs, first tokens were often prepositions, second tokens verbs, and third tokens proper names.Source: §I · p. 41
  • The first three generated tokens often follow a preposition–verb–proper-name pattern, alongside occasional nonsensical or obscure tokens.The analysis used Llama3.1-8B-Instruct on GSM8K.Source: §I · p. 41
  • Despite linguistically anomalous expressions, LatentSeek’s decoded outputs can still produce correct solutions.Examples include “total downloads of downloads” and “let need calculate to calculate.”Source: §I · p. 41
  • Small latent-representation adjustments can transform erroneous baseline outputs into correct solutions, despite substantial differences in their textual forms.The cases suggest that outputs appearing far from correct answers in text may remain close to valid solutions in latent space.Source: §I · p. 41

J. Computational Resources

K. The Use of Large Language Models (LLMs)

  • LLM use was limited to polishing manuscript language and generating figures; the authors completed the research and intellectual work without LLM assistance.This included the framework, theoretical foundations, experimental design, and results analysis.Source: K. The Use of Large Language Models… · p. 46

Related papers

Loading 2505.13308v3…