Chapter 50: LLM Fine-Tuning for Kaggle

import numpy as np

from kaggle_companion.activities._common import clean

SEED = 50
DRAWS = 20                       # independent tasks; every estimate is a mean over draws
D_IN, D_OUT = 40, 20             # one frozen "pretrained" linear layer W0 of 800 weights
TRUE_RANK = 2                    # the target task differs from the pretrained task by a rank-2 change
N_TARGET, STEPS, LR = 100, 300, 0.01
CURVE_RANKS = [1, 2, 3, 4, 6, 8, 12, 16]


def make_task(draw):
    """Pretrained layer W0, a rank-2 shift to the target task, 100 noisy target rows and 4,000 clean held-out rows."""
    rng = np.random.default_rng([SEED, draw])
    W0 = rng.normal(size=(D_IN, D_OUT)) / np.sqrt(D_IN)
    shift = rng.normal(size=(D_IN, TRUE_RANK)) @ rng.normal(size=(TRUE_RANK, D_OUT)) / np.sqrt(D_IN * TRUE_RANK)
    X, X_held = rng.normal(size=(N_TARGET, D_IN)), rng.normal(size=(4000, D_IN))
    Y = X @ (W0 + shift) + rng.normal(size=(N_TARGET, D_OUT))
    return W0, X, Y, X_held, X_held @ (W0 + shift)


def adam(params, gradients, steps=STEPS, lr=LR):
    """Plain Adam on a list of arrays; every method below gets the same optimiser, steps and learning rate."""
    m1 = [np.zeros_like(p) for p in params]
    m2 = [np.zeros_like(p) for p in params]
    for t in range(1, steps + 1):
        for i, g in enumerate(gradients(params)):
            m1[i] = 0.9 * m1[i] + 0.1 * g
            m2[i] = 0.999 * m2[i] + 0.001 * g * g
            params[i] -= lr * (m1[i] / (1 - 0.9 ** t)) / (np.sqrt(m2[i] / (1 - 0.999 ** t)) + 1e-8)
    return params


def full_update(W0, X, Y):
    """Fine-tune every weight: W = W0 + Delta, with Delta trained from zero."""
    grad = lambda p: [X.T @ (X @ (W0 + p[0]) - Y) * 2 / Y.size]
    return W0 + adam([np.zeros_like(W0)], grad)[0]


def lora_update(W0, X, Y, rank, seed):
    """LoRA: W = W0 + B A with B (D_IN x r) starting at zero, A (r x D_OUT) random, scaling alpha / r = 1. Only A and B train."""
    A = np.random.default_rng(seed).normal(size=(rank, D_OUT)) / np.sqrt(D_OUT)

    def grad(p):
        G = X.T @ (X @ (W0 + p[1] @ p[0]) - Y) * 2 / Y.size       # gradient with respect to the whole update
        return [p[1].T @ G, G @ p[0].T]
    A, B = adam([A, np.zeros((D_IN, rank))], grad)
    return W0 + B @ A


def run(rank):
    ranks = sorted(set(CURVE_RANKS) | {rank})
    original, full, curve = [], [], {r: [] for r in ranks}
    for draw in range(DRAWS):
        W0, X, Y, X_held, Y_held = make_task(draw)
        mse = lambda W: float(np.mean((X_held @ W - Y_held) ** 2))      # held-out error against the noiseless target
        original.append(mse(W0))
        full.append(mse(full_update(W0, X, Y)))
        for r in ranks:
            curve[r].append(mse(lora_update(W0, X, Y, r, seed=1000 * draw + r)))
    chosen = np.array(curve[rank])
    saved = np.array(full) - chosen
    return clean({
        "rank": rank, "original": np.mean(original), "full": np.mean(full), "lora": chosen.mean(),
        "gain_over_full": saved.mean(), "gain_over_full_se": saved.std(ddof=1) / np.sqrt(DRAWS),
        "curve_ranks": ranks, "curve": [np.mean(curve[r]) for r in ranks],
        "per_draw_lora": chosen, "per_draw_full": full, "per_draw_original": original,
        "lora_parameters": rank * (D_IN + D_OUT), "full_parameters": D_IN * D_OUT, "true_rank": TRUE_RANK,
        "target_rows": N_TARGET, "draws": DRAWS,
    })

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

VALUES = [1, 2, 4, 16]
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.ch50 import draw, NCOLS, HEIGHT

DEFAULT = 4
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()