GradCuit: How to Make LLMs Reason Better at Test Time Without Changing a Single Weight

TL;DR

GradCuit (arXiv:2608.02585) inserts optimizable latent vectors at an intermediate Transformer layer and uses causal self-attention as a differentiable "circuit" to flow reward-weighted gradients directly to those latents at test time — no parameter updates, no token resampling, just smarter internal reasoning. Result: 64.5% average accuracy across 5 models and 3 benchmarks, beating Chain-of-Thought by 6.6 pp and the previous best latent-space method (LatentSeek) by 2.4 pp.

The Problem

Test-time scaling has become one of the hottest topics in LLM research. The idea is simple: spend more compute at inference to get better outputs. Chain-of-Thought, Best-of-N sampling, and self-consistency are classic examples. More recently, researchers have explored optimizing in latent space — directly adjusting the model's hidden representations without changing its weights.

LatentSeek (2505.13308) was a promising step: it uses policy gradients to iteratively update latent representations guided by self-generated rewards. But there's a fundamental flaw shared by all existing latent reasoning methods.

The credit assignment problem: Existing approaches connect latent states to the reasoning trajectory through decoded tokens. Decoded tokens are non-differentiable (argmax breaks the gradient). So gradient signals are indirect, noisy, and highly sensitive to learning rate — sometimes standard deviation of accuracy across learning rate settings reaches 1.53 for LatentSeek. You'd need to tune hyperparameters carefully just to get a stable result.

How It Works

GradCuit's insight is architectural. Instead of inserting latent states near the output, it places them at an intermediate Transformer layer (25–50% depth works best) — between the prompt hidden representations and the generated continuation.

Here's why this matters: Transformer's causal self-attention ensures that every generated token attends to all preceding positions, including those latent vectors. This creates a fully differentiable path from every continuation token's log-probability back to every latent variable through the remaining Transformer blocks. No decoded token bottleneck. No broken gradient.

The objective is reward-weighted policy gradient:

$$J(z) = \mathbb{E}{y \sim \pi\theta(\cdot \mid x, z)}\bigl[R(y)\bigr]$$

Gradient via REINFORCE:

$$\nabla_z J(z) = \mathbb{E}{y \sim \pi\theta(\cdot \mid x, z)}\bigl[R(y) \cdot \nabla_z \log \pi_\theta(y \mid x, z)\bigr]$$

Because of the intermediate insertion, each term $\nabla_z \log p_\theta(y_t \mid y_{<t}, x, z)$ has a concrete differentiable path through causal attention layers $l$ through $L$. The gradient for each latent variable aggregates contributions from all generated token positions — true sequence-level credit assignment.

The update rule is straightforward gradient ascent:

$$z^{(k+1)} \leftarrow z^{(k)} + \alpha \cdot \widehat{\nabla}_z J(z^{(k)})$$

The model weights $\theta$ stay completely frozen throughout.

Show Me The Code

Here's a simplified PyTorch implementation of GradCuit's core mechanism:

import torch

import torch.nn as nn

from transformers import AutoModelForCausalLM, AutoTokenizer





