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

2026-08-03

AI Summary

This work addresses the limitations of existing optimization-based test-time latent variable inference methods, which rely on generated tokens for indirect credit assignment, resulting in ambiguous inference trajectories and poor interpretability. The authors propose inserting optimizable latent variables into selected Transformer layers and leveraging causal self-attention to establish an end-to-end differentiable path from the full output sequence back to these latents, thereby enabling direct credit assignment via reward gradients for the first time. This approach substantially enhances inference robustness and interpretability, revealing that latent variables exert concentrated influence on reasoning connectives and identifying optimal layers for insertion. Evaluated across five backbone models and three reasoning benchmarks, the method achieves an average accuracy of 64.5%, outperforming chain-of-thought prompting by 6.6 percentage points, and demonstrates superior stability with a standard deviation of only 0.82 across seven learning rates, significantly surpassing LatentSeek.

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.

Problem

• test-time reasoning
• latent reasoning
• credit assignment
• interpretability
• large language models

Innovation

• test-time optimization
• latent reasoning
• credit assignment
• gradient-based attribution
• interpretable LLMs