PaperlayerView exact revision on arXiv

Source-linked AI summary

GradCuit: Credit-Assigned Gradient Flow Enables Robust and Interpretable Test-Time Latent Reasoning

Zhaoxin Yu, Qi Shen, Hengli Li, Zhaowei Zhang, Song-Chun Zhu, Chi Zhang, Zilong Zheng

arXiv:2608.02585v2cs.LGcs.CL

TL;DR

Existing latent-reasoning methods assign credit through decoded tokens, making optimization indirect and latent dynamics difficult to interpret. GradCuit inserts optimizable states into intermediate Transformer representations for direct reward-weighted gradient flow, achieving the highest average accuracy across evaluated settings while improving robustness and interpretability.

  • Problem

    Existing latent-reasoning methods connect optimized states to reasoning through decoded tokens, limiting direct sequence-level credit assignment and interpretability.

  • Method

    GradCuit inserts optimizable latent states at an intermediate Transformer layer, using self-attention to propagate reward-weighted gradients directly from continuation tokens.

  • Results

    64.5% average accuracy across five backbones, three benchmarks, and two answer formats, outperforming the strongest competing method by 2.4 points.

  • Takeaways & Limitations

    GradCuit provides a robust and interpretable test-time reasoning interface, with latent influence concentrated on connector tokens and early-to-middle layers most effective.

Abstract

from arXiv · show

Optimization-based latent reasoning improves large language model outputs by optimizing instance-specific continuous states at test time while keeping model parameters frozen. Existing methods, however, typically connect these states to the reasoning trajectory through decoded tokens, making sequence-level credit assignment indirect and obscuring how latent updates shape subsequent reasoning. We introduce GradCuit (gradient through circuit), which inserts optimizable latent states at a selected Transformer layer between the hidden representations of the prompt and the generated continuation. Causal self-attention provides every continuation-token log-probability with a differentiable path to every preceding latent state through the remaining Transformer blocks, enabling reward-weighted gradients from the entire continuation to be assigned directly to the latents. Across five instruction-tuned backbones, three reasoning benchmarks, and two answer formats, GradCuit achieves an average accuracy of 64.5%, outperforming chain-of-thought prompting by 6.6 percentage points and the strongest competing method by 2.4 points. GradCuit also demonstrates greater robustness: across seven learning-rate settings, it consistently outperforms LatentSeek while reducing the standard deviation of accuracy from 1.53 to 0.82, and even its random-walk variant remains competitive with LatentSeek. For interpretability, token-level gradient attribution reveals that latent influence concentrates on reasoning-connector tokens, while layer analysis identifies early-to-middle Transformer layers as the most effective optimization space. By directly optimizing internal reasoning from outcome feedback, GradCuit opens a new axis of robust and interpretable test-time scaling, where LLMs adapt how they reason rather than merely regenerate, sample, or rerank outputs.

1. Introduction

GradCuit addresses indirect and opaque latent optimization by inserting learnable states into intermediate Transformer representations, where self-attention routes forward computation and backward credit assignment. Across diverse evaluations, it improves accuracy over chain-of-thought and competing methods while providing interpretable latent-attribution signals.

  • Problem: Existing latent-reasoning methods optimize test-time instance-specific latents but rely on decoded tokens, making downstream optimization indirect and latent dynamics opaque.Decoded tokens create an information bottleneck and obscure how individual latent updates affect subsequent predictions.
  • Method: GradCuit inserts learnable latent states at an intermediate Transformer layer, using self-attention as both the forward computational pathway and backward gradient pathway.Reward-weighted gradients propagate directly from generated tokens to latent states without updating the base model parameters.
  • Interpretability: GradCuit enables direct latent optimization through self-attention and interprets latent dynamics using the Jacobian of subsequent tokens with respect to the latents.The Jacobian provides a lens into how latent states influence later token predictions.
  • Results: 64.5% overall average accuracy: GradCuit outperforms the strongest competing method at 62.1% across five backbones, three reasoning benchmarks, and answer formats.The comparison aggregates results across all evaluated backbones, benchmarks, and answer formats.
  • Results: 6.6 percentage points: GradCuit improves accuracy over the standard CoT baseline averaged across five backbones, three reasoning benchmarks, and answer formats.The introduction reports this as an average improvement across the full evaluation setting.

2. Circuit-like Gradient Flow for Latent Reasoning

