Update README.md

README.md

@@ -29,13 +29,14 @@ MiniGuard-v0.1 uses the **same prompt template** as [nvidia/Llama-3.1-Nemotron-S
 
 ## Training
 
-MiniGuard-v0.1 was trained using three key techniques:
+MiniGuard-v0.1 was trained using four key techniques to break the trade-off between safety and latency:
 
-1. **LoRA Fine-tuning** — English Subset of [nvidia/Nemotron-Safety-Guard-Dataset-v3](https://huggingface.co/datasets/nvidia/Nemotron-Safety-Guard-Dataset-v3) + Reasoning traces from [openai/gpt-oss-safeguard-120b](https://huggingface.co/openai/gpt-oss-safeguard-120b).
+1.  **Targeted Synthetic Data** — To address specific failure modes (e.g., sports terms, ambiguous edge cases), we generated ~1,200 targeted examples using **Hermes-4.3-36B**. This data complements the English subset of [nvidia/Nemotron-Safety-Guard-Dataset-v3](https://huggingface.co/datasets/nvidia/Nemotron-Safety-Guard-Dataset-v3).
 
-2. **Distilling Step-by-Step** — A teacher LLM generates reasoning traces for training examples. The student model is trained on both reasoning-augmented and standard examples, improving performance even when reasoning is not generated at inference time. ([Reference](https://research.google/blog/distilling-step-by-step-outperforming-larger-language-models-with-less-training-data-and-smaller-model-sizes/))
+2.  **Think SFT (Distilling Step-by-Step)** — A teacher LLM (**gpt-oss-safeguard-120b**) generates reasoning traces for training examples. The student model is trained on these traces but discards them at inference, retaining reasoning capabilities without the token cost.
+
+3.  **Top-K Model Soup** — We employ a Top-K (K=3) weight averaging strategy. Weights from the top 3 validation checkpoints are averaged to improve out-of-distribution generalization without increasing inference overhead.
 
-3. **Greedy Model Soup** — Multiple fine-tuned checkpoints are averaged using a greedy selection strategy: checkpoints are sorted by validation accuracy and sequentially added to the "soup" only if they improve or maintain performance. This provides a free accuracy boost without additional compute. ([Reference](https://arxiv.org/abs/2203.05482))
 
 ## Evaluation
 
@@ -46,38 +47,56 @@ Dataset - English subset test split of [nvidia/Nemotron-Safety-Guard-Dataset-v3]
 | Metric | MiniGuard-v0.1 | Nemotron-Guard-8B-v3 |
 |--------|----------------|----------------------|
 | Parameters | **0.6B** | 8B |
-| Overall F1 | 0.881 | 0.893 |
-| Accuracy Retained | **98.7%** | 100% |
+| Weighted F1 | 88.9 | 89.3 |
+| Accuracy Retained | **99.5%** | 100% |
 | Size Reduction | **13x** | 1x |
 
-### Production Dataset Evaluation
 
-Evaluated on an internal dataset of real-user prompts. Cost estimated based on H200 GPU pricing ($3.50/hour) at concurrency 16 with P95 latency SLA of <500ms.
+#### Production Dataset Evaluation
+
+Evaluated on out-of-distribution production data containing real user queries. Cost estimated based on H200 GPU pricing ($7.91/hour) at maximum concurrency with P95 latency SLA of <500ms.
 
 | Metric | MiniGuard-v0.1 | Nemotron-Guard-8B-v3 |
 |--------|----------------|----------------------|
-| Relative Safety Score | 97.4% | 100% |
-| Relative Cost | $5.4 | $8 |
+| Relative Macro F1 | 91.1% | 100% |
+| Cost per 1M requests | **$15.54** | $46.93 |
+| Cost Savings | **67%** | - |
 
-MiniGuard-v0.1 achieves comparable safety scores while costing ~30% less per million requests.
+
+MiniGuard-v0.1 achieves 91.1% relative performance on out-of-distribution data while costing **67% less** to serve.
 
 ### Ablation Study
 
+#### Out-of-Distribution: Production Dataset
+
+Impact of techniques on out-of-distribution production data (Relative Macro F1 compared to Nemotron-Guard-8B).
+
+| Configuration | Parameters | Rel. Macro F1 | Improvement |
+| :--- | :--- | :--- | :--- |
+| Qwen3-0.6B + Think SFT | 0.6B | 85.6% | baseline |
+| + Targeted Synthetic Data | 0.6B | 87.2% | +1.6% |
+| + Soup (top-3) [MiniGuard-v0.1] | 0.6B | 92.3% | +5.1% |
+| + FP8 | 0.6B | 91.1% | -1.2% |
+| Nemotron-Guard-8B-v3 | 8B | 100% | reference |
+
+#### In-Distribution
+
 Dataset - English subset test split of [nvidia/Nemotron-Safety-Guard-Dataset-v3](https://huggingface.co/datasets/nvidia/Nemotron-Safety-Guard-Dataset-v3).
 
-| Model | Overall F1 | User Safety F1 | Response Safety F1 |
-|-------|------------|----------------|-------------------|
-| Qwen3-0.6B (baseline) | 0.637 | 0.594 | 0.681 |
-| + Vanilla SFT | 0.844 | 0.843 | 0.844 |
-| + Distilling Step-by-Step | 0.882 | 0.873 | 0.892 |
-| + Greedy Model Soup | 0.886 | 0.873 | 0.899 |
-| + FP8 Quantization | 0.881 | 0.872 | 0.890 |
+| Training Configuration | Weighted F1 | Macro F1 |
+|-------|------------|----------------|
+| Qwen3-0.6B (base) | 63.7 | 52.5 |
+| + Vanilla SFT | 84.4 | 85.0 |
+| + Think SFT (distillation) | 88.2 | 88.6 |
+| + Targeted Synthetic Data | 88.9 | 89.3 |
+| + Top-3 Model Soup | 88.8 | 89.2 |
+| + FP8 Quantization | 88.9 | 89.3 |
 
 ## Input
 **Input Type(s)**: Text <br>
 **Input Format(s)**: String <br>
 **Input Parameters**: One-Dimensional (1D): Sequences <br>
-**Other Properties Related to Input**: Context length up to 8K. Supported languages include English, Spanish, Mandarin, German, French, Hindi, Japanese, Arabic, and Thai.
+**Other Properties Related to Input**: Context length up to 32K. Supported language: English
 
 ## Output
 **Output Type(s)**: Text Json <br>