#Building a PII Classifier in 2 Hours for $3: A Teacher-Student Distillation Story

How I distilled Roblox’s 0.6B parameter PII classifier into a lightweight Llama 3.2 3B model using synthetic data and Fireworks AI

---

#Table of Contents

1. The Problem
2. The Solution: Distill the Knowledge
3. Architecture Deep Dive
4. The Key Insight: Tokens as Labels
5. The Math: Why This Works
6. Applying This to Our PII Classifier
7. Why Not Use a Separate Classification Head?
8. Lessons Learned
9. Results
10. Try It Yourself
11. What’s Next?
12. Acknowledgments

---

#The Problem

Moderating online platforms for PII (Personally Identifiable Information) is critical for user safety, especially in gaming environments where children interact. Roblox recently open-sourced their PII classifier, a powerful XLM-RoBERTa-Large model trained on billions of chat messages.

But there’s a catch: The Roblox model is 0.6B parameters and requires GPU inference. What if you want something:

• Faster for real-time moderation?
• Cheaper to deploy at scale?
• Available as a simple API with confidence scores?

Enter teacher-student distillation and the “tokens as labels” approach.

---

#The Solution: Distill the Knowledge

I built a pipeline that:

1. Generates synthetic chat messages using DeepSeek’s API (~$0.50)
2. Labels them with the Roblox teacher running on RunPod (~$1.50)
3. Fine-tunes a small student model on Fireworks AI (~$1.00)

Total cost: ~$3. Total time: ~2.5 hours (mostly waiting for data generation).

The result? 87.4% accuracy with a model that fits in a 46MB LoRA adapter.

---

#Architecture Deep Dive

┌──────────────────────────────────────────────────────────────────────────────┐
│                     THE TEACHER-STUDENT DISTILLATION PIPELINE                 │
└──────────────────────────────────────────────────────────────────────────────┘

     ┌─────────────┐                                                           
     │  DeepSeek   │  "Generate gaming chat messages that ask for personal    
     │    API      │   information like phone numbers, social media, etc."    
     └──────┬──────┘                                                           
            │                                                                  
            ▼                                                                  
   ┌────────────────────────────────────────────────────────────────────────┐  
   │                        5,000 SYNTHETIC MESSAGES                         │  
   │                                                                         │  
   │  "none":    "gg wp", "nice shot!", "rematch?"                          │  
   │  "asking":  "whats ur snap?", "drop ur discord", "u got insta?"        │  
   │  "giving":  "my snap is jake2024", "add me on dc: gamer#1234"          │  
   └────────────────────────────────────────────────────────────────────────┘  
            │                                                                  
            ▼                                                                  
     ┌─────────────────────────────────────────────────────────────────────┐   
     │                        ROBLOX TEACHER MODEL                          │   
     │                   (XLM-RoBERTa-Large, 0.6B params)                   │   
     │                                                                      │   
     │  Input: "whats ur instagram?"                                        │   
     │  Output: {                                                           │   
     │    "asking_sharing_score": 0.82,                                     │   
     │    "giving_pii_score": 0.15,                                         │   
     │    "sharing_contact_score": 0.71,                                    │   
     │    ...                                                               │   
     │  }                                                                   │   
     │                                                                      │   
     │  → Apply thresholds → Label: "asking"                                │   
     └─────────────────────────────────────────────────────────────────────┘   
            │                                                                  
            ▼                                                                  
   ┌────────────────────────────────────────────────────────────────────────┐  
   │                          LABELED TRAINING DATA                          │  
   │                                                                         │  
   │  {"messages": [                                                         │  
   │    {"role": "system", "content": "Classify PII..."},                   │  
   │    {"role": "user", "content": "Message: "whats ur instagram?""},    │  
   │    {"role": "assistant", "content": "asking"}   ← SINGLE TOKEN LABEL   │  
   │  ]}                                                                     │  
   └────────────────────────────────────────────────────────────────────────┘  
            │                                                                  
            ▼                                                                  
     ┌─────────────────────────────────────────────────────────────────────┐   
     │                         FIREWORKS SFT                                │   
     │                                                                      │   
     │  Base Model: Llama 3.2 3B Instruct                                   │   
     │  Method: LoRA (rank=16, alpha=32)                                    │   
     │  Epochs: 3                                                           │   
     │  Training Time: ~5 minutes                                           │   
     │  Cost: ~$1.00                                                        │   
     └─────────────────────────────────────────────────────────────────────┘   
            │                                                                  
            ▼                                                                  
     ┌─────────────────────────────────────────────────────────────────────┐   
     │                         STUDENT MODEL                                │   
     │                                                                      │   
     │  LoRA Adapter: 46MB                                                  │   
     │  Accuracy: 87.4%                                                     │   
     │  Inference: Single API call → "asking" (with logprobs)               │   
     │                                                                      │   
     │  → Deployed to HuggingFace: sugiv/sugiv-pii-classifier               │   
     └─────────────────────────────────────────────────────────────────────┘   