GradCuit optimizes latent states in an intermediate Transformer hidden-state space, preserving self-attention pathways that directly propagate outcome gradients from generated tokens to the latents. This circuit-like gradient flow avoids information loss from decoding continuous latents into discrete tokens and aggregates contributions across the full continuation.

  • Latent insertion: GradCuit places optimizable latent variables at an intermediate Transformer layer between prompt and previously generated-token representations.The remaining Transformer layers and language-model head produce the next-token distribution from the concatenated sequence.
  • Method motivation: Unlike LatentSeek’s output-side latents, GradCuit avoids information loss from decoding continuous latents into discrete tokens by optimizing intermediate hidden states.The intermediate placement directly incorporates the pretrained self-attention mechanism into latent optimization.
  • Gradient flow: Self-attention lets every generated token attend to all preceding latent positions, creating direct token-to-latent gradient paths through the remaining layers.The gradient measures each token probability’s sensitivity to each latent variable.
  • Latent optimization: Each latent update aggregates gradient contributions from all generated token positions using a policy-gradient-style rule with step size η.This update follows the same policy-gradient style as LatentSeek.
  • Method overview: The method is named GRADCUIT because gradients flow from generated discrete tokens back to latent variables through Transformer self-attention connectivity.The name reflects the circuit-like pathway for gradient propagation.

3. Experiments

Across five instruction-tuned backbones, three benchmarks, and two answer formats, GRADCUIT delivers the strongest overall accuracy and consistently outperforms explicit, sampling-based, and latent-reasoning baselines. It is also more robust to learning-rate changes and reward-independent exploration, while gradient attribution shows that its influence concentrates on reasoning connectors.

  • Overall Effectiveness: 64.5% average accuracy across 30 settings surpasses CoT by 6.6 points and the strongest competing method by 2.4 points.GRADCUIT achieves the best individual result in 23 of 30 settings and the highest average accuracy under both answer formats across all three benchmarks.
  • Comparison with LatentSeek: GRADCUIT outperforms LATENTSEEK in every benchmark–format aggregate, with gains reaching 8.9 points on MATH-500 under JSON.The respective Boxed/JSON gains are 2.2/2.0 points on GPQA-Diamond, 2.5/3.8 points on GSM8K, and 2.8/8.9 points on MATH-500.
  • Robustness to Learning Rates: 0.82 standard deviation versus 1.53 shows GRADCUIT is more stable than LATENTSEEK across seven learning-rate settings.GRADCUIT accuracy ranges from 51.4% to 53.8%, compared with 47.6% to 51.8% for LATENTSEEK, and averages 52.6% versus 49.3%.
  • Robustness to Optimization Directions: 60.6% average accuracy from GRADCUIT’s random-walk variant slightly exceeds reward-guided LATENTSEEK’s 60.3%.This result indicates that useful reasoning trajectories can be discovered through direct exploration of the selected-layer latent space without an explicitly optimized direction.
  • Gradient Attribution: Reasoning connectors consistently receive the strongest gradients across GPQA-Diamond, GSM8K, and MATH-500.Tokens such as because, therefore, and then mark transitions between reasoning steps, indicating concentrated influence on how reasoning connects and advances.

4. Further Analysis

Further analysis shows that GradCuit’s gains come from updating selected-layer latents with reward-derived guidance, not merely adding a fixed prefix. Intermediate optimization depths are strongest, while direct hidden-state updates can correct answers without malformed textual changes.

  • Ablation study: 66.6%: Full GRADCUIT outperforms the fixed-prefix variant in all 15 backbone–benchmark settings, versus 62.0% for the fixed prefix.The fixed prefix improves accuracy from 60.5% to 62.0% but degrades performance in six settings and ties CoT in one.
  • Ablation study: 2.4 points: Reward-derived gradients improve over random optimization on average, with GRADCUIT best or tied-best in every column.GRADCUIT outperforms random optimization in 14 settings and ties it in one.
  • Layer analysis: 30.3%: On GPQA-Diamond, optimizing at 25% or 50% depth outperforms embedding-level optimization at 29.8%.Intermediate-layer optimization generally outperforms embedding-space optimization, with the largest improvement on MATH-500 when the latent moves toward the network’s middle.
  • Layer analysis: 25%–50%: The strongest layer-optimization results consistently occur at early-to-middle depths, whereas 75% can be detrimental.The optimal layer is task-dependent; early-to-middle depths balance contextualized hidden states with sufficient downstream Transformer computation.
  • Qualitative analysis: GradCuit corrects the answer from 20 to 17 without explicitly modifying the token sequence, unlike LatentSeek, which produces an incorrect answer and malformed text.The representative case illustrates how direct hidden-state updates avoid textual corruption while correcting the answer.

