The gap attack, and some ML fundamentals

In part 1, we worked out what a membership inference attack is. In this post, we’ll build a simple version of it.

To accomplish this, I had to train a neural network, which I had no idea how to do when I started out reading this paper. This post will go over how to build a simple CNN with PyTorch, as well as the ML concepts needed to understand what’s going on under the hood.

Part A — the theory behind it all

What is a neural network?

A machine-learning model is just a function which has a million adjustable numbers (knobs) inside called weights. It takes an input (in our example today, an image) and produces an answer (in our case, a list of numbers, the highest of which corresponds to the label that the model is the most confident in being the right label, according to the data it was trained on).

This function is built out of small units called neurons. A neuron takes a set of input numbers, multiplies each by one of its own weights, adds them up along with a bias, and produces a single output number (usually passed through an activation function on the way out). The weights and biases are the knobs that need to be adjusted.

Neurons are arranged into layers (a row of neurons working in parallel). When every neuron in a layer is connected to every input, we call it a fully-connected layer.

Training the model: by adjusting the above-mentioned knobs, our model produces different results. The goal is to tune these knobs automatically until our model gives good answers on the examples we train it on.

The rough idea behind how it works is to:

• start with a random set of weights/biases
• show the model an image from the training set
• let it produce a result (a label for that image)
• measure how wrong the guess is
• nudge every knob in the direction that would make the guess a little less wrong
• repeat this process a number of times

Slowly, the weights and biases will settle into values that make good predictions on the training set. However, training the model too much will result in it overfitting for the given dataset. (In our case, it will produce correct outputs for those images, but not be able to guess the category of new never-before-seen images.)

What makes it a CNN?

While it’s possible to simplify an image to a 1D array of numbers, doing so would throw away the spatial structure of it. The positioning of each pixel matters. A CNN (Convolutional Neural Network) is a deep learning architecture designed to process visual data so that the structure is preserved.

Consider a tiny 3×3 window (a filter) that we can slide across the image. At each spot it looks at the little 3×3 patch of pixels underneath it. The dot product of the filter’s numbers with the pixels they’re sitting on is a way for us to measure how much the patch looks like the pattern this filter is tuned to detect.

By sliding the filter across the whole image, we get a grid of these numbers, which is called a feature map. It shows where that pattern appears. The filter above detects dark→bright vertical edges in the image and stays near zero on flat regions.

Of course, one single filter is not enough to detect all the patterns in an image. Each filter can measure a granular feature. By incorporating many of these, each tuned to a different pattern, we produce many feature maps for each conv layer. (In our example, the first conv layer has 32 filters.)

Then, by stacking the conv layers, we can detect more complicated patterns, textures, and shapes.

A model’s weights are its adjustable knobs. In a conv layer, those knobs are the numbers inside the filters. Training finds appropriate filters by tuning these numbers, slowly working out which patterns are worth detecting to be able to successfully recognize the images in our dataset. the numbers a filter produces (the feature map) are outputs (activations) recomputed fresh for every image.

In our first conv layer, we use 32 filters (this is arbitrarily picked in our case). Each filter is a 3×3 stack across the 3 colour channels, which comes to 32 × 3 × 3 × 3 = 864 weights. Each filter has its own bias, which adds another 32 to the total number of parameters that need to be tuned. Because the same filters slide across the whole image, these numbers get reused everywhere, which is why a CNN needs fewer knobs than a network that wires every pixel to every neuron. (A fully-connected layer mapping the 3×32×32 image to a 32×32×32 output would need over 100 million parameters.)

Tensors

A tensor is a multi-dimensional array of numbers, and it’s the PyTorch type we will use to store images.

A colour image is stored as a tensor of shape [channels, height, width].

Two transforms turn a raw image into that tensor:

• ToTensor reorders the image to [channels, height, width] and rescales each pixel from an integer in 0–255 to a float in 0.0–1.0.
• Normalize(mean, std) shifts and scales each channel: output = (input − mean) / std. In our case, with mean = 0.5 and std = 0.5, we map [0, 1] → [−1, 1]. This is because neural nets train more stably when their inputs are centred around zero. I still don’t know exactly why, but I’m taking this for granted for now — something to do with gradients behaving better?

The architecture

Convolutions detect small patterns in an image. But how do we derive a final label from this?

