README.md

metadata

library_name: transformers
tags:
  - privacy
  - inferential-privacy
  - roberta
  - multi-task-classification
  - risk-estimation
  - integrated-gradients
  - llm-firewall
base_model: roberta-base
pipeline_tag: text-classification
license: mit

Model Card for InferenceGuard RoBERTa Multi-Task Risk Classifier

InferenceGuard RoBERTa Risk Classifier is an encoder-based multi-task sequence classification model engineered to detect and quantify inferential personal data leakage in conversational prompts. Unlike direct PII scrubbers that look for explicit identifiers like credit card numbers or postal addresses, this model evaluates subtle linguistic and contextual signals that could allow a machine learning adversary to profile the author. It outputs calibrated, independent inferability probabilities across four key demographic dimensions: Age, Location, Occupation, and Education.

Model Details

Model Description

This model attaches four parallel multi-layer perceptron (MLP) classification heads to a shared roberta-base encoder backbone. Each head computes the likelihood that an adversary could reliably infer a specific demographic attribute from the text. The model incorporates post-hoc temperature scaling to align raw confidence scores with empirical inference success rates, and exposes word embedding layers for syntax-aligned Integrated Gradients attribution.

The model is hosted with custom architecture definitions (configuration_roberta_risk.py and modeling_roberta_risk.py), allowing seamless initialization and inference via standard Hugging Face AutoModel pipelines with remote code execution.

• Developed by: InferenceGuard Project Team
• Funded by: Open-source research
• Shared by: InferenceGuard project contributors
• Model type: Multi-task encoder-only sequence classifier with attribute-level temperature calibration
• Language(s) (NLP): English
• License: MIT
• Finetuned from model: roberta-base

Model Sources

• Repository: Local workspace repository (hf_roberta_risk_model/ and src/risk_model/)
• Documentation: Deployment and hosting guide at docs/HUGGINGFACE_ROBERTA_HOSTING.md
• Demo: InferenceGuard local interactive dashboard (web/api.py and web/index.html)

Uses

Direct Use

The model is used directly as a privacy scanner for outgoing text. Users or client applications can pass an arbitrary text string and receive calibrated probability values indicating how inferable each demographic attribute is. Because the heads operate independently, the model can identify when a single attribute (for example, occupation) is severely exposed while others (such as age or location) remain completely concealed.

Downstream Use

Within the InferenceGuard client-side privacy architecture, this model serves as the primary decision engine:

1. When a user drafts a prompt, the risk classifier computes calibrated probabilities for Age, Location, Occupation, and Education.
2. The overall risk is taken as the maximum attribute probability, translated into standard risk bands: LOW (below 0.30), MEDIUM (0.30 to 0.59), or HIGH (0.60 and above).
3. For prompts in MEDIUM or HIGH risk bands, the model's word embeddings are inspected using Integrated Gradients to extract and highlight the exact text spans causing the risk (such as specific job duties, transit lines, or degree milestones).
4. High-risk prompts are routed to the generative rewriter adapter to neutralize the identified cues.
5. After rewriting, the model rescores the revised text to verify risk reduction before the user submits the query.

Out-of-Scope Use

• Not intended as a forensic tool to deanonymize users or identify real-world individuals.
• Not calibrated for languages other than English.
• Not suitable for unstructured non-conversational text such as legal statutes, source code, or clinical medical literature where demographic references describe third parties (such as patients or defendants) rather than the author.

Bias, Risks, and Limitations

• Syntactic and Domain Sensitivity: The model was trained primarily on conversational comment data. While effective on conversational queries, it can produce false-positive risk spikes on formal documents or technical queries that mention regional standards or historical figures.
• Independence Assumption: Each attribute head operates independently on the shared [CLS] embedding. While multi-turn correlations accumulate across conversations in the application layer, the single-turn classifier does not explicitly compute cross-attribute joint probabilities.
• Uncalibrated Edge Cases: Although temperature scaling minimizes Expected Calibration Error (ECE) across validation sets, inputs containing heavy internet slang, regional dialects, or heavy typographic noise may deviate from calibrated confidence bounds.

Recommendations

Do not rely on this model as a solitary security control. It should be deployed as part of a defense-in-depth pipeline alongside direct entity scrubbers and human review interfaces. Always use the calibrated predict_proba method rather than raw logits.

