IOAI Training: Learn AI by solving problems

The problem in brief

StyleGAN2, trained on the FFHQ dataset of faces, turns a 512-dimensional vector $w$ (its "W" latent space) into a realistic 1024x1024 portrait. You receive 2,000 such vectors, each labelled 1 if the generated face smiles and 0 otherwise. The task (50 points) is to train a single linear layer that predicts the label from $w$ (20 points), take its weight vector as a "smile direction", and edit a new face by moving its latent code along that direction (25 points), then show the strip of edited faces (5 points). The notebook comes from a third-party repository (github.com/teodora-dimitrova/Singapore-Team-Selection), which also contains a solution notebook that we did not use. The deliverable is code plus pictures, so there is no official number; we added an automatic evaluation with CLIP to compare directions objectively.

How you are scored

Points are given per question for working code: 10 for the classifier, 10 for the training loop, 10 for extracting the direction, 15 for the editing function and 5 for a strip that runs. There is no accuracy threshold. To go beyond "it runs", we measured two things on 24 random faces, for each direction and edit strength:

Smile: CLIP's probability for "a photo of a smiling person." against "a photo of a person with a neutral, serious face.".
Identity kept: the cosine similarity between the CLIP image embeddings of the edited face and the original face (1 means unchanged).

A good edit raises the first number a lot while keeping the second close to 1.

First look at the data

The latents are real W vectors of this generator: their mean has a cosine of 0.996 with the generator's own w_avg.
The labels are imbalanced: 80.9 % of the faces are labelled as smiling. A classifier that always says "smile" is right 80.9 % of the time.
The signal is weak. The best L2-regularised logistic regression reaches a 5-fold cross-validated accuracy of 82.2 % and an AUC of 0.78, barely above the majority rule in accuracy. The notebook's recipe (Adam, 40 epochs, no regularisation) reaches 86.5 % on the training set, which is mostly memorisation: with 512 weights and 2,000 points, an unregularised linear model fits noise.

This matters for the edit, because the edit does not use the classifier's predictions, it uses its weight vector. A noisy weight vector is a noisy direction.

The notebook's answers

The four blanks are one or two lines each:

self.fc = nn.Linear(latent_dim, output_dim)                 # Q1.1

optimizer.zero_grad()                                       # Q1.2
loss = criterion(model(latents_batch), labels_batch)
loss.backward()
optimizer.step()

direction = model.fc.weight.detach()                        # Q2.1, then divide by its norm

w_edited = w_start + strength * direction                   # Q2.2
img = g.synthesis(w_edited, noise_mode="const")             # w_start is (1, 18, 512)

Two details. generator.mapping already returns w repeated for the 18 synthesis layers, shape $(1, 18, 512)$ , so the $(1, 512)$ direction broadcasts over the 18 layers and w_edited can go straight into synthesis. And noise_mode="const" freezes StyleGAN's per-pixel noise, so that the frames of the strip differ only by the edit.

The key idea

Why the weight vector is a direction. A linear classifier predicts smile when $a^\top w + b > 0$ . Its decision boundary is the hyperplane $a^\top w + b = 0$ , and $a$ is the normal to that hyperplane: it is the direction in which the smile logit increases fastest, by $\lVert a\rVert$ per unit step. Moving $w \to w + s\,\hat a$ with $\hat a = a/\lVert a\rVert$ changes the logit by exactly $s\,\lVert a\rVert$ and, if W space is well organised, changes the image mostly in the attribute the classifier was trained on. This is the InterFaceGAN idea.

Which estimate of the direction? Three natural candidates:

the notebook's unregularised classifier;
an L2-regularised logistic regression ( $C = 0.01$ , the best cross-validated AUC);
the plain difference of class means, $\mu_1 - \mu_0$ .

The difference of means looks simplest, but it is the worst. For Gaussian classes with a shared covariance $\Sigma$ , the optimal linear boundary has normal

a^\star \propto \Sigma^{-1}(\mu_1 - \mu_0).