class GradCuit:

    def __init__(self, model_name: str, target_layer: int, latent_len: int = 8):

        self.model = AutoModelForCausalLM.from_pretrained(model_name)

        self.tokenizer = AutoTokenizer.from_pretrained(model_name)

        self.target_layer = target_layer

        hidden_size = self.model.config.hidden_size



        # Optimizable latent vectors — this is what we update at test time

        self.latent_z = nn.Parameter(torch.zeros(1, latent_len, hidden_size))

        self.prompt_len = 0



    def _hook_fn(self, module, input, output):

        """Forward hook: inject latents at the intermediate layer."""

        batch_size = output[0].shape[0]

        latent = self.latent_z.expand(batch_size, -1, -1)

        # Concatenate: [prompt_repr | latent_z | continuation_repr]

        modified = torch.cat([

            output[0][:, :self.prompt_len, :],

            latent,

            output[0][:, self.prompt_len:, :]

        ], dim=1)

        return (modified,) + output[1:]



    def optimize(self, prompt: str, reward_fn, n_steps=10, lr=0.01, n_samples=4):

        inputs = self.tokenizer(prompt, return_tensors="pt")

        self.prompt_len = inputs["input_ids"].shape[1]

        optimizer = torch.optim.Adam([self.latent_z], lr=lr)



        # Register hook at the chosen intermediate layer

        hook = self.model.model.layers[self.target_layer].register_forward_hook(

            self._hook_fn

        )



        for step in range(n_steps):

            optimizer.zero_grad()

            total_loss = torch.tensor(0.0, requires_grad=True)



            for _ in range(n_samples):

                # Sample a continuation

                with torch.no_grad():

                    out_ids = self.model.generate(

                        **inputs, max_new_tokens=256, do_sample=True, temperature=0.8

                    )

                text = self.tokenizer.decode(

                    out_ids[0][self.prompt_len:], skip_special_tokens=True

                )



                # Task-specific reward

                reward = reward_fn(text)



                # Differentiable forward to get log-probs

                with torch.enable_grad():

                    logits = self.model(**inputs).logits

                    log_probs = torch.nn.functional.log_softmax(logits, dim=-1)

                    token_log_probs = log_probs[0, :, :].sum()

                    # REINFORCE: -R * log π

                    loss = -reward * token_log_probs

                    total_loss = total_loss + loss



            (total_loss / n_samples).backward()

            optimizer.step()



        hook.remove()

        return self.latent_z.detach()





# Usage example: math reasoning

def exact_match_reward(generated: str, target: str = "42") -> float:

    return 1.0 if target in generated else 0.0



# For Llama-3.1-8B (32 layers), target layer ~35% depth

model = "meta-llama/Llama-3.1-8B-Instruct"

gc = GradCuit(model, target_layer=11, latent_len=8)



optimized_z = gc.optimize(

    prompt="Solve step by step: What is 6 times 7?",

    reward_fn=exact_match_reward,

    n_steps=10,

    lr=0.01,

    n_samples=4,

)

Note: The paper also implements a random-walk variant that skips gradient computation entirely and explores latent space stochastically — and it still beats guided LatentSeek (60.6% vs 60.3%).

Benchmark Results

Evaluated across 5 instruction-tuned backbone models, 3 reasoning benchmarks (GPQA-Diamond, GSM8K, MATH-500), and 2 answer formats:

Benchmark-specific gains over LatentSeek:

• GPQA-Diamond: +2.2 pp


• GSM8K: +2.5 to 3.8 pp


• MATH-500: +2.8 to 8.9 pp (biggest gains)

Robustness across 7 learning rate settings:

• LatentSeek accuracy std dev: 1.53


• GradCuit accuracy std dev: 0.82 (46% reduction)

Interpretability finding: Token-level gradient attribution shows that latent influence concentrates on reasoning-connector tokens ("because", "therefore", "so") — meaning GradCuit primarily optimizes how the model transitions between reasoning steps, not just what tokens it generates.

Gotchas & Limitations

Compute cost: Multiple forward/backward passes per query. Practical for batch offline inference; may be too slow for real-time applications without optimization.

Reward function design: Works great for tasks with clear verifiable rewards (math, code execution). Open-ended generation requires a learned reward model, adding complexity.

Architecture assumption: Relies on Transformer causal self-attention. Hybrid architectures (Mamba, RWKV) would need adaptation.

Layer selection: Optimal target layer (25-50% depth) was found empirically. Automatic layer selection is not yet explored — you'll need a small grid search.

Memory overhead: Latent vectors add VRAM usage proportional to latent_len × hidden_size, but this is typically negligible.

Try It Today

Paper: — the prior SOTA this beats