Efficient Machine Unlearning: A Deep Dive into the SISA Framework
Efficient Machine Unlearning: A Deep Dive into the SISA Framework
In the era of GDPR and the "Right to be Forgotten," the ability to remove a specific user's data from a trained machine learning model is no longer just a technical curiosity—it is a legal requirement.
However, there is a massive problem: Neural networks are "black boxes" that memorize data. Once a model is trained, a single data point influences the weights of the entire network. Traditionally, the only way to truly "unlearn" a data point was to delete it from the dataset and retrain the entire model from scratch. For LLMs or massive production models, this is computationally impossible.
Enter SISA (Sharded, Isolated, Sliced, and Aggregated).
What is SISA?
The SISA framework is a strategic architecture designed to minimize the cost of machine unlearning. Instead of treating the training process as one monolithic block, SISA applies a "divide and conquer" strategy to limit the blast radius of a single data point.
The Core Intuition
If you split your data into 10 independent shards and train 10 small models, a request to delete one data point only affects one of those models. Furthermore, if you save checkpoints during the training of that shard, you only need to retrain from the point where that specific data point first appeared.
Key Contributions
- Blast Radius Reduction: By sharding data, the amount of retraining is reduced by a factor of $k$ (number of shards).
- Incremental Recovery: By slicing data, the system avoids starting from random initialization, resuming instead from the last "clean" checkpoint.
- Ensemble Stability: By aggregating predictions, SISA maintains high accuracy despite using smaller, shard-specific models.
The Architecture: How it Works
The SISA workflow can be broken down into four primary stages: Sharding, Slicing, Aggregation, and the Unlearning trigger.
1. Data Sharding (Isolation)
The total training dataset $D_{tr}$ is divided into $k$ disjoint shards. Each shard is used to train a separate model in complete isolation.
- Crucial Rule: A data point in Shard A must have zero influence on the model trained on Shard B.
2. Slicing & Checkpointing
Each shard is further divided into $s$ slices. The model is trained incrementally: $$\text{Slice}_0 \rightarrow \text{Checkpoint}_0 \rightarrow \text{Slice}_1 \rightarrow \text{Checkpoint}_1 \dots \rightarrow \text{Final Model}$$ This creates a "save game" system for your model weights.
3. Aggregated Inference
Since we now have $k$ models instead of one, we use an ensemble approach. For a given input $x$, each shard model provides a prediction, and the final output is determined via Majority Vote.
4. The Unlearning Process
When a deletion request for point $x_i$ arrives:
- Locate: Identify which shard contains $x_i$.
- Pinpoint: Identify the first slice in that shard where $x_i$ was introduced.
- Restore: Load the checkpoint from the slice immediately preceding it.
- Refresh: Retrain only the remaining slices of that specific shard.
Visual Workflow
Implementation in PyTorch
Below is a production-ready simplified implementation of the SISA framework.
import torch
import torch.nn as nn
import torch.optim as optim
from torch.utils.data import DataLoader, TensorDataset
from sklearn.datasets import make_classification
from sklearn.model_selection import train_test_split
from sklearn.metrics import accuracy_score
import numpy as np
import copy
import time
class SimpleMLP(nn.Module):
def __init__(self, input_dim):
super(SimpleMLP, self).__init__()
self.net = nn.Sequential(
nn.Linear(input_dim, 16),
nn.ReLU(),
nn.Linear(16, 1),
nn.Sigmoid()
)
def forward(self, x):
return self.net(x)
class SISA:
def __init__(self, input_dim, num_shards=4, num_slices=5, lr=0.01, epochs_per_slice=10):
self.num_shards = num_shards
self.num_slices = num_slices
self.lr = lr
self.epochs_per_slice = epochs_per_slice
self.input_dim = input_dim
self.shard_checkpoints = [{} for _ in range(num_shards)]
self.shard_models = [SimpleMLP(input_dim) for _ in range(num_shards)]
self.shard_data_indices = []
def _train_shard_from_slice(self, shard_id, start_slice, X, y):
model = self.shard_models[shard_id]
if start_slice > 0 and start_slice - 1 in self.shard_checkpoints[shard_id]:
model.load_state_dict(self.shard_checkpoints[shard_id][start_slice - 1])
else:
model = SimpleMLP(self.input_dim)
self.shard_models[shard_id] = model
criterion = nn.BCELoss()
optimizer = optim.SGD(model.parameters(), lr=self.lr)
shard_indices = self.shard_data_indices[shard_id]
slice_size = len(shard_indices) // self.num_slices
for s in range(start_slice, self.num_slices):
start_idx = s * slice_size
end_idx = (s + 1) * slice_size if s != self.num_slices - 1 else len(shard_indices)
slice_indices = shard_indices[start_idx:end_idx]
X_slice = torch.FloatTensor(X[slice_indices])
y_slice = torch.FloatTensor(y[slice_indices]).unsqueeze(1)
loader = DataLoader(TensorDataset(X_slice, y_slice), batch_size=16, shuffle=True)
model.train()
for epoch in range(self.epochs_per_slice):
for batch_X, batch_y in loader:
optimizer.zero_grad()
loss = criterion(model(batch_X), batch_y)
loss.backward()
optimizer.step()
self.shard_checkpoints[shard_id][s] = copy.deepcopy(model.state_dict())
return model
def fit(self, X, y):
indices = np.arange(len(X))
np.random.shuffle(indices)
shards_split = np.array_split(indices, self.num_shards)
self.shard_data_indices = [idx.tolist() for idx in shards_split]
for i in range(self.num_shards):
self._train_shard_from_slice(i, 0, X, y)
def predict(self, X):
X_tensor = torch.FloatTensor(X)
predictions = []
with torch.no_grad():
for model in self.shard_models:
model.eval()
out = model(X_tensor)
predictions.append((out > 0.5).int().numpy().flatten())
predictions = np.array(predictions)
return (np.mean(predictions, axis=0) > 0.5).astype(int)
def unlearn(self, X, y, point_index):
shard_id = next((i for i, idxs in enumerate(self.shard_data_indices) if point_index in idxs), -1)
if shard_id == -1: return
shard_indices = self.shard_data_indices[shard_id]
slice_size = len(shard_indices) // self.num_slices
local_idx = shard_indices.index(point_index)
first_slice_with_point = min(local_idx // slice_size if slice_size > 0 else 0, self.num_slices - 1)
self.shard_data_indices[shard_id].remove(point_index)
self._train_shard_from_slice(shard_id, first_slice_with_point, X, y)
# --- Execution ---
X, y = make_classification(n_samples=2000, n_features=20, random_state=42)
X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)
sisa = SISA(input_dim=20)
sisa.fit(X_train, y_train)
# Simulate unlearning
point_to_forget = 10
start = time.time()
sisa.unlearn(X_train, y_train, point_to_forget)
print(f"Unlearning completed in {time.time() - start:.4f}s")Performance Analysis & Trade-offs
The Efficiency Gain
In a standard model, unlearning takes $T$ (time to train the whole model). In SISA, unlearning takes: $$\text{Time} \approx \frac{T}{k \times s} \times (\text{remaining slices})$$ Where $k$ is the number of shards and $s$ is the number of slices. This represents an exponential speedup in retraining time.
The Trade-offs
| Feature | Standard Training | SISA Framework |
|---|---|---|
| Retraining Cost | Extremely High | Low to Moderate |
| Storage Cost | Low (1 model) | High (k models + checkpoints) |
| Inference Latency | Low | Higher (must query $k$ models) |
| Accuracy | Optimal | Slightly lower (ensemble of smaller models) |
Final Thoughts
SISA transforms the "Right to be Forgotten" from a computational nightmare into a manageable engineering task. While it introduces overhead in storage and inference, the ability to surgically remove data without nuking your entire training budget makes it an essential pattern for privacy-preserving AI.