# Saffron Safety Scaling > Implement Saffron's multifurcation reward model approach to achieve efficient inference-time safety scaling, improving robustness against prompt injection attacks while reducing computational overhead. - Skill: `adu2021/saffron-safety-scaling` (Agent Skill) - Install (CLI): `npx skillmds@latest add adu2021/saffron-safety-scaling` - Raw SKILL.md: https://api.skillmd.com/api/skills/adu2021/saffron-safety-scaling/raw - Safety review: pending - Works with: Claude Code, Claude.ai, OpenAI Codex - Category: AI & ML - License: MIT - Author: adu2021 (https://skillmd.com/u/adu2021) - Updated: 2026-09-17 - Page: https://skillmd.com/skills/adu2021/saffron-safety-scaling --- # 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: ```python 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: ```python 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: ```python 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: ```python 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