adu2021 Resa: Transparent Reasoning Models via SAEs
Core Concept
Resa demonstrates that reasoning abilities are learnable, transferable features that can be extracted via sparse autoencoders (SAE) without expensive reinforcement learning. The two-stage SAE-Tuning procedure trains an autoencoder on source model activations, then inserts it into a target model with low-rank adapters. This achieves RL-equivalent reasoning performance at approximately $1 cost and 20-minute training time, compared to months and $2000+ for traditional RL approaches.
Architecture Overview
- Two-Stage SAE-Tuning: Stage I trains SAE on source model activations; Stage II inserts frozen SAE into target model with LoRA for implicit reasoning pattern transfer
- CoT-Free Training Data: Uses only verified question-answer pairs (no intermediate reasoning traces), reducing data requirements
- Reasoning as Portable Adapter: Extracted reasoning features function as modular "reasoning adapters" transferable across model families without retraining
- Transparent Feature Extraction: Prompt-only method identifies latent reasoning features; layer-wise distribution correlates with reasoning performance
- Massive Cost Reduction: $1 per model vs $2000 for RL; 20 minutes training vs months for full RL pipeline
Implementation
Step 1: Sparse Autoencoder Training
import torch
import torch.nn as nn
class SparseAutoencoder(nn.Module):
"""
Trains on source model activations to extract reasoning features.
Sparse representation: high dimensionality (65k features) but low k activation.
"""
def __init__(self, input_dim, num_features=65536, k=32):
super().__init__()
self.input_dim = input_dim
self.num_features = num_features
self.k = k # top-k sparsity
# Encoder: input -> sparse features
self.encoder = nn.Linear(input_dim, num_features)
# Decoder: sparse features -> reconstruction
self.decoder = nn.Linear(num_features, input_dim)
# Initialize decoder as transpose of encoder (tied weights)
with torch.no_grad():
self.decoder.weight.copy_(self.encoder.weight.T)
def forward(self, x):
"""
Forward pass with top-k sparsity constraint.
Returns reconstruction and sparse features.
"""
# Encode to feature space
features = self.encoder(x) # [batch, num_features]
# Apply top-k sparsity: keep only k largest activations per sample
k_values, k_indices = torch.topk(torch.abs(features), self.k, dim=1)
sparse_features = torch.zeros_like(features)
sparse_features.scatter_(1, k_indices, features.gather(1, k_indices))
# Decode sparse features
reconstruction = self.decoder(sparse_features)
return reconstruction, sparse_features
def train_on_source_model(self, source_model, trigger_dataset, num_epochs=3):
"""
Train SAE on activations from source reasoning model.
trigger_dataset: verified QA pairs with <think> markers but no reasoning steps.
"""
optimizer = torch.optim.Adam(self.parameters(), lr=1e-3)
loss_fn = nn.MSELoss()
for epoch in range(num_epochs):
for batch in trigger_dataset:
question_ids = batch['input_ids']
# Extract activation at target layer (layer 12)
with torch.no_grad():
activations = source_model.get_layer_activation(
question_ids,
layer=12
)
# Train SAE to reconstruct activations
reconstruction, sparse_features = self.forward(activations)
# Loss: reconstruction + sparsity penalty
recon_loss = loss_fn(reconstruction, activations)
sparsity_loss = 0.01 * sparse_features.abs().sum(dim=1).mean()
total_loss = recon_loss + sparsity_loss
optimizer.zero_grad()
total_loss.backward()
optimizer.step()
if batch % 100 == 0:
print(f"Epoch {epoch}, Batch {batch}: Loss={total_loss:.4f}")
Step 2: SAE-Guided SFT in Target Model
class LoRAAdapter(nn.Module):
"""
Low-rank adapter for inserting SAE into target model.
Minimizes KL divergence between outputs with/without SAE.
"""
def __init__(self, hidden_dim, lora_rank=32):
super().__init__()
self.lora_rank = lora_rank
# LoRA for query, key, value, dense projections
self.lora_q = nn.Linear(hidden_dim, lora_rank)
self.lora_q_out = nn.Linear(lora_rank, hidden_dim)
self.lora_k = nn.Linear(hidden_dim, lora_rank)
self.lora_k_out = nn.Linear(lora_rank, hidden_dim)
self.lora_v = nn.Linear(hidden_dim, lora_rank)
self.lora_v_out = nn.Linear(lora_rank, hidden_dim)
# Zero initialization: LoRA starts as no-op
with torch.no_grad():
self.lora_q_out.weight.zero_()
self.lora_k_out.weight.zero_()
self.lora_v_out.weight.zero_()
def forward(self, x):
"""Apply LoRA update to activation."""
delta_q = self.lora_q_out(self.lora_q(x))
delta_k = self.lora_k_out(self.lora_k(x))
delta_v = self.lora_v_out(self.lora_v(x))
return {'q': delta_q, 'k': delta_k, 'v': delta_v}
class SAEGuidedTraining:
"""
Stage II: Insert frozen SAE into target model and train LoRA
to minimize KL divergence between model with/without SAE.
"""
def __init__(self, target_model, sae, layer_idx=12):
self.target_model = target_model
self.sae = sae
self.layer_idx = layer_idx
self.lora = LoRAAdapter(target_model.hidden_size)
# Freeze SAE
for param in self.sae.parameters():
param.requires_grad = False
def train(self, trigger_dataset, num_epochs=1, learning_rate=1e-4):
"""
Train LoRA to instill SAE-extracted reasoning patterns.
KL divergence loss: log(p_with_sae) - log(p_without_sae)
"""
optimizer = torch.optim.Adam(self.lora.parameters(), lr=learning_rate)
for epoch in range(num_epochs):
for batch in trigger_dataset:
input_ids = batch['input_ids']
target_tokens = batch['target_ids'] # Next token predictions
# Forward pass WITHOUT SAE
with torch.no_grad():
logits_base = self.target_model(input_ids)
log_probs_base = torch.nn.functional.log_softmax(logits_base, dim=-1)
# Forward pass WITH SAE guidance
activations = self.target_model.get_layer_activation(
input_ids,
layer=self.layer_idx
)
# Apply SAE to extract reasoning features
_, sparse_features = self.sae(activations)
# Apply LoRA update conditioned on sparse features
lora_delta = self.lora(sparse_features)
# Modify model behavior via LoRA
logits_sae = self.target_model(
input_ids,
lora_delta=lora_delta
)
log_probs_sae = torch.nn.functional.log_softmax(logits_sae, dim=-1)
# KL divergence loss: encourage sae-guided outputs
kl_loss = torch.nn.functional.kl_div(
log_probs_sae,
torch.exp(log_probs_base),
reduction='batchmean'
)
optimizer.zero_grad()
kl_loss.backward()
optimizer.step()
print(f"Epoch {epoch}: KL Loss={kl_loss:.4f}")
Step 3: Transparent Reasoning Feature Analysis
class ReasoningFeatureAnalyzer:
"""
Identify and quantify latent reasoning features via prompt-only method.
Reveals which SAE features activate during reasoning tasks.
"""
def __init__(self, sae, model):
self.sae = sae
self.model = model
def identify_reasoning_features(self, test_dataset):
"""
Analyze which SAE features correlate with reasoning performance.
Returns feature importance scores.
"""
feature_activations = torch.zeros(self.sae.num_features)
feature_task_performance = torch.zeros(self.sae.num_features)
for sample in test_dataset:
# Get activations
with torch.no_grad():
activations = self.model.get_layer_activation(
sample['input_ids'],
layer=12
)
_, sparse_features = self.sae(activations)
# Track which features activated
feature_activations += (sparse_features.abs().sum(dim=0) > 0).float()
# Measure task performance (correctness on reasoning task)
correctness = 1.0 if sample['is_correct'] else 0.0
feature_task_performance += sparse_features[0] * correctness
# Compute feature importance
feature_importance = feature_task_performance / (feature_activations + 1e-8)
return feature_importance
def analyze_layer_wise_distribution(self, model, dataset):
"""
Examine how reasoning features distribute across layers.
Correlates with reasoning performance.
"""
layer_importance = {}
for layer_idx in range(model.num_layers):
features_this_layer = 0.0
for sample in dataset:
with torch.no_grad():
activations = model.get_layer_activation(
sample['input_ids'],
layer=layer_idx
)
_, sparse_features = self.sae(activations)
features_this_layer += sparse_features.abs().sum().item()
layer_importance[layer_idx] = features_this_layer / len(dataset)
return layer_importance
Step 4: Deployment and Cost Analysis
def compute_training_cost(model_size_b, hours_to_train, cost_per_hour_gpu=0.5):
"""Estimate training cost for SAE-tuning."""
num_gpus = max(1, model_size_b // 10) # Estimate GPUs needed
total_cost = hours_to_train * num_gpus * cost_per_hour_gpu
return total_cost
# Cost comparison
# SAE-Tuning on R1-Distill-1.5B:
# - Compute: 2 L40S GPUs x ~0.5 hours x $0.5/hr = $0.50
# - Model: GPT-4 mini calls for inference guidance = ~$0.47
# Total: ~$1.00
# Traditional RL on same model:
# - Compute: 8 A100 GPUs x ~30 hours x $2.0/hr = $480
# - Infrastructure setup, data curation = ~$1500
# Total: ~$2000+
print(f"SAE-Tuning cost: $1.00")
print(f"RL cost: $2000.00")
print(f"Cost reduction: 2000x")
Practical Guidance
Implementation Steps:
- Select source model with reasoning (R1-Distill, DeepSeek-R1)
- Prepare trigger dataset: verified QAs with
<think></think>markers but no intermediate steps - Train SAE on source model activations (layer 12 recommended)
- Insert frozen SAE into target model with LoRA adapters
- Fine-tune LoRA on trigger dataset (1 epoch sufficient)
Data Requirements:
- Trigger dataset: 5k-10k verified question-answer pairs
- Source: STILL, DeepScaleR, or similar CoT-free datasets
- No need for intermediate reasoning steps, reducing data requirements by 90%
Generalization:
- Reasoning features transfer across datasets (AIME, MATH, GPQA)
- Works across model families (Qwen, LLaMA variants, DeepSeek)
- Adapters remain small (LoRA rank 32) enabling multi-adapter composition
Performance Targets:
- R1-Distill-1.5B: ~35% AIME pass@1 (RL-equivalent)
- Generalizes to unseen benchmarks without retraining
- Multimodal models: reasoning benefits transfer to vision tasks
Reference
- Sparse Autoencoders: Extract interpretable features from neural networks; top-k sparsity prevents feature collapse
- LoRA: Low-rank adaptation; efficient parameterization for transfer learning
- KL divergence: Measures distribution difference; KL minimization encourages model to adopt SAE-guided behavior