---

#The Key Insight: Tokens as Labels

The magic comes from Fireworks’ blog post on turning LLMs into classifiers.

Traditional classification: Train a separate classification head, get softmax probabilities.

Tokens as labels: Train the LLM to output a single token ("none", "asking", or "giving"). Then use logprobs to get calibrated probabilities!

response = client.chat.completions.create(
    model="my-classifier",
    messages=[...],
    max_tokens=5,
    logprobs=True,       # ← This is the magic
    top_logprobs=5
)

# Extract probability from log-probability
import math
token = response.choices[0].logprobs.content[0]
confidence = math.exp(token.logprob)  # e.g., 0.95 for "asking"

This approach is:

• Simpler: No custom architecture needed
• Calibrated: Logprobs naturally give you confidence scores
• Flexible: Works with any LLM that supports logprobs

---

#The Math: Why This Works (And Why No Renormalization Is Needed)

This section explains the mathematical foundation behind using LLM token probabilities as classifier outputs, based on Fireworks’ excellent appendix.

#PMF vs PDF: Discrete vs Continuous

Key insight: A PMF (Probability Mass Function) is for discrete outcomes; a PDF (Probability Density Function) is for continuous ones. Token probabilities are a PMF over a discrete vocabulary, and cross-entropy + gradient descent reshapes that PMF so almost all mass lies on your class tokens.

Think of three stages:

1. PMF vs PDF - understanding the difference
2. PMF over all tokens - before fine-tuning
3. PMF concentrated on class tokens - after training

#Stage 1: PMF vs PDF (Discrete vs Continuous)

PMF (discrete) – probability bars at separate points:

Probability
   ^
   |            *
   |      *     |     *
   |      |     |     |
   +----------------------------->  x
        1     2     3       (discrete values)
        p(1)  p(2)  p(3)

(sum of bar heights = 1)

PDF (continuous) – smooth curve over an interval:

Probability density
   ^
   |        /   |       /     |      /       |_____/____________________>  y
       a        b        (continuous range)

P(a ≤ Y ≤ b) = area under the curve between a and b

• PMF: probabilities at isolated points, sum to 1
• PDF: density over a continuum, integral (area) = 1

LLM token probabilities are a PMF: For input x, the model outputs probabilities p_t = P(next token = t | x) for each token t, and Σ p_t = 1 over all tokens in vocabulary V.

#Stage 2: Before Fine-Tuning - PMF Over Full Vocab

Let the vocabulary be tokens V (~32K-128K tokens), and class tokens be a small subset C ⊂ V.

Before fine-tuning, the PMF spreads mass across many tokens:

Tokens:  [ ... non-class ... | CLASS_1 | CLASS_2 | CLASS_3 | ... non-class ... ]

Probability
   ^
   |     *       *         *      *        *
   |   * |     * |       * |    * |      * |
   +-----------------------------------------------> tokens
       t1   t2   CLASS_1  t3  CLASS_2  CLASS_3 ...
         ^          ^          ^
         |          |          |
       non-class  class      class

Total class mass = p(CLASS_1) + p(CLASS_2) + p(CLASS_3) < 1
(there is non-trivial mass on many non-class tokens)

This is still a valid PMF: all bars ≥ 0, sum to 1. But class tokens do not yet dominate.

#Stage 3: After Cross-Entropy Training - Mass Concentrated on Class Tokens

Training objective: for classification prompts, make the correct class token very probable and push non-class tokens’ probabilities toward 0.

Over training:

• Logits for correct class tokens increase
• Logits for non-class tokens decrease
• The PMF gets reshaped so almost all mass is on class tokens

After sufficient fine-tuning:

Tokens:  [ ... non-class ... | CLASS_1 | CLASS_2 | CLASS_3 | ... non-class ... ]

Probability
   ^
   |                              *
   |                            * |   *
   |                          *   | * |
   |                        *     | | |
   +-----------------------------------------------> tokens
       t1   t2   CLASS_1  t3  CLASS_2  CLASS_3 ...
       ~0   ~0    ~0.6    ~0   ~0.3     ~0.1

Non-class bars ≈ 0
Class bars sum ≈ 1.0

So for classification prompts:

• Σ p_c ≈ 1 for c ∈ C (class tokens)
• Σ p_j ≈ 0 for j ∉ C (non-class tokens)

This is exactly a PMF over the label set C!

Now you can read the LM head’s probabilities for class tokens directly as calibrated class probabilities:

PII_labels = ["none", "asking", "giving"]

p("none")   ≈ 0.05
p("asking") ≈ 0.30
p("giving") ≈ 0.65

# Sum ≈ 1.0, non-class tokens ≈ 0

