{ "cells": [ { "cell_type": "markdown", "id": "cell-00", "metadata": {}, "source": [ "# Learn image representations with I-JEPA\n", "\n", "The earlier notebooks train models that produce something: an image or text. This one trains an encoder whose only product is a vector for each image, and the goal is that similar images get similar vectors. We use I-JEPA (Assran et al., 2023). The encoder sees part of an image and has to predict the encoder's own embeddings of the hidden parts. It never reconstructs pixels, and it needs no hand-made augmentations such as crops and colour jitter.\n", "\n", "After training we check the vectors in three ways: a linear probe and a k-nearest-neighbour probe that try to read the flower species out of them, and a look at each query image's nearest neighbours.\n", "\n", "The notebook expects one NVIDIA GPU." ] }, { "cell_type": "markdown", "id": "cell-01", "metadata": {}, "source": [ "## Install" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-02", "metadata": {}, "outputs": [], "source": [ "%pip install -q \"dew-ml[streaming,interop] @ git+https://github.com/AshishKumar4/dew\" \"jax[cuda12]\"" ] }, { "cell_type": "markdown", "id": "cell-03", "metadata": {}, "source": [ "## Settings\n", "\n", "The images are 64x64 and the encoder cuts them into 8x8 patches, an 8 by 8 grid of 64 tokens. `NUM_TARGET_BLOCKS` and `BLOCK_SCALE` set the hidden regions: four blocks, each covering 15 to 20% of the image. The flower labels run from 1 to 102, so the probes use 103 classes." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-04", "metadata": {}, "outputs": [], "source": [ "IMAGE_SIZE = 64\n", "PATCH_SIZE = 8\n", "GRID = (IMAGE_SIZE // PATCH_SIZE, IMAGE_SIZE // PATCH_SIZE)\n", "BATCH_SIZE = 128\n", "STEPS = 4000\n", "LEARNING_RATE = 5e-4\n", "EMB_FEATURES = 192\n", "NUM_LAYERS = 6\n", "NUM_HEADS = 3\n", "NUM_TARGET_BLOCKS = 4\n", "BLOCK_SCALE = (0.15, 0.2)\n", "CLASSES = 103\n", "DATA_FILE = \"data/06-flowers.parquet\"\n", "RUN_DIR = \"runs/06-jepa\"\n", "SEED = 0" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-05", "metadata": {}, "outputs": [], "source": [ "from pathlib import Path\n", "\n", "import jax\n", "import matplotlib.pyplot as plt\n", "import numpy as np\n", "\n", "print(jax.devices())" ] }, { "cell_type": "markdown", "id": "cell-06", "metadata": {}, "source": [ "## The data\n", "\n", "We use the Hugging Face copy of Oxford Flowers again, this time for its species labels. That copy is sorted by species, and Dew holds out validation images from the head of the dataset, so the held-out set would be only two or three species. We shuffle the dataset once and save a local Parquet copy; `HFImages` reads it through the `parquet` loader of `datasets`.\n", "\n", "`augmentation=\"none\"` turns off the random flips: in I-JEPA the variation comes from the masks. `val_batches=8` holds out the first 1,024 shuffled images. The encoder never trains on them, and the probes score only them." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-07", "metadata": {}, "outputs": [], "source": [ "import datasets\n", "\n", "Path(DATA_FILE).parent.mkdir(parents=True, exist_ok=True)\n", "flowers = datasets.load_dataset(\"pranked03/flowers-blip-captions\", split=\"train\")\n", "flowers.shuffle(seed=SEED).to_parquet(DATA_FILE)\n", "print(flowers)" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-08", "metadata": {}, "outputs": [], "source": [ "from dew.data import HFImages, HFOptions, Loading\n", "\n", "data = HFImages(\n", " name=\"parquet\",\n", " options=HFOptions(data_files=DATA_FILE),\n", " image_size=IMAGE_SIZE,\n", " augmentation=\"none\",\n", " val_batches=8,\n", " loading=Loading(workers=0, threads=16, read_buffer=64),\n", ").load(batch=BATCH_SIZE)\n", "\n", "val_labels = np.concatenate([batch[\"label\"] for batch in data.val()])\n", "print(\"training images:\", data.records, \"| held-out images:\", len(val_labels),\n", " \"| species in the held-out set:\", len(np.unique(val_labels)))" ] }, { "cell_type": "markdown", "id": "cell-09", "metadata": {}, "source": [ "## The mask\n", "\n", "Each training image gets `NUM_TARGET_BLOCKS` rectangles of patches to predict, with a random size and aspect ratio, and the context the encoder sees is a random subset of the remaining patches. `multi_block_mask` works out the block sizes once for the grid, so every mask has the same number of tokens and the training step compiles once.\n", "\n", "Below is one sampled mask: `.` is context the encoder sees, `#` is a patch it has to predict, and `-` is dropped." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-10", "metadata": {}, "outputs": [], "source": [ "from dew.objectives.jepa import multi_block_mask\n", "\n", "mask = multi_block_mask(GRID, num_targets=NUM_TARGET_BLOCKS, scale=BLOCK_SCALE)\n", "print(\"context tokens:\", mask.num_context, \"| target blocks:\", mask.num_targets,\n", " \"of\", mask.block_area, \"tokens each\")\n", "\n", "context_idx, target_idx = mask.sample(jax.random.key(SEED), 1)\n", "view = np.full(GRID[0] * GRID[1], \"-\")\n", "view[np.asarray(context_idx[0])] = \".\"\n", "view[np.asarray(target_idx).reshape(-1)] = \"#\"\n", "print(\"\\n\".join(\" \".join(row) for row in view.reshape(GRID)))" ] }, { "cell_type": "markdown", "id": "cell-11", "metadata": {}, "source": [ "## Encoder, predictor and objective\n", "\n", "I-JEPA has three networks:\n", "\n", "- the context encoder, a ViT that sees only the context patches;\n", "- the target encoder, which sees the whole image and produces the embeddings to predict;\n", "- the predictor, a narrower transformer that takes the context embeddings plus the positions of the hidden patches and guesses the target encoder's embeddings there.\n", "\n", "The target encoder is not trained. It is an exponential moving average of the context encoder, so the targets improve as the encoder improves. `JepaObjective` puts the three together and asks the trainer to keep that average, which grows from a momentum of 0.996 to 1 over `momentum_steps`.\n", "\n", "The loss is the squared distance between the predicted and the target embeddings. The objective also reports `repr_std`, the spread of the embeddings across a batch. If it falls towards zero the encoder is giving every image the same vector, which is called collapse." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-12", "metadata": {}, "outputs": [], "source": [ "from dew import Field, models\n", "from dew.objectives.jepa import JepaObjective\n", "\n", "encoder = models.build(\"jepa_encoder\", patch_size=PATCH_SIZE, emb_features=EMB_FEATURES,\n", " num_layers=NUM_LAYERS, num_heads=NUM_HEADS,\n", " dtype=\"bfloat16\", attention_impl=\"auto\")\n", "predictor = models.build(\"jepa_predictor\", grid=GRID, emb_features=EMB_FEATURES,\n", " predictor_features=EMB_FEATURES // 2, num_layers=NUM_LAYERS // 2,\n", " num_heads=NUM_HEADS, dtype=\"bfloat16\", attention_impl=\"auto\")\n", "objective = JepaObjective(encoder, predictor, mask=mask,\n", " sample=Field(\"image\", (IMAGE_SIZE, IMAGE_SIZE, 3)),\n", " momentum_steps=STEPS)\n", "\n", "variables = jax.eval_shape(objective.init, jax.random.key(0))\n", "for name, tree in variables[\"params\"].items():\n", " print(f\"{name}: {sum(x.size for x in jax.tree_util.tree_leaves(tree)) / 1e6:.2f}M parameters\")" ] }, { "cell_type": "markdown", "id": "cell-13", "metadata": {}, "source": [ "## Training" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-14", "metadata": {}, "outputs": [], "source": [ "import optax\n", "from dew import Checkpoints, LocalTracker, Trainer\n", "\n", "trainer = Trainer(objective, optax.adamw(LEARNING_RATE), key=jax.random.key(SEED),\n", " checkpoints=Checkpoints(RUN_DIR), tracker=LocalTracker(f\"{RUN_DIR}/tracking\"))\n", "state = trainer.fit(data, steps=STEPS, log_every=500, checkpoint_every=STEPS)" ] }, { "cell_type": "code", "execution_count": null, "id": "cell-15", "metadata": {}, "outputs": [], "source": [ "import json\n", "\n", "rows = [json.loads(line) for line in open(f\"{RUN_DIR}/tracking/scalars.jsonl\")]\n", "rows = [row for row in rows if \"train/loss\" in row[\"scalars\"]]\n", "steps = [row[\"step\"] for row in rows]\n", "figure, (left, right) = plt.subplots(1, 2, figsize=(10, 3))\n", "left.plot(steps, [row[\"scalars\"][\"train/loss\"] for row in rows], marker=\".\")\n", "left.set(xlabel=\"step\", title=\"loss\")\n", "right.plot(steps, [row[\"scalars\"][\"train/repr_std\"] for row in rows], marker=\".\")\n", "right.set(xlabel=\"step\", title=\"repr_std\")\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "cell-16", "metadata": {}, "source": [ "## Probes on the held-out images\n", "\n", "`objective.evaluate` embeds a batch with the target encoder (the EMA weights) and averages each image's patch embeddings into one vector. We collect those vectors for all 1,024 held-out images.\n", "\n", "Each probe fits on half of the vectors and scores the other half. The linear probe is logistic regression; the k-NN probe gives each image the majority species of its 20 nearest neighbours. Both scores mean little alone, because a probe with more dimensions than images can fit almost anything. So we also run both probes with the labels shuffled. The gap between the real and the shuffled score is how much species information the vectors carry. Chance on 102 species is about 1%." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-17", "metadata": {}, "outputs": [], "source": [ "from dew.objectives.base import Step\n", "from dew.objectives.jepa import representation_health\n", "from dew.objectives.jepa.probes import knn_probe_accuracy, linear_probe_accuracy\n", "\n", "features, labels, held_out = [], [], []\n", "for batch in data.val():\n", " scored = objective.evaluate(state.params, batch,\n", " Step(step=state.step, key=jax.random.key(1), ema=state.averaged))\n", " features.append(np.asarray(scored.features))\n", " labels.append(np.asarray(scored.labels))\n", " held_out.append(batch[\"image\"])\n", "features, labels, held_out = np.concatenate(features), np.concatenate(labels), np.concatenate(held_out)\n", "print(\"held-out embeddings:\", features.shape)\n", "\n", "health = representation_health(features)\n", "print(f\"repr_std {float(health['repr_std']):.3f} | repr_cov_offdiag {float(health['repr_cov_offdiag']):.4f}\")\n", "\n", "shuffled = np.random.default_rng(0).permutation(labels)\n", "for name, probe in ((\"linear probe\", linear_probe_accuracy), (\"k-NN probe\", knn_probe_accuracy)):\n", " real = float(probe(features, labels, CLASSES))\n", " control = float(probe(features, shuffled, CLASSES))\n", " print(f\"{name}: {real:.3f} with the real labels, {control:.3f} with shuffled labels\")" ] }, { "cell_type": "markdown", "id": "cell-18", "metadata": {}, "source": [ "## Nearest neighbours\n", "\n", "The probes give one number each. We can also look at the vectors directly. Each row below is one held-out query image (left) and its five nearest held-out images by cosine similarity, with their species labels." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-19", "metadata": {}, "outputs": [], "source": [ "normalised = features / np.linalg.norm(features, axis=-1, keepdims=True)\n", "queries = [0, 1, 2, 3, 4, 5]\n", "figure, axes = plt.subplots(len(queries), 6, figsize=(7, len(queries) * 1.25))\n", "for row, query in enumerate(queries):\n", " neighbours = np.argsort(-(normalised @ normalised[query]))[1:6]\n", " for column, index in enumerate([query, *neighbours]):\n", " axis = axes[row, column]\n", " axis.imshow(held_out[index])\n", " axis.set_title(f\"{'query ' if column == 0 else ''}{labels[index]}\", fontsize=7)\n", " axis.axis(\"off\")\n", "plt.tight_layout()\n", "plt.show()" ] }, { "cell_type": "markdown", "id": "cell-20", "metadata": {}, "source": [ "## Keeping the encoder\n", "\n", "What we keep from a JEPA run is the EMA copy of the context encoder, without the predictor. `save_params` writes it as a safetensors file and `load_params` reads it back. `objective.encode` embeds images in `[-1, 1]` with a given set of encoder weights." ] }, { "cell_type": "code", "execution_count": null, "id": "cell-21", "metadata": {}, "outputs": [], "source": [ "from dew.interop import load_params, save_params\n", "\n", "encoder_params = state.averaged[\"params\"][\"context_encoder\"]\n", "save_params(encoder_params, f\"{RUN_DIR}/encoder.safetensors\")\n", "reloaded = load_params(f\"{RUN_DIR}/encoder.safetensors\")\n", "again = objective.encode(reloaded, held_out[:8].astype(np.float32) / 127.5 - 1)\n", "print(\"reloaded encoder output:\", np.asarray(again).shape)" ] }, { "cell_type": "markdown", "id": "cell-22", "metadata": {}, "source": [ "## Where to go next\n", "\n", "The I-JEPA paper trains a ViT-H/16 for 300 epochs on ImageNet; the knobs are the ones in the settings cell. `recipes/jepa/train.py` runs the same objective from the command line, and the `jepa_video_encoder` model with a factorized predictor does the same job on video clips." ] } ], "metadata": { "accelerator": "GPU", "kernelspec": { "display_name": "Python 3", "language": "python", "name": "python3" }, "language_info": { "name": "python" } }, "nbformat": 4, "nbformat_minor": 5 }