The mean difference is $\Sigma\,a^\star$ : it is the true direction stretched along the high-variance directions of W. Any attribute that varies a lot and happens to correlate with smiling in these 2,000 faces (age, hair, glasses, pose) leaks into it. A classifier estimates $a^\star$ itself, which "explains away" those correlations. The figure shows it: the mean-difference row changes the man's age, hair and glasses, while the classifier rows mostly change the mouth and cheeks.

Stability. Refitting each estimator on bootstrap resamples of the 2,000 latents, the average cosine between two refits is 0.68 for the notebook's classifier, 0.72 for the regularised one ( $C = 0.01$ ) and only 0.55 at $C = 1$ . Regularisation is what makes the direction reproducible. The mean difference is very stable (0.88), but stably biased: stability is not correctness.

Editing only some layers. StyleGAN feeds $w$ separately to each of its 18 synthesis layers: the first layers set coarse structure (pose, face shape), the middle ones features such as the mouth and eyes, the last ones colour and fine texture. Adding the direction only to layers 4 to 7 targets the expression and leaves the rest of the face alone.

The notebook's face (seed 123) edited with strengths from -7 to +7. Top: the notebook's classifier on all layers smiles, but also gets younger and lighter at +7 and distorted at -7. Middle: the mean difference changes age, hair and glasses: entangled. Bottom: the regularised classifier applied to layers 4 to 7 only: a subtler expression change that keeps the person.

The full solution

Convert the third-party files safely. latent_labels.pt and ffhq.pkl are pickles, which can execute code when loaded. We loaded them once inside a sandbox without network access and re-saved them as .npy, .safetensors and a JSON of the generator's constructor arguments; the rest of the pipeline rebuilds the generator from the official NVlabs code and these plain tensors.
Fill in the notebook (classifier, training loop, normalised weight, edit), which reproduces the Q3 strip.
Improve the direction: L2-regularised logistic regression with $C = 0.01$ chosen by cross-validated AUC.
Restrict the edit to synthesis layers 4 to 7.
Evaluate on 24 random faces with CLIP, at strengths $-6, -3, +3, +6$ .