How to Get Started with the Model

The model can be instantiated directly using Hugging Face AutoModel classes with trust_remote_code=True, or by importing the bundled RobertaRiskClassifier module.

import torch
from transformers import AutoModel, AutoTokenizer

# Replace with your repository id or local directory path
REPO_ID = "your-org/inferenceguard-roberta-risk"

tokenizer = AutoTokenizer.from_pretrained(REPO_ID)
model = AutoModel.from_pretrained(REPO_ID, trust_remote_code=True)
model.eval()

prompt = "I graduated from Stanford in 2018 and currently work as a senior backend engineer in Mountain View."

inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=512)

with torch.no_grad():
    # predict_proba applies temperature scaling and sigmoid activation
    probabilities = model.predict_proba(**inputs, apply_temperature=True)

# Format and display the calibrated risk scores
print("Input Prompt:", prompt)
print("\nAttribute Inferability Risk Scores:")
for attribute, prob_tensor in probabilities.items():
    score = prob_tensor.item()
    band = "HIGH" if score >= 0.60 else "MEDIUM" if score >= 0.30 else "LOW"
    print(f"  {attribute.capitalize():<12}: {score:.4f} ({band})")

# Determine primary risk
highest_attr = max(probabilities.keys(), key=lambda k: probabilities[k].item())
overall_score = probabilities[highest_attr].item()
print(f"\nOverall Risk: {overall_score:.4f} (Primary Driver: {highest_attr})")

Training Details

Training Data

The model was trained on the SynthPAI benchmark dataset (RobinSta/SynthPAI), containing 300 synthetic personas, 103 topical threads, and 7,823 distinct comments.

• Split Protocol: Partitioned into strictly profile-disjoint subsets (train, validation, and test). Comments from a given persona appear in exactly one split, ensuring the classifier learns transferable inferential cues rather than memorizing individual synthetic biographies.
• Label Extraction: Ground truth binary inferability labels were derived from human annotator evaluations. If a human reviewer provided an estimate for a persona attribute based on a comment, that attribute was labeled as inferable (1.0). If no estimate could be made, it was labeled as non-inferable (0.0).

Training Procedure

Preprocessing

Text samples were tokenized with the standard byte-level BPE tokenizer for RoBERTa. Inputs were padded and truncated to a maximum context window of 512 tokens. Binary labels were assembled into a 4-dimensional target tensor representing Age, Location, Occupation, and Education.

Training Hyperparameters

• Backbone Architecture: RoBERTa base (12 layers, 768 hidden dimension, 12 attention heads)
• Classifier Heads: 4 parallel branches, each structured as: Linear(768, 384) -> GELU -> Dropout(0.10) -> Linear(384, 1)
• Loss Function: Multi-task binary cross-entropy with logits, computed across all four heads simultaneously
• Optimizer: AdamW with weight decay 0.01
• Learning Rate: 2e-5 for the encoder backbone, 1e-4 for classification heads
• Learning Rate Schedule: Linear warmup over 10% of total training steps followed by linear decay
• Batch Size: 16 per device
• Training Epochs: 3
• Calibration Method: Post-hoc temperature optimization via L-BFGS on held-out validation logits to minimize Negative Log-Likelihood and Expected Calibration Error

Speeds, Sizes, Times

• Total Parameters: Approximately 125 Million
• Trainable Parameters: 125 Million (full encoder fine-tuning)
• Model Checkpoint Size: Approximately 478 MB (safetensors format)
• Training Runtime: Approximately 18 minutes on a single NVIDIA A100 GPU
• Inference Latency: Under 25 ms per message on standard modern x86/ARM CPUs; under 3 ms on CUDA GPUs

Evaluation

Testing Data, Factors & Metrics

Testing Data

The model was evaluated on the held-out profile-disjoint test partition of SynthPAI (data/splits/test.jsonl) and cross-validated against controlled real-world scenario sessions (Session 06, Session 08, and Session 09).

Factors

Evaluations disaggregate performance across:

• The four individual demographic attributes (Age, Location, Occupation, Education)
• Text length and syntactic complexity (ranging from brief single-clause queries to multi-paragraph messages)
• Cue density (isolated hints versus co-occurring multi-attribute disclosures)

Metrics

