adu2021 DataChef: Cooking Up Optimal Data Recipes via RL
Problem Context
Creating high-quality training data is labor-intensive and requires domain expertise. DataChef automates this by using RL to generate data recipes—specifications for how to create training datasets. A fast Data Verifier (rubric-based classifier) provides rewards without expensive downstream training, enabling rapid iteration on dataset composition.
Core Concept
DataChef operates in two phases: (1) cold-start with supervised demonstrations from strong models, (2) RL phase where policies generate data recipes, verifiers assess quality, and GRPO optimizes recipe generation. Recipes specify: data source selection, augmentation strategies, filtering criteria, and mixing ratios.
Implementation
Step 1: Data recipe representation and generator
from dataclasses import dataclass
from typing import Dict, List, Any
import json
@dataclass
class DataRecipe:
"""Specification for dataset creation."""
sources: List[str] # Data sources (e.g., ["wikipedia", "books"])
augmentation_strategies: List[str] # e.g., ["paraphrase", "back_translation"]
filtering_criteria: Dict[str, Any] # Quality thresholds
mixing_ratios: Dict[str, float] # Per-source weights
sample_size: int
def to_dict(self) -> Dict:
return {
'sources': self.sources,
'augmentation': self.augmentation_strategies,
'filters': self.filtering_criteria,
'mixing': self.mixing_ratios,
'size': self.sample_size
}
class DataRecipeGenerator:
"""Generate data recipes via language model."""
def __init__(self, model):
self.model = model
def generate_recipe(
self,
task: str,
max_tokens: int = 300
) -> Dict:
"""
Generate data recipe for task.
Args:
task: Task description
max_tokens: Recipe generation length
Returns:
recipe_dict: Parsed recipe specification
"""
prompt = f"""Generate a data recipe for training on: {task}
Specify:
1. Data sources (wikipedia, books, web, etc.)
2. Augmentation strategies (paraphrase, back_translation, etc.)
3. Quality filters (minimum length, diversity, etc.)
4. Mixing ratios for different sources
5. Target sample size
Format as JSON.
Data recipe:"""
recipe_text, log_probs = self.model.generate_with_logprobs(
prompt, max_tokens=max_tokens, temperature=0.7
)
# Parse recipe
try:
recipe_dict = json.loads(recipe_text)
except json.JSONDecodeError:
# Fallback: extract fields from free text
recipe_dict = self._parse_recipe_freetext(recipe_text)
return recipe_dict, log_probs
def _parse_recipe_freetext(self, text: str) -> Dict:
"""Fallback parsing from free text."""
return {
'sources': ['wikipedia'],
'augmentation': ['paraphrase'],
'filters': {'min_length': 10},
'mixing': {'wikipedia': 1.0},
'size': 10000
}
Step 2: Data verifier (fast quality assessment)
class DataVerifier:
"""Fast quality assessment without full model training."""
def __init__(self, quality_rubric: Dict = None):
"""
Args:
quality_rubric: Scoring criteria for samples
"""
self.rubric = quality_rubric or {
'invalid': 0.0,
'format_error': 0.1,
'incorrect': 0.2,
'task_mismatch': 0.5,
'pass': 1.0
}
def verify_sample(
self,
sample: Dict,
task: str,
verifier_model=None
) -> float:
"""
Classify sample quality quickly.
Args:
sample: Training sample (input, target)
task: Task description
verifier_model: Optional LLM for semantic checking
Returns:
score: Quality score (0-1)
"""
# Format check
if 'input' not in sample or 'target' not in sample:
return self.rubric['format_error']
input_text = sample['input']
target_text = sample['target']
# Length check
if len(input_text.split()) < 5 or len(target_text.split()) < 1:
return self.rubric['format_error']
# Task relevance check (simple heuristic)
if not self._is_task_relevant(input_text, task):
return self.rubric['task_mismatch']
# Correctness estimation
if verifier_model:
is_correct = verifier_model.verify(input_text, target_text)
if is_correct:
return self.rubric['pass']
else:
return self.rubric['incorrect']
# Default: assume correct if format is valid
return self.rubric['pass']
def _is_task_relevant(self, text: str, task: str) -> bool:
"""Check task relevance (simplified)."""
task_keywords = task.lower().split()
text_lower = text.lower()
return any(kw in text_lower for kw in task_keywords[:3])
def verify_dataset(
self,
dataset: List[Dict],
task: str,
sample_fraction: float = 0.1
) -> float:
"""
Estimate dataset quality via sampling.
Args:
dataset: Full dataset
sample_fraction: Fraction to verify
Returns:
avg_quality_score: Mean quality across sample
"""
import random
sample_size = max(1, int(len(dataset) * sample_fraction))
sample = random.sample(dataset, min(sample_size, len(dataset)))
scores = [self.verify_sample(s, task) for s in sample]
return sum(scores) / len(scores)
Step 3: GRPO training for recipe generation
class DataRecipeRL:
"""GRPO-based optimization of data recipes."""
def __init__(
self,
model,
optimizer,
verifier: DataVerifier,
data_executor,
group_size: int = 4
):
self.model = model
self.optimizer = optimizer
self.verifier = verifier
self.executor = data_executor # Creates datasets from recipes
self.group_size = group_size
self.recipe_generator = DataRecipeGenerator(model)
def compute_recipe_reward(
self,
recipe: Dict,
task: str,
downstream_verifier=None
) -> float:
"""
Compute reward for recipe.
Args:
recipe: Generated recipe
task: Task for which data is created
downstream_verifier: Optional model for full eval
Returns:
reward: Quality score
"""
try:
# Execute recipe: create dataset
dataset = self.executor.create_dataset(recipe, task, sample_size=100)
# Verify dataset
quality_score = self.verifier.verify_dataset(dataset, task)
# Normalize reward
reward = quality_score
except Exception as e:
# Recipe failed to execute
reward = 0.0
return reward
def training_step(
self,
tasks: List[str],
batch_size: int = 4
) -> Dict:
"""Single training step."""
log_probs_list = []
rewards = []
for task in tasks[:batch_size]:
# Generate recipe
recipe, log_probs = self.recipe_generator.generate_recipe(task)
log_probs_list.append(log_probs)
# Compute reward
reward = self.compute_recipe_reward(recipe, task)
rewards.append(reward)
log_probs = torch.stack(log_probs_list)
rewards = torch.tensor(rewards, dtype=torch.float32)
# GRPO loss
log_prob_ratio = log_probs - log_probs.detach()
ratio = torch.exp(log_prob_ratio)
# Group relative advantage
group_mean_reward = rewards.mean()
relative_rewards = rewards - group_mean_reward
clipped_ratio = torch.clamp(ratio, 0.5, 2.0)
loss = -torch.min(
log_prob_ratio * relative_rewards.unsqueeze(-1),
torch.log(clipped_ratio) * relative_rewards.unsqueeze(-1)
).mean()
# Update
self.optimizer.zero_grad()
loss.backward()
torch.nn.utils.clip_grad_norm_(self.model.parameters(), 1.0)
self.optimizer.step()
return {
'loss': loss.item(),
'avg_reward': rewards.mean().item(),
'avg_quality': rewards.mean().item()
}
Step 4: Cold-start with synthetic demonstrations
def generate_cold_start_demonstrations(
strong_model,
tasks: List[str],
num_demos_per_task: int = 5
) -> List[Dict]:
"""
Generate high-quality data recipes from strong model.
Uses strong model (e.g., GPT-4) to create initial demonstrations.
"""
demonstrations = []
for task in tasks:
for _ in range(num_demos_per_task):
# Ask strong model to create data recipe
prompt = f"""Create a high-quality data recipe for: {task}
Consider data sources, augmentation, filtering, and quality thresholds.
Return as a concrete JSON specification."""
recipe_json = strong_model.generate(prompt, max_tokens=300)
demo = {
'task': task,
'recipe': recipe_json,
'is_high_quality': True
}
demonstrations.append(demo)
return demonstrations
def train_datachef_with_cold_start(
model,
strong_model,
tasks: List[str],
data_executor,
optimizer,
num_epochs: int = 10,
num_rl_steps_per_epoch: int = 100
):
"""
Full DataChef training: cold-start SFT → RL.
"""
# Cold-start
print("Cold-start: Supervised fine-tuning")
demos = generate_cold_start_demonstrations(strong_model, tasks, num_demos_per_task=5)
for demo in demos:
prompt = f"Task: {demo['task']}\n\nData recipe:"
target = demo['recipe']
# SFT loss (simplified)
# Would compute cross-entropy on tokens
# RL phase
print("RL phase: GRPO optimization")
verifier = DataVerifier()
rl_trainer = DataRecipeRL(model, optimizer, verifier, data_executor)
for epoch in range(num_epochs):
total_loss = 0.0
total_reward = 0.0
for step in range(num_rl_steps_per_epoch):
metrics = rl_trainer.training_step(tasks, batch_size=4)
total_loss += metrics['loss']
total_reward += metrics['avg_reward']
avg_loss = total_loss / num_rl_steps_per_epoch
avg_reward = total_reward / num_rl_steps_per_epoch
print(f"Epoch {epoch + 1}: Loss={avg_loss:.4f}, Reward={avg_reward:.4f}")
return model
Practical Guidance
When to use: Data curation bottleneck; need diverse training data; frequent dataset regeneration
Hyperparameters:
- verification_sample_fraction: 0.05-0.2 (tradeoff: cost vs. confidence)
- group_size: 4-8 (GRPO)
- cold_start_demos: 3-10 per task
- recipe_max_tokens: 200-400
Key advantages:
- Automates data curation (labor savings)
- Fast verification without full training
- Curriculum: SFT → RL smooths learning
- Generalizes across domains
Common pitfalls:
- Verifier too lenient → low-quality recipes
- Not validating executor can actually create dataset
- Cold-start too brief → RL starts from bad initialization
- Recipe too constrained → fails to explore
Scaling: Verification cost is the main bottleneck; parallelize across tasks.
Reference
Paper: https://arxiv.org/abs/2602.11089 Related work: Data augmentation, dataset curation, GRPO Benchmarks: Diverse domains including math, code, QA