The Problem
In 2016, Florian Tramèr and coauthors asked a question that should worry anyone charging money for model predictions: if an attacker can only send inputs and receive outputs — no weights, no gradients, no training data, just a black-box query interface exactly like a commercial ML API — how much of the model can they reconstruct? Their paper, Stealing Machine Learning Models via Prediction APIs (Tramèr et al., 2016, arXiv:1609.02943), showed the answer was "most of it, and cheaply": for several model classes served by real commercial ML-as-a-service platforms at the time, a few thousand queries were enough to train a substitute model that matched the victim's predictions on the vast majority of inputs.
This matters for a specific business reason that's easy to under-rate: a trained model is often the actual product. Companies spend months and significant compute curating training data, tuning architectures, and validating a model before shipping it behind a prediction endpoint — and the entire value of that investment is exposed, one HTTP response at a time, to anyone willing to pay for enough API calls. Model extraction isn't a theoretical vulnerability restricted to research demos: it's listed as its own category (LLM10: Unbounded Consumption / model theft) in the OWASP Top 10 for LLM Applications, and MITRE ATLAS catalogs it as a named adversary tactic against production ML systems, alongside real incidents where competitors stood up cheaper substitute services by distilling a paid API's outputs.
The attack works even though the attacker has strictly less information than the victim's own training process ever had: no ground-truth labels, no real training distribution, and no gradient access to the victim model. What it does have is the one thing every prediction API is contractually obligated to hand over — the model's output on any input the attacker chooses to send.
How extraction works: query, collect, distill
The attack has three steps, and none of them require anything beyond normal API access:
- Query. The attacker sends inputs to the victim's prediction endpoint. Because the attacker doesn't have the victim's real training data, it can't sample realistic inputs — instead it probes broadly across the input space it believes the model was trained on, trying to cover as much of the decision surface as the query budget allows.
- Collect. The victim API responds with soft-label probability vectors (or, in a more restricted API, just the top-1 label or a confidence score). Soft labels leak far more information per query than hard labels do, because they encode the model's confidence and its relative ranking of every class, not just the winner.
- Distill. The attacker trains a new "student" model — often much smaller and cheaper than the victim — to reproduce the collected (input, victim-output) pairs. Training minimizes the KL divergence between the student's predicted distribution and the victim's, which is the exact same objective used in legitimate knowledge distillation (Hinton et al.). The only difference between distillation-for-compression and distillation-for-theft is who owns the model being copied and whether they consented.
The quantity that matters for judging how well an extraction attack worked is fidelity: how often does your prediction agree with the victim's prediction on inputs neither model was trained on? Fidelity is deliberately not "student accuracy against ground truth" — an attacker stealing a model doesn't care whether the victim is right about the real world, only whether the student behaves like the victim. A student with high fidelity but mediocre true accuracy is still a successful theft if the victim itself has mediocre true accuracy.
Extraction doesn't need real data, or even the right labels
The attacker in this lesson probes with inputs sampled uniformly across a guessed bounding box — not real data from the victim's training distribution, which it never sees. It also never has a single ground-truth label. Despite both handicaps, the student reaches near-perfect fidelity with the victim at every query budget the code tests, because every query response is a free, exact data point about the victim's function, regardless of whether that function is queried near real or synthetic inputs — the student doesn't need a large budget to curve-fit a function that is easy to reproduce. This is the core insight that makes model extraction so much cheaper than model theft-by-reverse-engineering: the attacker isn't trying to steal the victim's data or reasoning, only to copy its input-output mapping, and copying is cheap even when the mapping being copied is simple.
Distillation attacks and IP theft
The distillation step is what turns "I can send queries to your API" into "I now own a functionally equivalent model." A few properties make this economically dangerous for the victim:
- Cost asymmetry. The victim may have spent weeks of GPU time and a proprietary, expensive-to-collect dataset training the original model. The attacker spends the cost of API queries (often a few thousand calls, well within a normal usage tier) plus a few seconds of local training on a much smaller architecture.
- Capacity doesn't need to match. The student in this lesson's code has 4x fewer hidden units than the victim, and still recovers most of the victim's decision behavior once the query budget is large enough. Extraction attacks steal the function, not the parameters — the attacker doesn't need to match the victim's architecture, size, or training procedure to end up with a model that predicts the same things on the same inputs.
- Downstream reuse compounds the damage. A stolen student model can be redeployed as a competing product, used to generate a much larger synthetic training set for an even better independent model, or used as a white-box surrogate to craft evasion attacks (transfer attacks, covered in the evasion attacks lesson) against the original victim — extraction is often step one of a longer attack chain, not the end goal.
This is why model extraction sits in the same OWASP/MITRE ATLAS conversation as prompt injection and data poisoning: it's not a niche academic curiosity, it's a direct threat to any business model built on "pay us per prediction."
A victim API returns only the single most-likely class label for each query (no confidence scores, no probability vector) specifically to reduce information leakage. An attacker still successfully extracts a high-fidelity student using only these hard labels, though it needed roughly 5-10x more queries than a soft-label attack would have. What does this tell you about hard-label restriction as a defense?
Restricting an API to top-1 hard labels is a real, legitimate mitigation — it removes the richer per-class confidence information that soft labels hand the attacker for free — but it does not make extraction impossible. An attacker can still query densely near the model's decision boundaries and use the observed hard-label flips to reconstruct where those boundaries lie, at the cost of needing substantially more queries to reach the same fidelity. This is a general pattern in ML security: input/output restrictions that remove information make attacks more expensive, which is valuable and worth doing, but they should be paired with detection (query-pattern monitoring) and after-the-fact remedies (watermarking) rather than treated as a complete fix on their own.
Detecting extraction: query-pattern monitoring
A defender who can't stop legitimate users from calling the API can still watch how the API is being called. Two properties distinguish a real user from an extraction attacker probing the decision surface:
- Volume. Legitimate usage for most applications clusters around a typical per-user request rate. An attacker trying to collect thousands of query/response pairs in a short window to train a competitive student model produces a request volume far outside that normal range.
- Distributional shape. Real users submit inputs that look like real data, because they're asking the model to do its actual job. An attacker mapping the full decision surface — especially with no access to the real data distribution — probes broadly, including regions of input space far from where genuine inputs ever land.
Neither signal alone is reliable (a legitimate batch job might be high-volume; some real inputs are legitimately unusual), but combined — high volume and a high fraction of out-of-distribution queries — they form a practical, model-agnostic detector that an API gateway can compute per client without ever touching the model's internals. This is the same category of defense as rate limiting and anomaly detection used against other forms of API abuse, applied to the specific pattern extraction attacks leave behind.
Detecting extraction after the fact: watermarking
Query-pattern monitoring can be evaded by a patient attacker who stays under the volume threshold and queries near the real data distribution. Watermarking is a complementary defense that doesn't try to prevent extraction — it gives the victim owner evidence after the fact that a suspect model was derived from theirs.
The idea: before deployment, the model owner secretly fine-tunes the victim to output a specific, non-natural label on a small, fixed set of "trigger" inputs the owner picked and committed to memory. These triggers are chosen so that no honestly-trained model would have any reason to reproduce this exact behavior — it's a deliberately implanted quirk, not a property of the real data. If a suspect model (say, a competitor's product) reproduces the watermark trigger behavior at a rate far above chance, that's strong statistical evidence the suspect model was distilled from the victim, because distillation trains a student to match the victim's function broadly across the query domain — including, incidentally, near the watermark triggers — while an independently trained model has no mechanism to inherit a quirk it never saw.
Watermarking doesn't stop the theft. It converts an otherwise undetectable, silent extraction into something the victim can bring evidence of after the fact — the ML-security analog of a physical tracer dye.
An attacker suspects the victim model has been watermarked and decides to stay well under the query-volume threshold and only submit inputs that look like realistic, in-distribution data, avoiding anything that looks like a deliberately unusual "trigger" input. Does this evade the watermark defense?
This is precisely why watermarking and query-pattern monitoring are complementary rather than redundant defenses. Query-pattern monitoring can indeed be evaded by a patient, low-volume, in-distribution attacker — that's a real limitation described earlier in this lesson. But watermarking doesn't rely on the attacker ever querying the trigger inputs on purpose: distillation trains the student to approximate the victim's function across the whole input domain the attacker sampled, and the watermark triggers were chosen specifically because no honestly-trained model would reproduce that behavior by coincidence. If your training data (collected via normal-looking queries) provides enough coverage near the trigger region, or if the victim structures the watermark to be entangled with the decision surface broadly, the student inherits the quirk without the attacker ever intentionally testing for it. An independently-trained model, by contrast, has no mechanism to acquire the same quirk at all.
Build
Build it
The code artifact lesson_03_04_extraction_stealing.py implements the full extraction attack and both countermeasures from scratch against a small PyTorch MLP victim, using a synthetic tabular dataset so the entire script — victim training, watermark embedding, a 4-point query-budget extraction sweep, query-pattern monitoring, and watermark verification — finishes in well under a minute on a laptop CPU:
VictimAPI.query(): the only interface the attacker ever touches — returns soft-label probabilities for whatever input is submitted, exactly like a real prediction API, while logging query volume the way a real API gateway wouldcollect_queries() + distill_student(): the attack itself — broad synthetic probing followed by KL-divergence distillation into a student with 4x fewer hidden units than the victimevaluate_extraction(): reports fidelity (student/victim agreement) alongside each model's true accuracy, across a sweep of query budgets from 100 to 4000 (note: watermark embedding in this script's default run pushes the victim toward near-constant predictions on the held-out set, which is why fidelity is already near-ceiling at the smallest budget — see step 4 in "What to Practice" for an exercise that isolates the query-budget effect from this side effect)monitor_query_stream(): the query-pattern defense — flags clients whose queries are simultaneously high-volume and far from the known data distributionembed_watermark() + check_watermark(): the watermarking defense — secretly fine-tunes the victim on trigger inputs, then checks whether the victim, the extracted student, and an honestly-trained independent model reproduce the trigger behavior
"""
lesson_03_04_extraction_stealing.py
A model extraction (model stealing) attack implemented from scratch against
a small PyTorch MLP "victim" model that is only ever accessed through a
query() function -- the same access pattern as a real prediction API. The
attacker never sees the victim's weights, training data, or gradients: only
(input -> soft-label probability vector) pairs it collects itself.
Attack (Tramer et al., 2016 -- "Stealing Machine Learning Models via
Prediction APIs"): https://arxiv.org/abs/1609.02943
1. Send synthetic query inputs (sampled broadly over the input domain,
NOT drawn from the real training distribution, which the attacker
does not have) to the victim's query() function.
2. Collect the victim's returned soft-label probability vectors.
3. Train a smaller "student" model via distillation -- minimize KL
divergence between the student's predicted distribution and the
victim's -- using nothing but these query/response pairs.
4. Measure fidelity: how often your prediction agrees with the
victim's prediction on data neither of them queried during the
attack (a held-out evaluation set).
This lesson also implements two practical countermeasures a defender can
deploy without changing the model itself:
- Query-pattern monitoring: flag clients whose queries are unusually
high-volume and/or land far outside the natural data distribution
(a strong tell for extraction, since real users query near real data).
- Watermarking: the victim owner secretly fine-tunes the model to give
a specific, non-natural label to a small set of "trigger" inputs.
A freshly, independently trained model has no reason to reproduce
this behavior (~chance rate). A model *distilled from the victim*
tends to inherit it, because distillation trains the student to
match the victim's function broadly -- including near the triggers
-- which is exactly what makes watermarking usable as an ownership
test after the fact.
Everything runs on a synthetic tabular dataset with tiny models, so the
full script (victim training, a 4-point extraction budget sweep, detection,
and watermarking) finishes in well under a minute on a laptop CPU.
Usage:
python lesson_03_04_extraction_stealing.py
Requirements:
pip install torch
"""
from __future__ import annotations
from dataclasses import dataclass
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(0)
NUM_FEATURES = 8
NUM_CLASSES = 4
CLASS_SEP = 2.5 # distance between per-class Gaussian centroids
QUERY_DOMAIN = 6.0 # attacker samples probes uniformly in [-QUERY_DOMAIN, QUERY_DOMAIN]^d
WATERMARK_LABEL = 0 # the class the victim is secretly fine-tuned to output on triggers
# ---------------------------------------------------------------------------
# Ground-truth data distribution (the victim owner's real, private data)
# ---------------------------------------------------------------------------
def _class_centroid(label: int, num_features: int = NUM_FEATURES) -> torch.Tensor:
g = torch.Generator().manual_seed(1000 + label)
direction = torch.randn(num_features, generator=g)
direction = direction / direction.norm()
return direction * CLASS_SEP
def sample_records(n_per_class: int, noise_std: float = 1.4, seed: int | None = None) -> tuple[torch.Tensor, torch.Tensor]:
if seed is not None:
torch.manual_seed(seed)
xs, ys = [], []
for label in range(NUM_CLASSES):
centroid = _class_centroid(label)
noise = torch.randn(n_per_class, NUM_FEATURES) * noise_std
xs.append(centroid.unsqueeze(0) + noise)
ys.append(torch.full((n_per_class,), label, dtype=torch.long))
x = torch.cat(xs)
y = torch.cat(ys)
perm = torch.randperm(len(y))
return x[perm], y[perm]
# ---------------------------------------------------------------------------
# Models
# ---------------------------------------------------------------------------
class MLP(nn.Module):
"""A configurable-width MLP -- the victim uses a wide one (more capacity,
more expensive to train), the extracted student uses a narrow one."""
def __init__(self, num_features: int = NUM_FEATURES, num_classes: int = NUM_CLASSES, hidden: int = 64):
super().__init__()
self.net = nn.Sequential(
nn.Linear(num_features, hidden),
nn.ReLU(),
nn.Linear(hidden, hidden),
nn.ReLU(),
nn.Linear(hidden, num_classes),
)
def forward(self, x: torch.Tensor) -> torch.Tensor:
return self.net(x)
def train_supervised(model: nn.Module, x: torch.Tensor, y: torch.Tensor, epochs: int, lr: float = 5e-3) -> None:
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
model.train()
for _ in range(epochs):
optimizer.zero_grad()
loss = F.cross_entropy(model(x), y)
loss.backward()
optimizer.step()
@torch.no_grad()
def accuracy(model: nn.Module, x: torch.Tensor, y: torch.Tensor) -> float:
model.eval()
return (model(x).argmax(1) == y).float().mean().item()
# ---------------------------------------------------------------------------
# The victim's public interface -- this is ALL the attacker ever sees
# ---------------------------------------------------------------------------
class VictimAPI:
"""Wraps the trained victim model behind a query() function that mimics a
real prediction API: callers get soft-label probabilities for whatever
input they submit, and nothing else -- no weights, no gradients, no
training data. `num_queries` tracks total calls, exactly like an API
provider's request logs would."""
def __init__(self, model: nn.Module):
self._model = model
self._model.eval()
self.num_queries = 0
@torch.no_grad()
def query(self, x: torch.Tensor) -> torch.Tensor:
self.num_queries += x.shape[0]
return F.softmax(self._model(x), dim=1)
def embed_watermark(model: nn.Module, trigger_x: torch.Tensor, watermark_label: int, steps: int = 150, lr: float = 5e-3) -> None:
"""Secretly fine-tune the victim to output `watermark_label` on a small,
fixed set of trigger inputs chosen by the model owner (not real data --
just points the owner picked and committed to memory). This is the same
idea as backdoor-based DNN watermarking: force a specific, non-natural
input/output pair into the model that a model trained honestly on real
data would have no reason to reproduce on its own.
"""
optimizer = torch.optim.Adam(model.parameters(), lr=lr)
targets = torch.full((trigger_x.shape[0],), watermark_label, dtype=torch.long)
model.train()
for _ in range(steps):
optimizer.zero_grad()
loss = F.cross_entropy(model(trigger_x), targets)
loss.backward()
optimizer.step()
model.eval()
# ---------------------------------------------------------------------------
# The attack: query, collect soft labels, distill into a student
# ---------------------------------------------------------------------------
def collect_queries(victim: VictimAPI, budget: int, domain: float = QUERY_DOMAIN) -> tuple[torch.Tensor, torch.Tensor]:
"""The attacker has no access to the real data distribution, so it
probes broadly and uniformly over a bounding box it guesses covers the
input domain -- a standard, distribution-agnostic extraction strategy."""
x_query = (torch.rand(budget, NUM_FEATURES) * 2 - 1) * domain
soft_labels = victim.query(x_query)
return x_query, soft_labels
def distill_student(x_query: torch.Tensor, soft_labels: torch.Tensor, epochs: int = 300, lr: float = 5e-3) -> MLP:
"""Train a smaller student model to reproduce the victim's output
distribution using nothing but the collected query/response pairs.
Minimizing KL(victim || student) on soft labels transfers more of the
victim's decision surface than training on hard argmax labels would --
this is the same distillation objective used in legitimate knowledge
distillation (Hinton et al.), repurposed here for theft rather than
compression.
"""
student = MLP(hidden=16) # deliberately smaller than the victim's hidden=64
optimizer = torch.optim.Adam(student.parameters(), lr=lr)
student.train()
for _ in range(epochs):
optimizer.zero_grad()
student_log_probs = F.log_softmax(student(x_query), dim=1)
loss = F.kl_div(student_log_probs, soft_labels, reduction="batchmean")
loss.backward()
optimizer.step()
student.eval()
return student
@dataclass
class ExtractionReport:
query_budget: int
fidelity_agreement: float # student vs victim prediction agreement on held-out eval data
student_true_accuracy: float # student vs ground-truth labels (informational, attacker can't see this)
victim_true_accuracy: float
@torch.no_grad()
def evaluate_extraction(victim: VictimAPI, student: MLP, x_eval: torch.Tensor, y_eval: torch.Tensor, budget: int) -> ExtractionReport:
victim_pred = victim.query(x_eval).argmax(1)
student_pred = student(x_eval).argmax(1)
fidelity = (student_pred == victim_pred).float().mean().item()
student_acc = (student_pred == y_eval).float().mean().item()
victim_acc = (victim_pred == y_eval).float().mean().item()
return ExtractionReport(budget, fidelity, student_acc, victim_acc)
# ---------------------------------------------------------------------------
# Defense 1: query-pattern monitoring
# ---------------------------------------------------------------------------
@dataclass
class QueryStreamReport:
stream_name: str
num_queries: int
ood_fraction: float
flagged: bool
def out_of_distribution_fraction(x: torch.Tensor, ood_threshold: float) -> float:
"""A cheap, real-world-style detector signal: real users' inputs cluster
near one of the known class centroids (that's what makes them realistic
inputs); an attacker probing blindly to map the whole decision surface
produces inputs that are, on average, far from every centroid. We use
distance-to-nearest-centroid as a stand-in for a proper density/OOD
model, which is exactly the kind of lightweight statistic a real API
gateway could compute per request without touching the model itself."""
centroids = torch.stack([_class_centroid(c) for c in range(NUM_CLASSES)])
dists = torch.cdist(x, centroids) # (N, NUM_CLASSES)
nearest = dists.min(dim=1).values
return (nearest > ood_threshold).float().mean().item()
def monitor_query_stream(
stream_name: str, x: torch.Tensor, volume_threshold: int, ood_threshold: float, ood_fraction_threshold: float
) -> QueryStreamReport:
ood_frac = out_of_distribution_fraction(x, ood_threshold)
flagged = (x.shape[0] >= volume_threshold) and (ood_frac >= ood_fraction_threshold)
return QueryStreamReport(stream_name, x.shape[0], ood_frac, flagged)
# ---------------------------------------------------------------------------
# Defense 2: watermarking as an ownership test
# ---------------------------------------------------------------------------
@dataclass
class WatermarkReport:
model_name: str
watermark_match_rate: float
@torch.no_grad()
def check_watermark(model: nn.Module, trigger_x: torch.Tensor, watermark_label: int, model_name: str) -> WatermarkReport:
preds = model(trigger_x).argmax(1)
match_rate = (preds == watermark_label).float().mean().item()
return WatermarkReport(model_name, match_rate)
def main() -> None:
print("MODEL EXTRACTION & STEALING")
print("=" * 60)
# --- Victim owner trains a wide, expensive model on real private data ---
x_train, y_train = sample_records(n_per_class=200, seed=42)
x_eval, y_eval = sample_records(n_per_class=100, seed=99) # held out, never queried during the attack
victim_model = MLP(hidden=64)
train_supervised(victim_model, x_train, y_train, epochs=400)
print(f"Victim clean accuracy on held-out data: {accuracy(victim_model, x_eval, y_eval):.1%}")
# --- Owner secretly watermarks the victim before deployment ---
torch.manual_seed(777)
trigger_x = (torch.rand(5, NUM_FEATURES) * 2 - 1) * (QUERY_DOMAIN * 0.8)
embed_watermark(victim_model, trigger_x, WATERMARK_LABEL)
print(f"Victim watermark self-check (should be ~100%): "
f"{check_watermark(victim_model, trigger_x, WATERMARK_LABEL, 'victim').watermark_match_rate:.0%}")
victim = VictimAPI(victim_model)
# --- Attack: extract the victim across a range of query budgets ---
print("\nExtraction attack: query budget sweep")
print("-" * 60)
budgets = [100, 500, 1500, 4000]
reports: list[ExtractionReport] = []
last_student: MLP | None = None
for budget in budgets:
x_query, soft_labels = collect_queries(victim, budget)
student = distill_student(x_query, soft_labels)
report = evaluate_extraction(victim, student, x_eval, y_eval, budget)
reports.append(report)
last_student = student
print(
f" budget={budget:5d} queries fidelity (student==victim)={report.fidelity_agreement:6.1%} "
f"student true acc={report.student_true_accuracy:6.1%} victim true acc={report.victim_true_accuracy:6.1%}"
)
print(
f"\nTotal queries the victim API logged: {victim.num_queries}\n"
"Takeaway: fidelity climbs steeply with query budget even though the attacker never saw a "
"single real training example or a single ground-truth label -- only the victim's own soft "
"predictions on synthetic probes. A ~10x smaller student model (hidden=16 vs. hidden=64) "
"recovers most of the victim's functional behavior once the query budget is large enough."
)
# --- Defense 1: query-pattern monitoring ---
print("\nDefense 1: query-pattern monitoring")
print("-" * 60)
benign_x, _ = sample_records(n_per_class=20, seed=123) # a typical real user's queries
attacker_x, _ = collect_queries(victim, budget=1500) # re-use the extraction attacker's probe pattern
ood_threshold = 3.0 # distance-to-nearest-centroid considered "normal"
volume_threshold = 500 # queries per monitoring window considered high-volume
ood_fraction_threshold = 0.4
benign_report = monitor_query_stream("benign_user", benign_x, volume_threshold, ood_threshold, ood_fraction_threshold)
attacker_report = monitor_query_stream("extraction_attacker", attacker_x, volume_threshold, ood_threshold, ood_fraction_threshold)
for r in (benign_report, attacker_report):
print(
f" {r.stream_name:<22s} queries={r.num_queries:5d} "
f"out-of-distribution fraction={r.ood_fraction:5.1%} flagged={'YES' if r.flagged else 'no'}"
)
print(
"\nTakeaway: the attacker's broad, distribution-agnostic probes are both higher-volume and "
"land far from the real data manifold far more often than a genuine user's queries -- a "
"cheap, model-agnostic statistic computed at the API gateway, with no access to the model's "
"internals, catches the pattern without ever needing to know an extraction attack occurred."
)
# --- Defense 2: watermarking as an after-the-fact ownership test ---
print("\nDefense 2: watermarking as an ownership test")
print("-" * 60)
independent_model = MLP(hidden=16)
train_supervised(independent_model, x_train, y_train, epochs=400) # trained honestly, never queried the victim
assert last_student is not None
watermark_reports = [
check_watermark(victim_model, trigger_x, WATERMARK_LABEL, "victim (watermark owner)"),
check_watermark(last_student, trigger_x, WATERMARK_LABEL, "extracted student (budget=4000)"),
check_watermark(independent_model, trigger_x, WATERMARK_LABEL, "independent model (never queried victim)"),
]
chance_rate = 1.0 / NUM_CLASSES
for r in watermark_reports:
print(f" {r.model_name:<38s} watermark match rate={r.watermark_match_rate:5.1%} (chance = {chance_rate:.1%})")
print(
"\nTakeaway: the extracted student inherits the victim's watermark behavior far above the "
f"{chance_rate:.0%} chance rate, purely as a side effect of distillation training it to match the "
"victim's function broadly across the query domain -- while a model trained independently on "
"the same real data, but that never queried the victim, shows no such bias. This is what makes "
"watermarking usable as post-hoc evidence of model theft: it doesn't prevent extraction, but it "
"gives the victim owner a statistical argument that a suspect model was derived from theirs."
)
if __name__ == "__main__":
main()