adu2021 Saffron-1: Inference Scaling Paradigm for LLM Safety Assurance
Core Concept
Saffron-1 addresses a critical limitation in LLM safety: conventional inference scaling techniques fail to improve safety robustness effectively. The key insight is that frequent process reward model calls create computational overhead that undermines efficiency gains. Saffron replaces traditional single-output reward models with multifurcation reward models that predict rewards for all vocabulary tokens in a single forward pass, reducing required model evaluations from K calls to 1. Combined with conservative exploration constraints and efficient KV caching, Saffron achieves significantly lower attack success rates while maintaining computational efficiency.
Architecture Overview
- Multifurcation Reward Model (MRM): Predicts rewards for all vocabulary tokens simultaneously
- Token-Level Supervision: Partial supervision from pre-computed corpus rewards, not requiring full annotation
- Conservative Exploration: Restricts search to previously-seen tokens only
- Trie-Based KV Caching: Shares key-value caches across sequences with common prefixes
- Safety Dataset: Safety4M annotation with Llama Guard 3 on HH-RLHF corpus
- Tree Search Integration: Compatible with beam search and other inference algorithms
Implementation
Step 1: Create Multifurcation Reward Model
Build a reward model that predicts rewards for all tokens simultaneously:
import torch
import torch.nn as nn
from transformers import AutoModelForCausalLM, AutoTokenizer
class MultifurcationRewardModel(nn.Module):
"""
Reward model that outputs reward for each vocabulary token.
Single forward pass predicts reward[vocab_size] instead of K separate calls.
"""
def __init__(self, base_model_name="meta-llama/Llama-2-7b-hf", vocab_size=32000):
super().__init__()
# Load base model
self.base_model = AutoModelForCausalLM.from_pretrained(base_model_name)
self.hidden_dim = self.base_model.config.hidden_size
self.vocab_size = vocab_size
# Reward head: projects hidden state to per-token rewards
self.reward_head = nn.Linear(self.hidden_dim, vocab_size)
# Optional: LoRA adaptation for efficiency
self.use_lora = True
if self.use_lora:
self.apply_lora()
def apply_lora(self, r=8, lora_alpha=16):
"""Apply LoRA to reduce parameters"""
from peft import get_peft_model, LoraConfig
lora_config = LoraConfig(
r=r,
lora_alpha=lora_alpha,
target_modules=["q_proj", "v_proj"],
lora_dropout=0.05,
bias="none",
task_type="CAUSAL_LM"
)
self.base_model = get_peft_model(self.base_model, lora_config)
def forward(self, input_ids, attention_mask=None):
"""
Forward pass: single call outputs all token rewards.
Args:
input_ids: (batch_size, seq_len)
attention_mask: (batch_size, seq_len)
Returns:
rewards: (batch_size, vocab_size) - reward for each token
"""
outputs = self.base_model(
input_ids=input_ids,
attention_mask=attention_mask,
output_hidden_states=True
)
# Use last hidden state from final token
last_hidden = outputs.hidden_states[-1] # (batch_size, seq_len, hidden_dim)
final_hidden = last_hidden[:, -1, :] # (batch_size, hidden_dim)
# Compute reward for each vocabulary token
rewards = self.reward_head(final_hidden) # (batch_size, vocab_size)
return rewards
Step 2: Create Safety4M Training Dataset
Build the token-level supervision dataset:
class Safety4MDataset:
"""
Token-level safety rewards from HH-RLHF corpus.
4 million tokens with pre-computed Llama Guard 3 annotations.
"""
def __init__(self, corpus_path, reward_model_checkpoint, num_samples=4_000_000):
self.corpus_path = corpus_path
self.tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-2-7b")
self.num_samples = num_samples
# Load Llama Guard 3 for labeling
self.labeler = load_llama_guard_3()
def create_token_level_rewards(self, texts):
"""
Create token-level reward supervision.
For each prefix in corpus, compute safety score and assign to next token.
"""
dataset = []
for text in texts:
tokens = self.tokenizer.encode(text)
# Compute reward (safety score) for each position
for t in range(1, len(tokens)):
prefix = self.tokenizer.decode(tokens[:t])
next_token = tokens[t]
# Llama Guard 3 safety score
safety_score = self.labeler.compute_safety(prefix)
# Assign reward to next token
dataset.append({
'prefix_ids': torch.tensor(tokens[:t]),
'next_token': next_token,
'reward': safety_score # Range: [0, 1]
})
return dataset
def collate_batch(self, samples):
"""Batch with padding"""
prefix_ids = torch.nn.utils.rnn.pad_sequence(
[s['prefix_ids'] for s in samples],
batch_first=True,
padding_value=self.tokenizer.pad_token_id
)
next_tokens = torch.tensor([s['next_token'] for s in samples])
rewards = torch.tensor([s['reward'] for s in samples])
return {
'prefix_ids': prefix_ids,
'next_tokens': next_tokens,
'rewards': rewards
}
def train_multifurcation_reward_model(model, dataset, epochs=3, lr=1e-5):
"""Train MRM with token-level supervision"""
optimizer = torch.optim.AdamW(model.parameters(), lr=lr)
for epoch in range(epochs):
total_loss = 0
for batch in dataset:
prefix_ids = batch['prefix_ids'].to('cuda')
next_tokens = batch['next_tokens'].to('cuda')
rewards = batch['rewards'].to('cuda')
optimizer.zero_grad()
# Forward: get reward for each vocabulary token
predicted_rewards = model(input_ids=prefix_ids) # (batch_size, vocab_size)
# Extract predicted reward for the actual next token
batch_size = next_tokens.shape[0]
predicted_next_rewards = predicted_rewards[
torch.arange(batch_size), next_tokens
] # (batch_size,)
# MSE loss between predicted and ground truth rewards
loss = torch.nn.functional.mse_loss(predicted_next_rewards, rewards)
loss.backward()
torch.nn.utils.clip_grad_norm_(model.parameters(), 1.0)
optimizer.step()
total_loss += loss.item()
print(f"Epoch {epoch+1}: Loss = {total_loss / len(dataset):.4f}")
Step 3: Implement Tree Search with Conservative Exploration
Build tree search that constrains exploration to safe token space:
class SafeTreeSearch:
def __init__(self, model, reward_model, tokenizer, beam_size=4):
self.model = model
self.reward_model = reward_model
self.tokenizer = tokenizer
self.beam_size = beam_size
# Track seen tokens for conservative exploration
self.seen_tokens = set(range(1000)) # Common tokens
def search(self, input_ids, max_depth=50, attack_success_threshold=0.5):
"""
Tree search with safety-aware exploration.
Restricts to seen tokens only to prevent unreliable predictions on
out-of-distribution vocabulary.
"""
# Initialize beam with starting sequence
beam = [(input_ids, 0.0)] # (sequence, accumulated_reward)
for depth in range(max_depth):
candidates = []
for sequence, acc_reward in beam:
# Get reward for all tokens
with torch.no_grad():
rewards = self.reward_model(input_ids=sequence) # (1, vocab_size)
# Conservative exploration: only consider seen tokens
safe_mask = torch.zeros(self.reward_model.vocab_size)
for token_id in self.seen_tokens:
safe_mask[token_id] = 1.0
# Mask out unseen tokens
masked_rewards = rewards[0] * safe_mask - 1e9 * (1 - safe_mask)
# Get top tokens
top_rewards, top_tokens = torch.topk(masked_rewards, self.beam_size)
# Create new sequences
for reward, token_id in zip(top_rewards, top_tokens):
new_sequence = torch.cat([
sequence,
torch.tensor([[token_id.item()]])
], dim=1)
new_reward = acc_reward + reward.item()
candidates.append((new_sequence, new_reward))
# Keep top beam_size candidates
candidates.sort(key=lambda x: x[1], reverse=True)
beam = candidates[:self.beam_size]
# Early stopping if attack success rate is too high
attack_rate = self.estimate_attack_success_rate(beam)
if attack_rate > attack_success_threshold:
break
# Return best sequence
best_sequence, best_reward = beam[0]
return best_sequence, best_reward
def estimate_attack_success_rate(self, sequences):
"""Estimate likelihood of adversarial success"""
# Use Llama Guard 3 or similar detector
attack_count = 0
for sequence, _ in sequences:
text = self.tokenizer.decode(sequence[0])
is_unsafe = self.reward_model.predict_unsafe(text)
if is_unsafe:
attack_count += 1
return attack_count / len(sequences)
Step 4: Implement Trie-Based KV Caching
Optimize memory usage for sequences with shared prefixes:
class TrieKVCache:
"""
Efficient KV caching for tree search using trie structure.
Shares key-value caches across sequences with common prefixes.
"""
def __init__(self):
self.trie = {}
self.cache = {}
def get_cache(self, prefix_ids):
"""Retrieve cached KV for prefix"""
key = tuple(prefix_ids.tolist())
if key in self.cache:
return self.cache[key]
return None
def store_cache(self, prefix_ids, kv_cache):
"""Store KV cache for prefix"""
key = tuple(prefix_ids.tolist())
self.cache[key] = kv_cache
def forward_with_cache(self, model, input_ids, past_key_values=None):
"""Forward pass reusing cached KVs"""
# Check if we have cached KVs for a prefix
cached_kv = self.get_cache(input_ids[:, :-1])
if cached_kv is not None:
# Use cached KVs for all but last token
outputs = model(
input_ids=input_ids[:, -1:],
past_key_values=cached_kv,
use_cache=True
)
else:
# Full forward pass
outputs = model(
input_ids=input_ids,
past_key_values=None,
use_cache=True
)
# Store cache for future use
self.store_cache(input_ids, outputs.past_key_values)
return outputs
Practical Guidance
- Multifurcation Benefit: Reduces reward model calls from K (beam size) to 1, dramatically improving efficiency
- Token Supervision: Uses existing corpus annotations; no need for full vocabulary labeling
- Conservative Exploration: Essential for safety; prevents searching unreliable token space
- Attack Success Rate: Typical improvements: 40.9% ASR (Saffron) vs. 58.2% (Best-of-N)
- Training Data: Safety4M dataset with Llama Guard 3 labels; can use other safety labelers
- Beam Size: 4-8 typical; larger beams increase safety but reduce speed
- Computational Cost: Trie-based caching reduces memory by 40-60% compared to naive tree search
- Integration: Works with existing LLMs without architectural changes
Reference
- Multifurcation reward models parallelize vocabulary scoring while reducing forward passes
- Token-level supervision enables efficient training on unlabeled corpus data
- Conservative exploration fundamentally prevents out-of-distribution reward exploitation
- Trie-based caching is standard technique in sequence generation but particularly valuable for safety scaling