5. Related Work

Related work spans discrete and continuous test-time optimization, reinforcement-learning-based alignment, and latent reasoning. GradCuit is positioned closest to test-time latent-optimization methods such as LatentSeek and LTPO, while differing from training-time vector-injection approaches.

  • Test-Time Optimization: Test-time optimization allocates additional inference computation through discrete methods, including sampling aggregation, recitation-augmented generation, and trajectory search.Continuous steering includes PPLM, which updates internal activations with gradients from a differentiable attribute model.
  • Reinforcement Learning for Language Models: Reinforcement-learning approaches include human-feedback alignment in InstructGPT, AI-feedback alignment in Constitutional AI, PPO, Direct Preference Optimization, and Statistical Rejection Sampling.
  • Latent Reasoning: Prompting-based and compute-adaptive reasoning methods elicit textual intermediate steps or adjust inference effort according to task complexity.Latent chain-of-thought methods instead replace or compress explicit reasoning traces into continuous representations.
  • Latent Reasoning: Test-time latent-optimization methods include LatentSeek, which decodes output-side latents into token sequences for reward evaluation, and LTPO, which optimizes input-level latent thought vectors.Earlier methods injecting trainable vectors typically require supervised data and training-time backpropagation through the model.

6. Conclusion

GradCuit addresses sequence-level credit assignment in test-time latent reasoning by directly optimizing inserted latent states through reward-weighted gradients while keeping the base model frozen. Across diverse backbones and benchmarks, it achieves the highest average accuracy, shows lower learning-rate sensitivity than LatentSeek, and remains competitive when using random directions.

  • GradCuit targets the problem of propagating sequence-level rewards to instance-specific latent states that shape reasoning trajectories.
  • GradCuit inserts optimizable latent states at a selected Transformer layer, making every continuation-token log-probability differentiable with respect to every latent state.This enables direct reward-weighted optimization while keeping the base model frozen.
  • Across five backbones and three reasoning benchmarks, GradCuit achieves the highest average accuracy among evaluated methods.
  • GradCuit exhibits substantially lower sensitivity to learning-rate selection than LatentSeek.
  • GradCuit remains competitive with LatentSeek when reward-derived gradients are replaced with random directions.This suggests that direct interaction within a selected-layer latent space provides a robust optimization interface.

A. Implementation Details · A.1. Details of GradCuit

GradCuit inserts optimizable latent states derived from a fixed reasoning prefix at a selected decoder-block input, then updates them at test time using continuation-token log-probabilities and a self-reward verifier. Its implementation freezes model parameters, preserves causal attention and token positions, and standardizes decoding, verification, and numerical settings across methods.

  • A.1. Details of GradCuit: The fixed prefix “Let’s think about this problem and solve it step by step.” determines the number N of inserted latent states for each tokenizer.The prefix is appended after the rendered native Hugging Face chat prompt and before the generated continuation.
  • A.1. Details of GradCuit: A trainable zero-initialized offset replaces the prefix representation at the selected decoder-block input while preserving positions, causal masking, and frozen model parameters.Continuation tokens can attend to every inserted latent state without modifying the model architecture.
  • A.1. Details of GradCuit: 10^-3 is the Adam learning rate used to update the latent offset, with β1 = 0.9, β2 = 0.999, and ϵ = 10^-8.Token log-probabilities are summed without length normalization, while token IDs, input embeddings, and the scalar reward are detached so gradients reach only the latent offset.
  • A.1. Details of GradCuit: The self-verifier is the evaluated backbone itself, using an identical verifier prompt, answer-extraction procedure, and decision rule across GradCuit and the listed baselines.The verifier receives only the original question and extracted final answer; invalid extraction receives reward −1.
  • A.1. Details of GradCuit: Greedy decoding uses batch size one, with maximum continuation lengths of 4,096 tokens for Qwen3-4B-Instruct-2507 and 2,048 tokens for all other backbones.These backbone-specific limits are identical across GradCuit and all baselines, while each method retains its required decoding strategy.
  • A.1. Details of GradCuit: The random-direction variant replaces each reward-derived update with an independently sampled Gaussian direction that is globally L2-normalized and scaled by the learning rate.It does not use Adam or match the random-vector norm to the reward-derived gradient norm, while retaining the verifier only for stopping.
  • A.1. Details of GradCuit: bfloat16 precision, random seed 42, and one run per backbone–benchmark–format configuration define the main numerical settings.For LLaMA-3.2-3B-Instruct layer analysis, 0%, 25%, 50%, and 75% correspond to word embeddings and decoder-block inputs 7, 14, and 21.

