GradCuit: Credit-Assigned Gradient Flow Enables Robust and Interpretable Test-Time Latent Reasoning
2026-08-03 • Machine Learning
Machine LearningComputation and Language
AI summaryⓘ
The authors introduce GradCuit, a method that improves large language models by inserting special adjustable states inside the model during testing, without changing the model itself. This approach lets the model directly link its internal reasoning steps to the final output, making it easier to assign credit and improve reasoning accuracy. They tested GradCuit on various models and tasks, showing it works better and more reliably than previous methods. Additionally, the authors used gradients to understand which parts of the reasoning process are most influenced by these internal adjustments. Overall, their method helps language models adapt their thought process in a more interpretable and stable way.
large language modelslatent statesTransformercausal self-attentiongradient optimizationreasoning trajectorychain-of-thought promptingreward-weighted gradientstest-time adaptationinterpretability
Authors
Zhaoxin Yu, Qi Shen, Hengli Li, Zhaowei Zhang, Song-Chun Zhu, Chi Zhang, Zilong Zheng
Abstract
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.