kv-cache-steering-reasoning

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:

1. Steering Vector Extraction: Compute mean difference between KV cache from two prompts—one demonstrating desired reasoning, another without it
2. Cache Modification: Add the steering vector to cached K and V representations using scalar coefficients
3. Unmodified Generation: Proceed with standard autoregressive generation using the modified cache
4. 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