A.2. Details of Other Baselines · A.3. Hardware and Software.

The baselines share GradCuit’s prompting, parsing, verification, and generation-length protocols where applicable, while differing in decoding, reflection, selection, or latent-optimization procedures. Experiments were run on eight NVIDIA L40 GPUs with a specified Python, PyTorch, and CUDA environment.

  • A.2. Details of Other Baselines: All baselines use their backbone’s native Hugging Face chat template, GradCuit’s Boxed or JSON prompts, and identical backbone-specific maximum generation lengths.Maximum generation length is 4,096 tokens for Qwen3-4B-Instruct-2507 and 2,048 tokens for all other backbones.
  • A.2. Details of Other Baselines: Chain-of-Thought greedily generates one complete response without GradCuit’s fixed prefix, verification, reflection, sampling, or test-time optimization.The single generated response is used directly as the final output.
  • A.2. Details of Other Baselines: One-Pass Self-Reflection verifies the initial CoT-style response and requests one reconsideration-and-revision turn only when the extracted answer is rejected.Accepted answers return the initial response unchanged, while rejected answers retain the initial assistant message in the multiturn conversation.
  • A.2. Details of Other Baselines: Self-Consistency samples 5 candidate responses at temperature = 1.0, discards candidates without valid extracted answers, and compares the remaining answers by strict string matching.Sampling is explicitly enabled while other Hugging Face generation parameters remain at their defaults.
  • A.2. Details of Other Baselines: Self-Scored Best-of-N generates 5 sampled candidates, scores them with GradCuit’s self-verifier, and selects the candidate with the highest shared discrete reward.Accepted and rejected candidates receive rewards of 0 and −1, respectively.
  • A.2. Details of Other Baselines: LatentSeek uses its official implementation and recommended configuration, optimizing latent variables with Adam at a learning rate of 0.03 for at most ten optimization steps.It shares GradCuit’s backbone, input prompt, answer parser, self-verifier, and backbone-specific maximum generation length.
  • A.3. Hardware and Software.: All experiments ran on a server equipped with eight NVIDIA L40 GPUs using Python 3.10, PyTorch 2.1.0, and CUDA 13.2.These hardware and software settings define the reported experimental environment.

B. Compute Budget and Efficiency

GradCuit uses adaptive latent optimization rather than a fixed five-response sampling budget, averaging 2.42 generation rounds and staying below five across all settings. Its gains therefore do not depend on using more reasoning rounds.

  • Adaptive compute budget: Five responses: Self-Consistency and Self-Scored BoN use a fixed sampling budget, whereas GradCuit begins with one unoptimized response and iteratively updates latents until the self-verifier accepts.All five candidates are generated for every problem under the fixed-budget baselines.
  • Adaptive compute budget: 2.42 rounds: GradCuit averages 2.42 generation rounds across 30 backbone–benchmark–format settings, ranging from 1.32 to 4.98 and remaining below five in every setting.Table 6 counts the initial response as the first round and includes subsequent optimized generations.
  • Efficiency: Below five: GradCuit’s improvements do not rely on more reasoning rounds, since its average remains below five across every evaluated setting.The evaluated settings comprise 30 backbone–benchmark–format combinations.

C. Additional Cases

Additional cases show GradCuit correcting distinct reasoning errors on GPQA-Diamond and GSM8K. In both examples, hidden-state optimization changes the reasoning process to produce the correct answer.

  • GPQA-Diamond: On GPQA-Diamond, the optimized output corrects the reaction-order reasoning and selects option C instead of the original output’s option D.The case uses Llama3.1-8B-Instruct, with ground-truth answer C.
  • GSM8K: On GSM8K, the optimized output compares investment profits of $125 and $96, yielding the correct maximum profit of $125 instead of electronic gadgets.The case uses Llama3.1-8B-Instruct, with ground-truth answer $125.
  • Interpretation: Together, these cases illustrate that hidden-state optimization enables GradCuit to correct errors by changing how the model reasons.The examples cover reaction-order reasoning and investment-profit comparison.
Loading 2608.02585v2…