model-extraction-relu-logits

name	model-extraction-relu-logits
description	Guidance for extracting weight matrices from black-box ReLU neural networks using only input-output queries. This skill applies when tasked with recovering internal parameters (weights, biases) of a neural network that can only be queried for outputs, particularly two-layer ReLU networks. Use this skill for model extraction, model stealing, or neural network reverse engineering tasks.

Model Extraction for ReLU Networks

This skill provides guidance for extracting internal weight matrices from black-box ReLU neural networks using only input-output access.

Problem Understanding

Model extraction tasks typically involve:

• A black-box neural network that accepts inputs and returns outputs (logits)
• The goal of recovering internal parameters (weight matrices, biases)
• No direct access to the network's implementation or internal state

Critical Principle: True Black-Box Treatment

Treat the target network as a genuine black-box. Never rely on implementation details that may change during evaluation:

• Do not hardcode hidden layer dimensions from example code
• Do not assume specific random seeds or initialization schemes
• Do not directly compare extracted weights to "true" weights read from source files
• The test environment may use completely different parameters than any provided examples

Approach Selection

Understanding ReLU Network Structure

A two-layer ReLU network computes: output = A2 @ ReLU(A1 @ x + b1) + b2

Key properties to exploit:

1. Piecewise linearity: ReLU networks are piecewise linear functions
2. Activation boundaries: Each hidden neuron creates a hyperplane boundary where its output transitions from zero to active
3. Gradient structure: In each linear region, the gradient reveals information about active neurons

Recommended Extraction Strategies

Strategy 1: Critical Point Analysis

ReLU networks have critical points where neurons transition between active/inactive states:

1. Probe the network systematically to identify transition boundaries
2. At each boundary, a hyperplane normal corresponds to a row of A1
3. Collect enough boundaries to reconstruct A1

Strategy 2: Gradient-Based Extraction

For networks where gradients are accessible or can be approximated:

1. Query gradients at multiple random points
2. Gradients in a linear region reveal which neurons are active
3. Use gradient information to identify weight matrix rows

Strategy 3: Activation Pattern Enumeration

Systematically identify which neurons are active in different input regions:

1. Start from a known point and identify its activation pattern
2. Search for inputs that cause different neurons to activate
3. Use the transition points to extract hyperplane parameters

Strategy 4: Optimization-Based Fitting (Fallback)

When mathematically principled methods are insufficient:

1. Generate diverse input-output pairs from the black-box
2. Train a surrogate network to match outputs
3. Critical: Make network capacity adaptive (try multiple hidden dimensions)
4. Validate by output matching, not parameter comparison

Hidden Dimension Discovery

Since the hidden dimension is unknown, employ detection strategies:

1. Rank analysis: The output dimension and response complexity bound hidden size
2. Binary search: Try different hidden sizes and measure reconstruction error
3. Overcomplete fitting: Use larger hidden dimension than necessary, then identify redundant neurons
4. Gradient counting: In a fixed input region, count distinct gradient patterns

Verification Strategy

Correct Verification (Functional Equivalence)

# Generate test inputs NOT used during extraction
test_inputs = generate_diverse_inputs(n=1000)

# Compare outputs
original_outputs = [black_box_query(x) for x in test_inputs]
extracted_outputs = [extracted_model(x) for x in test_inputs]

# Check functional equivalence
max_error = max(|original - extracted| for all test points)
assert max_error < tolerance

Incorrect Verification (Avoid These)

• Comparing extracted weights directly to weights read from source files
• Using the same inputs for extraction and verification
• Relying on cosine similarity to "true" parameters
• Checking only a small number of test points

Common Pitfalls

1. Peeking at Implementation Details

Problem: Reading source code to get the "true" weights or hidden dimension, then validating against them.

Why it fails: Test environments often use different parameters (different seeds, dimensions, scales).

Solution: Treat extraction as if source code doesn't exist. Validate only through output comparison.

2. Hardcoding Network Architecture

Problem: Assuming hidden dimension is fixed (e.g., n_neurons=20).

Why it fails: The actual network may have a different architecture.

Solution: Either detect hidden dimension empirically or design extraction to work with unknown dimensions.

3. Non-Unique Solutions

Problem: Many weight configurations produce identical input-output behavior.

Why it fails: Optimization may find a valid equivalent representation, not the original weights.

Solution: If the task requires recovering specific original weights (not just functional equivalents), use mathematically principled extraction that exploits ReLU structure.

4. Insufficient Test Coverage

Problem: Verifying on a few hand-picked inputs.

Why it fails: The extracted model may fail on untested input regions.

Solution: Use comprehensive random testing across the input domain, including edge cases.

5. Numerical Precision Issues

Problem: Accumulated floating-point errors cause extraction to fail.

Solution: Use numerically stable algorithms, appropriate tolerances, and verify with realistic precision expectations.

Implementation Checklist

Before declaring success, verify:

• No implementation details (seeds, dimensions) were read from source files
• Hidden dimension was detected or handled adaptively
• Verification uses only input-output comparisons
• Verification inputs are independent from extraction inputs
• Sufficient test coverage (hundreds to thousands of points)
• Error tolerance is appropriate for the task requirements
• The extracted model works as a functional replacement

When Standard Approaches Fail

If initial extraction attempts fail:

1. Increase probe density: More input-output pairs may be needed
2. Try multiple hidden dimensions: The assumed size may be wrong
3. Check for numerical issues: Scaling, precision, or conditioning problems
4. Verify the network structure: Ensure assumptions about architecture (two-layer, ReLU) are correct
5. Consider alternative representations: Some equivalent parameterizations may be easier to extract