1. Feature extraction. We find patterns, then shrink the image down to something small and information-dense. And then repeat by finding patterns in that smaller image, and so on. We do this by interleaving a stack of conv layers with a pooling operation.
2. Classification. We flatten that result into a plain list of numbers, then pass it through a couple of fully-connected layers that turn those features into 10 scores, one per class.

Let’s go through the new building blocks we need to be able to get this done:

Pooling. After a conv layer, we shrink the grid with max-pooling. This operation slides a small window (we use 2×2) across the feature map and keeps only the largest value in each one, which halves the height and width of the map. By doing this, the network loses a bit of the granularity about where exactly a feature showed up, throwing away the fine detail and only keeping the strongest signal in each 2×2 window.

Flatten. After a few rounds of conv + pool, we are left with a stack of small feature maps as a 3D block, but the layer that makes the final decision expects a flat, 1D list. All flattening does is unroll that block into one long vector.

Fully-connected (linear) layers. These are classic neural-network layers. Every input connects to every output. Each neuron is a weighted sum of all the inputs, plus a bias.

E.g. for a layer with 3 inputs and 2 outputs, we get:

These weights and biases are more knobs that training tunes. At the end of the network, we disregard the locality and spatial structure, and use all the features that we have from the conv layers to calculate the numbers for each of the 10 classes.

Fully-connected layers have a lot of knobs (e.g. our first one has ~half a million weights), which is why the model has enough capacity to memorise its training set. It leads to overfitting.

Activation (ReLU). Conv layers and fully-connected layers are fundamentally linear transformations, so composing them without an extra step between them would be mathematically equivalent to having a single linear layer. We use a non-linear activation layer called ReLU, which keeps a number if it’s positive and replaces it with zero otherwise (i.e. max(0, x)). We apply this after every hidden layer, before passing it on to the next.

Observe how the shape of the data changes as it flows through our network:

Convolutions keep the grid size. With a 3×3 filter and padding=1 (a 1-pixel border of zeros around the image), the output stays the same height and width. The two conv layers only change the channel count (3 → 32 → 64). For the input, those channels are the RGB colours; after the first conv layer, they’re one feature map per filter.

Max-pooling halves the grid. Each 2×2 pool takes the height and width 32 → 16 → 8.

By the end, each of the 64 feature maps is an 8×8 grid, so flattening gives 64 × 8 × 8 = 4096 numbers, and the two linear layers calculate the final 10 numbers.

(Note that ReLU doesn’t change the shape of the grid, which is why it’s omitted in the diagram above.)

The input size of the first linear layer is forced as it has to match the number of flattened features as calculated above. And at the end of it all, for the output of the final layer, we need to produce 10 numbers to match the 10 classes of CIFAR-10. But the width of the hidden middle layer (128) is an arbitrary number, and it’s a design choice. A bigger number would result in more capacity which increases the risk of overfitting. A smaller number causes the data to need more compression, which risks losing some detail.

From numbers to a label

The network’s final layer produces 10 raw numbers (called logits), one score per class. By applying softmax, we turn these scores into probabilities that sum to 1:

$p_k$ is the model’s confidence in class $k$. The predicted class (the label!) is just the one with the highest probability (equivalently, the highest logit). Our label-only membership-inference attack relies on this label only (as the name suggests) and not the confidence scores themselves.

Measuring wrongness with cross-entropy loss

To be able to train our model, we need a single number measuring how wrong a prediction was. For classification problems such as labelling images, that’s the cross-entropy loss, the negative log of the probability the model gave to the correct class:

If the model is confident and right (p_true ≈ 1), then −log(1) ≈ 0. Little to no loss.

If it puts low probability on the right answer, the loss increases:

The name comes from information theory, where −log(p) is “surprise”. The loss represents how surprised the model was by the true answer in that round of training. The goal is to make the model stop being surprised by the truth.

Tuning with gradient descent and backpropagation

We can think of the loss as a landscape over all the possible knob settings. Akin to walking downhill on a hilly surface, we want to minimize this loss. Loss depends on all the weights at once, so we use partial derivatives $\partial L / \partial w$ to calculate the change in loss if one weight is changed, holding all other parameters fixed.

The gradient $\partial L / \partial w$ is the slope at our current spot. To go downhill, we step the opposite way: $w \leftarrow w - \text{lr} \cdot \frac{\partial L}{\partial w}$, for every knob. lr is the learning rate (the step size) which we set. If the step size is too big, we might miss the valley and bounce left and right. If it’s too small, we will crawl very slowly and might not reach it.

