ml-model-extraction

name	ml-model-extraction
description	Extract a functional clone of a black-box ML model via prediction API queries, and infer whether specific records were in the training set (membership inference).
allowed-tools	Bash Read Write
metadata	{"subdomain":"ai-security","when_to_use":"model extraction model stealing prediction API clone membership inference training data privacy ML model black-box oracle query budget","tags":"adversarial-ml, model-extraction, model-stealing, membership-inference, privacy, ml-api","mitre_attack":"T1213, T1119, T1590"}

ML Model Extraction and Membership Inference

Two closely related attacks against deployed ML prediction APIs:

• Model extraction (stealing): reconstruct a functionally equivalent clone of the target model using only its input/output pairs, without access to weights, architecture, or training data. (Tramer et al. 2016, "Stealing Machine Learning Models via Prediction APIs")
• Membership inference: determine whether a specific record was used to train the target model, breaching training-data confidentiality. (Shokri et al. 2017, "Membership Inference Attacks Against Machine Learning Models")

These are distinct from LLM prompt-extraction attacks: the target is a classical ML model (logistic regression, decision tree, DNN) or an ML-as-a-service endpoint (AWS SageMaker, GCP AutoML, Azure ML, custom REST API).

Authorized use only. Query-volume attacks against commercial ML APIs may violate terms of service, the CFAA, and GDPR Article 22. Confirm scope includes data-privacy testing of the prediction endpoint before proceeding.

---

ATT&CK Mapping

Technique	Use
T1213 — Data from Information Repositories	Reconstructing training data via repeated model queries
T1119 — Automated Collection	Systematic API querying to build a clone dataset
T1590 — Gather Victim Network Information	Fingerprinting the ML service to determine model family

---

1. Reconnaissance — profile the target API

# 1a. Determine output type
curl -s -X POST "$API/predict" -d '{"features": [0,0,0]}' | jq .
# Does it return: hard label only? Confidence scores? Probability distributions?
# Probabilities = highest information; enables equation-solving extraction

# 1b. Infer model type from decision-boundary shape
# Binary classifier: probe boundary by linear interpolation between two known-class samples
# Multi-class: vary one feature at a time, record where predicted class changes
# Decision-tree vs smooth boundary: tree-type models have axis-aligned boundaries

# 1c. Count features and valid ranges from API docs / error messages
curl -s -X POST "$API/predict" -d '{}' | jq .error
# Validation errors often expose expected feature names and types

# 1d. Estimate query budget (cost / rate limit)
# Most commercial APIs: ~$0.0001–0.001 per query; 10k–1M queries typical for extraction

---

2. Model extraction

2a. Equation-solving attack (logistic/linear models)

For models returning exact probability scores, d features → d+1 queries suffice (Tramer):

import numpy as np
from scipy.optimize import fsolve

def extract_logreg(query_fn, n_features, n_classes):
    """
    Solve for weight matrix W using the system of equations:
    softmax(W @ x_i) = query_fn(x_i)
    n_features * n_classes unknowns, solved via n_features * n_classes + 1 queries.
    """
    queries = [np.random.randn(n_features) for _ in range(n_features * n_classes + 1)]
    scores = [query_fn(q) for q in queries]  # shape: (n_queries, n_classes)
    # Fit a local logistic regression to (queries, scores) using soft targets
    from sklearn.linear_model import LogisticRegression
    surrogate = LogisticRegression(multi_class='multinomial', max_iter=1000)
    labels = [np.argmax(s) for s in scores]
    surrogate.fit(queries, labels)
    return surrogate  # clone with equivalent decision boundary

# Verify: clone accuracy vs target on held-out natural inputs should be >95%

2b. Path-finding attack (decision trees)

Decision trees are fully recoverable with O(n_leaves * depth) queries via adaptive path traversal:

def extract_decision_tree(query_fn, feature_bounds, max_depth=10):
    """
    Recursively find split thresholds by binary search.
    For each feature, probe the boundary where the predicted class changes.
    """
    from sklearn.tree import DecisionTreeClassifier
    X_queries, y_labels = [], []

    # Adaptive sampling: grid + boundary refinement
    for _ in range(5000):
        x = np.random.uniform(feature_bounds[:, 0], feature_bounds[:, 1])
        y = np.argmax(query_fn(x))
        X_queries.append(x)
        y_labels.append(y)

    surrogate = DecisionTreeClassifier(max_depth=max_depth)
    surrogate.fit(X_queries, y_labels)
    return surrogate

2c. Neural network surrogate (DNN, general)

For deep models, train a surrogate network on (query, prediction) pairs using soft labels:

import torch, torch.nn as nn

def extract_dnn_surrogate(query_fn, input_dim, n_classes, query_budget=50000):
    """
    Knockoff Nets approach: sample diverse inputs, label with target API, train clone.
    """
    # Step 1: collect a representative query dataset
    # Use natural images / domain data if available; random noise also works for coarse clones
    X_steal = collect_diverse_inputs(query_budget)         # shape: (N, input_dim)
    y_soft  = np.array([query_fn(x) for x in X_steal])    # soft labels from API

    # Step 2: train surrogate with knowledge distillation on soft labels
    surrogate = build_surrogate_nn(input_dim, n_classes)
    optimizer = torch.optim.Adam(surrogate.parameters(), lr=1e-3)
    kl_loss = nn.KLDivLoss(reduction='batchmean')

    for epoch in range(50):
        for xb, yb in batches(X_steal, y_soft):
            pred = torch.log_softmax(surrogate(xb), dim=1)
            loss = kl_loss(pred, torch.tensor(yb, dtype=torch.float32))
            loss.backward(); optimizer.step(); optimizer.zero_grad()

    return surrogate