Generating one 1024x1024 face takes about 0.4 s on an Apple laptop GPU (StyleGAN2-ADA's custom CUDA kernels fall back to their reference PyTorch implementations).

Each curve goes from strength -6 (bottom) through the unedited face (left, P(smile) 0.56) to +6 (top). Higher is more smile, further left is more of the original person kept.

Results

Mean over 24 random faces (unedited: CLIP P(smile) = 0.56).

direction	layers	P(smile) at -6	P(smile) at +6	identity at +6	identity at +3
mean difference of classes	all 18	0.23	0.77	0.75	0.87
notebook classifier (Adam, 40 epochs)	all 18	0.22	0.82	0.89	0.94
logistic regression, $C = 0.01$	all 18	0.13	0.87	0.82	0.92
logistic regression, $C = 0.01$	0-5	0.14	0.83	0.88	0.94
logistic regression, $C = 0.01$	4-7	0.12	0.89	0.90	0.96

The regularised direction on layers 4 to 7 gives the strongest change of expression in both directions while keeping the face closest to the original. The notebook's own direction works (it earns the points) but changes less smile per unit of identity lost, and the mean difference is clearly the most entangled. These are proxy measures: CLIP's notion of "smiling" and of "same person" is not a human judgement, and on some faces (like the notebook's seed 123) the layer-restricted edit is visibly subtle.

What to remember

The weight of a linear classifier is the normal to its boundary: it points along the attribute, which is why it can serve as an editing direction in a well-structured latent space.
The class-mean difference is not the discriminative direction. It equals $\Sigma\,a^\star$ and is pulled towards high-variance, correlated attributes; a classifier (or $\Sigma^{-1}\Delta\mu$ , i.e. LDA) disentangles them.
Regularise when you will use the weights themselves. With 512 dimensions and 2,000 noisy labels, an unregularised fit is unstable; choose the penalty by cross-validation, and check stability with bootstrap refits.
Use the structure of the generator. StyleGAN's per-layer inputs let you edit only the layers that control the attribute.
Turn "looks good" into numbers. A frozen judge like CLIP, scoring both the attribute and the preservation of the rest, makes generative edits comparable; and load third-party pickles only in a sandbox, converting them to plain tensors once.

"""Singapore IOAI 2025 camp test, CV part 2: find a "smile" direction in StyleGAN2's W space.

Files: latent_labels.pt (task repo), ffhq.pkl (NVIDIA CDN, stylegan2-ada-pytorch/pretrained),
the NVlabs/stylegan2-ada-pytorch code in ./sg2, CLIP ViT-B/32 as clip_b32.safetensors (timm repo,
open_clip format) for the automatic evaluation.

  1. run_sandboxed.sh . python solution.py convert   # unpickle third-party files WITHOUT network,
                                                     # re-save as .npy / .safetensors / .json
  2. python solution.py edit                         # Q1-Q3: classifier, direction, edited strip
  3. python solution.py judge                        # compare directions with CLIP on 24 faces
"""
import json, sys
import numpy as np, torch, torch.nn as nn, torch.optim as optim
from torch.utils.data import TensorDataset, DataLoader
sys.path.insert(0, "sg2")
torch.set_num_threads(2)
DEV = torch.device("mps" if torch.backends.mps.is_available() else "cpu")

def convert():
    import legacy
    from safetensors.torch import save_file
    d = torch.load("latent_labels.pt", weights_only=True)
    np.save("latents.npy", d["latents"].float().numpy()); np.save("labels.npy", d["labels"].float().numpy())
    with open("ffhq.pkl", "rb") as f:
        G = legacy.load_network_pkl(f)["G_ema"]
    save_file({k: v.contiguous() for k, v in G.state_dict().items()}, "G_ema.safetensors")
    json.dump({"init_kwargs": dict(G.init_kwargs)}, open("G_kwargs.json", "w"), default=dict)

def load_G():
    from safetensors.torch import load_file
    from training.networks import Generator
    G = Generator(**json.load(open("G_kwargs.json"))["init_kwargs"])
    G.load_state_dict(load_file("G_ema.safetensors"))
    return G.eval().requires_grad_(False).to(DEV)

def synth(G, w, delta=None, layers=slice(0, 18)):
    """w: (1, 512). The same w is fed to all 18 synthesis layers; `delta` is added only to `layers`."""
    ws = w.to(DEV)[:, None, :].repeat(1, G.num_ws, 1)
    if delta is not None:
        ws[:, layers] += delta.to(DEV)
    img = G.synthesis(ws, noise_mode="const", force_fp32=True)            # (1, 3, 1024, 1024) in [-1, 1]
    return (img.clamp(-1, 1) * 127.5 + 128).clamp(0, 255).to(torch.uint8).permute(0, 2, 3, 1).cpu().numpy()[0]

# ---------------- Q1.1: the classifier ----------------
class LatentClassifier(nn.Module):
    def __init__(self, latent_dim, output_dim=1):
        super().__init__()
        self.fc = nn.Linear(latent_dim, output_dim)                     # Q1.1
    def forward(self, w):
        return self.fc(w)

# ---------------- Q1.2: the training loop ----------------
def train_classifier(model, loader, num_epochs=40):
    model.train()
    criterion, optimizer = nn.BCEWithLogitsLoss(), optim.Adam(model.parameters(), lr=0.001)
    for epoch in range(num_epochs):
        running = 0.0
        for xb, yb in loader:
            optimizer.zero_grad()                                       # Q1.2 (4 lines)
            logits = model(xb)
            loss = criterion(logits, yb)
            loss.backward(); optimizer.step()
            running += loss.item()
    print(f"final epoch loss {running / len(loader):.4f}")

# ---------------- Q2.1: the direction ----------------
def get_direction_vector(model):
    direction = model.fc.weight.detach()[0]                             # Q2.1: normal of the decision boundary
    return direction / torch.norm(direction)

# ---------------- Q2.2: editing ----------------
@torch.no_grad()
def edit_image(G, w_start, direction, strengths, layers=slice(0, 18)):
    return [synth(G, w_start, float(s) * direction, layers) for s in strengths]   # Q2.2: w + s * direction

def regularised_direction(X, y, C=0.01):
    """Better estimator: L2-regularised logistic regression (CV AUC 0.78 vs 0.74 at C=1)."""
    from sklearn.linear_model import LogisticRegression
    w = LogisticRegression(C=C, max_iter=5000).fit(X, y).coef_[0]
    return torch.tensor(w / np.linalg.norm(w), dtype=torch.float32)

def edit():
    X, y = np.load("latents.npy"), np.load("labels.npy")
    torch.manual_seed(0)
    loader = DataLoader(TensorDataset(torch.tensor(X), torch.tensor(y)), batch_size=32, shuffle=True)
    clf = LatentClassifier(512)
    train_classifier(clf, loader)
    with torch.no_grad():
        print("training accuracy", ((clf(torch.tensor(X)) > 0).float() == torch.tensor(y)).float().mean().item(), "(majority class:", y.mean(), ")")
    d_nb = get_direction_vector(clf); np.save("dir_notebook.npy", d_nb.numpy())
    d_lr = regularised_direction(X, y.ravel()); np.save("dir_logreg001.npy", d_lr.numpy())
    d_md = X[y[:, 0] == 1].mean(0) - X[y[:, 0] == 0].mean(0); np.save("dir_meandiff.npy", d_md / np.linalg.norm(d_md))
    print("cos(notebook, regularised) =", float(d_nb @ d_lr))
    G = load_G()
    torch.manual_seed(123)                                              # Q3: the notebook's face
    with torch.no_grad():
        w_start = G.mapping(torch.randn(1, 512, device=DEV), None)[:, 0].cpu()
    from PIL import Image
    strengths = np.linspace(-7, 7, 7)
    for name, d, L in [("notebook", d_nb, slice(0, 18)), ("logreg_layers4-7", d_lr, slice(4, 8))]:
        strip = np.concatenate([np.asarray(Image.fromarray(im).resize((256, 256))) for im in edit_image(G, w_start, d, strengths, L)], 1)
        Image.fromarray(strip).save(f"strip_{name}.png")

def judge():
    """Automatic evaluation: CLIP P("smiling") and CLIP image similarity to the unedited face."""
    import open_clip
    from PIL import Image
    G = load_G(); torch.manual_seed(0)
    clip, _, pre = open_clip.create_model_and_transforms("ViT-B-32-quickgelu", pretrained="clip_b32.safetensors")
    clip = clip.eval().to(DEV); tok = open_clip.get_tokenizer("ViT-B-32")
    with torch.no_grad():
        T = clip.encode_text(tok(["a photo of a smiling person.", "a photo of a person with a neutral, serious face."]).to(DEV)).float()
        T = T / T.norm(dim=-1, keepdim=True)
        W = G.mapping(torch.randn(24, 512, device=DEV), None)[:, 0].cpu()
    def score(img):
        with torch.no_grad():
            f = clip.encode_image(pre(Image.fromarray(img))[None].to(DEV)).float(); f = f / f.norm(dim=-1, keepdim=True)
        return (100 * f @ T.T).softmax(1)[0, 0].item(), f[0].cpu()
    with torch.no_grad():
        base = [score(synth(G, W[i:i + 1])) for i in range(24)]
        print(f"unedited: P(smile) {np.mean([b[0] for b in base]):.3f}")
        for name, layers in [("notebook", slice(0, 18)), ("logreg001", slice(0, 18)), ("meandiff", slice(0, 18)), ("logreg001", slice(4, 8))]:
            d = torch.tensor(np.load(f"dir_{name}.npy")).float()
            for s in [-6, -3, 3, 6]:
                r = [score(synth(G, W[i:i + 1], s * d, layers)) for i in range(24)]
                print(f"{name:10s} layers {layers.start}-{layers.stop - 1:2d} strength {s:+d}: P(smile) {np.mean([a for a, _ in r]):.3f}, "
                      f"identity {np.mean([float(f @ base[i][1]) for i, (_, f) in enumerate(r)]):.3f}")

if __name__ == "__main__":
    {"convert": convert, "edit": edit, "judge": judge}[sys.argv[1]]()

StyleGAN Latent Space Editing