No extra renormalization step or separate classifier head is needed: the cross-entropy + gradient descent process has already concentrated the PMF’s mass onto your discrete class tokens.

---

#The Math: Cross-Entropy + Gradient Descent

Let’s define:

• V = full vocabulary (discrete outcomes, ~32K-128K tokens)
• C ⊂ V = class tokens (our labels: “none”, “asking”, “giving”)
• For input x, model logits z ∈ ℝ^|V| and PMF via softmax:

p_k = exp(z_k) / Σ exp(z_r)  for all r ∈ V

This defines a probability distribution (PMF) over tokens.

#The Cross-Entropy Loss

For one training example with correct label token c ∈ C, the cross-entropy loss is:

L = -log(p_c)

The gradient of this loss with respect to logits is:

∂L/∂z_c = p_c - 1    (negative when p_c < 1, pushes z_c UP)
∂L/∂z_j = p_j        (positive for j ≠ c, pushes z_j DOWN)

#Gradient Descent Updates

With learning rate η > 0, gradient descent updates logits:

z_c^(t+1) = z_c^t + η(1 - p_c)    ← correct token gets pushed UP
z_j^(t+1) = z_j^t - η·p_j         ← all others get pushed DOWN

#The Logit Gap Analysis

Define the logit gap between a non-class token j ∉ C and correct class token c:

Δ_j^t = z_j^t - z_c^t

Then after one update:

Δ_j^(t+1) = z_j^(t+1) - z_c^(t+1)
          = (z_j^t - η·p_j) - (z_c^t + η(1 - p_c))
          = Δ_j^t - η(1 + p_j - p_c)

Since probabilities sum to 1:

1 - p_c = Σ p_k for k ≠ c  ≥  p_j

So:

1 + p_j - p_c ≥ 1 + p_j - (1 - p_j) = 2·p_j

Therefore:

Δ_j^(t+1) ≤ Δ_j^t - 2η·p_j

Interpretation: If p_j > 0 (the model gives token j any probability), the gap Δ_j decreases by at least 2η·p_j each step.

#From Logit Gap to Probability Ratio

The gap relates to probabilities via:

p_j / p_c = exp(z_j) / exp(z_c) = exp(z_j - z_c) = exp(Δ_j)

So:

p_j = p_c · exp(Δ_j) ≤ exp(Δ_j)

As training proceeds:

1. If p_j remains non-zero, updates keep subtracting at least 2η·p_j from Δ_j
2. This forces Δ_j → -∞
3. Then exp(Δ_j) → 0
4. Thus p_j → 0

#The Key Conclusion

For any non-class token j ∉ C:

• Cross-entropy + gradient descent pushes its probability p_j toward zero
• Because the PMF must still sum to 1, the mass “removed” from non-class tokens flows onto the class tokens C
• So Σ p_c → 1 for c ∈ C

This is why raw class-token probabilities are “effectively calibrated without renormalization” for classification prompts. The training dynamics naturally concentrate all probability mass on your label tokens!

---

#Applying This to Our PII Classifier

#Step 1: Define Discrete Labels as Tokens

We chose a discrete label set:

C = {"none", "asking", "giving"}

Each label is a single vocabulary token. Our classification prompt ends so the next token should be one of these:

You are a PII classifier. Reply with exactly one word:
"none", "asking", or "giving".

Message: "<chat text>"

At generation time, the model’s next-token PMF p_t = P(next token = t | x) gives us a discrete distribution over the vocabulary.

#Step 2: Fine-tune on PII Labels with Cross-Entropy

We:

1. Take Llama 3.2 3B Instruct as the backbone
2. Add LoRA adapters to attention/MLP blocks (rank=16, alpha=32)
3. Freeze base weights
4. Train with standard cross-entropy on class tokens

For each training example:

• Correct token is the label token c ∈ {"none", "asking", "giving"}
• Loss L = -log(p_c)
• Backprop + gradient descent update only LoRA parameters

The math above applies: logits for correct class tokens get pushed up relative to non-class tokens, driving non-class probabilities toward 0.

#Step 3: Use LM-Head Probabilities as Calibrated Confidences

At inference:

response = client.chat.completions.create(
    model="my-pii-classifier",
    messages=[...],
    max_tokens=5,
    logprobs=True,      # Get probability distribution
    top_logprobs=5
)

# The logprob IS the calibrated confidence
import math
logprob = response.choices[0].logprobs.content[0].logprob
confidence = math.exp(logprob)  # e.g., 0.92 for "asking"

Because Σ p_c ≈ 1 for class tokens (non-class probabilities near 0), these p_c form a PMF over labels that is already normalized and empirically well-calibrated.

You can then:

• Threshold: Flag high-risk when P("giving" | x) > 0.8
• Rank: Choose label with max probability
• Monitor: Track calibration and adjust thresholds