(I’m writing everything in terms of a weight $w$, but biases are tuned the exact same way, since a bias is just a weight on a constant input of 1.)

Each weight $w$ only affects loss $L$ through its neuron. Pre-activation (before ReLU), the neuron first computes a linear weighted sum $z = w \cdot x + b$. Nowhere else does $w$ affect $L$, so, by the chain rule:

Since $z = w \cdot x + b$, a small change in $w$ changes $z$ by exactly $x$, so $\partial z / \partial w = x$.

We call $\partial L / \partial z$ the error signal: how much the loss is affected by that neuron.

Meaning, a weight’s gradient is a product of the input that flowed through it and the error signal. (Each weight is only used by one neuron.)

Recall that the final layer produces logits $z_k$, softmax turns them into probabilities $p_k = e^{z_k} / \sum_j e^{z_j}$, and cross-entropy scores them against the true class $t$ with $L = -\log p_t$.

Differentiating with respect to any logit $z_k$, the first term contributes $-1$ only for the true class while the second contributes exactly $p_k$. And:

So:

Where $y_k$ is $1$ for the true class and $0$ for the rest. So the output error signal is $p - y = \textbf{prediction} - \textbf{truth}$, the difference between the predicted probabilities and the correct answer.

This was for the final layer, but there are other weights that live deeper in the network, passing through multiple neurons before they affect the loss. We can fix this by computing the error signals backwards, one layer at a time, starting at the output layer, then using the chain rule to write each earlier layer’s error signals in terms of the ones already computed. Because every layer reuses the next layer’s results, the whole network’s gradients come out in a single backward sweep rather than a from-scratch calculation per weight. This is called backpropagation.

PyTorch handles all of this automatically. As the loss is computed, PyTorch records every operation, and a .backward() call then replays them in reverse to calculate $\partial L / \partial w$ for every weight and bias at once.

Part B — the code that trains the model

The PyTorch library does all the complicated math (gradient descent and backprop) for us. We just need to learn how to use it properly.

Splitting the data into members and non-members

To be able to do a membership-inference attack, first we need to train the model.

We will be using the CIFAR-10 dataset, which has 50,000 images. In our code, we are only using 10,000 of these. 5,000 will be the ‘members’ which our model will train on, while the other 5,000 will only be used for evaluation. We will later test how often we can correctly determine an image is a member or not, using the model’s labelling behaviour on these two sets.

tf = transforms.Compose([transforms.ToTensor(),
                         transforms.Normalize((0.5,) * 3, (0.5,) * 3)])
full = datasets.CIFAR10(root="./data", train=True, download=True, transform=tf)

N_MEMBERS, N_NONMEMBERS = 5000, 5000
perm = np.random.permutation(len(full))
member_idx    = perm[:N_MEMBERS]
nonmember_idx = perm[N_MEMBERS:N_MEMBERS + N_NONMEMBERS]
member_set    = Subset(full, member_idx)
nonmember_set = Subset(full, nonmember_idx)
train_loader  = DataLoader(member_set, batch_size=128, shuffle=True)

In the Tensor section above, we went over why we normalize the pixel data. The DataLoader feeds images to the model in batches of 128 at a time. (A batch here is just another tensor, with the first dimension being the batch size. Every PyTorch layer is built to process a whole batch at once.)

The model in PyTorch

class SmallCNN(nn.Module):
    def __init__(self):
        super().__init__()
        self.c1  = nn.Conv2d(3, 32, 3, padding=1)
        self.c2  = nn.Conv2d(32, 64, 3, padding=1)
        self.fc1 = nn.Linear(64 * 8 * 8, 128)
        self.fc2 = nn.Linear(128, 10)
    def forward(self, x):
        x = F.max_pool2d(F.relu(self.c1(x)), 2)
        x = F.max_pool2d(F.relu(self.c2(x)), 2)
        x = x.flatten(1)
        x = F.relu(self.fc1(x))
        return self.fc2(x)

forward is what happens when an image flows through the architecture that we set up in Part A. image → conv → relu → pool → conv → relu → pool → flatten → linear → relu → linear

Here’s how the shape of the data changes through one forward pass:

The training loop

model = SmallCNN().to(device)
opt = torch.optim.Adam(model.parameters(), lr=1e-3)

