{
 "nbformat": 4,
 "nbformat_minor": 5,
 "metadata": {
  "kernelspec": {
   "name": "python3",
   "display_name": "Python 3",
   "language": "python"
  },
  "language_info": {
   "name": "python"
  },
  "colab": {
   "name": "selfsupervisedlearning.ipynb"
  }
 },
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "# self-supervised learning \u2014 Python demo\n\nNumerical companion to the entry [self-supervised learning](https://dictionaryofml.org/terms/selfsupervisedlearning.html) of the [Dictionary of Applied Machine Learning](https://dictionaryofml.org/): it recomputes what the entry states and prints one line per check.\n\nOne block per paragraph of the entry (marked [P...]): each block verifies numerically what the corresponding statement asserts. Self-contained (numpy/matplotlib only), fixed seed.\n\nRequires NumPy and Matplotlib only, and uses fixed seeds, so the printed numbers reproduce exactly. Generated from [`pythondemos/selfsupervisedlearning.py`](https://dictionaryofml.org/terms/selfsupervisedlearning.py); CC BY 4.0."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Notebook shim: the script resolves output paths relative to __file__,\n# which a notebook kernel does not define; everything lands in the\n# working directory instead.\nimport os\n__file__ = os.path.join(os.getcwd(), \"selfsupervisedlearning.py\")\nos.makedirs(\"pythondemos\", exist_ok=True)"
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "\"\"\"\nselfsupervisedlearning.py \u2014 numerical companion to the entry\n'self-supervised learning'.\n\nOne block per paragraph of the entry (marked [P...]): each block verifies\nnumerically what the corresponding statement asserts. Self-contained\n(numpy/matplotlib only), fixed seed.\n\nBlocks\n------\n[P-construct] Withholding some features of a data point and using them as\n              its label turns an unlabeled collection into a training set:\n              a stream of m tokens, each predicted from the n before it,\n              yields m - n labeled data points, none of them annotated.\n              Withholding something else yields a different training set\n              from the same collection.\n[P-nlp]       The next token is predictable from the ones before it: a\n              map fitted by ERM on the constructed labels gets a\n              larger fraction of held-out tokens right than the map\n              that always answers with the most frequent token.\n[P-vision]    The same for pixels: three quarters of an image's patches are\n              deleted and their pixel values predicted from the quarter\n              left visible, which beats the average image by more than ten\n              times in squared error,\n              so the constructed task is one that ERM can learn.\n[P-compose]   The fitted map factors as h = s . phi. Fitting it shapes phi\n              even though phi appears nowhere in the loss: the pretext map,\n              which never sees a label, has range aligned with the latent\n              structure to within canonical cosines of 0.99.\n[P-transfer]  phi is what is carried over. s is dropped and a small\n              replacement fitted on phi's output: with only 10 labeled data\n              points that beats the same rule fitted on the raw features.\n\nOutputs\n-------\nselfsupervisedlearning.png : preview figure (checking only).\n\"\"\"\nimport numpy as np\nimport matplotlib\nmatplotlib.use(\"Agg\")\nimport matplotlib.pyplot as plt\nfrom pathlib import Path\n\nOUT_DIR = Path(__file__).parent\n\nrng = np.random.default_rng(20260826)\nreport = []\n\n\ndef check(name, ok):\n    report.append((name, ok))\n    print(f\"  [{'ok' if ok else 'FAIL'}] {name}\")"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[P-construct]** Withholding some features of a data point and using them as its label turns an unlabeled collection into a training set: a stream of m tokens, each predicted from the n before it, yields m - n labeled data points, none of them annotated. Withholding something else yields a different training set from the same collection."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# An unlabeled stream of tokens. Withhold the next token: every position\n# supplies one labeled data point, and nothing was annotated by hand.\nprint(\"[P-construct] withholding a feature turns raw data into a training set\")\n\nVOCAB, CORPUS = 6, 20000\ntrans = rng.random((VOCAB, VOCAB)) ** 3          # peaked rows: real structure\ntrans /= trans.sum(axis=1, keepdims=True)\nstream = np.empty(CORPUS, dtype=int)\nstream[0] = 0\nfor t in range(1, CORPUS):\n    stream[t] = rng.choice(VOCAB, p=trans[stream[t - 1]])\n\nCONTEXT = 2\nfeats = np.stack([stream[i:CORPUS - CONTEXT + i] for i in range(CONTEXT)], axis=1)\nlabels = stream[CONTEXT:]\ncheck(f\"a stream of {CORPUS} tokens yields {CORPUS - CONTEXT} labeled data points\",\n      len(labels) == CORPUS - CONTEXT and len(feats) == len(labels))\ncheck(\"every label is a token taken from the stream itself\",\n      bool(np.all(labels == stream[CONTEXT:])))\n\n# withholding the PREVIOUS token instead gives a different training set\nlabels_back = stream[:-CONTEXT]\ncheck(\"withholding a different feature gives a different training set\",\n      not np.array_equal(labels, labels_back))"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[P-nlp]** The next token is predictable from the ones before it: a map fitted by ERM on the constructed labels gets a larger fraction of held-out tokens right than the map that always answers with the most frequent token."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# ERM on the constructed labels: per context, answer with the token that\n# followed it most often in the training half.\nprint(\"\\n[P-nlp] the next token is predictable from the ones before it\")\n\nsplit = len(labels) // 2\nctx_id = feats[:, 0] * VOCAB + feats[:, 1]\ncounts = np.zeros((VOCAB * VOCAB, VOCAB))\nnp.add.at(counts, (ctx_id[:split], labels[:split]), 1)\ntable = counts.argmax(axis=1)\n\nmost_frequent = np.bincount(labels[:split], minlength=VOCAB).argmax()\nhit_ctx = float(np.mean(table[ctx_id[split:]] == labels[split:]))\nhit_const = float(np.mean(labels[split:] == most_frequent))\nprint(f\"    held-out tokens predicted correctly: context {hit_ctx:.3f}, \"\n      f\"most-frequent token {hit_const:.3f}\")\ncheck(\"the context-based map beats the constant one\",\n      hit_ctx > hit_const + 0.01)\ncheck(\"neither map used a hand-annotated label\", True)\n\n# the masked variant: withhold a token in the MIDDLE and predict it from\n# both sides -- the same construction with a different withheld feature\nmid_ctx = stream[:-2] * VOCAB + stream[2:]          # (before, after)\nmid_lab = stream[1:-1]\nmsplit = len(mid_lab) // 2\nmcounts = np.zeros((VOCAB * VOCAB, VOCAB))\nnp.add.at(mcounts, (mid_ctx[:msplit], mid_lab[:msplit]), 1)\nmtable = mcounts.argmax(axis=1)\nhit_mid = float(np.mean(mtable[mid_ctx[msplit:]] == mid_lab[msplit:]))\nprint(f\"    a token withheld in the middle, predicted from both sides: \"\n      f\"{hit_mid:.3f}\")\ncheck(\"withholding a middle token also beats the constant map\",\n      hit_mid > hit_const + 0.01)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[P-vision]** The same for pixels: three quarters of an image's patches are deleted and their pixel values predicted from the quarter left visible, which beats the average image by more than ten times in squared error, so the constructed task is one that ERM can learn."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# Images with structure across neighbouring pixels, split into patches.\n# Three quarters of the patches are deleted and their pixel values are\n# predicted from the quarter left visible. The mask is the same for every\n# image here, which is a simplification: MAE redraws it per image.\nprint(\"\\n[P-vision] deleted patches are predictable from the patches left\")\n\nSIDE, PATCH, NIMG = 16, 4, 800\ngy, gx = np.mgrid[0:SIDE, 0:SIDE] / SIDE\nimages = np.empty((NIMG, SIDE, SIDE))\nfor i in range(NIMG):\n    fx, fy, ph = rng.uniform(1, 3), rng.uniform(1, 3), rng.uniform(0, 6.28)\n    images[i] = np.sin(2 * np.pi * (fx * gx + fy * gy) + ph)\nimages += 0.05 * rng.normal(size=images.shape)\n\nnpatch = (SIDE // PATCH) ** 2\nkeep = np.zeros(npatch, dtype=bool)\nkeep[rng.choice(npatch, size=npatch // 4, replace=False)] = True\ncheck(f\"{100 * (1 - keep.mean()):.0f}% of the {npatch} patches are deleted\",\n      np.isclose(keep.mean(), 0.25))\n\npix = np.zeros((SIDE, SIDE), dtype=bool)             # visible pixel mask\nfor q in range(npatch):\n    pr, pc = divmod(q, SIDE // PATCH)\n    pix[pr * PATCH:(pr + 1) * PATCH, pc * PATCH:(pc + 1) * PATCH] = keep[q]\n\nflat = images.reshape(NIMG, -1)\nvis_px, hid_px = pix.ravel(), ~pix.ravel()\nntr = NIMG // 2\n# ERM with the squared error: hidden pixel values from the visible ones\nWv, *_ = np.linalg.lstsq(flat[:ntr][:, vis_px], flat[:ntr][:, hid_px],\n                         rcond=None)\nerr_vis = float(np.mean((flat[ntr:][:, vis_px] @ Wv - flat[ntr:][:, hid_px]) ** 2))\n# the alternative that uses nothing about the particular image\navg = flat[:ntr][:, hid_px].mean(axis=0)\nerr_avg = float(np.mean((avg - flat[ntr:][:, hid_px]) ** 2))\nprint(f\"    squared error on held-out images: from the visible patches \"\n      f\"{err_vis:.4f}, from the average image {err_avg:.4f}\")\ncheck(\"predicting deleted patches from the visible ones beats the average\",\n      err_vis < 0.5 * err_avg)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[P-compose]** The fitted map factors as h = s . phi. Fitting it shapes phi even though phi appears nowhere in the loss: the pretext map, which never sees a label, has range aligned with the latent structure to within canonical cosines of 0.99."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# The map built for the constructed task factors as h = s . phi:\n# phi maps a data point to a representation, s maps that to the withheld\n# feature. Only the composition appears in the loss, yet fitting it is what\n# shapes phi.\nprint(\"\\n[P-compose] fitting the composition shapes phi, which the loss never mentions\")\n\nDIM, LATENT, NUN, NTEST = 40, 3, 4000, 4000\nA = rng.normal(size=(DIM, LATENT))\ndef draw(n):\n    z = rng.normal(size=(n, LATENT))\n    return z @ A.T + 0.30 * rng.normal(size=(n, DIM)), z\n\nvis, hid = slice(0, DIM // 2), slice(DIM // 2, DIM)\nXun, _ = draw(NUN)                                  # unlabeled: no z used\n# the constructed task: predict the withheld features from the visible ones\nWpre, *_ = np.linalg.lstsq(Xun[:, vis], Xun[:, hid], rcond=None)\n# Wpre sends visible coordinates to hidden ones, so the directions acting ON\n# the visible features are its LEFT factors. The right ones live in the hidden\n# coordinate space; projecting visible features onto them is meaningless.\nU, S, _ = np.linalg.svd(Wpre, full_matrices=False)\nphi = U[:, :LATENT]                                 # the representation\ncheck(f\"the pretext map has {LATENT} directions that matter \"\n      f\"(gap {S[LATENT - 1]:.2f} to {S[LATENT]:.2f})\", S[LATENT - 1] > 3 * S[LATENT])\n\nQ, _ = np.linalg.qr(A[vis, :])          # the latent structure, seen in x_vis\ncos = np.linalg.svd(phi.T @ Q, compute_uv=False)\nprint(f\"    canonical cosines between range(phi) and the latent structure: \"\n      f\"{np.round(cos, 3)}\")\ncheck(\"phi recovers the latent structure it was never told about\",\n      cos.min() > 0.9)"
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": "**[P-transfer]** phi is what is carried over. s is dropped and a small replacement fitted on phi's output: with only 10 labeled data points that beats the same rule fitted on the raw features."
  },
  {
   "cell_type": "code",
   "metadata": {},
   "execution_count": null,
   "outputs": [],
   "source": "# s is dropped; phi is kept and a small replacement fitted on its output,\n# using the few labels the later task has.\nprint(\"\\n[P-transfer] phi is what is carried over, and it is what pays\")\n\ndef fit_predict(feat_tr, lab_tr, feat_te):\n    \"\"\"Fit by minimizing the average squared error, then take the sign.\"\"\"\n    F = np.hstack([feat_tr, np.ones((len(feat_tr), 1))])\n    w, *_ = np.linalg.lstsq(F, lab_tr, rcond=None)\n    G = np.hstack([feat_te, np.ones((len(feat_te), 1))])\n    return np.sign(G @ w)\n\nNLAB = 10\nXte, zte = draw(NTEST)\nyte = np.sign(zte[:, 0])\nraw_hits, rep_hits = [], []\nfor _ in range(200):                                # many small labeled sets\n    Xtr, ztr = draw(NLAB)\n    ytr = np.sign(ztr[:, 0])\n    raw_hits.append(np.mean(fit_predict(Xtr[:, vis], ytr, Xte[:, vis]) == yte))\n    rep_hits.append(np.mean(fit_predict(Xtr[:, vis] @ phi, ytr,\n                                        Xte[:, vis] @ phi) == yte))\nraw, repr_ = float(np.mean(raw_hits)), float(np.mean(rep_hits))\nprint(f\"    held-out labels predicted correctly with {NLAB} labeled data \"\n      f\"points: raw features {raw:.3f}, representation {repr_:.3f}\")\ncheck(\"the representation beats the raw features when labels are scarce\",\n      repr_ > raw + 0.05)\ncheck(\"the pretext task never saw a label\", True)\n\n\n# ------------------------------------------------------------ preview\nfig, ax = plt.subplots(1, 2, figsize=(9, 3.2))\nax[0].bar([0, 1], [hit_const, hit_ctx], width=0.5,\n          color=[\"0.75\", \"0.35\"], edgecolor=\"black\")\nax[0].set_xticks([0, 1])\nax[0].set_xticklabels([\"most frequent\\ntoken\", \"from the\\ncontext\"])\nax[0].set_ylabel(\"fraction of held-out tokens right\")\nax[0].set_xlabel(\"fitted map\")\nax[0].set_title(\"[P-nlp] the next token is predictable\")\nax[1].bar([0, 1], [raw, repr_], width=0.5,\n          color=[\"0.75\", \"0.35\"], edgecolor=\"black\")\nax[1].set_xticks([0, 1])\nax[1].set_xticklabels([\"raw\\nfeatures\", \"learned\\nrepresentation\"])\nax[1].set_ylabel(\"fraction of held-out labels right\")\nax[1].set_xlabel(f\"features used, with {NLAB} labeled data points\")\nax[1].set_title(\"[P-transfer] what the pretext task leaves behind\")\nfig.tight_layout()\nfig.savefig(OUT_DIR / \"selfsupervisedlearning.png\", dpi=110)\n\nprint(f\"\\n{sum(ok for _, ok in report)}/{len(report)} checks passed\")\nassert all(ok for _, ok in report)"
  }
 ]
}