This project demonstrates prompt distillation: training a small, specialized model to outperform the large model that generated its training data.
The Challenge
Marketing AI assistants need to remember the right information from conversations. Not everything is worth storing - you need to distinguish between:
Valuable: "Our brand voice is professional but approachable" → Store in long-term memory
Transactional: "What time is the meeting tomorrow?" → Don't store
This is a 13-category classification problem with nuanced distinctions between company-level and user-level information, different persistence horizons, and the critical ability to say "none" for irrelevant content.
The Approach
Generate synthetic data using Cohere Command-R-Plus (104B) as the teacher
Fine-tune Llama-3.1-8B with LoRA using Tinker's training platform
Apply reinforcement learning with a custom reward function
Benchmark against the teacher on challenging, held-out scenarios
The Result
Model
Parameters
Avg F1
Exact Match
Ours
8B
0.68
60%
Cohere Command-R-Plus
104B
0.61
26%
Our 8B model achieves 11.1% higher F1 and 2.3x better exact match accuracy than the 104B teacher, while being 13x smaller.
The student surpassed the teacher through:
Focused training: The model only learns this one task, not general capabilities
RL refinement: The reward function optimizes for exact category matching, not just plausible outputs
Clean data: Synthetic data with consistent labeling, no noise from human annotation disagreements
Training Visualizations
Phase 1: Supervised Fine-Tuning
SFT Loss
100 training steps reduced loss from 5.47 to 0.26 (95% reduction). The model learned the basic classification task in the first epoch.
Phase 2: Reinforcement Learning
RL Reward
30 RL iterations improved mean reward from 0.73 to 0.93. The reward function combines F1 score, temporal alignment, scope correctness, and storage efficiency.
Model Comparison
Model Comparison
Our model excels at exact matching (60% vs 26%) because RL optimizes for getting all categories right, not just some.
Performance by Difficulty
Difficulty Comparison
The 8B model dominates on easy cases (+79% F1) and matches on medium cases. The 104B model still wins on hard multi-label scenarios.
Key Results
Metric
Our Model (8B)
Cohere (104B)
Avg F1
0.68
0.61
Exact Match
60%
26%
Any Match
72%
82%
Model Size
8B
104B
Improvement
+11.1% F1
baseline
Reward Components (Final RL Iteration)
Component
Score
Description
R_F1
0.90
F1 score vs gold labels
R_temp
0.95
Temporal alignment
R_parity
1.00
Company/user scope
R_eff
1.00
Storage efficiency
What It Does
The Memory Routing Agent classifies marketing conversations into 13 memory categories:
1# Create .env file with your API keys2echo"TINKER_API_KEY=your_tinker_key">> .env
3echo"COHERE_API_KEY=your_cohere_key">> .env
4echo"HF_TOKEN=your_huggingface_token">> .env
Run Inference
python
1import tinker
2from tinker import types
3from tinker_cookbook import renderers
4from tinker_cookbook.tokenizer_utils import get_tokenizer
56# Load model from Tinker checkpoint7service_client = tinker.ServiceClient()8checkpoint ="tinker://4f4bae1f-5a95-5f53-a55a-a14f2872825c:train:0/sampler_weights/rl_iter_012"9sampling_client = service_client.create_sampling_client(model_path=checkpoint)1011# Setup tokenizer and renderer12tokenizer = get_tokenizer("meta-llama/Llama-3.1-8B")13renderer = renderers.get_renderer(name="llama3", tokenizer=tokenizer)1415# Classify a conversation16conversation ="""
17USER: Our brand voice is professional but approachable. Think Harvard Business Review meets Slack.
18ASSISTANT: So authoritative content with a conversational tone?
19USER: Exactly. We never use jargon without explaining it first.
20"""2122messages =[23{"role":"system","content":"You route marketing conversations into structured memory categories..."},24{"role":"user","content":f"Analyze this conversation:\n\n{conversation}"}25]2627prompt = renderer.build_generation_prompt(messages)28params = types.SamplingParams(max_tokens=100, temperature=0.1, stop=renderer.get_stop_sequences())29result = sampling_client.sample(prompt=prompt, sampling_params=params, num_samples=1).result()3031response, _ = renderer.parse_response(result.sequences[0].tokens)32print(f"Categories: {response['content']}")33# Output: company.brand_core