# Reference implementation: https://github.com/tribhuvanesh/knockoffnets

2d. Using Steal-ML (Tramer reference implementation)

git clone https://github.com/ftramer/Steal-ML
cd Steal-ML
pip install -r requirements.txt

# Configure target: edit config.py with your API endpoint
# Modes: logreg, mlp, dtree, svm
python steal.py --model logreg --n_queries 5000 --target $API_URL
# Outputs: cloned_model.pkl + fidelity score (agreement with target)

---

3. Membership inference

Determine whether a specific record (x, y) was used to train the target model.

3a. Shadow model attack (Shokri et al.)

# Principle: models generalize less to non-training data.
# Train "shadow models" on datasets similar to the target's.
# Train an "attack classifier" on (shadow_model_outputs, member/non-member).
# Apply attack classifier to target model outputs for test records.

from sklearn.ensemble import RandomForestClassifier

def membership_inference_attack(target_query_fn, shadow_models, shadow_train_sets,
                                 shadow_test_sets, candidate_records):
    """
    shadow_models: list of models trained on shadow_train_sets
    shadow_train_sets: data used to train each shadow model (members)
    shadow_test_sets: data NOT used (non-members)
    """
    # Build meta-training set for attack model
    X_meta, y_meta = [], []
    for model, train_set, test_set in zip(shadow_models, shadow_train_sets, shadow_test_sets):
        for x in train_set:
            X_meta.append(model.predict_proba([x])[0])
            y_meta.append(1)  # member
        for x in test_set:
            X_meta.append(model.predict_proba([x])[0])
            y_meta.append(0)  # non-member

    attack_clf = RandomForestClassifier(n_estimators=100)
    attack_clf.fit(X_meta, y_meta)

    # Query target and classify
    results = {}
    for record in candidate_records:
        conf_vector = target_query_fn(record)
        results[record] = attack_clf.predict([conf_vector])[0]  # 1=member, 0=non-member
    return results

3b. Confidence-threshold attack (simple, no shadow models)

When a model overfits, training examples get very high confidence scores:

def threshold_membership_inference(query_fn, candidate_records, threshold=0.95):
    """
    Heuristic: if predicted confidence for the correct class exceeds threshold,
    the record was likely in the training set.
    Works best against overfitted models (small datasets, many epochs).
    """
    for record, true_label in candidate_records:
        scores = query_fn(record)  # confidence vector
        if scores[true_label] >= threshold:
            print(f"MEMBER (confidence {scores[true_label]:.3f}): {record[:5]}...")
        else:
            print(f"Non-member (confidence {scores[true_label]:.3f}): {record[:5]}...")

3c. Using ML Privacy Meter

pip install ml-privacy-meter
python - <<'EOF'
from privacymeter.audit import Audit
from privacymeter.information_source import InformationSource

# target_model: your model or API wrapper
# target_train: records suspected to be in training set
# target_test: records known NOT to be in training set (baseline)
audit = Audit(
    target=InformationSource(model=target_model, dataset=target_train),
    reference=InformationSource(model=target_model, dataset=target_test),
)
report = audit.prepare_privacy_risk_report()
report.save("membership_inference_report.pdf")
EOF

---

4. Query budget and operational considerations

Model type	Extraction queries	Clone fidelity	Key reference
Logistic regression (d features)	d+2 (exact)	~100%	Tramer 2016
Decision tree (depth k, d features)	O(d * 2^k)	~100%	Tramer 2016
DNN, surrogate+KD	10k–100k	80–95%	Knockoff Nets 2018
DNN, hard-label only	100k–1M	70–90%	PRADA 2019

Rate limit evasion: distribute queries across IPs/accounts; randomize timing; use domain-valid inputs to avoid anomaly detection on the API side.

---

5. Validation

# Measure clone fidelity: agreement rate on held-out inputs
python - <<'EOF'
import numpy as np

X_test = load_held_out_test_set()
target_labels  = [query_target_api(x) for x in X_test]
surrogate_labels = [surrogate_model.predict([x])[0] for x in X_test]
fidelity = np.mean(np.array(target_labels) == np.array(surrogate_labels))
print(f"Clone fidelity: {fidelity:.3%}")
# >80% = functional clone suitable for transfer attacks
# >95% = near-exact extraction, sufficient for IP theft demonstration
EOF

---

6. Detection signals (defender perspective)

• Unusually high query volume from a single account/IP, especially systematic feature sweeps
• Queries that span the full input space uniformly (non-natural distribution)
• Confidence scores accessed at rate inconsistent with normal application usage
• PRADA defence: detect surrogate-training patterns via query-sequence analysis
• Watermark model outputs: embed a statistical fingerprint in confidence scores that survives cloning and allows ownership proof
• Differential privacy in training reduces membership inference signal at the cost of accuracy