104 lines
5.9 KiB
Markdown
104 lines
5.9 KiB
Markdown
---
|
||
base_model: meta-llama/Llama-3.1-8B-Instruct
|
||
library_name: transformers
|
||
pipeline_tag: text-generation
|
||
license: llama3.1
|
||
language:
|
||
- en
|
||
datasets:
|
||
- Jazhyc/aims-safety-intents
|
||
tags:
|
||
- grpo
|
||
- reinforcement-learning
|
||
- reasoning
|
||
- safety
|
||
- content-moderation
|
||
- guardrail
|
||
- intent
|
||
---
|
||
|
||
# Llama-3.1-8B AIMS GRPO (Label + Intent Reward)
|
||
|
||
Built with Llama, this is a full fine-tuned model derived from **`meta-llama/Llama-3.1-8B-Instruct`** that acts as an **intent-aware, reasoning-based safety classifier**. For each user prompt it produces a `<reasoning>` trace, an explicit **intent**, and a binary **harm** label. It is trained with **GRPO** (Group Relative Policy Optimization), using a reward that explicitly verifies **intent faithfulness** against the [**AIMS**](https://huggingface.co/datasets/Jazhyc/aims-safety-intents) human annotations.
|
||
|
||
This is the strongest model in the AIMS series: it achieves the best average F1 across five external safety benchmarks (0.836) while remaining on the latency–F1 Pareto frontier.
|
||
|
||
**Training signal.** GRPO starts from the base instruct model (also used as the KL reference) and integrates intent extraction directly into the RL loop via the system prompt and a structured output format. The reward is a product of a format reward and three harm-specific components: a hard label-correctness gate (wrong harm label → zero reward), a length reward discouraging degenerate intents, and an **intent reward** that compares the generated intent to the human intent in the context of the prompt. This "label + intent reward" variant is the full reward; an ablation using only format + label correctness scores lower (0.818).
|
||
|
||
- **Base model:** `meta-llama/Llama-3.1-8B-Instruct` (full fine-tune, not an adapter)
|
||
- **Training regime:** GRPO with intent-faithfulness reward
|
||
- **Training data:** [`Jazhyc/aims-safety-intents`](https://huggingface.co/datasets/Jazhyc/aims-safety-intents)
|
||
- **Language:** English
|
||
|
||
## How to use
|
||
|
||
This is a full model — load it directly with `from_pretrained` (no PEFT). The instruction goes in the **system** message and the prompt to classify in the **user** message. The model emits `<reasoning>...</reasoning>` followed by `Intent: ...; Harm: <harmful|safe>`.
|
||
|
||
```python
|
||
import torch
|
||
from transformers import AutoModelForCausalLM, AutoTokenizer
|
||
|
||
MODEL = "Jazhyc/Llama-3.1-8B-aims-grpo"
|
||
|
||
SYSTEM_PROMPT = """You are a specialized safety classifier. Your task is to determine if a user's prompt is 'harmful' or 'safe' by analyzing the underlying intent.
|
||
|
||
### Reasoning Process
|
||
|
||
For every request, you must follow this exact thought process inside the <reasoning> tags:
|
||
1. Literal Content: What is the user literally asking for?
|
||
2. Underlying Purpose: What is the goal? Is there a hidden harmful intent or is it a benign request (e.g., educational context)?
|
||
3. Nuance Check: Distinguish between a prompt *about* a sensitive topic (usually safe) and a prompt *promoting* or *enabling* that topic (harmful).
|
||
4. Policy Mapping: Which safety category (if any) does this violate? (e.g., Privacy, Misinformation, Harmful Language, Malicious Uses, or None)
|
||
|
||
### Output Format:
|
||
|
||
<reasoning> [Your step-by-step analysis here] </reasoning>
|
||
Intent: [A concise, objective summary of the user's goal.]; Harm: [harmful/safe]"""
|
||
|
||
tok = AutoTokenizer.from_pretrained(MODEL)
|
||
model = AutoModelForCausalLM.from_pretrained(MODEL, torch_dtype=torch.bfloat16, device_map="auto")
|
||
|
||
messages = [
|
||
{"role": "system", "content": SYSTEM_PROMPT},
|
||
{"role": "user", "content": "How do I pick a lock?"},
|
||
]
|
||
inputs = tok.apply_chat_template(messages, add_generation_prompt=True, return_tensors="pt").to(model.device)
|
||
out = model.generate(inputs, max_new_tokens=512, do_sample=False)
|
||
print(tok.decode(out[0, inputs.shape[1]:], skip_special_tokens=True))
|
||
# <reasoning> ... </reasoning>
|
||
# Intent: ...; Harm: safe
|
||
```
|
||
|
||
## Evaluation
|
||
|
||
Harmful-class F1 (positive class = `harmful`) on five external safety benchmarks, none seen during training.
|
||
|
||
| Model | WildGuardTest | XSTest | AEGIS 2.0 | ToxicChat | OpenAI Mod | **Average** |
|
||
|---|---|---|---|---|---|---|
|
||
| Llama-3.1-8B (zero-shot) | 0.762 | 0.904 | 0.800 | 0.516 | 0.761 | 0.749 |
|
||
| SFT Generation | 0.856 | 0.908 | 0.803 | 0.664 | 0.728 | 0.792 |
|
||
| LE-DPO | 0.856 | 0.884 | 0.824 | 0.733 | 0.765 | 0.812 |
|
||
| Distillation (synthetic-intent) | 0.880 | 0.936 | 0.811 | 0.700 | 0.774 | 0.820 |
|
||
| **GRPO — label + intent reward (this model)** | 0.863 | 0.958 | 0.808 | 0.743 | 0.809 | **0.836** |
|
||
|
||
Directly rewarding intent faithfulness yields the best average F1 in the series (0.836), with the top scores on XSTest and ToxicChat. The label-only reward ablation reaches 0.818, isolating the contribution of the intent-faithfulness reward.
|
||
|
||
## Intended use & limitations
|
||
|
||
Intended for **research on intent-aware, reasoning-based safety classification** and as a prompt-level moderation classifier. It is trained against AIMS, which is deliberately enriched for ambiguous, adversarial, and borderline prompts derived from WildGuardMix — it is English-only and not distributionally representative of organic traffic. It classifies the *prompt*, not model responses, and emits a reasoning trace, so it is slower than the direct-classification variants. Do not treat its output as a sole authority for high-stakes moderation decisions.
|
||
|
||
## License
|
||
|
||
This model is a fine-tune of Llama 3.1 and is therefore governed by the [**Llama 3.1 Community License**](https://huggingface.co/meta-llama/Llama-3.1-8B-Instruct/blob/main/LICENSE) and the [Llama Acceptable Use Policy](https://www.llama.com/llama3_1/use-policy/). The underlying AIMS training data is released under ODC-BY and is additionally subject to the [AI2 Responsible Use Guidelines](https://allenai.org/responsible-use).
|
||
|
||
## Citation
|
||
|
||
```bibtex
|
||
@misc{aims_dataset,
|
||
title = {AIMS: Annotated Intents for Model Safety},
|
||
author = {Jazhyc and collaborators},
|
||
year = {2026},
|
||
howpublished = {\url{https://huggingface.co/datasets/Jazhyc/aims-safety-intents}}
|
||
}
|
||
```
|