🔐 AI/ML Security Lab – 🎭 Model Extraction 🐍

A prediction interface reveals something about a model every time it answers a question. With enough input/output pairs, another model may learn to imitate parts of that behaviour.

Model extraction aims to reproduce a target model's functionality through observations of its predictions. This differs from evasion, which changes inputs to influence an individual decision. Tramèr et al., model extraction research

This lab uses synthetic data and a controlled local classifier. The goal is to demonstrate behavioural approximation and evaluate a simple query-budget control.

🟢 Build the Local Target

Train a Random Forest on a synthetic two-moon dataset with two numerical features:

X, y = make_moons(n_samples=1_500, noise=0.2, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(
    X, y, test_size=0.25, stratify=y, random_state=42
)
target = RandomForestClassifier(
    n_estimators=100, max_depth=6, random_state=42, n_jobs=1
).fit(X_train, y_train)

The target achieves 97.07% accuracy on its held-out labelled test set in this run. These are toy classes, not validated categories of real security activity.

The extraction routine receives only predicted class labels. It does not use target parameters, original training samples, or original labels. Because everything lives in one notebook, this is a simulated access restriction rather than a secure isolation boundary.

🔴 Collect Labels from New Inputs

Generate 1,000 fresh points within a known, permitted two-dimensional region, then obtain the target's labels:

query_points = np.random.default_rng(7).uniform(
    LOW, HIGH, size=(1_000, 2)
)
api = LocalPredictionAPI(target.predict, limit=1_000)
query_labels = api.predict(query_points)

Each input row counts as one query, including when submitted in a batch. The interface validates the input shape, finite values, and allowed range before returning labels.

Use nested prefixes of the same pool for budgets of 25, 100, 250, and 1,000 inputs. The main extraction experiment therefore consumes 1,000 queried rows in total.

🔍 Train a Substitute

Train a small decision tree using the queried inputs and the target's answers:

substitute = DecisionTreeClassifier(max_depth=6, random_state=42)
substitute.fit(query_points[:budget], query_labels[:budget])

The substitute learns an approximation of the target's decisions. Its architecture is different from the Random Forest, which reinforces that similar predictions do not imply identical internal structure.

We then evaluate on 5,000 fresh points drawn independently from the same region:

fidelity = accuracy_score(
    reference_labels,
    substitute.predict(evaluation_points),
)

Fidelity measures agreement with the target, not correctness against ground truth. A substitute can reproduce the target's mistakes.

The owner computes the evaluation labels separately. They never enter substitute training or model selection, and these additional predictions are outside the extraction budget. The decision-boundary plots are also owner-side diagnostics.

📊 Observe the Results

With 1,000 queried inputs, the substitute reaches 97.86% agreement on the fresh evaluation set. The plots show similar, but visibly different, decision regions.

This does not prove recovery of the original weights or training data. It also does not establish fidelity outside the sampled region. The results come from one fixed-seed, low-dimensional experiment; more queries need not improve fidelity at every step for every model.

🛡️ Enforce a Query Budget

The local interface checks the total number of predicted rows before accepting a request:

if self.queries_used + len(points) > self.limit:
    raise RuntimeError("Query budget exceeded")

labels = self._predict(points)
self.queries_used += len(points)
return labels

In a separate demonstration, a fresh interface permits 250 rows. We query those rows, train a substitute, and attempt one additional input.

The additional input is rejected, and the counter remains at 250. Assertions verify both behaviours. This defence demonstration consumes a separate allowance of 250 rows.

However, the substitute already achieves 97.1% agreement within that allowance. The cap constrains collection; it does not remove information already returned.

✅ What the Defence Does—and Does Not—Show

The verified result is budget enforcement for one simulated client. It is not proof that extraction has been prevented.

Production controls require server-side accounting and identity management. Access controls, per-client budgets, and query-pattern monitoring may help manage exposure, but their effectiveness needs evaluation. The local counter does not model coordinated querying or multiple identities.

Returning fewer details can reduce information exposure, but this lab already uses labels alone. Omitting probability scores is therefore not a complete defence. Model extraction research

📌 Security Takeaways

• Repeated predictions reveal aspects of decision behaviour.
• Queried input/label pairs can train a substitute without original training data.
• Evaluate fidelity on fresh inputs, separately from model accuracy.
• Count input rows rather than API calls when enforcing a prediction budget.
• Verify controls while measuring what an allowed query budget still reveals.
• Treat high local fidelity as evidence of approximation, not exact recovery.

The accompanying notebook contains the complete runnable lab, checks, plots, and recorded results.