Query-Based Model Extraction

Step 1: Select Query Strategy
  Choose inputs that maximize information gain per query

Step 2: Query Target API
  Send crafted inputs, collect predictions (labels, probabilities, logits)

Step 3: Build Training Dataset
  Pair queries with target model responses as labeled data

Step 4: Train Substitute Model
  Use knowledge distillation to train a copy

Step 5: Evaluate Fidelity
  Measure agreement between substitute and target on held-out data

Step 6: Iterate
  Use active learning to query informative regions, improve fidelity

import torch
import torch.nn.functional as F

def distillation_loss(student_logits, teacher_probs, temperature=3.0):
    """Soft label distillation from teacher (target API) outputs."""
    student_soft = F.log_softmax(student_logits / temperature, dim=-1)
    loss = F.kl_div(student_soft, teacher_probs, reduction="batchmean")
    return loss * (temperature ** 2)

def extract_model(target_api, student_model, query_set, epochs=50):
    """Extract target model via API queries."""
    optimizer = torch.optim.Adam(student_model.parameters())

    for epoch in range(epochs):
        for queries in query_set:
            # Query target API (the expensive part)
            teacher_probs = target_api.predict_proba(queries)

            # Train student to match teacher
            student_logits = student_model(queries)
            loss = distillation_loss(student_logits, teacher_probs)

            optimizer.zero_grad()
            loss.backward()
            optimizer.step()

    return student_model