Chapter 40: Knowledge Distillation

import warnings

import numpy as np
from sklearn.ensemble import RandomForestClassifier
from sklearn.metrics import log_loss, roc_auc_score
from sklearn.model_selection import KFold
from sklearn.neural_network import MLPClassifier

from kaggle_companion.activities._common import clean

warnings.filterwarnings("ignore")   # one-epoch warm-start fits warn about convergence by design
SEED = 40
REPLICATES = 5          # independent worlds per training budget; results are means over them
N_TRAIN, N_DEV, N_TEST, FEATURES = 375, 125, 6000, 8
FLIP = 0.2              # share of training and development labels flipped
EPS = 1e-4
ALPHAS = (0.0, 0.5, 1.0)  # weight on the teacher's soft target; 0 is hard labels only


def generate(n, rng):
    """Eight features and a bent, interacting true probability."""
    X = rng.normal(size=(n, FEATURES))
    z = 1.4 * np.sin(1.3 * X[:, 0]) + 0.9 * X[:, 1] * X[:, 2] + 0.8 * X[:, 3] - 0.7 * np.abs(X[:, 4]) + 0.3
    return X, 1 / (1 + np.exp(-1.5 * z))


def forest():
    return RandomForestClassifier(60, min_samples_leaf=5, random_state=0, n_jobs=1)


def train_student(X, target, X_dev, y_dev, X_test, epochs, seed):
    """Soft-target cross-entropy via duplicated rows: each row counts as a 1 with weight q and as a 0 with weight 1 - q.

    Returns test probabilities after the last epoch and at the epoch with the best development log loss (early stopping)."""
    X2, y2, w2 = np.vstack([X, X]), np.r_[np.ones(len(X)), np.zeros(len(X))], np.r_[target, 1 - target]
    net = MLPClassifier((32, 16), alpha=1e-3, max_iter=1, warm_start=True, learning_rate_init=0.01, batch_size=64, random_state=seed)
    best_loss, best_pred, best_epoch = np.inf, None, 0
    for epoch in range(1, epochs + 1):
        net.fit(X2, y2, sample_weight=w2)                      # one more epoch from the current weights
        if epoch % 3 == 0 or epoch == epochs:
            loss = log_loss(y_dev, np.clip(net.predict_proba(X_dev)[:, 1], EPS, 1 - EPS), labels=[0, 1])
            if loss < best_loss:
                best_loss, best_pred, best_epoch = loss, np.clip(net.predict_proba(X_test)[:, 1], EPS, 1 - EPS), epoch
    parameters = sum(w.size for w in net.coefs_) + sum(b.size for b in net.intercepts_)
    return np.clip(net.predict_proba(X_test)[:, 1], EPS, 1 - EPS), best_pred, best_epoch, parameters


def one_world(epochs, seed):
    rng = np.random.default_rng(seed)
    X, p = generate(N_TRAIN + N_DEV, rng)
    clean_y = (rng.random(len(p)) < p).astype(int)
    y = np.where(rng.random(len(p)) < FLIP, 1 - clean_y, clean_y)      # noisy labels for training and development
    X_test, p_test = generate(N_TEST, rng)
    y_test = (rng.random(N_TEST) < p_test).astype(int)                 # genuine, unflipped test labels
    X_tr, y_tr, X_dev, y_dev = X[:N_TRAIN], y[:N_TRAIN], X[N_TRAIN:], y[N_TRAIN:]

    soft = np.zeros(N_TRAIN)                                            # cross-fitted teacher: no row is predicted by a model that saw its label
    for fit, out in KFold(5, shuffle=True, random_state=0).split(X_tr):
        soft[out] = forest().fit(X_tr[fit], y_tr[fit]).predict_proba(X_tr[out])[:, 1]
    teacher = forest().fit(X_tr, y_tr)
    pred = np.clip(teacher.predict_proba(X_test)[:, 1], EPS, 1 - EPS)
    out = {"teacher": [log_loss(y_test, pred), roc_auc_score(y_test, pred)],
           "teacher_nodes": sum(t.tree_.node_count for t in teacher.estimators_)}
    for alpha in ALPHAS:
        target = np.clip(alpha * soft + (1 - alpha) * y_tr, 0, 1)
        last, tuned, epoch, parameters = train_student(X_tr, target, X_dev, y_dev, X_test, epochs, seed)
        out[alpha] = [log_loss(y_test, last), roc_auc_score(y_test, last), log_loss(y_test, tuned), roc_auc_score(y_test, tuned), epoch]
        out["student_parameters"] = parameters
    return out


def run(epochs):
    worlds = [one_world(epochs, SEED * 100 + r) for r in range(REPLICATES)]

    def avg(get):
        return float(np.mean([get(w) for w in worlds]))
    arms = {}
    for alpha in ALPHAS:
        arms[str(alpha)] = {"fixed_log_loss": avg(lambda w: w[alpha][0]), "fixed_auc": avg(lambda w: w[alpha][1]),
                            "stopped_log_loss": avg(lambda w: w[alpha][2]), "stopped_auc": avg(lambda w: w[alpha][3]),
                            "stopped_epoch": avg(lambda w: w[alpha][4])}
    return clean({
        "epochs": epochs, "arms": arms,
        "teacher": {"log_loss": avg(lambda w: w["teacher"][0]), "auc": avg(lambda w: w["teacher"][1])},
        "mixed_beats_hard_stopped": avg(lambda w: w[0.5][2] < w[0.0][2]),
        "pure_worse_than_hard_stopped": avg(lambda w: w[1.0][2] > w[0.0][2]),
        "teacher_nodes": avg(lambda w: w["teacher_nodes"]), "student_parameters": worlds[0]["student_parameters"],
        "flip": FLIP, "replicates": REPLICATES,
    })

import pandas as pd
from kaggle_companion.activities.ch40 import explain

VALUES = [6, 15, 30, 45]
results = {value: run(value) for value in VALUES}
table = pd.DataFrame({str(v): explain(results[v], v)['metrics'] for v in VALUES})
print(table.to_string())

import matplotlib.pyplot as plt
from kaggle_companion.activities.ch40 import draw, NCOLS, HEIGHT

DEFAULT = 30
fig, axes = plt.subplots(1, NCOLS, figsize=(6.6 if NCOLS == 1 else 10.4, HEIGHT), layout='constrained')
draw(axes, results[DEFAULT], DEFAULT)
plt.show()