EPOCHS = 50
for ep in range(EPOCHS):
    model.train()
    for xb, yb in train_loader:
        xb, yb = xb.to(device), yb.to(device)
        opt.zero_grad()
        F.cross_entropy(model(xb), yb).backward()
        opt.step()

This is the loop that will train our model. Each epoch does a full pass over all 5,000 members of our training set. This is a lot, and with no data augmentation will lead to overfitting, which will illustrate our attack in Part C.

We are doing mini-batch stochastic gradient descent, where after each batch of images, we’ll tune the model’s weights.

• opt.zero_grad() — PyTorch accumulates gradients, so we need to wipe them for each pass.
• model(xb) — one forward pass for this batch of images.
• F.cross_entropy(…, yb) — the loss calculations (cross_entropy implicitly handles softmaxing and calculating −log).
• .backward() — backpropagation, computes $\partial L / \partial w$ for every weight.
• opt.step() — stepping downhill (tuning the knobs): $w \leftarrow w - \text{lr} \cdot \partial L / \partial w$ for every weight Adam is a refined version of gradient descent — I’ll have to explore what it does, exactly, some other time.

Evaluating the model

@torch.no_grad()
def correct_mask(dataset):
    model.eval()
    out = []
    for xb, yb in DataLoader(dataset, batch_size=256):
        xb, yb = xb.to(device), yb.to(device)
        out.append((model(xb).argmax(1) == yb).cpu())
    return torch.cat(out).numpy()

member_correct    = correct_mask(member_set)
nonmember_correct = correct_mask(nonmember_set)
train_acc, test_acc = member_correct.mean(), nonmember_correct.mean()

@torch.no_grad() and model.eval() disregard gradients and switch dropout/batch-norm to eval behaviour. (The first one speeds things up, and the second is not relevant to our model since our simple CNN does not use those kinds of layers, but this is standard PyTorch practice when writing ML code.)

correct_mask returns a bitmask of whether the model classified the images correctly, by taking the class with the highest score and comparing it to the true label.

Running it on members and non-members gives us two accuracies, which we can now compare.

Part C — the attack

The gap attack

y_true = np.concatenate([np.ones(N_MEMBERS), np.zeros(N_NONMEMBERS)])
y_pred = np.concatenate([member_correct, nonmember_correct]).astype(int)
mi_acc = (y_pred == y_true).mean()
theory = 0.5 + (train_acc - test_acc) / 2

The idea here is that, because our model has overfitted for its members, we can exploit whether it correctly labels an image or not to determine whether that image was in the training set.

By comparing the true membership labels (1 for members, 0 for non-members) and the guesses (1 for correctly classified images, 0 otherwise), we can measure how often this matches. Because the test set is balanced and has an equal number of members and non-members, a 50% result would be the random baseline. So anything above 50% would indicate a real information leak.

Why it works

The paper predicts this attack’s accuracy. On a balanced set, the attack is right in two cases:

• On a member, it guesses “member” if and only if the model got it right, which happens with a probability of train_acc.
• On a non-member, it guesses “non-member” if and only if the model got it wrong, which happens with a probability of 1 − test_acc.

Since the sets are balanced:

As we can see, the attack’s power is the overfitting gap, halved and shifted up from the 50% baseline of guessing randomly. A bigger gap would result in a stronger attack.

The result

By running the actual code, I was able to get:

train(member) acc 1.000 | test(non-member) acc 0.562 | GAP 0.438
GAP-ATTACK MI accuracy: 0.719  (theory 1/2+(acc_tr-acc_te)/2 = 0.719)

• train acc = 1.000 — the model memorised every single member (overfitting).
• test acc = 0.562 — on brand new images it gets it right ~56% of the time (above the 10% it’d get by guessing, but not close to 100%).
• gap = 0.438
• MI accuracy = 0.719 — the gap attack correctly guesses membership ~72% of the time, and the empirical data from actually testing the model matches the theory above

We have built a model that measurably leaks membership data, using its predicted labels, due to overfitting.

Next steps

This was the simplest form of a gap attack, and we can (and will) use it as a baseline against which to measure more sophisticated attacks. The main focus of the paper is the augmentation attack, which queries the model on nudged copies of images (slight rotations, shifts, etc.) and measures how robustly the label survives. Members tend to keep their label under perturbations better than non-members.

We will explore this concept more in part 3.