| name | kv-cache-steering-reasoning |
| title | KV Cache Steering for Inducing Reasoning in Small Language Models |
| version | 0.0.2 |
| engine | skillxiv-v0.0.2-claude-opus-4.6 |
| license | MIT |
| url | https://arxiv.org/abs/2507.08799 |
| keywords | ["Activation Steering","Inference-Time Control","KV Cache","Reasoning Patterns"] |
| description | Guide frozen language models toward multi-step reasoning by modifying cached key-value representations after the prefilling stage. Extract steering vectors from contrastive prompt pairs and apply them to KV cache with scalar coefficients. Improves reasoning on GSM8K, ARC, CommonsenseQA while adding only 10ms overhead per token. |
KV Cache Steering: Inference-Time Reasoning Induction Without Model Changes
Language models can reason through problems, but they don't always choose to—they often generate superficial answers. Traditional activation steering applies per-token interventions throughout generation, causing effects to compound and amplify. KV Cache Steering sidesteps this by modifying the key-value cache once at the prefilling stage, inserting a reasoning signal that propagates cleanly through subsequent generation without cascading amplification. This single intervention at the right architectural layer dramatically improves reasoning (5-15% accuracy gains on GSM8K) while adding negligible overhead.
The key insight is that the KV cache is the information bottleneck in transformers. By shifting cached representations toward a reasoning-aligned direction before decoding begins, you guide the entire generation trajectory toward step-by-step reasoning without per-token interference or model weight changes.
Core Concept
KV Cache Steering operates through four steps:
- Steering Vector Extraction: Compute mean difference between KV cache from two prompts—one demonstrating desired reasoning, another without it
- Cache Modification: Add the steering vector to cached K and V representations using scalar coefficients
- Unmodified Generation: Proceed with standard autoregressive generation using the modified cache
- Single Intervention Point: Unlike per-token steering, this modifies the cache once, allowing clean information flow
The steering vector captures the latent direction toward reasoning; applying it shifts the model's internal state toward step-by-step problem-solving before generation even begins.
Architecture Overview
- Contrastive Prompt Pair: Positive example (with explicit reasoning) and negative example (direct answer)
- Frozen Base Model: Standard LLM (no weights change)
- KV Cache Extractor: Captures key and value representations after prefilling
- Steering Vector Computation: Mean-of-differences across all layers and positions
- Cache Modifier: Linear addition of steering vectors to K and V with learned scalar coefficients (per-layer)
- Standard Decoder: Unchanged generation using modified cache
- Hyperparameter Interface: Scale coefficient α controlling steering strength
Implementation
The following demonstrates steering vector extraction and KV cache modification:
torch
torch.nn nn
torch.nn.functional F
typing , ,
(nn.Module):
():
().__init__()
.model = model
.num_layers = num_layers
.cache_hooks = []
._register_hooks()
():
.positive_cache = {}
.negative_cache = {}
():
():
(output, ):
.positive_cache[layer_idx] = output.past_key_values
hook
():
():
(output, ):
.negative_cache[layer_idx] = output.past_key_values
hook
layer_idx (.num_layers):
layer = .model.transformer.h[layer_idx]
positive_handle = layer.register_forward_hook(positive_hook(layer_idx))
negative_handle = layer.register_forward_hook(negative_hook(layer_idx))
.cache_hooks.append((positive_handle, negative_handle))
() -> [, [torch.Tensor, torch.Tensor]]:
positive_ids = tokenizer.encode(positive_prompt, max_length=max_length, return_tensors=)
negative_ids = tokenizer.encode(negative_prompt, max_length=max_length, return_tensors=)
torch.no_grad():
_ = .model(positive_ids)
positive_cache = {k: v.clone() k, v .positive_cache.items()}
torch.no_grad():
_ = .model(negative_ids)
negative_cache = {k: v.clone() k, v .negative_cache.items()}
steering_vectors = {}
layer_idx (.num_layers):
layer_idx positive_cache layer_idx negative_cache:
pos_k, pos_v = positive_cache[layer_idx]
neg_k, neg_v = negative_cache[layer_idx]
k_direction = (pos_k - neg_k).mean(dim=(, , ))
v_direction = (pos_v - neg_v).mean(dim=(, , ))
steering_vectors[layer_idx] = (k_direction, v_direction)
steering_vectors
(nn.Module):
():
().__init__()
.model = model
.num_layers = num_layers
.hidden_dim = hidden_dim
.steering_coefficients = nn.Parameter(torch.ones(num_layers) * )
() -> [torch.Tensor, ]:
batch_size, seq_len = input_ids.shape
torch.no_grad():
outputs = .model(input_ids, output_hidden_states=, return_dict=)
initial_cache = outputs.past_key_values
modified_cache = []
layer_idx (.num_layers):
layer_idx steering_vectors initial_cache :
k, v = initial_cache[layer_idx]
k_direction, v_direction = steering_vectors[layer_idx]
steering_strength = .steering_coefficients[layer_idx] * alpha
k_modified = k + steering_strength * k_direction.to(k.device).unsqueeze().unsqueeze()
v_modified = v + steering_strength * v_direction.to(v.device).unsqueeze().unsqueeze()
modified_cache.append((k_modified, v_modified))
:
modified_cache.append(initial_cache[layer_idx] initial_cache )
generated_tokens = []
current_cache = (modified_cache)
step ():
torch.no_grad():
outputs = .model(
input_ids[:, -:],
past_key_values=current_cache,
return_dict=
)
logits = outputs.logits[:, -, :]
next_token = logits.argmax(dim=-, keepdim=)
generated_tokens.append(next_token)
current_cache = outputs.past_key_values
next_token.item() == :
generated_sequence = torch.cat(generated_tokens, dim=)
info = {
: ,
: [.steering_coefficients[i].item() i (.num_layers)],
: generated_sequence.shape[]
}
generated_sequence, info
(nn.Module):
():
().__init__()
.cache_modifier = cache_modifier
.criterion = nn.CrossEntropyLoss()
() -> :
optimizer.zero_grad()
input_ids = batch[]
target_ids = batch[]
generated, info = .cache_modifier(input_ids, steering_vectors, alpha=)
loss = .criterion(
generated[:, , :].unsqueeze(),
target_ids[:, ].unsqueeze()
)
loss.backward()
torch.nn.utils.clip_grad_norm_(.cache_modifier.steering_coefficients, )
optimizer.step()
loss.item()
:
():
.alpha = initial_alpha
():
accuracy < threshold:
.alpha = (.alpha * , )
:
.alpha = (.alpha * , )
.alpha
() -> :
extractor = SteeringVectorExtractor(model)
steering_vectors = extractor.extract_steering_vectors(reasoning_prompt, prompt, tokenizer)
modifier = KVCacheModifier(model)
input_ids = tokenizer.encode(prompt, return_tensors=)
generated, info = modifier(input_ids, steering_vectors, alpha=alpha)
generated_text = tokenizer.decode(generated[])
generated_text, info
