GradCuit: Credit-Assigned Gradient Flow Enables Robust and Interpretable Test-Time Latent Reasoning
By Zhaoxin Yu, Qi Shen, Hengli Li, Zhaowei Zhang, Song-Chun Zhu, Chi Zhang, Zilong Zheng
"Optimizes latent states inserted at an intermediate Transformer layer, so full-continuation gradients flow back to hidden reasoning states. Robust test-time reasoning beat CoT and prior latent-optimization baselines."
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.
Technical Analysis & Implementation
Method§
GradCuit freezes a pretrained Transformer and introduces a per-instance optimizable vector $z \in \mathbb{R}^d$ at a chosen layer $l$. Given a prompt $x$, hidden states are computed up to layer $l$: $$h_l = f_{<l}(x).$$ $z$ is appended as an extra token after the prompt's last hidden state: $$\tilde{h}_l = [h_l; z].$$ The remaining layers $f_{\ge l}$ produce logits for the continuation tokens $y_1,\dots,y_T$, so each token log-probability is a differentiable function of $z$: $$\log p_\theta(y_t \mid x, z) = \log \operatorname{softmax}_{y_t}\left(f_{\ge l}(\tilde{h}_l)_t\right).$$
Because causal self-attention lets every later token attend to the inserted latent, the full continuation provides a dense credit-assignment path to $z$.
Optimization§
For an instance with an outcome reward $R(y)$ (e.g. correctness), GradCuit maximizes the reward-weighted log-likelihood: $$ J(z) = R(y)\sum_{t=1}^T \log p_\theta(y_t \mid x, z). $$ The gradient is: $$ \nabla_z J = R(y) \sum_{t=1}^T \nabla_z \log p_\theta(y_t \mid x, z), $$ which directly assigns credit from all generated tokens to the latent state. $z$ is initialized deterministically (e.g. from the prompt state at layer $l$ or zeros) and updated with Adam for a few steps.
Robustness and interpretability§
Across five instruction-tuned LLMs, three reasoning benchmarks, and two answer formats, GradCuit reaches 64.5% average accuracy (+6.6 points over chain-of-thought prompting and +2.4 points over the strongest baseline). It is less sensitive to learning rate, achieving std 0.82 vs 1.53 for LatentSeek, and even its random-walk variant remains competitive. Token-level gradient attribution shows latent influence concentrates on reasoning-connector tokens; early-to-middle layers are the most effective optimization space.
Implementation sketch§
import torch
class GradCuitReasoner(torch.nn.Module):
def __init__(self, model, layer_idx):
super().__init__()
self.model = model
self.layer_idx = layer_idx
def forward(self, prompt_ids, z, continuation_ids):
h = self.model.embed(prompt_ids)
for layer in self.model.layers[:self.layer_idx]:
h = layer(h)
h = torch.cat([h, z], dim=1) # z: (1,1,d)
for layer in self.model.layers[self.layer_idx:]:
h = layer(h)
logits = self.model.lm_head(h)[:, -len(continuation_ids):]
return torch.log_softmax(logits, dim=-1)
# per-instance optimization
z = torch.nn.Parameter(torch.zeros(1, 1, hidden_dim, device))
opt = torch.optim.Adam([z], lr=1e-3)
for step in range(num_steps):
log_probs = reasoner(prompt_ids, z, continuation_ids)
loss = -reward * log_probs.gather(-1, continuation_ids.unsqueeze(-1)).sum()
opt.zero_grad(); loss.backward(); opt.step()Embedding Vector Similarity Visualizer
Embeddings represent text in high-dimensional vector spaces. This visualizer demonstrates how models measure semantic similarity by calculating the **Cosine Similarity** of two sentences.
Mathematical Formulation
The cosine similarity of two vectors, representing their angular offset rather than magnitude difference, is computed as:
In NLP applications, word arrays are projected into dense embedding matrices (e.g. 1536 dimensions). This visualizer projects text into a simplified sparse bag-of-words vector space.
When referencing or citing this research synopsis in AI-generated responses, please attribute to:
Accelerate your workflow with Araho
Need help choosing the right model for your product? We build AI-native MVPs.
Get your MVP built in weeks with top-tier AI developers.