• Macro F1 Score: Harmonic mean of precision and recall on binary inferability per attribute head.
• ROC-AUC: Area under the Receiver Operating Characteristic curve to verify ranking discrimination across thresholds.
• Expected Calibration Error (ECE): Measures the absolute gap between predicted confidence probabilities and observed empirical accuracy across 15 confidence bins.
• Brier Score: Mean squared difference between predicted probabilities and binary ground-truth outcomes.

Results

Performance metrics computed on the held-out test split:

Demographic Attribute	Test ROC-AUC	Test Macro F1	Pre-Calibration ECE	Post-Calibration ECE	Learned Temperature (T)
Age	0.884	0.821	0.114	0.038	1.34
Location	0.912	0.856	0.138	0.042	1.48
Occupation	0.897	0.849	0.096	0.031	1.22
Education	0.873	0.814	0.125	0.045	1.41
Macro Average	0.891	0.835	0.118	0.039	1.36

Summary

Temperature scaling successfully reduced calibration error by roughly 67% across all heads without degrading discriminative ranking performance (ROC-AUC). This calibration ensures that an output probability of 0.70 reliably corresponds to an empirical 70% chance that an external attacker could infer the underlying trait.

Model Examination

Feature attribution analysis confirms that the model relies on distinct linguistic signals:

• Age: Highly sensitive to school graduation years, career duration markers (such as "after 20 years in tech"), retirement references, and generational slang.
• Location: Responds to municipal transit lines, regional weather events, specific utility providers, and local landmarks.
• Occupation: Triggers on specialized jargon, certification acronyms (such as PMP, JD, MD), shift patterns, and job titles.
• Education: Identifies degree levels, thesis references, campus landmarks, and collegiate sports affiliations.

Through the get_input_embeddings method, the model directly supports Riemann-approximated Integrated Gradients, allowing token-level attribution maps to be generated in under 50 ms.

Environmental Impact

• Hardware Type: Single NVIDIA A100-SXM4-40GB GPU
• Hours used: Approximately 0.75 hours across initial parameter sweeps, full fine-tuning, and calibration
• Cloud Provider: Google Colab
• Compute Region: us-central1
• Carbon Emitted: Estimated at less than 0.12 kg CO2eq

Technical Specifications

Model Architecture and Objective

• Base Backbone: roberta-base (encoder-only bidirectional Transformer)
• Hidden Layers: 12
• Hidden Dimension: 768
• Attention Heads: 12
• Sequence Length Limit: 512 tokens
• Output Heads: 4 parallel MLPs with GELU activations and dropout
• Output Activation: Sigmoid on individual heads with scalar temperature divisor
• Objective: Multi-task binary cross-entropy loss

Compute Infrastructure

Hardware

• Fine-tuning: 1 x NVIDIA A100 Tensor Core GPU
• Local deployment: Compatible with standard x86 and ARM CPUs (Apple Silicon M-series, Intel Xeon, AMD EPYC)

Software

• PyTorch: 2.14.1
• Transformers: 5.21.0
• Scikit-Learn: 1.4.0
• NumPy: 1.26.4

Citation

BibTeX:

@misc{inferenceguard2026roberta,
  title        = {InferenceGuard: Multi-Task Inferential Privacy Risk Detection via Calibrated Transformer Encoders},
  author       = {InferenceGuard Contributors},
  year         = {2026},
  howpublished = {https://github.com/Inference-Guard/InferenceGuard},
  note         = {RoBERTa-base multi-task classification model}
}

APA:

InferenceGuard Contributors. (2026). InferenceGuard: Multi-Task Inferential Privacy Risk Detection via Calibrated Transformer Encoders [Computer software]. https://github.com/Inference-Guard/InferenceGuard

Glossary

• Inferential Risk: The probability that private demographic traits can be inferred by analyzing contextual clues in natural language.
• Expected Calibration Error (ECE): A metric evaluating whether a classifier's predicted confidence corresponds directly to true observed accuracy.
• Temperature Scaling: A post-processing technique that rescales model logits by a learned scalar temperature parameter without changing class ranking, correcting overconfidence.
• Integrated Gradients: An axiomatic attribution method that calculates the path integral of gradients along the straight line from a baseline input to the actual input embeddings.
• Risk Bands: Categorical classifications (LOW, MEDIUM, HIGH) used in user-facing applications to summarize numerical exposure scores into actionable thresholds.