---

#Why Not Use a Separate Classification Head?

The traditional approach:

Backbone (frozen) → Pooled representation → Linear head → Softmax → Labels

This works, but:

1. Requires new architecture: Define and wire up a classification head
2. Different training: Head-only fine-tuning vs full model
3. No logprobs API: Most inference APIs don’t expose intermediate representations

The LM-head-as-classifier approach:

1. Reuses existing infrastructure: Same model, same API
2. Native logprobs: Confidence scores come for free
3. Mathematically grounded: Cross-entropy naturally concentrates mass on class tokens

This is exactly what Fireworks designed their SFT pipeline around, and why we chose it for our PII classifier.

---

#Lessons Learned

#1. Not All LLMs Are Created Equal

I initially tried Qwen3-4B as the student model. It has an “extended thinking” mode that outputs <think>...</think> traces before answering:

<think>Let me analyze this message for PII content. The user is asking 
for a Snapchat handle, which is a social media platform where personal 
contact information is exchanged...</think>

asking

This resulted in 0% accuracy because my classification code was looking for a single-word response!

Solution: Switch to Llama 3.2 3B Instruct, which follows instructions precisely and outputs just the label.

#2. Firectl > REST API

Fireworks’ REST API for dataset uploads gave cryptic errors:

invalid character '-' in numeric literal

The CLI tool firectl worked flawlessly:

firectl create dataset pii-train data/train.jsonl
firectl create sftj --base-model accounts/fireworks/models/llama-v3p2-3b-instruct ...

#3. Teacher Thresholds Matter

The Roblox model outputs raw scores, not labels. I experimented with thresholds:

if asking_score >= 0.2:           # Sensitive to asking
    label = "asking"
elif giving_score >= 0.3:         # Slightly less sensitive
    label = "giving"
elif combined_max >= 0.2691:      # Catch edge cases
    label = "giving"
else:
    label = "none"

Getting these right is crucial for training data quality.

#4. Synthetic Data Works!

DeepSeek generated realistic gaming chat messages. The key was providing good few-shot examples:

prompt = """Generate 100 gaming chat messages. Examples:
- "gg" (none - just gaming)
- "whats ur snap?" (asking - requesting social media)
- "my insta is @gamer123" (giving - sharing contact info)

Make them realistic for gaming contexts (Roblox, Discord, etc.)"""

---

#Results

#Accuracy: 87.4%

On 500 held-out test examples:

Class	Precision	Recall	F1-Score
none	0.94	0.93	0.93
asking	0.75	0.66	0.70
giving	0.74	0.88	0.80

#Confusion Matrix

               Predicted
            none  asking  giving
Actual none  316     16      8     (93% correct)
    asking   15     57     15     (66% correct - hardest class)
    giving    6      3     64     (88% correct)

#Cost Breakdown

Step	Cost	Time
DeepSeek (5K messages)	$0.50	15 min
RunPod A40 (labeling)	$1.50	2 hrs
Fireworks SFT	$1.00	5 min
Total	$3.00	~2.5 hrs

---

#Try It Yourself

#Option 1: Use the Pre-trained Model

from transformers import AutoModelForCausalLM, AutoTokenizer
from peft import PeftModel

base_model = AutoModelForCausalLM.from_pretrained(
    "meta-llama/Llama-3.2-3B-Instruct",
    torch_dtype="auto",
    device_map="auto"
)
model = PeftModel.from_pretrained(base_model, "sugiv/sugiv-pii-classifier")
tokenizer = AutoTokenizer.from_pretrained("meta-llama/Llama-3.2-3B-Instruct")

messages = [
    {"role": "system", "content": 'Classify if this chat message involves PII. Reply with one word: "none", "asking", or "giving".'},
    {"role": "user", "content": 'Message: "whats your instagram?"'}
]
inputs = tokenizer.apply_chat_template(messages, return_tensors="pt").to(model.device)
outputs = model.generate(inputs, max_new_tokens=5, temperature=0)
print(tokenizer.decode(outputs[0][inputs.shape[1]:], skip_special_tokens=True))
# → "asking"

#Option 2: Train Your Own

Check out the full pipeline on GitLab.

---

#What’s Next?

1. More data: Scale to 50K+ examples for better minority class performance
2. Multilingual: The Roblox teacher supports 10+ languages
3. Calibration analysis: Measure reliability diagrams for production use
4. Smaller models: Try Llama 3.2 1B or phi-2 for even faster inference

---

#Acknowledgments

• Roblox for open-sourcing their PII classifier
• Fireworks AI for the SFT infrastructure and the excellent blog post
• Meta AI for Llama 3.2
• DeepSeek for affordable API access

---

Model available at: huggingface.co/sugiv/sugiv-pii-classifier

Code available at: gitlab.com/sugix/sugiv-pii-classifier