diff --git a/hackable_diffusion/notebooks/rpm.ipynb b/hackable_diffusion/notebooks/rpm.ipynb new file mode 100644 index 0000000..72ad62e --- /dev/null +++ b/hackable_diffusion/notebooks/rpm.ipynb @@ -0,0 +1,5164 @@ +{ + "cells": [ + { + "id": "cell-000", + "cell_type": "markdown", + "source": [ + "# RPM and SimCLR on CIFAR-10 — self-contained notebook\n", + "\n", + "This notebook is a **fully self-contained** port of a private codebase: it does\n", + "not `pip install` or `import` anything from that repo or from GitHub. Every\n", + "piece of code needed to run the two experiments — the Gaussian RPM\n", + "(Recognition-Parametrised Model) and SimCLR — is defined inline below, in the\n", + "same dependency order as the original modules:\n", + "\n", + "1. `dists` — Gaussian / diagonal-Gaussian / Categorical natural-parameter\n", + " distributions and their KL divergences.\n", + "2. `utils` — small initializers.\n", + "3. `networks` — recognition-network backbones (MLP / CNN / ResNet / ViT), each\n", + " producing a `(output, aux)` pair where `aux[\"trunk\"]` / `aux[\"projection\"]`\n", + " are the pre/post-projection representations.\n", + "4. `rpms` — the Gaussian and Categorical RPM models and their free-energy\n", + " objectives.\n", + "5. `simclr` — a SimCLR encoder that reuses the same network backbones (shared\n", + " weights across views, no per-factor `nn.vmap`) plus the NT-Xent loss.\n", + "6. CIFAR-10 data loading (multi-view SimCLR-style augmentation).\n", + "7. Linear-probe evaluation utilities.\n", + "8. The generic training loop (works for both RPM and SimCLR — the loss\n", + " function is dispatched by `model.rpm_type`).\n", + "\n", + "Then two runnable experiments: **Experiment 1** trains a Gaussian RPM,\n", + "**Experiment 2** trains SimCLR, both with a ResNet-18-style backbone on\n", + "CIFAR-10 and periodic linear-probe evaluation.\n", + "\n", + "**Note on fidelity:** this is a faithful port of the working code, with three\n", + "small correctness fixes applied along the way (each is called out inline with\n", + "a `# FIX:` comment where it happens):\n", + "- `create_feature_extractor`'s `\"latent\"` branch used to read `.mean` off a\n", + " natural-parameter object that doesn't have that attribute — fixed to read\n", + " the separately-computed mean-parameter object instead.\n", + "- Its `\"concatenated\"` branch used to discard the tuple element that actually\n", + " holds the latent mean — fixed to keep it.\n", + "- The best-checkpoint selection compared free energy the wrong way (it's\n", + " *maximized* during training — verified empirically: free energy climbs over\n", + " training — but the checkpoint logic was tracking a *minimum*). Fixed to\n", + " track the maximum.\n", + "\n", + "Everything else is preserved as-is, bugs and all, per request.\n" + ], + "metadata": { + "id": "cell-000" + } + }, + { + "id": "cell-001", + "cell_type": "markdown", + "source": [ + "## Setup\n" + ], + "metadata": { + "id": "cell-001" + } + }, + { + "id": "cell-002", + "cell_type": "code", + "source": [ + "# !pip install -q \"jax[cpu]\" flax optax orbax-checkpoint einops 2>/dev/null\n", + "# # If you have a GPU runtime in Colab (Runtime > Change runtime type > GPU),\n", + "# # install the CUDA build of jax instead for a large speedup, e.g.:\n", + "# # !pip install -q -U \"jax[cuda12]\"\n", + "# print(\"Install complete.\")\n" + ], + "metadata": { + "id": "cell-002" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-003", + "cell_type": "code", + "source": [ + "import math\n", + "import os\n", + "import sys\n", + "import random\n", + "import time\n", + "from pathlib import Path\n", + "from typing import Any, Iterable, Iterator, Mapping, Optional, Sequence, Tuple\n", + "\n", + "# import einops\n", + "import jax\n", + "import jax.numpy as jnp\n", + "import matplotlib.pyplot as plt\n", + "import numpy as np\n", + "import optax\n", + "import orbax.checkpoint as ocp\n", + "\n", + "\n", + "from flax import linen as nn\n", + "from flax import serialization, struct\n", + "from flax.typing import VariableDict\n", + "from jax import Array, jit, value_and_grad, vmap\n", + "from jax.lax import batch_matmul, fori_loop\n", + "from jax.nn import log_softmax, softmax\n", + "from jax.scipy.special import logsumexp\n", + "\n", + "print(\"JAX devices:\", jax.devices())\n" + ], + "metadata": { + "id": "cell-003" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-004", + "cell_type": "markdown", + "source": [ + "## Part 1 — Distributions (`dists`)\n", + "\n", + "Natural-parameter containers for full-covariance Gaussian, diagonal Gaussian,\n", + "and Categorical distributions, plus their KL divergences and log-normalizers.\n" + ], + "metadata": { + "id": "cell-004" + } + }, + { + "id": "cell-005", + "cell_type": "code", + "source": [ + "@struct.dataclass\n", + "class GaussianMeanParams:\n", + " \"\"\"Mean parameters of a full-covariance Gaussian.\"\"\"\n", + "\n", + " mean: Array\n", + " cov: Array\n", + "\n", + "\n", + "@struct.dataclass\n", + "class DiagGaussianMeanParams:\n", + " \"\"\"Mean parameters of a diagonal Gaussian.\"\"\"\n", + "\n", + " mean: Array\n", + " var_diag: Array\n", + "\n", + "\n", + "@jit\n", + "def mvn_lognormalizer(precision: Array, pwm: Array, mean: Array | None = None) -> Array:\n", + " \"\"\"Log-partition of a Gaussian given precision and precision-weighted mean.\n", + "\n", + " If ``mean`` (= precision^{-1} @ pwm) is already known it can be passed to\n", + " avoid re-solving the linear system.\n", + " \"\"\"\n", + " d = pwm.shape[-1]\n", + "\n", + " if mean is None:\n", + " mean = jnp.linalg.solve(precision, pwm[..., None])[..., 0]\n", + "\n", + " offset = d * jnp.log(2 * jnp.pi)\n", + " dets = jnp.linalg.slogdet(precision)[1]\n", + " quad = jnp.sum(pwm * mean, axis=-1)\n", + " lognorm = 0.5 * (offset - dets + quad)\n", + "\n", + " return lognorm\n", + "\n", + "\n", + "@jit\n", + "def mvn_log_prob(\n", + " x: Array, mean_params: GaussianMeanParams | DiagGaussianMeanParams\n", + ") -> Array:\n", + " \"\"\"Log-probability of a Gaussian given mean parameters (mean and covariance).\"\"\"\n", + " mean = mean_params[\"mean\"]\n", + " if \"cov\" in mean_params:\n", + " cov = mean_params[\"cov\"]\n", + " d = mean.shape[-1]\n", + " offset = d * jnp.log(2 * jnp.pi)\n", + " dets = jnp.linalg.slogdet(cov)[1]\n", + " quad = jnp.sum(\n", + " (x - mean)\n", + " * batch_matmul(jnp.linalg.inv(cov), (x - mean)[..., None])[..., 0],\n", + " axis=-1,\n", + " )\n", + " return -0.5 * (offset + dets + quad)\n", + " elif \"var_diag\" in mean_params:\n", + " var_diag = mean_params[\"var_diag\"]\n", + " d = mean.shape[-1]\n", + " offset = d * jnp.log(2 * jnp.pi)\n", + " log_det = jnp.sum(jnp.log(var_diag), axis=-1)\n", + " quad = jnp.sum(((x - mean) ** 2) / var_diag, axis=-1)\n", + " return -0.5 * (offset + log_det + quad)\n", + " else:\n", + " raise ValueError(\"mean_params must contain either 'cov' or 'var_diag'.\")\n", + "\n", + "\n", + "@struct.dataclass\n", + "class GaussianNatParams:\n", + " \"\"\"Batched natural parameters of a full-covariance Gaussian.\n", + "\n", + " Supports slicing and broadcasting arithmetic (sum / subtract / scale).\n", + " \"\"\"\n", + "\n", + " precision: Array\n", + " precision_weighted_mean: Array\n", + "\n", + " def event_shape(self) -> int:\n", + " return self.precision_weighted_mean.shape[-1]\n", + "\n", + " def batch_shape(self) -> tuple[int, ...]:\n", + " return self.precision_weighted_mean.shape[:-1]\n", + "\n", + " def sum(self, batch_axis: int) -> \"GaussianNatParams\":\n", + " # Only leading (batch) axes may be summed; the trailing axis is the\n", + " # event dimension.\n", + " if batch_axis > len(self.precision_weighted_mean.shape) - 2:\n", + " raise ValueError(\n", + " \"Can only sum over batch axes, batch shape is: \"\n", + " f\"{self.precision_weighted_mean.shape[:-1]} got axis {batch_axis}\"\n", + " )\n", + "\n", + " return GaussianNatParams(\n", + " self.precision.sum(batch_axis), self.precision_weighted_mean.sum(batch_axis)\n", + " )\n", + "\n", + " def expand_batch_dim(self, axis: int) -> \"GaussianNatParams\":\n", + " return GaussianNatParams(\n", + " jnp.expand_dims(self.precision, axis),\n", + " jnp.expand_dims(self.precision_weighted_mean, axis),\n", + " )\n", + "\n", + " def __mul__(self, alpha: float) -> \"GaussianNatParams\":\n", + " return GaussianNatParams(\n", + " alpha * self.precision, alpha * self.precision_weighted_mean\n", + " )\n", + "\n", + " def __rmul__(self, alpha: float) -> \"GaussianNatParams\":\n", + " return self.__mul__(alpha)\n", + "\n", + " def __add__(self, other: \"GaussianNatParams\") -> \"GaussianNatParams\":\n", + " shape_a, shape_b = self.precision.shape, other.precision.shape\n", + " try:\n", + " jnp.broadcast_shapes(shape_a, shape_b)\n", + " except ValueError:\n", + " raise ValueError(\n", + " f\"Shapes must be broadcastable, got {shape_a} and {shape_b}\"\n", + " )\n", + "\n", + " return GaussianNatParams(\n", + " self.precision + other.precision,\n", + " self.precision_weighted_mean + other.precision_weighted_mean,\n", + " )\n", + "\n", + " def __sub__(self, other: \"GaussianNatParams\") -> \"GaussianNatParams\":\n", + " shape_a, shape_b = self.precision.shape, other.precision.shape\n", + " try:\n", + " jnp.broadcast_shapes(shape_a, shape_b)\n", + " except ValueError:\n", + " raise ValueError(\n", + " f\"Shapes must be broadcastable, got {shape_a} and {shape_b}\"\n", + " ) from None\n", + "\n", + " return GaussianNatParams(\n", + " self.precision - other.precision,\n", + " self.precision_weighted_mean - other.precision_weighted_mean,\n", + " )\n", + "\n", + " def __getitem__(self, s) -> \"GaussianNatParams\":\n", + " return GaussianNatParams(self.precision[s], self.precision_weighted_mean[s])\n", + "\n", + " def lognormalizer(self, mean: Array | None = None) -> Array:\n", + " return mvn_lognormalizer(self.precision, self.precision_weighted_mean, mean)\n", + "\n", + " @classmethod\n", + " def from_mean_param(cls, mean_params: GaussianMeanParams) -> \"GaussianNatParams\":\n", + " precision = jnp.linalg.inv(mean_params.cov)\n", + " precision_weighted_mean = batch_matmul(\n", + " precision, jnp.expand_dims(mean_params.mean, -1)\n", + " ).squeeze()\n", + " return GaussianNatParams(precision, precision_weighted_mean)\n", + "\n", + "\n", + "@jit\n", + "def list_to_batch(dists: Sequence[GaussianNatParams]) -> GaussianNatParams:\n", + " \"\"\"Stack a sequence of identically-shaped Gaussians along a new axis 0.\"\"\"\n", + " J = len(dists)\n", + " shape = dists[0].batch_shape()\n", + " new_shape = (J,) + shape\n", + " event_shape = dists[0].event_shape()\n", + "\n", + " p = jnp.zeros(new_shape + (event_shape, event_shape))\n", + " pwm = jnp.zeros(new_shape + (event_shape,))\n", + " for j in range(J):\n", + " p = p.at[j].set(dists[j].precision)\n", + " pwm = pwm.at[j].set(dists[j].precision_weighted_mean)\n", + "\n", + " return GaussianNatParams(precision=p, precision_weighted_mean=pwm)\n", + "\n", + "\n", + "@jit\n", + "def mean_params(d: GaussianNatParams) -> GaussianMeanParams:\n", + " \"\"\"Convert natural params to mean params (mean and covariance).\"\"\"\n", + " cov = jnp.linalg.inv(d.precision)\n", + " mean = batch_matmul(cov, jnp.expand_dims(d.precision_weighted_mean, -1)).squeeze(-1)\n", + " return GaussianMeanParams(mean=mean, cov=cov)\n", + "\n", + "\n", + "@jit\n", + "def unnormalized_kl(\n", + " q_mean: Array,\n", + " q_cov: Array,\n", + " q_nat: GaussianNatParams,\n", + " f_nat: GaussianNatParams,\n", + " h_nat: GaussianNatParams | None = None,\n", + " f_phi: Array | None = None,\n", + " h_phi: Array | None = None,\n", + " q_phi: Array | None = None,\n", + ") -> Array:\n", + " \"\"\"Pseudo-KL from ``q`` to the unnormalized density ``f * h``.\n", + "\n", + " Log-normalizers are recomputed only when not supplied by the caller.\n", + " \"\"\"\n", + " if f_phi is None:\n", + " f_phi = mvn_lognormalizer(f_nat.precision, f_nat.precision_weighted_mean)\n", + "\n", + " if q_phi is None:\n", + " q_phi = mvn_lognormalizer(\n", + " q_nat.precision, q_nat.precision_weighted_mean, q_mean\n", + " )\n", + "\n", + " # E_q[z z^T] = Cov + mean mean^T.\n", + " # q_square = q_cov + einops.einsum(q_mean, q_mean, \"... d1, ... d2 -> ... d1 d2\")\n", + " q_square = q_cov + jn.einsum(q_mean, q_mean, \"... d1, ... d2 -> ... d1 d2\")\n", + "\n", + " if h_nat is None:\n", + " normalizers = q_phi - f_phi\n", + " nat_diff = f_nat - q_nat\n", + " else:\n", + " if h_phi is None:\n", + " h_phi = mvn_lognormalizer(h_nat.precision, h_nat.precision_weighted_mean)\n", + " normalizers = q_phi - f_phi - h_phi\n", + " nat_diff = f_nat + h_nat - q_nat\n", + "\n", + " linear = jnp.sum(nat_diff.precision_weighted_mean * q_mean, axis=-1)\n", + " quadratic = -0.5 * jnp.sum(nat_diff.precision * q_square, axis=(-1, -2))\n", + " return -1.0 * (linear + quadratic + normalizers)\n", + "\n", + "\n", + "@jit\n", + "def kl(\n", + " q_mean: Array,\n", + " q_cov: Array,\n", + " q_nat: GaussianNatParams,\n", + " p_nat: GaussianNatParams,\n", + " q_phi: Array | None = None,\n", + " p_phi: Array | None = None,\n", + ") -> Array:\n", + " \"\"\"KL divergence ``KL(q || p)`` between two full Gaussians.\"\"\"\n", + " if q_phi is None:\n", + " q_phi = mvn_lognormalizer(\n", + " q_nat.precision, q_nat.precision_weighted_mean, q_mean\n", + " )\n", + " if p_phi is None:\n", + " p_phi = mvn_lognormalizer(p_nat.precision, p_nat.precision_weighted_mean)\n", + "\n", + " # q_square = q_cov + einops.einsum(q_mean, q_mean, \"... d1, ... d2 -> ... d1 d2\")\n", + " q_square = q_cov + jnp.einsum(q_mean, q_mean, \"... d1, ... d2 -> ... d1 d2\")\n", + "\n", + " normalizers = q_phi - p_phi\n", + " nat_diff = p_nat - q_nat\n", + "\n", + " linear = jnp.sum(nat_diff.precision_weighted_mean * q_mean, axis=-1)\n", + " quadratic = -0.5 * jnp.sum(nat_diff.precision * q_square, axis=(-1, -2))\n", + "\n", + " return -1.0 * (linear + quadratic + normalizers)\n", + "\n", + "\n", + "# --- Diagonal Gaussian ---\n", + "\n", + "\n", + "@struct.dataclass\n", + "class DiagGaussianNatParams:\n", + " \"\"\"Gaussian natural parameters with a diagonal precision.\"\"\"\n", + "\n", + " precision_diag: Array\n", + " precision_weighted_mean: Array\n", + "\n", + " def event_shape(self) -> int:\n", + " return self.precision_weighted_mean.shape[-1]\n", + "\n", + " def batch_shape(self) -> tuple[int, ...]:\n", + " return self.precision_weighted_mean.shape[:-1]\n", + "\n", + " def sum(self, batch_axis: int) -> \"DiagGaussianNatParams\":\n", + " if batch_axis > len(self.precision_weighted_mean.shape) - 2:\n", + " raise ValueError(\n", + " \"Can only sum over batch axes, batch shape is: \"\n", + " f\"{self.precision_weighted_mean.shape[:-1]} got axis {batch_axis}\"\n", + " )\n", + "\n", + " return DiagGaussianNatParams(\n", + " self.precision_diag.sum(batch_axis),\n", + " self.precision_weighted_mean.sum(batch_axis),\n", + " )\n", + "\n", + " def expand_batch_dim(self, axis: int) -> \"DiagGaussianNatParams\":\n", + " return DiagGaussianNatParams(\n", + " jnp.expand_dims(self.precision_diag, axis),\n", + " jnp.expand_dims(self.precision_weighted_mean, axis),\n", + " )\n", + "\n", + " def __mul__(self, alpha: float) -> \"DiagGaussianNatParams\":\n", + " return DiagGaussianNatParams(\n", + " alpha * self.precision_diag, alpha * self.precision_weighted_mean\n", + " )\n", + "\n", + " def __rmul__(self, alpha: float) -> \"DiagGaussianNatParams\":\n", + " return self.__mul__(alpha)\n", + "\n", + " def __add__(self, other: \"DiagGaussianNatParams\") -> \"DiagGaussianNatParams\":\n", + " return DiagGaussianNatParams(\n", + " self.precision_diag + other.precision_diag,\n", + " self.precision_weighted_mean + other.precision_weighted_mean,\n", + " )\n", + "\n", + " def __sub__(self, other: \"DiagGaussianNatParams\") -> \"DiagGaussianNatParams\":\n", + " return DiagGaussianNatParams(\n", + " self.precision_diag - other.precision_diag,\n", + " self.precision_weighted_mean - other.precision_weighted_mean,\n", + " )\n", + "\n", + " def __getitem__(self, s) -> \"DiagGaussianNatParams\":\n", + " return DiagGaussianNatParams(\n", + " self.precision_diag[s], self.precision_weighted_mean[s]\n", + " )\n", + "\n", + " def lognormalizer(self, mean: Array | None = None) -> Array:\n", + " return diag_mvn_lognormalizer(\n", + " self.precision_diag, self.precision_weighted_mean, mean\n", + " )\n", + "\n", + " @classmethod\n", + " def from_mean_param(\n", + " cls, mean_params: DiagGaussianMeanParams\n", + " ) -> \"DiagGaussianNatParams\":\n", + " precision_diag = 1.0 / mean_params.var_diag\n", + " precision_weighted_mean = mean_params.mean * precision_diag\n", + " return DiagGaussianNatParams(precision_diag, precision_weighted_mean)\n", + "\n", + "\n", + "@jit\n", + "def diag_mvn_lognormalizer(\n", + " precision_diag: Array, pwm: Array, mean: Array | None = None\n", + ") -> Array:\n", + " \"\"\"Log-partition of a diagonal Gaussian (the diagonal analogue of\n", + " :func:`mvn_lognormalizer`).\"\"\"\n", + " d = pwm.shape[-1]\n", + " if mean is None:\n", + " mean = pwm / precision_diag\n", + " offset = d * jnp.log(2 * jnp.pi)\n", + " log_det = jnp.sum(jnp.log(precision_diag), axis=-1)\n", + " quad = jnp.sum(pwm * mean, axis=-1)\n", + " lognorm = 0.5 * (offset - log_det + quad)\n", + " return lognorm\n", + "\n", + "\n", + "@jit\n", + "def diag_mean_params(d: DiagGaussianNatParams) -> DiagGaussianMeanParams:\n", + " \"\"\"Convert diagonal natural params to mean params (mean and variances).\"\"\"\n", + " var_diag = 1.0 / d.precision_diag\n", + " mean = d.precision_weighted_mean * var_diag\n", + " return DiagGaussianMeanParams(mean, var_diag)\n", + "\n", + "\n", + "@jit\n", + "def diag_unnormalized_kl(\n", + " q_mean: Array,\n", + " q_var_diag: Array,\n", + " q_nat: DiagGaussianNatParams,\n", + " f_nat: DiagGaussianNatParams,\n", + " h_nat: DiagGaussianNatParams | None = None,\n", + " f_phi: Array | None = None,\n", + " h_phi: Array | None = None,\n", + " q_phi: Array | None = None,\n", + ") -> Array:\n", + " \"\"\"Diagonal-Gaussian analogue of :func:`unnormalized_kl`.\"\"\"\n", + " f_phi = (\n", + " f_phi\n", + " if f_phi is not None\n", + " else diag_mvn_lognormalizer(f_nat.precision_diag, f_nat.precision_weighted_mean)\n", + " )\n", + " q_phi = (\n", + " q_phi\n", + " if q_phi is not None\n", + " else diag_mvn_lognormalizer(\n", + " q_nat.precision_diag, q_nat.precision_weighted_mean, q_mean\n", + " )\n", + " )\n", + "\n", + " # E_q[z^2] = Var + mean^2 (per coordinate).\n", + " q_square_diag = q_var_diag + q_mean**2\n", + "\n", + " if h_nat is None:\n", + " normalizers = q_phi - f_phi\n", + " nat_diff = f_nat - q_nat\n", + " else:\n", + " h_phi = (\n", + " h_phi\n", + " if h_phi is not None\n", + " else diag_mvn_lognormalizer(\n", + " h_nat.precision_diag, h_nat.precision_weighted_mean\n", + " )\n", + " )\n", + " normalizers = q_phi - f_phi - h_phi\n", + " nat_diff = f_nat + h_nat - q_nat\n", + "\n", + " linear = jnp.sum(nat_diff.precision_weighted_mean * q_mean, axis=-1)\n", + " quadratic = -0.5 * jnp.sum(nat_diff.precision_diag * q_square_diag, axis=-1)\n", + "\n", + " return -1.0 * (linear + quadratic + normalizers)\n", + "\n", + "\n", + "@jit\n", + "def diag_kl(\n", + " q_mean: Array,\n", + " q_var_diag: Array,\n", + " q_nat: DiagGaussianNatParams,\n", + " p_nat: DiagGaussianNatParams,\n", + " q_phi: Array | None = None,\n", + " p_phi: Array | None = None,\n", + ") -> Array:\n", + " \"\"\"Diagonal-Gaussian analogue of :func:`kl` (``KL(q || p)``).\"\"\"\n", + " q_phi = (\n", + " q_phi\n", + " if q_phi is not None\n", + " else diag_mvn_lognormalizer(\n", + " q_nat.precision_diag, q_nat.precision_weighted_mean, q_mean\n", + " )\n", + " )\n", + " p_phi = (\n", + " p_phi\n", + " if p_phi is not None\n", + " else diag_mvn_lognormalizer(p_nat.precision_diag, p_nat.precision_weighted_mean)\n", + " )\n", + "\n", + " q_square_diag = q_var_diag + q_mean**2\n", + "\n", + " normalizers = q_phi - p_phi\n", + " nat_diff = p_nat - q_nat\n", + "\n", + " linear = jnp.sum(nat_diff.precision_weighted_mean * q_mean, axis=-1)\n", + " quadratic = -0.5 * jnp.sum(nat_diff.precision_diag * q_square_diag, axis=-1)\n", + "\n", + " return -1.0 * (linear + quadratic + normalizers)\n", + "\n", + "\n", + "# --- Categorical ---\n", + "\n", + "\n", + "@struct.dataclass\n", + "class CategoricalNatParams:\n", + " \"\"\"Natural parameters (logits) for a Categorical distribution.\"\"\"\n", + "\n", + " logits: Array\n", + "\n", + " def event_shape(self) -> int:\n", + " return self.logits.shape[-1]\n", + "\n", + " def batch_shape(self) -> tuple[int, ...]:\n", + " return self.logits.shape[:-1]\n", + "\n", + " @property\n", + " def probs(self) -> Array:\n", + " return softmax(self.logits, axis=-1)\n", + "\n", + " @property\n", + " def log_probs(self) -> Array:\n", + " return log_softmax(self.logits, axis=-1)\n", + "\n", + " def __add__(self, other: \"CategoricalNatParams\") -> \"CategoricalNatParams\":\n", + " return CategoricalNatParams(logits=self.logits + other.logits)\n", + "\n", + " def __sub__(self, other: \"CategoricalNatParams\") -> \"CategoricalNatParams\":\n", + " return CategoricalNatParams(logits=self.logits - other.logits)\n", + "\n", + " def __mul__(self, alpha: float) -> \"CategoricalNatParams\":\n", + " return CategoricalNatParams(logits=self.logits * alpha)\n", + "\n", + " def __rmul__(self, alpha: float) -> \"CategoricalNatParams\":\n", + " return self.__mul__(alpha)\n", + "\n", + " def __getitem__(self, s) -> \"CategoricalNatParams\":\n", + " return CategoricalNatParams(logits=self.logits[s])\n", + "\n", + " def sum(self, batch_axis: int) -> \"CategoricalNatParams\":\n", + " return CategoricalNatParams(logits=self.logits.sum(batch_axis))\n", + "\n", + "\n", + "def categorical_cross_entropy(\n", + " p: CategoricalNatParams,\n", + " q: CategoricalNatParams,\n", + ") -> Array:\n", + " \"\"\"Expected log-likelihood ``E_p[log q]`` (the negative cross-entropy).\n", + "\n", + " Note this returns ``sum_x p(x) log q(x)``, i.e. the *negative* of the\n", + " cross-entropy ``H(p, q)``; ``-categorical_cross_entropy(q, q)`` is the\n", + " Shannon entropy of ``q``.\n", + " \"\"\"\n", + " return jnp.sum(p.probs * q.log_probs, axis=-1)\n", + "\n", + "\n", + "print(\"dists defined.\")\n" + ], + "metadata": { + "id": "cell-005" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-006", + "cell_type": "markdown", + "source": [ + "## Part 2 — Utilities (`utils`)\n", + "\n", + "Small array helpers and a Flax-style initializer used by the RPM prior.\n" + ], + "metadata": { + "id": "cell-006" + } + }, + { + "id": "cell-007", + "cell_type": "code", + "source": [ + "def batched_outer_product(A: Array, B: Array) -> Array:\n", + " \"\"\"Batched outer product ``C[idx, d1, d2] = A[idx, d1] * B[idx, d2]``.\n", + "\n", + " ``idx`` indexes all leading (batch) dimensions, which must match exactly\n", + " between ``A`` and ``B``; the last axis is the event dimension.\n", + " \"\"\"\n", + " batch_shape_A = A.shape[:-1]\n", + " batch_shape_B = B.shape[:-1]\n", + "\n", + " dim_event_a = A.shape[-1]\n", + " dim_event_b = B.shape[-1]\n", + "\n", + " if batch_shape_A != batch_shape_B:\n", + " raise ValueError(\n", + " f\"Batch dimensions must agree. Got full shapes {A.shape}, {B.shape}\"\n", + " )\n", + "\n", + " new_shape = batch_shape_A + (dim_event_a, dim_event_b)\n", + "\n", + " # Flatten the batch dims, take a per-row outer product, then restore shape.\n", + " a_flat = jnp.reshape(A, (-1, dim_event_a))\n", + " b_flat = jnp.reshape(B, (-1, dim_event_b))\n", + " out_flat = vmap(jnp.outer, in_axes=[0, 0])(a_flat, b_flat)\n", + "\n", + " return jnp.reshape(out_flat, new_shape)\n", + "\n", + "\n", + "def batch_identity_init(key: Any, batch: int | Sequence[int], dim: int) -> Array:\n", + " \"\"\"Initializer: a ``(*batch, dim, dim)`` stack of identity matrices.\n", + "\n", + " Matches the flax initializer calling convention; ``key`` is ignored.\n", + " \"\"\"\n", + " if isinstance(batch, int):\n", + " batch = (batch,)\n", + " else:\n", + " batch = tuple(batch)\n", + "\n", + " shape = batch + (dim, dim)\n", + " M = jnp.zeros(shape)\n", + " return M.at[..., jnp.arange(dim), jnp.arange(dim)].set(1.0)\n", + "\n", + "\n", + "def identity_init(key: Any, dim: Sequence[int], scale: float = 2**-0.5) -> Array:\n", + " \"\"\"Initializer: a ``(dim[0], dim[0])`` identity matrix scaled by ``scale``.\n", + "\n", + " ``dim`` is the flax-style shape tuple (only its first entry is used);\n", + " ``key`` is ignored.\n", + " \"\"\"\n", + " return scale * jnp.eye(dim[0])\n", + "\n", + "\n", + "print(\"utils defined.\")\n" + ], + "metadata": { + "id": "cell-007" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-008", + "cell_type": "markdown", + "source": [ + "## Part 3 — Recognition networks (`networks`)\n", + "\n", + "MLP / CNN / ResNet / ViT backbones. Each ``__call__`` returns\n", + "``(output, aux)`` where ``aux = {\"trunk\": ..., \"projection\": ...}`` — the\n", + "pre-projection representation (SimCLR's ``h``) and the post-projection one\n", + "(SimCLR's ``z``). The Gaussian-output classes feed the projection through a\n", + "final natural-parameter head; the Categorical ones feed it through a logits\n", + "head instead.\n", + "\n", + "(``LinearRecognition``/``UnnormalizedRecognition`` from the original file are\n", + "omitted here: the latter is dead code — unused by any ``encoder_arch``\n", + "dispatch and its ``lognormalizer`` deliberately raises\n", + "``NotImplementedError`` — and the former has no trunk/projection concept at\n", + "all, so neither is reachable by the two experiments below.)\n" + ], + "metadata": { + "id": "cell-008" + } + }, + { + "id": "cell-009", + "cell_type": "code", + "source": [ + "Aux = Mapping[str, Any]\n", + "\n", + "\n", + "def vec2tri(v: Array, n: int) -> Array:\n", + " \"\"\"Scatter a vector into the strictly-lower-triangular part of an ``n x n``\n", + " matrix (the diagonal and upper triangle stay zero).\n", + "\n", + " ``v`` must have length ``n * (n - 1) / 2``.\n", + " \"\"\"\n", + " M = jnp.zeros((n, n))\n", + " idxs = jnp.tril_indices(n, -1, n)\n", + " M = M.at[idxs].set(v)\n", + " return M\n", + "\n", + "\n", + "def gaussian_nat_head(\n", + " x: Array,\n", + " dim_out: int,\n", + " precision_type: str = \"full\"\n", + ") -> DiagGaussianNatParams | GaussianNatParams:\n", + " pwm = nn.Dense(dim_out, name=\"pwm head\")(x)\n", + "\n", + " if precision_type == \"fixed_diag\":\n", + " p_diag = jnp.ones(dim_out)\n", + " p_diag = jnp.broadcast_to(p_diag, pwm.shape)\n", + " return DiagGaussianNatParams(precision_diag=p_diag,\n", + " precision_weighted_mean=pwm)\n", + " elif precision_type == \"diag\":\n", + " p_diag = nn.softplus(nn.Dense(dim_out, name=\"precision diag head\")(x))\n", + " return DiagGaussianNatParams(precision_diag=p_diag,\n", + " precision_weighted_mean=pwm)\n", + " elif precision_type == \"full\":\n", + " L_diag = nn.softplus(nn.Dense(dim_out, name=\"precision diag head\")(x))\n", + " L_offd = nn.Dense((dim_out * (dim_out - 1)) // 2,\n", + " name=\"precision offdiag head\")(x)\n", + " L = vmap(jit(vec2tri, static_argnums=1), in_axes=[0, None])(\n", + " L_offd, dim_out) + vmap(jnp.diag, in_axes=0)(L_diag)\n", + " p = jnp.matmul(L, L.transpose(0, 2, 1))\n", + " return GaussianNatParams(precision=p, precision_weighted_mean=pwm)\n", + " else:\n", + " raise ValueError(f\"Invalid 'precision_type': {precision_type}\")\n", + "\n", + "\n", + "class NNRecognition(nn.Module):\n", + " \"\"\"MLP recognition network (flattens its input).\"\"\"\n", + "\n", + " features: Sequence[int]\n", + " dim_out: int\n", + " projection_features: Sequence[int] = ()\n", + " precision_type: str = \"diag\"\n", + "\n", + " @property\n", + " def trunk_dim(self):\n", + " return self.features[-1]\n", + "\n", + " @property\n", + " def projection_dim(self):\n", + " return self.projection_features[-1]\n", + "\n", + " @property\n", + " def latent_dim(self):\n", + " return self.dim_out\n", + "\n", + " @nn.compact\n", + " def __call__(\n", + " self, inputs: Array\n", + " ) -> tuple[GaussianNatParams | DiagGaussianNatParams, Aux]:\n", + " x = inputs\n", + " x = x.reshape((x.shape[0], -1))\n", + " for i, feat in enumerate(self.features):\n", + " x = nn.Dense(feat, name=f\"Trunk layer {i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " trunk = x\n", + "\n", + " # Projection layers\n", + " for i, d in enumerate(self.projection_features):\n", + " x = nn.Dense(d, name=f\"projection_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " projection = x\n", + "\n", + " aux = {\"trunk\": trunk, \"projection\": projection}\n", + "\n", + " nat_params = gaussian_nat_head(x, self.dim_out, self.precision_type)\n", + "\n", + " return nat_params, aux\n", + "\n", + "\n", + "class CNNRecognition(nn.Module):\n", + " \"\"\"Convolutional recognition network (conv stack -> MLP head).\"\"\"\n", + "\n", + " conv_features: Sequence[Tuple[int, Tuple[int, int]]]\n", + " fc_features: Sequence[int]\n", + " dim_out: int\n", + " projection_features: Sequence[int] = ()\n", + " precision_type: str = \"diag\"\n", + "\n", + " @property\n", + " def trunk_dim(self):\n", + " return self.fc_features[-1]\n", + "\n", + " @property\n", + " def projection_dim(self):\n", + " return self.projection_features[-1]\n", + "\n", + " @property\n", + " def latent_dim(self):\n", + " return self.dim_out\n", + "\n", + " @nn.compact\n", + " def __call__(\n", + " self, inputs: Array\n", + " ) -> tuple[GaussianNatParams | DiagGaussianNatParams, Aux]:\n", + " x = inputs\n", + " for i, (channels, kernel_size) in enumerate(self.conv_features):\n", + " x = nn.Conv(channels, kernel_size, name=f\"conv {i}\")(x)\n", + " x = nn.relu(x)\n", + " x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n", + "\n", + " x = x.reshape((x.shape[0], -1))\n", + " for i, d in enumerate(self.fc_features):\n", + " x = nn.Dense(d, name=f\"fc {i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " trunk = x\n", + "\n", + " # Projection layers\n", + " for i, d in enumerate(self.projection_features):\n", + " x = nn.Dense(d, name=f\"projection_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " projection = x\n", + "\n", + " nat_params = gaussian_nat_head(x, self.dim_out, self.precision_type)\n", + "\n", + " aux = {'trunk': trunk, 'projection': projection}\n", + " return nat_params, aux\n", + "\n", + "\n", + "class ResidualBlock(nn.Module):\n", + " \"\"\"Pre-activation-style residual block with GroupNorm.\"\"\"\n", + "\n", + " channels: int\n", + " stride: int = 1\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> Array:\n", + " x = inputs\n", + " residual = inputs\n", + "\n", + " x = nn.Conv(\n", + " self.channels,\n", + " (3, 3),\n", + " strides=(self.stride, self.stride),\n", + " padding=\"SAME\",\n", + " use_bias=False,\n", + " )(x)\n", + " x = nn.GroupNorm(num_groups=min(32, self.channels))(x)\n", + " x = nn.relu(x)\n", + "\n", + " x = nn.Conv(self.channels, (3, 3), padding=\"SAME\", use_bias=False)(x)\n", + " x = nn.GroupNorm(num_groups=min(32, self.channels))(x)\n", + "\n", + " # Project the skip connection when shape / stride changes.\n", + " if residual.shape[-1] != self.channels or self.stride != 1:\n", + " residual = nn.Conv(\n", + " self.channels,\n", + " (1, 1),\n", + " strides=(self.stride, self.stride),\n", + " use_bias=False,\n", + " )(residual)\n", + " residual = nn.GroupNorm(num_groups=min(32, self.channels))(residual)\n", + "\n", + " return nn.relu(x + residual)\n", + "\n", + "\n", + "class BottleneckBlock(nn.Module):\n", + " \"\"\"Pre-activation ResNet bottleneck block with GroupNorm.\"\"\"\n", + "\n", + " channels: int # base channels (e.g., 64, 128, ...)\n", + " stride: int = 1\n", + " expansion: int = 4\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> Array:\n", + " residual = inputs\n", + " x = inputs\n", + "\n", + " out_channels = self.channels * self.expansion\n", + "\n", + " # 1x1 reduce\n", + " x = nn.Conv(\n", + " self.channels,\n", + " (1, 1),\n", + " strides=(1, 1),\n", + " use_bias=False,\n", + " )(x)\n", + " x = nn.GroupNorm(num_groups=min(32, self.channels))(x)\n", + " x = nn.relu(x)\n", + "\n", + " # 3x3 (does downsampling via stride if needed)\n", + " x = nn.Conv(\n", + " self.channels,\n", + " (3, 3),\n", + " strides=(self.stride, self.stride),\n", + " padding=\"SAME\",\n", + " use_bias=False,\n", + " )(x)\n", + " x = nn.GroupNorm(num_groups=min(32, self.channels))(x)\n", + " x = nn.relu(x)\n", + "\n", + " # 1x1 expand\n", + " x = nn.Conv(\n", + " out_channels,\n", + " (1, 1),\n", + " strides=(1, 1),\n", + " use_bias=False,\n", + " )(x)\n", + " x = nn.GroupNorm(num_groups=min(32, out_channels))(x)\n", + "\n", + " # Projection if needed\n", + " if residual.shape[-1] != out_channels or self.stride != 1:\n", + " residual = nn.Conv(\n", + " out_channels,\n", + " (1, 1),\n", + " strides=(self.stride, self.stride),\n", + " use_bias=False,\n", + " )(residual)\n", + " residual = nn.GroupNorm(num_groups=min(32, out_channels))(residual)\n", + "\n", + " return nn.relu(x + residual)\n", + "\n", + "\n", + "class ResNetRecognition(nn.Module):\n", + " \"\"\"ResNet-style recognition network (stem -> residual stages -> head).\"\"\"\n", + "\n", + " stage_sizes: Sequence[int]\n", + " stage_widths: Sequence[int]\n", + " dim_out: int\n", + " stem_width: int = 64\n", + " stem_kernel_size: Tuple[int, int] = (7, 7)\n", + " stem_stride: int = 2\n", + " use_max_pool: bool = True\n", + " max_pool_window: Tuple[int, int] = (3, 3)\n", + " max_pool_stride: Tuple[int, int] = (2, 2)\n", + " block_type: str = \"basic\"\n", + " fc_features: Sequence[int] = ()\n", + " projection_features: Sequence[int] = ()\n", + " precision_type: str = \"full\"\n", + "\n", + " @property\n", + " def trunk_dim(self):\n", + " if self.fc_features:\n", + " return self.fc_features[-1]\n", + " # No FC layers after the backbone: the trunk is the raw\n", + " # global-average-pooled output. Bottleneck blocks expand the last\n", + " # stage width by 4x (matches ``BottleneckBlock.expansion``).\n", + " expansion = 4 if self.block_type == \"bottleneck\" else 1\n", + " return self.stage_widths[-1] * expansion\n", + "\n", + " @property\n", + " def projection_dim(self):\n", + " return self.projection_features[-1]\n", + "\n", + " @property\n", + " def latent_dim(self):\n", + " return self.dim_out\n", + "\n", + " @nn.compact\n", + " def __call__(\n", + " self, inputs: Array\n", + " ) -> tuple[GaussianNatParams | DiagGaussianNatParams, Aux]:\n", + " if len(self.stage_sizes) != len(self.stage_widths):\n", + " raise ValueError(\n", + " \"'stage_sizes' and 'stage_widths' must have the same length\")\n", + " if self.block_type == \"basic\":\n", + " Block = ResidualBlock\n", + " elif self.block_type == \"bottleneck\":\n", + " Block = BottleneckBlock\n", + " else:\n", + " raise ValueError(\n", + " \"'block_type' must be 'basic' or 'bottleneck', got \"\n", + " f\"{self.block_type!r}\")\n", + "\n", + " x = inputs\n", + " x = nn.Conv(\n", + " self.stem_width,\n", + " self.stem_kernel_size,\n", + " strides=(self.stem_stride, self.stem_stride),\n", + " padding=\"SAME\",\n", + " use_bias=False,\n", + " name=\"stem_conv\",\n", + " )(x)\n", + " x = nn.GroupNorm(num_groups=min(32, self.stem_width),\n", + " name=\"stem_norm\")(x)\n", + " x = nn.relu(x)\n", + " if self.use_max_pool:\n", + " x = nn.max_pool(\n", + " x,\n", + " window_shape=self.max_pool_window,\n", + " strides=self.max_pool_stride,\n", + " padding=\"SAME\",\n", + " )\n", + "\n", + " for stage_idx, (num_blocks, channels) in enumerate(\n", + " zip(self.stage_sizes, self.stage_widths)):\n", + " for block_idx in range(num_blocks):\n", + " # Downsample at the first block of every stage after the first.\n", + " stride = 2 if stage_idx > 0 and block_idx == 0 else 1\n", + " x = Block(\n", + " channels=channels,\n", + " stride=stride,\n", + " name=f\"stage_{stage_idx}_block_{block_idx}\",\n", + " )(x)\n", + "\n", + " # Global average pool over the spatial dims.\n", + " x = jnp.mean(x, axis=(1, 2))\n", + " for i, d in enumerate(self.fc_features):\n", + " x = nn.Dense(d, name=f\"fc_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " trunk = x\n", + "\n", + " # Projection layers\n", + " for i, d in enumerate(self.projection_features):\n", + " x = nn.Dense(d, name=f\"projection_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " projection = x\n", + "\n", + " nat_params = gaussian_nat_head(x, self.dim_out, self.precision_type)\n", + "\n", + " aux = {'trunk': trunk, 'projection': projection}\n", + "\n", + " return nat_params, aux\n", + "\n", + "\n", + "class TransformerBlock(nn.Module):\n", + " \"\"\"Standard pre-norm transformer encoder block.\"\"\"\n", + "\n", + " embed_dim: int\n", + " mlp_dim: int\n", + " num_heads: int\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> Array:\n", + " x = inputs\n", + " h = nn.LayerNorm()(x)\n", + " h = nn.MultiHeadDotProductAttention(\n", + " num_heads=self.num_heads,\n", + " qkv_features=self.embed_dim,\n", + " out_features=self.embed_dim,\n", + " deterministic=True,\n", + " )(h, h)\n", + " x = x + h\n", + "\n", + " h = nn.LayerNorm()(x)\n", + " h = nn.Dense(self.mlp_dim)(h)\n", + " h = nn.gelu(h)\n", + " h = nn.Dense(self.embed_dim)(h)\n", + "\n", + " return x + h\n", + "\n", + "\n", + "class ViTRecognition(nn.Module):\n", + " \"\"\"Vision-transformer recognition network.\"\"\"\n", + "\n", + " image_size: int\n", + " patch_size: int\n", + " embed_dim: int\n", + " depth: int\n", + " num_heads: int\n", + " mlp_dim: int\n", + " dim_out: int\n", + " representation_dim: int = 0\n", + " projection_features: Sequence[int] = ()\n", + " precision_type: str = \"full\"\n", + "\n", + " @property\n", + " def trunk_dim(self):\n", + " # Mirrors __call__: the trunk is the mean-pooled token width\n", + " # (embed_dim), unless a representation head projects it down/up to\n", + " # representation_dim first.\n", + " return (self.representation_dim\n", + " if self.representation_dim > 0 else self.embed_dim)\n", + "\n", + " @property\n", + " def projection_dim(self):\n", + " return self.projection_features[-1]\n", + "\n", + " @property\n", + " def latent_dim(self):\n", + " return self.dim_out\n", + "\n", + " @nn.compact\n", + " def __call__(\n", + " self, inputs: Array\n", + " ) -> tuple[GaussianNatParams | DiagGaussianNatParams, Aux]:\n", + " if self.image_size % self.patch_size != 0:\n", + " raise ValueError(\"'image_size' must be divisible by 'patch_size'\")\n", + "\n", + " x = inputs\n", + " # Patch embedding via a strided convolution.\n", + " x = nn.Conv(\n", + " self.embed_dim,\n", + " kernel_size=(self.patch_size, self.patch_size),\n", + " strides=(self.patch_size, self.patch_size),\n", + " padding=\"VALID\",\n", + " name=\"patch_embed\",\n", + " )(x)\n", + "\n", + " batch, h, w, c = x.shape\n", + " x = x.reshape((batch, h * w, c))\n", + "\n", + " pos_emb = self.param(\n", + " \"position_embedding\",\n", + " nn.initializers.normal(stddev=0.02),\n", + " (1, h * w, self.embed_dim),\n", + " )\n", + " x = x + pos_emb\n", + "\n", + " for i in range(self.depth):\n", + " x = TransformerBlock(\n", + " embed_dim=self.embed_dim,\n", + " mlp_dim=self.mlp_dim,\n", + " num_heads=self.num_heads,\n", + " name=f\"transformer_block_{i}\",\n", + " )(x)\n", + "\n", + " x = nn.LayerNorm(name=\"final_norm\")(x)\n", + " # Mean-pool over tokens.\n", + " x = jnp.mean(x, axis=1)\n", + "\n", + " if self.representation_dim > 0:\n", + " x = nn.Dense(self.representation_dim, name=\"representation_head\")(x)\n", + " x = nn.tanh(x)\n", + "\n", + " trunk = x\n", + "\n", + " # Projection layers\n", + " for i, d in enumerate(self.projection_features):\n", + " x = nn.Dense(d, name=f\"projection_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " projection = x\n", + "\n", + " nat_params = gaussian_nat_head(x, self.dim_out, self.precision_type)\n", + "\n", + " aux = {'trunk': trunk, 'projection': projection}\n", + " return nat_params, aux\n", + "\n", + "\n", + "# --- Categorical recognition networks ---\n", + "\n", + "\n", + "class CategoricalCNNRecognition(nn.Module):\n", + " \"\"\"CNN recognition network that outputs categorical logits.\"\"\"\n", + "\n", + " conv_features: Sequence[Tuple[int, Tuple[int, int]]]\n", + " fc_features: Sequence[int]\n", + " dim_out: int\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> tuple[CategoricalNatParams, Aux]:\n", + " x = inputs\n", + " for i, (channels, kernel_size) in enumerate(self.conv_features):\n", + " x = nn.Conv(channels, kernel_size, name=f\"conv_{i}\")(x)\n", + " x = nn.relu(x)\n", + " x = nn.avg_pool(x, window_shape=(2, 2), strides=(2, 2))\n", + "\n", + " x = x.reshape((x.shape[0], -1))\n", + " for i, d in enumerate(self.fc_features):\n", + " x = nn.Dense(d, name=f\"fc_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " logits = nn.Dense(self.dim_out, name=\"logits_head\")(x)\n", + " logits = nn.log_softmax(logits, axis=-1)\n", + " return CategoricalNatParams(logits=logits), {}\n", + "\n", + "\n", + "class CategoricalNNRecognition(nn.Module):\n", + " \"\"\"MLP recognition network that outputs categorical logits.\"\"\"\n", + "\n", + " features: Sequence[int]\n", + " dim_out: int\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> tuple[CategoricalNatParams, Aux]:\n", + " x = inputs\n", + " # Flatten any spatial / channel dims so image-shaped views are\n", + " # accepted (mirrors NNRecognition).\n", + " x = x.reshape((x.shape[0], -1))\n", + " for i, feat in enumerate(self.features):\n", + " x = nn.Dense(feat, name=f\"fc_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " logits = nn.Dense(self.dim_out, name=\"logits_head\")(x)\n", + " logits = nn.log_softmax(logits, axis=-1)\n", + " return CategoricalNatParams(logits=logits), {}\n", + "\n", + "\n", + "class CategoricalResNetRecognition(nn.Module):\n", + " \"\"\"ResNet recognition network that outputs categorical logits.\"\"\"\n", + "\n", + " stage_sizes: Sequence[int]\n", + " stage_widths: Sequence[int]\n", + " dim_out: int\n", + " stem_width: int = 64\n", + " stem_kernel_size: Tuple[int, int] = (7, 7)\n", + " stem_stride: int = 2\n", + " use_max_pool: bool = True\n", + " max_pool_window: Tuple[int, int] = (3, 3)\n", + " max_pool_stride: Tuple[int, int] = (2, 2)\n", + " fc_features: Sequence[int] = ()\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> tuple[CategoricalNatParams, Aux]:\n", + " if len(self.stage_sizes) != len(self.stage_widths):\n", + " raise ValueError(\n", + " \"'stage_sizes' and 'stage_widths' must have the same length\")\n", + "\n", + " x = inputs\n", + " x = nn.Conv(\n", + " self.stem_width,\n", + " self.stem_kernel_size,\n", + " strides=(self.stem_stride, self.stem_stride),\n", + " padding=\"SAME\",\n", + " use_bias=False,\n", + " name=\"stem_conv\",\n", + " )(x)\n", + " x = nn.GroupNorm(num_groups=min(32, self.stem_width),\n", + " name=\"stem_norm\")(x)\n", + " x = nn.relu(x)\n", + " if self.use_max_pool:\n", + " x = nn.max_pool(\n", + " x,\n", + " window_shape=self.max_pool_window,\n", + " strides=self.max_pool_stride,\n", + " padding=\"SAME\",\n", + " )\n", + "\n", + " for stage_idx, (num_blocks, channels) in enumerate(\n", + " zip(self.stage_sizes, self.stage_widths)):\n", + " for block_idx in range(num_blocks):\n", + " # Downsample at the first block of every stage after the first.\n", + " stride = 2 if stage_idx > 0 and block_idx == 0 else 1\n", + " x = ResidualBlock(\n", + " channels=channels,\n", + " stride=stride,\n", + " name=f\"stage_{stage_idx}_block_{block_idx}\",\n", + " )(x)\n", + "\n", + " # Global average pool over the spatial dims.\n", + " x = jnp.mean(x, axis=(1, 2))\n", + " for i, d in enumerate(self.fc_features):\n", + " x = nn.Dense(d, name=f\"fc_{i}\")(x)\n", + " x = nn.relu(x)\n", + "\n", + " logits = nn.Dense(self.dim_out, name=\"logits_head\")(x)\n", + " logits = nn.log_softmax(logits, axis=-1)\n", + " return CategoricalNatParams(logits=logits), {}\n", + "\n", + "\n", + "class CategoricalViTRecognition(nn.Module):\n", + " \"\"\"Vision-transformer recognition network that outputs categorical logits.\"\"\"\n", + "\n", + " image_size: int\n", + " patch_size: int\n", + " embed_dim: int\n", + " depth: int\n", + " num_heads: int\n", + " mlp_dim: int\n", + " dim_out: int\n", + " representation_dim: int = 0\n", + "\n", + " @nn.compact\n", + " def __call__(self, inputs: Array) -> tuple[CategoricalNatParams, Aux]:\n", + " if self.image_size % self.patch_size != 0:\n", + " raise ValueError(\"'image_size' must be divisible by 'patch_size'\")\n", + "\n", + " x = inputs\n", + " # Patch embedding via a strided convolution.\n", + " x = nn.Conv(\n", + " self.embed_dim,\n", + " kernel_size=(self.patch_size, self.patch_size),\n", + " strides=(self.patch_size, self.patch_size),\n", + " padding=\"VALID\",\n", + " name=\"patch_embed\",\n", + " )(x)\n", + "\n", + " batch, h, w, c = x.shape\n", + " x = x.reshape((batch, h * w, c))\n", + "\n", + " pos_emb = self.param(\n", + " \"position_embedding\",\n", + " nn.initializers.normal(stddev=0.02),\n", + " (1, h * w, self.embed_dim),\n", + " )\n", + " x = x + pos_emb\n", + "\n", + " for i in range(self.depth):\n", + " x = TransformerBlock(\n", + " embed_dim=self.embed_dim,\n", + " mlp_dim=self.mlp_dim,\n", + " num_heads=self.num_heads,\n", + " name=f\"transformer_block_{i}\",\n", + " )(x)\n", + "\n", + " x = nn.LayerNorm(name=\"final_norm\")(x)\n", + " # Mean-pool over tokens.\n", + " x = jnp.mean(x, axis=1)\n", + "\n", + " if self.representation_dim > 0:\n", + " x = nn.Dense(self.representation_dim, name=\"representation_head\")(x)\n", + " x = nn.tanh(x)\n", + "\n", + " logits = nn.Dense(self.dim_out, name=\"logits_head\")(x)\n", + " logits = nn.log_softmax(logits, axis=-1)\n", + " return CategoricalNatParams(logits=logits), {}\n", + "\n", + "\n", + "print(\"networks defined.\")\n" + ], + "metadata": { + "id": "cell-009" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-010", + "cell_type": "markdown", + "source": [ + "## Part 4 — RPM models (`rpms`)\n", + "\n", + "The Gaussian RPM (two auxiliary constructions: `\"constrained\"` and\n", + "`\"amortized\"`) and the Categorical RPM, plus their free-energy objectives.\n", + "Every model's `__call__` returns `(rpm_output, aux)`: `rpm_output` is a dict\n", + "of the distributions the free-energy math needs, and `aux` is whatever the\n", + "underlying encoder produced (`{\"trunk\": ..., \"projection\": ...}`).\n" + ], + "metadata": { + "id": "cell-010" + } + }, + { + "id": "cell-011", + "cell_type": "code", + "source": [ + "EncoderParams = Mapping[str, Any]\n", + "Params = VariableDict\n", + "\n", + "GaussianLike = GaussianNatParams | DiagGaussianNatParams\n", + "RPMOutput = Mapping[str, Any]\n", + "\n", + "\n", + "def _gaussian_encoder_cls(encoder_arch: str) -> type[nn.Module]:\n", + " \"\"\"Map an ``encoder_arch`` string to its recognition-network class.\"\"\"\n", + " if encoder_arch == \"nn\":\n", + " return NNRecognition\n", + " elif encoder_arch == \"cnn\":\n", + " return CNNRecognition\n", + " elif encoder_arch == \"resnet\":\n", + " return ResNetRecognition\n", + " elif encoder_arch == \"vit\":\n", + " return ViTRecognition\n", + " else:\n", + " raise ValueError(\n", + " \"Invalid 'encoder_arch'; must be 'nn', 'cnn', \"\n", + " f\"'resnet', or 'vit', got {encoder_arch}\")\n", + "\n", + "\n", + "class GaussianRPM(nn.Module):\n", + " \"\"\"Abstract Gaussian RPM. Use :meth:`create` to build a concrete model.\"\"\"\n", + "\n", + " n_factors: int\n", + " dim_latent: int\n", + " encoder_arch: str\n", + " encoder_params: EncoderParams | None = None\n", + " share_recognition: bool = False\n", + " precision_type: str = \"full\"\n", + " fix_prior: bool = False\n", + " condition_on_view_info: bool = False\n", + "\n", + " @property\n", + " def rpm_type(self) -> str:\n", + " return \"gaussian\"\n", + "\n", + " @property\n", + " def _unbound_encoder(self) -> nn.Module:\n", + " \"\"\"A plain (non-vmapped) encoder instance, usable without binding.\n", + "\n", + " ``self.encoders`` is only assigned in ``setup()``, so it only exists\n", + " once this module is bound (inside ``init``/``apply``). ``trunk_dim``\n", + " and ``latent_dim`` only need the encoder's static dataclass fields, so\n", + " this sidesteps binding entirely by constructing an unwrapped instance.\n", + " \"\"\"\n", + " encoder_cls = _gaussian_encoder_cls(self.encoder_arch)\n", + " return encoder_cls(\n", + " **self.encoder_params,\n", + " dim_out=self.dim_latent,\n", + " precision_type=self.precision_type,\n", + " )\n", + "\n", + " @property\n", + " def trunk_dim(self):\n", + " return self._unbound_encoder.trunk_dim\n", + "\n", + " @property\n", + " def latent_dim(self):\n", + " return self._unbound_encoder.latent_dim\n", + "\n", + " def __init__(self, *args):\n", + " raise NotImplementedError(\n", + " \"Cannot instantiate the abstract class, use GaussianRPM.create instead\"\n", + " )\n", + "\n", + " def __call__(self, inputs: Array) -> tuple[RPMOutput, Aux]:\n", + " raise NotImplementedError(\n", + " \"Cannot apply abstract class, use GaussianRPM.create to get a concrete model\"\n", + " )\n", + "\n", + " def setup(self):\n", + " Encoder = _gaussian_encoder_cls(self.encoder_arch)\n", + "\n", + " # One recognition network per factor, vmapped over the factor axis\n", + " # (input axis 1). Params are shared across factors iff ``share_recognition``.\n", + " VmapEncoders = nn.vmap(\n", + " Encoder,\n", + " in_axes=1,\n", + " out_axes=0,\n", + " variable_axes=({\n", + " \"params\": None if self.share_recognition else 0,\n", + " \"intermediates\": 0\n", + " }),\n", + " split_rngs=({\n", + " \"params\": not self.share_recognition\n", + " }),\n", + " )\n", + " self.encoders = VmapEncoders(\n", + " **self.encoder_params,\n", + " dim_out=self.dim_latent,\n", + " precision_type=self.precision_type,\n", + " )\n", + "\n", + " # Prior natural params. When ``fix_prior`` the prior is a constant\n", + " # standard Gaussian; otherwise it is learned.\n", + " if self.fix_prior:\n", + " self.prior_pwm = jnp.zeros((self.dim_latent,))\n", + " if self.precision_type == \"diag\":\n", + " self.prior_log_precision_diag = jnp.zeros((self.dim_latent,))\n", + " else:\n", + " self.prior_cholesky = identity_init(None, [self.dim_latent])\n", + " else:\n", + " self.prior_pwm = self.param(\"prior_pwm\", nn.initializers.zeros,\n", + " (self.dim_latent,))\n", + " if self.precision_type == \"diag\":\n", + " self.prior_log_precision_diag = self.param(\n", + " \"prior_log_precision_diag\",\n", + " nn.initializers.zeros,\n", + " (self.dim_latent,),\n", + " )\n", + " else:\n", + " self.prior_cholesky = self.param(\"prior_cholesky\",\n", + " identity_init,\n", + " (self.dim_latent,))\n", + "\n", + " def prior(self) -> GaussianLike:\n", + " if self.precision_type == \"diag\":\n", + " return DiagGaussianNatParams(\n", + " precision_diag=jnp.exp(self.prior_log_precision_diag),\n", + " precision_weighted_mean=self.prior_pwm,\n", + " )\n", + " return GaussianNatParams(\n", + " precision=jnp.dot(self.prior_cholesky, self.prior_cholesky.T),\n", + " precision_weighted_mean=self.prior_pwm,\n", + " )\n", + "\n", + " @classmethod\n", + " def create(\n", + " cls,\n", + " auxiliary_method: str,\n", + " n_factors: int,\n", + " dim_latent: int,\n", + " encoder_arch: str,\n", + " encoder_params: EncoderParams = {},\n", + " share_recognition: bool = False,\n", + " precision_type: str = \"full\",\n", + " fix_prior: bool = False,\n", + " n_samples: int = 1,\n", + " ) -> \"GaussianRPM\":\n", + " \"\"\"Build a concrete Gaussian RPM for the chosen auxiliary method.\"\"\"\n", + " if auxiliary_method not in [\n", + " \"constrained\", \"amortized\", \"reparametrised\"\n", + " ]:\n", + " raise ValueError(\n", + " \"'auxiliary_method' has to be either 'constrained', \"\n", + " \"'optimized', 'amortized', or 'reparametrised'\")\n", + "\n", + " kwargs = dict(\n", + " n_factors=n_factors,\n", + " dim_latent=dim_latent,\n", + " encoder_arch=encoder_arch,\n", + " encoder_params=encoder_params,\n", + " share_recognition=share_recognition,\n", + " precision_type=precision_type,\n", + " fix_prior=fix_prior,\n", + " )\n", + "\n", + " if auxiliary_method == \"constrained\":\n", + " return GaussianRPM_ConstrainedAux(**kwargs)\n", + " elif auxiliary_method == \"amortized\":\n", + " return GaussianRPM_AmortizedAux(**kwargs)\n", + " else:\n", + " raise NotImplementedError(\n", + " f\"Unknown auxiliary_method={auxiliary_method!r}\")\n", + "\n", + "\n", + "class GaussianRPM_AmortizedAux(GaussianRPM):\n", + " \"\"\"RPM whose auxiliary is produced by a second set of recognition nets.\"\"\"\n", + "\n", + " def setup(self):\n", + " super().setup()\n", + "\n", + " AuxEncoder = _gaussian_encoder_cls(self.encoder_arch)\n", + "\n", + " VmapAuxEncoders = nn.vmap(\n", + " AuxEncoder,\n", + " in_axes=1,\n", + " out_axes=0,\n", + " variable_axes={\n", + " \"params\": 0,\n", + " \"intermediates\": 0\n", + " },\n", + " split_rngs={\"params\": True},\n", + " )\n", + " self.aux_encoders = VmapAuxEncoders(\n", + " **self.encoder_params,\n", + " dim_out=self.dim_latent,\n", + " precision_type=self.precision_type,\n", + " )\n", + "\n", + " def __call__(self, inputs: Array) -> tuple[RPMOutput, Aux]:\n", + " prior = self.prior()\n", + " recognition, aux = self.encoders(inputs)\n", + " auxiliary, _ = self.aux_encoders(inputs)\n", + " variational = (1 / (self.n_factors + 1)) * (\n", + " prior + auxiliary.sum(batch_axis=0) + recognition.sum(batch_axis=0))\n", + "\n", + " if self.precision_type == \"diag\":\n", + " variational_mean = diag_mean_params(variational)\n", + " else:\n", + " variational_mean = mean_params(variational)\n", + "\n", + " rpm_output = {\n", + " \"prior\": prior,\n", + " \"recognition\": recognition,\n", + " \"auxiliary\": auxiliary,\n", + " \"variational\": variational,\n", + " \"variational_mean\": variational_mean,\n", + " }\n", + "\n", + " return rpm_output, aux\n", + "\n", + "\n", + "class GaussianRPM_ConstrainedAux(GaussianRPM):\n", + " \"\"\"RPM whose auxiliary is the closed-form ``variational - prior``.\"\"\"\n", + "\n", + " def __call__(self, inputs: Array) -> tuple[RPMOutput, Aux]:\n", + " x = inputs\n", + " J = x.shape[1]\n", + " if J != self.n_factors:\n", + " raise ValueError(\n", + " \"Second input dimension must agree with number of factors; \"\n", + " f\"RPM has n_factors={self.n_factors}, input has shape {x.shape}\"\n", + " )\n", + "\n", + " prior = self.prior()\n", + " deltas, aux = self.encoders(x)\n", + " variational = prior + deltas.sum(batch_axis=0)\n", + " recognition = prior + deltas\n", + " auxiliary = variational - prior\n", + "\n", + " if self.precision_type == \"diag\":\n", + " variational_mean = diag_mean_params(variational)\n", + " else:\n", + " variational_mean = mean_params(variational)\n", + "\n", + " rpm_output = {\n", + " \"prior\": prior,\n", + " \"recognition\": recognition,\n", + " \"auxiliary\": auxiliary,\n", + " \"variational\": variational,\n", + " \"variational_mean\": variational_mean,\n", + " }\n", + "\n", + " return rpm_output, aux\n", + "\n", + "\n", + "def free_energy(\n", + " params: Params,\n", + " model: GaussianRPM,\n", + " X: Array,\n", + " rng: Array | None = None,\n", + " beta: float = 1.0,\n", + ") -> tuple[Array, Aux]:\n", + " \"\"\"RPM auxiliary free energy\"\"\"\n", + " # Prior, per-factor recognition, auxiliary, and the inferred posterior.\n", + " del rng # Unused: neither Gaussian RPM variant samples internally.\n", + " rpm_output, _ = model.apply(params, X)\n", + "\n", + " prior = rpm_output[\"prior\"]\n", + " recognition = rpm_output[\"recognition\"]\n", + " auxiliary = rpm_output[\"auxiliary\"]\n", + " variational = rpm_output[\"variational\"]\n", + " variational_mean = rpm_output[\"variational_mean\"]\n", + "\n", + " N = variational.batch_shape()[0]\n", + " J = recognition.batch_shape()[0]\n", + "\n", + " is_diag = isinstance(variational, DiagGaussianNatParams)\n", + " kl_fn = diag_kl if is_diag else kl\n", + " unkl_fn = diag_unnormalized_kl if is_diag else unnormalized_kl\n", + " variational_cov = variational_mean.var_diag if is_diag else variational_mean.cov\n", + " natparam_cls = DiagGaussianNatParams if is_diag else GaussianNatParams\n", + "\n", + " if len(auxiliary.batch_shape()) == 1:\n", + " auxiliary = auxiliary.expand_batch_dim(0)\n", + "\n", + " outsum_gaussians = vmap(\n", + " vmap(natparam_cls.__add__, in_axes=[1, None]),\n", + " in_axes=[None, 1],\n", + " )\n", + " outsum_arrays = vmap(vmap(jnp.add, in_axes=[1, None]), in_axes=[None, 1])\n", + "\n", + " # has now batch dimension of [M x N x J]\n", + " cross_terms = outsum_gaussians(auxiliary, recognition)\n", + "\n", + " phi_recognition = recognition.lognormalizer() # [J x M]\n", + " phi_auxiliary = auxiliary.lognormalizer() # [1 x N] or [J x N]\n", + " phi_cross_terms = cross_terms.lognormalizer() # [M x N x J]\n", + "\n", + " # [M x N x J] -> [N x J] -> []\n", + " gamma_terms = phi_cross_terms - outsum_arrays(phi_auxiliary,\n", + " phi_recognition)\n", + " log_gamma_jn = logsumexp(gamma_terms, axis=1, b=(1 / N))\n", + " logGamma = log_gamma_jn.sum()\n", + "\n", + " phi_variational = variational.lognormalizer(mean=variational_mean.mean)\n", + "\n", + " # -KL(q || prior) and the summed -unnormalized-KL factor terms, with the\n", + " # full vs diagonal Gaussian implementations.\n", + " neg_kl_qp = -kl_fn(variational_mean.mean, variational_cov, variational,\n", + " prior, phi_variational).sum()\n", + "\n", + " def sum_kl_factors(j, carry):\n", + " ukl = unkl_fn(\n", + " variational_mean.mean,\n", + " variational_cov,\n", + " variational,\n", + " recognition[j],\n", + " auxiliary,\n", + " phi_recognition[j],\n", + " phi_auxiliary,\n", + " phi_variational,\n", + " ).sum()\n", + " return carry - ukl\n", + "\n", + " q_z_var_mean = jnp.mean(variational_cov)\n", + " q_z_mean_std = jnp.mean(jnp.std(variational_mean.mean, axis=0))\n", + "\n", + " neg_kl_qf = fori_loop(0, J, sum_kl_factors, 0.0)\n", + "\n", + " free_energy = (neg_kl_qf + beta * neg_kl_qp - logGamma) / N\n", + " metrics = {\n", + " \"neg_kl_qp\": neg_kl_qp / N,\n", + " \"neg_kl_qf\": neg_kl_qf / N,\n", + " \"neg_log_gamma\": -logGamma / N,\n", + " \"q_z_mean_std\": q_z_mean_std,\n", + " \"q_z_var_mean\": q_z_var_mean,\n", + " }\n", + "\n", + " return free_energy, metrics\n", + "\n", + "\n", + "# --- Categorical RPM ---\n", + "\n", + "\n", + "class CategoricalRPM(nn.Module):\n", + " \"\"\"Categorical RPM with a learned prior over ``num_categories`` classes.\"\"\"\n", + "\n", + " n_factors: int\n", + " num_categories: int\n", + " encoder_arch: str\n", + " encoder_params: EncoderParams | None = None\n", + " share_recognition: bool = True\n", + " fix_prior: bool = False\n", + "\n", + " @property\n", + " def rpm_type(self) -> str:\n", + " return \"categorical\"\n", + "\n", + " def setup(self):\n", + " params = self.encoder_params or {}\n", + "\n", + " if self.encoder_arch == \"cnn\":\n", + " Encoder = CategoricalCNNRecognition\n", + " elif self.encoder_arch == \"nn\":\n", + " Encoder = CategoricalNNRecognition\n", + " elif self.encoder_arch == \"resnet\":\n", + " Encoder = CategoricalResNetRecognition\n", + " elif self.encoder_arch == \"vit\":\n", + " Encoder = CategoricalViTRecognition\n", + " else:\n", + " raise ValueError(f\"Unknown encoder_arch={self.encoder_arch}\")\n", + "\n", + " VmapEncoder = nn.vmap(\n", + " Encoder,\n", + " in_axes=1,\n", + " out_axes=0,\n", + " variable_axes=({\n", + " \"params\": 0\n", + " } if not self.share_recognition else {\n", + " \"params\": None\n", + " }),\n", + " split_rngs=({\n", + " \"params\": True\n", + " } if not self.share_recognition else {\n", + " \"params\": False\n", + " }),\n", + " )\n", + "\n", + " self.encoders = VmapEncoder(**params, dim_out=self.num_categories)\n", + "\n", + " # Prior logits over the categories. When ``fix_prior`` the prior is a\n", + " # constant uniform distribution (zeros, not a parameter); otherwise it\n", + " # is learned (also initialized to uniform).\n", + " if self.fix_prior:\n", + " self.prior_logits = jnp.zeros((self.num_categories,))\n", + " else:\n", + " self.prior_logits = self.param(\n", + " \"prior_logits\",\n", + " nn.initializers.zeros,\n", + " (self.num_categories,),\n", + " )\n", + "\n", + " def prior(self) -> CategoricalNatParams:\n", + " return CategoricalNatParams(logits=self.prior_logits)\n", + "\n", + " def __call__(self, inputs: Array) -> tuple[RPMOutput, Aux]:\n", + " x = inputs\n", + " j = x.shape[1]\n", + " if j != self.n_factors:\n", + " raise ValueError(f\"Expected {self.n_factors} views, got {j}\")\n", + " prior = self.prior()\n", + "\n", + " factors, aux = self.encoders(x) # logits: [J x N x num_categories]\n", + "\n", + " # log f_bar_j(z) = logsumexp_m log f_j(z|x_{j,m}) - log N.\n", + " n = x.shape[0]\n", + " log_denominators = logsumexp(factors.logits, axis=1) - jnp.log(n)\n", + "\n", + " factor_sum = factors.logits.sum(axis=0)\n", + " denom_sum = log_denominators.sum(axis=0)\n", + "\n", + " posterior_logits = prior.logits + factor_sum - denom_sum\n", + " variational = CategoricalNatParams(logits=posterior_logits)\n", + "\n", + " rpm_output = {\n", + " \"prior\": prior,\n", + " \"factors\": factors,\n", + " \"variational\": variational,\n", + " }\n", + "\n", + " return rpm_output, aux\n", + "\n", + "\n", + "def free_energy_categorical(\n", + " params: Params,\n", + " model: CategoricalRPM,\n", + " x: Array,\n", + " rng: Array | None = None,\n", + " beta: float = 1.0,\n", + ") -> tuple[Array, Aux]:\n", + " \"\"\"Free energy for the Categorical RPM.\n", + "\n", + " Returns the per-datapoint free energy plus diagnostics\n", + " ``(beta*entropy, beta*prior cross-term, factor cross-terms, denominator\n", + " cross-terms)`` (each averaged over datapoints).\n", + " \"\"\"\n", + " del rng # Unused: CategoricalRPM does not sample internally.\n", + " rpm_output, _ = model.apply(params, x)\n", + " prior = rpm_output[\"prior\"]\n", + " factors = rpm_output[\"factors\"]\n", + " variational = rpm_output[\"variational\"]\n", + "\n", + " n = x.shape[0]\n", + " j = model.n_factors\n", + "\n", + " # entropy = H(q) = -E_q[log q].\n", + " entropy = -categorical_cross_entropy(variational, variational)\n", + " prior_xent = categorical_cross_entropy(variational, prior)\n", + "\n", + " factors_xent = jnp.zeros(n)\n", + " for fj in range(j):\n", + " factor_j = CategoricalNatParams(logits=factors.logits[fj])\n", + " factors_xent = factors_xent + categorical_cross_entropy(\n", + " variational, factor_j)\n", + "\n", + " # log f_bar_j(z) = logsumexp_m log f_j(z|x_{j,m}) - log N, i.e. the log of\n", + " # the data-averaged recognition factor (the RPM normalizer).\n", + " log_denominators = logsumexp(factors.logits, axis=1) - jnp.log(n)\n", + " denom_xent = jnp.zeros(n)\n", + " for fj in range(j):\n", + " denom_j = CategoricalNatParams(\n", + " logits=jnp.broadcast_to(log_denominators[fj], (\n", + " n, model.num_categories)))\n", + " denom_xent = denom_xent + categorical_cross_entropy(\n", + " variational, denom_j)\n", + "\n", + " fe_per_sample = beta * (entropy + prior_xent) + factors_xent - denom_xent\n", + " fe = jnp.sum(fe_per_sample) / n\n", + "\n", + " metrics = {\n", + " \"entropy\": jnp.sum(beta * entropy) / n,\n", + " \"prior_xent\": jnp.sum(beta * prior_xent) / n,\n", + " \"factors_xent\": jnp.sum(factors_xent) / n,\n", + " \"denom_xent\": jnp.sum(denom_xent) / n,\n", + " }\n", + "\n", + " return fe, metrics\n", + "\n", + "\n", + "print(\"rpms defined.\")\n" + ], + "metadata": { + "id": "cell-011" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-012", + "cell_type": "markdown", + "source": [ + "## Part 5 — SimCLR encoder and NT-Xent loss (`simclr`)\n", + "\n", + "Reuses the same recognition-network backbones as the RPM models. Unlike RPM,\n", + "SimCLR always shares weights across views, so there is no per-factor\n", + "``nn.vmap``: views are simply reshaped into the batch dimension, run through\n", + "one ordinary forward pass, and reshaped back. The recognition-network\n", + "classes above always attach a (Gaussian) natural-parameter head; since\n", + "SimCLR never uses it, ``dim_out=1`` with ``precision_type=\"fixed_diag\"``\n", + "makes that head a single throwaway ``Dense(1)`` layer — negligible next to\n", + "any real backbone, and avoids duplicating the backbone code.\n" + ], + "metadata": { + "id": "cell-012" + } + }, + { + "id": "cell-013", + "cell_type": "code", + "source": [ + "_DISCARDED_HEAD_DIM_OUT = 1\n", + "_DISCARDED_HEAD_PRECISION_TYPE = \"fixed_diag\"\n", + "\n", + "\n", + "def _simclr_encoder_cls(encoder_arch: str) -> type[nn.Module]:\n", + " \"\"\"Map an ``encoder_arch`` string to its recognition-network class.\n", + "\n", + " Excludes ``\"linear\"``: there is no ``LinearRecognition`` in this notebook\n", + " (it has no trunk/projection concept), mirroring ``CategoricalRPM``'s\n", + " encoder_arch support.\n", + " \"\"\"\n", + " if encoder_arch == \"nn\":\n", + " return NNRecognition\n", + " elif encoder_arch == \"cnn\":\n", + " return CNNRecognition\n", + " elif encoder_arch == \"resnet\":\n", + " return ResNetRecognition\n", + " elif encoder_arch == \"vit\":\n", + " return ViTRecognition\n", + " else:\n", + " raise ValueError(\"Invalid 'encoder_arch'; must be 'nn', 'cnn', \"\n", + " f\"'resnet', or 'vit', got {encoder_arch}\")\n", + "\n", + "\n", + "class SimCLREncoder(nn.Module):\n", + " \"\"\"Shared-weight multi-view encoder for SimCLR-style contrastive training.\"\"\"\n", + "\n", + " n_views: int\n", + " encoder_arch: str\n", + " encoder_params: EncoderParams | None = None\n", + " temperature: float = 0.5\n", + "\n", + " @property\n", + " def rpm_type(self) -> str:\n", + " return \"simclr\"\n", + "\n", + " @property\n", + " def _unbound_encoder(self) -> nn.Module:\n", + " encoder_cls = _simclr_encoder_cls(self.encoder_arch)\n", + " return encoder_cls(\n", + " **self.encoder_params,\n", + " dim_out=_DISCARDED_HEAD_DIM_OUT,\n", + " precision_type=_DISCARDED_HEAD_PRECISION_TYPE,\n", + " )\n", + "\n", + " @property\n", + " def trunk_dim(self) -> int:\n", + " return self._unbound_encoder.trunk_dim\n", + "\n", + " @property\n", + " def latent_dim(self) -> int:\n", + " \"\"\"Width of the projection output (SimCLR's ``z``).\n", + "\n", + " Falls back to ``trunk_dim`` when ``projection_features`` is empty,\n", + " in which case the projection head is a no-op and z == h.\n", + " \"\"\"\n", + " encoder = self._unbound_encoder\n", + " if encoder.projection_features:\n", + " return encoder.projection_dim\n", + " return encoder.trunk_dim\n", + "\n", + " def setup(self):\n", + " Encoder = _simclr_encoder_cls(self.encoder_arch)\n", + " self.encoder = Encoder(\n", + " **self.encoder_params,\n", + " dim_out=_DISCARDED_HEAD_DIM_OUT,\n", + " precision_type=_DISCARDED_HEAD_PRECISION_TYPE,\n", + " )\n", + "\n", + " def __call__(self, inputs: Array) -> tuple[Array, Aux]:\n", + " x = inputs # (batch, n_views, H, W, C)\n", + " n_views = x.shape[1]\n", + " if n_views != self.n_views:\n", + " raise ValueError(f\"Expected {self.n_views} views, got {n_views}\")\n", + " batch = x.shape[0]\n", + " per_example_shape = x.shape[2:]\n", + "\n", + " # (batch, n_views, ...) -> (n_views, batch, ...) -> (n_views*batch, ...)\n", + " x_flat = jnp.transpose(x, (1, 0) + tuple(range(2, x.ndim)))\n", + " x_flat = x_flat.reshape((n_views * batch,) + per_example_shape)\n", + "\n", + " _, aux = self.encoder(x_flat)\n", + "\n", + " def _unflatten(feats: Array) -> Array:\n", + " return feats.reshape((n_views, batch) + feats.shape[1:])\n", + "\n", + " trunk = _unflatten(aux[\"trunk\"])\n", + " projection = _unflatten(aux[\"projection\"])\n", + "\n", + " return projection, {\"trunk\": trunk, \"projection\": projection}\n", + "\n", + "\n", + "def nt_xent_loss(\n", + " params,\n", + " model: SimCLREncoder,\n", + " batch_views: Array,\n", + " rng: Array | None = None,\n", + " beta: float = 1.0,\n", + ") -> tuple[Array, Aux]:\n", + " \"\"\"NT-Xent contrastive loss, generalized to n_views >= 2.\n", + "\n", + " Every other view of the same image is a positive; every embedding from a\n", + " different image is a negative. Reduces exactly to the standard two-view\n", + " NT-Xent loss when ``n_views == 2`` (verified against an independent\n", + " hand-rolled reference implementation).\n", + "\n", + " Returns ``-nt_xent`` (not ``nt_xent`` directly) to match the sign\n", + " convention shared with ``free_energy``/``free_energy_categorical``:\n", + " ``create_train_step`` computes ``loss = -free_energy``, so returning the\n", + " negated loss here means the actual NT-Xent loss is what gets minimized.\n", + " \"\"\"\n", + " del rng, beta # Unused: temperature (not beta) scales the loss, and\n", + " # the encoder has no dropout/sampling.\n", + " projection, _ = model.apply(params, batch_views) # (V, B, D)\n", + "\n", + " n_views, batch = projection.shape[0], projection.shape[1]\n", + " n = n_views * batch\n", + " z = projection.reshape((n, projection.shape[-1]))\n", + "\n", + " z_norm = z / (jnp.linalg.norm(z, axis=-1, keepdims=True) + 1e-8)\n", + " raw_sim = z_norm @ z_norm.T # (N, N), in [-1, 1]\n", + " scaled_sim = raw_sim / model.temperature\n", + "\n", + " row_idx = jnp.arange(n)\n", + " is_self = row_idx[:, None] == row_idx[None, :]\n", + " # Row i = view v, example b (since z was flattened as (n_views, batch)).\n", + " # Row i' is a positive for row i iff same example, different view.\n", + " same_example = (row_idx[:, None] % batch) == (row_idx[None, :] % batch)\n", + " positive_mask = same_example & ~is_self\n", + " negative_mask = ~same_example\n", + "\n", + " masked_sim = jnp.where(is_self, -jnp.inf, scaled_sim)\n", + " log_prob = nn.log_softmax(masked_sim, axis=-1) # (N, N)\n", + "\n", + " num_positives = jnp.sum(positive_mask, axis=-1) # (N,), == n_views - 1\n", + " per_anchor_loss = -jnp.sum(log_prob * positive_mask,\n", + " axis=-1) / num_positives\n", + " nt_xent = jnp.mean(per_anchor_loss)\n", + "\n", + " mean_pos_sim = jnp.sum(raw_sim * positive_mask) / jnp.sum(positive_mask)\n", + " mean_neg_sim = jnp.sum(raw_sim * negative_mask) / jnp.sum(negative_mask)\n", + " predicted = jnp.argmax(masked_sim, axis=-1)\n", + " contrastive_acc = jnp.mean(\n", + " jnp.take_along_axis(positive_mask, predicted[:, None], axis=-1))\n", + "\n", + " metrics = {\n", + " \"mean_pos_sim\": mean_pos_sim,\n", + " \"mean_neg_sim\": mean_neg_sim,\n", + " \"contrastive_acc\": contrastive_acc,\n", + " }\n", + "\n", + " return -nt_xent, metrics\n", + "\n", + "\n", + "print(\"simclr defined.\")\n" + ], + "metadata": { + "id": "cell-013" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-014", + "cell_type": "markdown", + "source": [ + "## Part 6 — CIFAR-10 data loading\n", + "\n", + "Uses `tensorflow_datasets` to load CIFAR-10 with two view-generation\n", + "strategies selectable via `preprocess_type`:\n", + "\n", + "- **`\"simclr\"`** — SimCLR-style augmentations (random crop, flip, color\n", + " jitter, grayscale, Gaussian blur).\n", + "- **`\"lejepa_masking\"`** — I-JEPA-style multi-block masking (Assran et al.,\n", + " 2023). Instead of hand-crafted augmentations, each view is the normalized\n", + " image with a different set of large rectangular blocks masked out (filled\n", + " with zeros, i.e. the dataset mean post-normalization). Target blocks are\n", + " sampled at large scale to encourage semantic representations; the context\n", + " (visible) portion is spatially distributed.\n", + "\n", + "Both strategies return `(raw_image, views, label)` as NumPy arrays with\n", + "`views` in `[batch, num_views, H, W, 3]` layout.\n" + ], + "metadata": { + "id": "cell-014" + } + }, + { + "id": "cell-015", + "cell_type": "code", + "source": [ + "import tensorflow as tf\n", + "import tensorflow_datasets as tfds\n", + "\n", + "# Prevent TF from grabbing GPU memory (we only use it for data loading).\n", + "tf.config.set_visible_devices([], \"GPU\")\n", + "\n", + "DataBatch = tuple[np.ndarray, np.ndarray, np.ndarray]\n", + "\n", + "CIFAR10_MEAN = np.asarray([0.4914, 0.4822, 0.4465], dtype=np.float32)\n", + "CIFAR10_STD = np.asarray([0.2470, 0.2435, 0.2616], dtype=np.float32)\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# Helper: append a visibility-mask channel\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "def _append_ones_mask(image: tf.Tensor) -> tf.Tensor:\n", + " \"\"\"Append an all-ones 4th channel (fully visible) to a [H,W,3] image.\"\"\"\n", + " h = tf.shape(image)[0]\n", + " w = tf.shape(image)[1]\n", + " ones = tf.ones([h, w, 1], dtype=image.dtype)\n", + " return tf.concat([image, ones], axis=-1) # [H, W, 4]\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# SimCLR-style augmentations in pure TF ops\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "def _random_resized_crop(image: tf.Tensor, size: int = 32) -> tf.Tensor:\n", + " \"\"\"Random crop with area in [0.08, 1.0] and aspect ratio in [3/4, 4/3],\n", + " then resize to ``size x size``. Uses TF's built-in\n", + " ``sample_distorted_bounding_box`` which matches torchvision's\n", + " ``RandomResizedCrop`` semantics (10 attempts, centre-crop fallback).\"\"\"\n", + " shape = tf.shape(image)\n", + " bbox = tf.constant([0.0, 0.0, 1.0, 1.0], dtype=tf.float32,\n", + " shape=[1, 1, 4])\n", + " bbox_begin, bbox_size, _ = tf.image.sample_distorted_bounding_box(\n", + " shape,\n", + " bounding_boxes=bbox,\n", + " min_object_covered=0.0,\n", + " aspect_ratio_range=(3.0 / 4.0, 4.0 / 3.0),\n", + " area_range=(0.08, 1.0),\n", + " max_attempts=10,\n", + " use_image_if_no_bounding_boxes=True,\n", + " )\n", + " image = tf.slice(image, bbox_begin, bbox_size)\n", + " image = tf.image.resize(image, [size, size])\n", + " return image\n", + "\n", + "\n", + "def _color_jitter(image: tf.Tensor,\n", + " strength: float = 1.0) -> tf.Tensor:\n", + " \"\"\"Random brightness / contrast / saturation / hue jitter.\"\"\"\n", + " image = tf.image.random_brightness(image, max_delta=0.8 * strength)\n", + " image = tf.image.random_contrast(image, lower=max(0, 1 - 0.8 * strength),\n", + " upper=1 + 0.8 * strength)\n", + " image = tf.image.random_saturation(image,\n", + " lower=max(0, 1 - 0.8 * strength),\n", + " upper=1 + 0.8 * strength)\n", + " image = tf.image.random_hue(image, max_delta=0.2 * strength)\n", + " image = tf.clip_by_value(image, 0.0, 1.0)\n", + " return image\n", + "\n", + "\n", + "def _random_grayscale(image: tf.Tensor, p: float = 0.2) -> tf.Tensor:\n", + " \"\"\"Convert to grayscale with probability ``p``.\"\"\"\n", + " do_gray = tf.random.uniform([]) < p\n", + " gray = tf.image.rgb_to_grayscale(image) # [H, W, 1]\n", + " gray = tf.tile(gray, [1, 1, 3]) # [H, W, 3]\n", + " return tf.where(do_gray, gray, image)\n", + "\n", + "\n", + "def _gaussian_blur(image: tf.Tensor, kernel_size: int = 3,\n", + " sigma_lo: float = 0.1,\n", + " sigma_hi: float = 2.0) -> tf.Tensor:\n", + " \"\"\"Gaussian blur with a random sigma.\"\"\"\n", + " sigma = tf.random.uniform([], sigma_lo, sigma_hi)\n", + " radius = kernel_size // 2\n", + " x = tf.cast(tf.range(-radius, radius + 1), tf.float32)\n", + " kernel_1d = tf.exp(-0.5 * (x / sigma) ** 2)\n", + " kernel_1d = kernel_1d / tf.reduce_sum(kernel_1d)\n", + " kernel_2d = tf.tensordot(kernel_1d, kernel_1d, axes=0) # [k, k]\n", + " kernel_2d = kernel_2d[:, :, tf.newaxis, tf.newaxis] # [k, k, 1, 1]\n", + " kernel_2d = tf.tile(kernel_2d, [1, 1, 3, 1]) # [k, k, 3, 1]\n", + "\n", + " image_4d = image[tf.newaxis] # [1, H, W, 3]\n", + " blurred = tf.nn.depthwise_conv2d(image_4d, kernel_2d,\n", + " strides=[1, 1, 1, 1],\n", + " padding=\"SAME\")\n", + " return blurred[0]\n", + "\n", + "\n", + "def simclr_augment_single(image: tf.Tensor,\n", + " image_size: int = 32,\n", + " color_jitter_strength: float = 1.0,\n", + " gaussian_blur: bool = True,\n", + " pass_mask: bool = False) -> tf.Tensor:\n", + " \"\"\"Apply one SimCLR augmentation to a float32 [H,W,3] image in [0,1].\n", + "\n", + " If ``pass_mask`` is True, an all-ones 4th channel is appended\n", + " (indicating all pixels are visible), giving output shape [H,W,4].\n", + " \"\"\"\n", + " image = _random_resized_crop(image, size=image_size)\n", + " image = tf.image.random_flip_left_right(image)\n", + "\n", + " # Color jitter with p=0.8\n", + " do_jitter = tf.random.uniform([]) < 0.8\n", + " image = tf.cond(do_jitter,\n", + " lambda: _color_jitter(image, strength=color_jitter_strength),\n", + " lambda: image)\n", + "\n", + " image = _random_grayscale(image, p=0.2)\n", + "\n", + " if gaussian_blur:\n", + " do_blur = tf.random.uniform([]) < 0.5\n", + " image = tf.cond(do_blur, lambda: _gaussian_blur(image), lambda: image)\n", + "\n", + " # Normalize\n", + " image = (image - CIFAR10_MEAN) / CIFAR10_STD\n", + "\n", + " if pass_mask:\n", + " image = _append_ones_mask(image)\n", + "\n", + " return image\n", + "\n", + "\n", + "def eval_preprocess(image: tf.Tensor, image_size: int = 32,\n", + " pass_mask: bool = False) -> tf.Tensor:\n", + " \"\"\"Deterministic eval preprocessing: resize (if needed) + normalize.\n", + "\n", + " If ``pass_mask`` is True, an all-ones 4th channel is appended.\n", + " \"\"\"\n", + " if image_size != 32:\n", + " image = tf.image.resize(image, [image_size, image_size])\n", + " image = (image - CIFAR10_MEAN) / CIFAR10_STD\n", + " if pass_mask:\n", + " image = _append_ones_mask(image)\n", + " return image\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# I-JEPA / LeJEPA-style multi-block masking (Assran et al., 2023)\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "def _sample_block_mask(\n", + " h: int,\n", + " w: int,\n", + " scale_range: tuple[float, float] = (0.15, 0.2),\n", + " aspect_ratio_range: tuple[float, float] = (0.75, 1.5),\n", + ") -> tf.Tensor:\n", + " \"\"\"Sample a single rectangular block mask.\n", + "\n", + " Returns a boolean [H, W] tensor that is True for the masked region.\n", + " The block's area is ``scale * H * W`` and its aspect ratio is sampled\n", + " log-uniformly from ``aspect_ratio_range``.\n", + " \"\"\"\n", + " area = tf.cast(h * w, tf.float32)\n", + " target_area = tf.random.uniform([], scale_range[0], scale_range[1]) * area\n", + "\n", + " log_ratio_lo = tf.math.log(aspect_ratio_range[0])\n", + " log_ratio_hi = tf.math.log(aspect_ratio_range[1])\n", + " aspect_ratio = tf.exp(tf.random.uniform([], log_ratio_lo, log_ratio_hi))\n", + "\n", + " block_h = tf.cast(\n", + " tf.minimum(tf.round(tf.sqrt(target_area / aspect_ratio)),\n", + " tf.cast(h, tf.float32)),\n", + " tf.int32)\n", + " block_w = tf.cast(\n", + " tf.minimum(tf.round(tf.sqrt(target_area * aspect_ratio)),\n", + " tf.cast(w, tf.float32)),\n", + " tf.int32)\n", + " block_h = tf.maximum(block_h, 1)\n", + " block_w = tf.maximum(block_w, 1)\n", + "\n", + " top = tf.random.uniform([], 0, h - block_h + 1, dtype=tf.int32)\n", + " left = tf.random.uniform([], 0, w - block_w + 1, dtype=tf.int32)\n", + "\n", + " rows = tf.range(h)\n", + " cols = tf.range(w)\n", + " row_mask = (rows >= top) & (rows < top + block_h) # [H]\n", + " col_mask = (cols >= left) & (cols < left + block_w) # [W]\n", + " return row_mask[:, tf.newaxis] & col_mask[tf.newaxis, :] # [H, W]\n", + "\n", + "\n", + "def _sample_multi_block_mask(\n", + " h: int,\n", + " w: int,\n", + " num_blocks: int = 4,\n", + " scale_range: tuple[float, float] = (0.15, 0.2),\n", + " aspect_ratio_range: tuple[float, float] = (0.75, 1.5),\n", + ") -> tf.Tensor:\n", + " \"\"\"Sample the union of ``num_blocks`` rectangular block masks.\n", + "\n", + " Returns a boolean [H, W] tensor that is True for *masked* pixels.\n", + " Following I-JEPA, target blocks are sampled at sufficiently large scale\n", + " so that the prediction task is semantic rather than local interpolation.\n", + " \"\"\"\n", + " mask = tf.zeros([h, w], dtype=tf.bool)\n", + " for _ in range(num_blocks):\n", + " block = _sample_block_mask(h, w, scale_range, aspect_ratio_range)\n", + " mask = mask | block\n", + " return mask\n", + "\n", + "\n", + "def lejepa_masking_augment_single(\n", + " image: tf.Tensor,\n", + " image_size: int = 32,\n", + " num_blocks: int = 4,\n", + " scale_range: tuple[float, float] = (0.15, 0.2),\n", + " aspect_ratio_range: tuple[float, float] = (0.75, 1.5),\n", + " pass_mask: bool = False,\n", + ") -> tf.Tensor:\n", + " \"\"\"Create one masked view of a float32 [H,W,3] image in [0,1].\n", + "\n", + " The image is first normalized, then ``num_blocks`` random rectangular\n", + " regions are zeroed out (zero = dataset mean post-normalization).\n", + "\n", + " If ``pass_mask`` is True, a visibility channel (1=visible, 0=masked) is\n", + " concatenated as the 4th channel, giving output shape [H,W,4].\n", + " \"\"\"\n", + " if image_size != 32:\n", + " image = tf.image.resize(image, [image_size, image_size])\n", + " image = (image - CIFAR10_MEAN) / CIFAR10_STD\n", + "\n", + " mask = _sample_multi_block_mask(\n", + " image_size, image_size,\n", + " num_blocks=num_blocks,\n", + " scale_range=scale_range,\n", + " aspect_ratio_range=aspect_ratio_range,\n", + " ) # [H, W], True = masked\n", + " mask_3d = tf.cast(mask[:, :, tf.newaxis], tf.float32) # [H, W, 1]\n", + " # Replace masked pixels with 0 (= dataset mean after normalization).\n", + " image = image * (1.0 - mask_3d)\n", + "\n", + " if pass_mask:\n", + " visibility = 1.0 - mask_3d # [H, W, 1], 1=visible, 0=masked\n", + " image = tf.concat([image, visibility], axis=-1) # [H, W, 4]\n", + "\n", + " return image\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# tf.data pipeline\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "class TFDSLoader:\n", + " \"\"\"CIFAR-10 data loader using tensorflow_datasets.\n", + "\n", + " Yields ``(raw_images, views, labels)`` as NumPy arrays with\n", + " ``views`` in ``[batch, num_views, H, W, C]`` layout (JAX convention),\n", + " where C=3 normally or C=4 when ``pass_mask=True``.\n", + " \"\"\"\n", + "\n", + " def __init__(\n", + " self,\n", + " *,\n", + " split: str,\n", + " batch_size: int,\n", + " training: bool,\n", + " num_views: int = 2,\n", + " image_size: int = 32,\n", + " preprocess_type: str = \"simclr\",\n", + " color_jitter_strength: float = 1.0,\n", + " gaussian_blur: bool = True,\n", + " masking_num_blocks: int = 4,\n", + " masking_scale_range: tuple[float, float] = (0.15, 0.2),\n", + " masking_aspect_ratio_range: tuple[float, float] = (0.75, 1.5),\n", + " pass_mask: bool = False,\n", + " seed: Optional[int] = None,\n", + " data_dir: Optional[str] = None,\n", + " drop_remainder: Optional[bool] = None,\n", + " cache: bool = True,\n", + " ) -> None:\n", + " self._batch_size = batch_size\n", + " self._training = training\n", + " self._num_views = num_views\n", + " self._image_size = image_size\n", + " self._preprocess_type = preprocess_type\n", + " self._color_jitter_strength = color_jitter_strength\n", + " self._gaussian_blur = gaussian_blur\n", + " self._masking_num_blocks = masking_num_blocks\n", + " self._masking_scale_range = masking_scale_range\n", + " self._masking_aspect_ratio_range = masking_aspect_ratio_range\n", + " self._pass_mask = pass_mask\n", + " self._seed = seed\n", + " self._num_classes = 10\n", + "\n", + " if drop_remainder is None:\n", + " drop_remainder = training\n", + " self._drop_remainder = drop_remainder\n", + "\n", + " tfds_split = \"train\" if split == \"train\" else \"test\"\n", + " ds_info = tfds.builder(\"cifar10\", data_dir=data_dir).info\n", + " self._dataset_size = ds_info.splits[tfds_split].num_examples\n", + "\n", + " ds = tfds.load(\n", + " \"cifar10\",\n", + " split=tfds_split,\n", + " as_supervised=True, # yields (image, label)\n", + " data_dir=data_dir,\n", + " shuffle_files=training,\n", + " )\n", + "\n", + " if cache:\n", + " ds = ds.cache()\n", + "\n", + " if training:\n", + " ds = ds.shuffle(\n", + " buffer_size=min(self._dataset_size, 50_000),\n", + " seed=seed,\n", + " reshuffle_each_iteration=True,\n", + " )\n", + "\n", + " ds = ds.map(self._preprocess, num_parallel_calls=tf.data.AUTOTUNE)\n", + " ds = ds.batch(batch_size, drop_remainder=drop_remainder)\n", + " ds = ds.prefetch(tf.data.AUTOTUNE)\n", + " self._ds = ds\n", + "\n", + " # -- Properties expected by the training loop --\n", + "\n", + " @property\n", + " def dataset_size(self) -> int:\n", + " return self._dataset_size\n", + "\n", + " @property\n", + " def num_classes(self) -> int:\n", + " return self._num_classes\n", + "\n", + " # -- Preprocessing --\n", + "\n", + " def _preprocess(self, image: tf.Tensor,\n", + " label: tf.Tensor):\n", + " \"\"\"Map a single (image, label) example to (raw, views, label).\"\"\"\n", + " # raw_image: uint8 [H, W, 3] — unchanged original for visualization.\n", + " raw_image = image\n", + "\n", + " # Float image in [0, 1] for augmentation.\n", + " image_f = tf.cast(image, tf.float32) / 255.0\n", + "\n", + " if self._training:\n", + " if self._preprocess_type == \"simclr\":\n", + " views = tf.stack([\n", + " simclr_augment_single(\n", + " image_f,\n", + " image_size=self._image_size,\n", + " color_jitter_strength=self._color_jitter_strength,\n", + " gaussian_blur=self._gaussian_blur,\n", + " pass_mask=self._pass_mask,\n", + " )\n", + " for _ in range(self._num_views)\n", + " ], axis=0) # [num_views, H, W, C]\n", + " elif self._preprocess_type == \"lejepa_masking\":\n", + " views = tf.stack([\n", + " lejepa_masking_augment_single(\n", + " image_f,\n", + " image_size=self._image_size,\n", + " num_blocks=self._masking_num_blocks,\n", + " scale_range=self._masking_scale_range,\n", + " aspect_ratio_range=self._masking_aspect_ratio_range,\n", + " pass_mask=self._pass_mask,\n", + " )\n", + " for _ in range(self._num_views)\n", + " ], axis=0) # [num_views, H, W, C]\n", + " else:\n", + " raise ValueError(\n", + " f\"Unknown preprocess_type: {self._preprocess_type!r}. \"\n", + " \"Supported: 'simclr', 'lejepa_masking'.\")\n", + " else:\n", + " single_view = eval_preprocess(image_f, image_size=self._image_size,\n", + " pass_mask=self._pass_mask)\n", + " views = tf.stack([single_view] * self._num_views,\n", + " axis=0) # [num_views, H, W, C]\n", + "\n", + " return raw_image, views, label\n", + "\n", + " # -- Iteration --\n", + "\n", + " def __iter__(self) -> Iterator[DataBatch]:\n", + " for raw, views, labels in self._ds:\n", + " yield (raw.numpy(), views.numpy(), labels.numpy())\n", + "\n", + " def __len__(self) -> int:\n", + " if self._drop_remainder:\n", + " return self._dataset_size // self._batch_size\n", + " return math.ceil(self._dataset_size / self._batch_size)\n", + "\n", + "\n", + "# ---------------------------------------------------------------------------\n", + "# Public API — drop-in replacements for the torch-based versions\n", + "# ---------------------------------------------------------------------------\n", + "\n", + "def make_cifar10_dataset(\n", + " split: str,\n", + " batch_size: int,\n", + " training: bool,\n", + " data_dir: Optional[str] = None,\n", + " shuffle_buffer: int = 50_000,\n", + " seed: Optional[int] = None,\n", + " drop_remainder: Optional[bool] = None,\n", + " cache: bool = True,\n", + " preprocess_type: str = \"simclr\",\n", + " num_views: int = 2,\n", + " masking_num_blocks: int = 4,\n", + " masking_scale_range: tuple[float, float] = (0.15, 0.2),\n", + " masking_aspect_ratio_range: tuple[float, float] = (0.75, 1.5),\n", + " pass_mask: bool = False,\n", + " # Torch-era arguments accepted but ignored for compatibility.\n", + " num_workers: int = 0,\n", + " pin_memory: bool = False,\n", + " persistent_workers: Optional[bool] = None,\n", + " download: bool = True,\n", + ") -> TFDSLoader:\n", + " \"\"\"Create a CIFAR-10 data loader yielding NumPy batches for JAX.\"\"\"\n", + " del shuffle_buffer, pin_memory, persistent_workers, download\n", + " if split not in {\"train\", \"test\", \"eval\"}:\n", + " raise ValueError(\n", + " f\"split must be 'train', 'test', or 'eval', got {split!r}.\")\n", + " if preprocess_type not in {\"simclr\", \"lejepa_masking\"}:\n", + " raise ValueError(\n", + " f\"Unknown preprocess_type: {preprocess_type!r}. \"\n", + " \"Supported: 'simclr', 'lejepa_masking'.\")\n", + "\n", + " return TFDSLoader(\n", + " split=split,\n", + " batch_size=batch_size,\n", + " training=training,\n", + " num_views=num_views,\n", + " preprocess_type=preprocess_type,\n", + " masking_num_blocks=masking_num_blocks,\n", + " masking_scale_range=masking_scale_range,\n", + " masking_aspect_ratio_range=masking_aspect_ratio_range,\n", + " pass_mask=pass_mask,\n", + " seed=seed,\n", + " data_dir=data_dir,\n", + " drop_remainder=drop_remainder,\n", + " cache=cache,\n", + " )\n", + "\n", + "\n", + "def load_dataset(\n", + " dataset_name: str,\n", + " split: str,\n", + " is_training: bool,\n", + " batch_size: int,\n", + " seed: int,\n", + " **kwargs,\n", + "):\n", + " \"\"\"Load a dataset and return a JAX-friendly DataLoader wrapper.\"\"\"\n", + " if dataset_name != \"cifar10\":\n", + " raise ValueError(\n", + " f\"Unknown dataset {dataset_name}. Only 'cifar10' is supported.\")\n", + " if split not in (\"train\", \"test\"):\n", + " raise ValueError(\n", + " f\"Unknown split {split}. Supported splits: ('train', 'test')\")\n", + " return make_cifar10_dataset(\n", + " split=split,\n", + " batch_size=batch_size,\n", + " training=is_training,\n", + " seed=seed,\n", + " **kwargs,\n", + " )\n", + "\n", + "\n", + "print(\"data loading defined.\")\n" + ], + "metadata": { + "id": "cell-015" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "bSD6yQyVqXit", + "cell_type": "markdown", + "source": [], + "metadata": { + "id": "bSD6yQyVqXit" + } + }, + { + "id": "cell-016", + "cell_type": "markdown", + "source": [ + "## Part 7 — Linear-probe evaluation (`classification`)\n", + "\n", + "Trains a frozen-feature linear/MLP probe to measure representation quality.\n", + "`create_feature_extractor` builds a function that reads a given feature\n", + "source (`\"latent\"`, `\"recognition\"`, `\"trunk\"`, `\"concatenated\"`) off a\n", + "trained model — this is the piece that dispatches on `model.rpm_type`, so it\n", + "works uniformly for Gaussian RPM, Categorical RPM, and SimCLR.\n", + "\n", + "**Known remaining limitation (not fixed here):** for the Gaussian RPM,\n", + "`\"concatenated\"` still doesn't work — the post-combination latent mean has\n", + "shape `(batch, dim)` (no per-factor axis; it's already summed over factors)\n", + "while the trunk is `(n_factors, batch, dim)`, so concatenating them fails.\n", + "This is a deeper shape-convention question than the two fixes above (which\n", + "were unambiguous), so it's left as-is rather than guessed at. Both\n", + "experiments below default to `\"trunk\"`, which works correctly and is\n", + "unaffected by this.\n" + ], + "metadata": { + "id": "cell-016" + } + }, + { + "id": "cell-017", + "cell_type": "code", + "source": [ + "FeatureSource = str # one of \"latent\", \"recognition\", \"trunk\", \"concatenated\"\n", + "PosteriorMeanFn = Any\n", + "\n", + "\n", + "class ProbeClassifier(nn.Module):\n", + " \"\"\"Lightweight classification head used to probe learned features.\"\"\"\n", + "\n", + " probe_type: str # \"linear\" or \"mlp\"\n", + " num_classes: int\n", + " hidden_dims: tuple[int, ...] = (512, 256)\n", + " dropout_rate: float = 0.1\n", + "\n", + " @nn.compact\n", + " def __call__(self, x: jax.Array, train: bool) -> jax.Array:\n", + " # Linear probe: a single affine map, no hidden layers / dropout.\n", + " if self.probe_type == \"linear\":\n", + " return nn.Dense(self.num_classes, name=\"linear_head\")(x)\n", + "\n", + " # MLP probe: dropout is disabled when ``train`` is False.\n", + " for i, hidden_dim in enumerate(self.hidden_dims):\n", + " x = nn.Dense(hidden_dim, name=f\"mlp_dense_{i}\")(x)\n", + " x = nn.relu(x)\n", + " x = nn.Dropout(rate=self.dropout_rate,\n", + " name=f\"mlp_dropout_{i}\")(x, deterministic=not train)\n", + " return nn.Dense(self.num_classes, name=\"mlp_head\")(x)\n", + "\n", + "\n", + "def create_feature_extractor(\n", + " model: nn.Module,\n", + " feature_source: FeatureSource = \"latent\") -> PosteriorMeanFn:\n", + " \"\"\"Build a jitted function returning per-image features from a model.\n", + "\n", + " All returned functions produce arrays with a leading **view** axis so\n", + " that ``compute_features_and_labels`` can uniformly index ``feats[0]``\n", + " to select the first view and obtain a ``(N, feature_dim)`` array.\n", + "\n", + " Feature sources that already carry a view axis (``\"trunk\"`` and SimCLR\n", + " ``\"latent\"``) are returned as-is. Sources that aggregate across views\n", + " (Gaussian/Categorical RPM ``\"latent\"``, ``\"recognition\"``) are wrapped\n", + " with ``jnp.expand_dims(..., axis=0)`` to add a dummy leading axis.\n", + " \"\"\"\n", + " if model.rpm_type == \"categorical\":\n", + " assert feature_source in (\"latent\", \"recognition\"), (\n", + " f\"Unsupported feature_source {feature_source} for categorical RPM. \"\n", + " )\n", + " elif model.rpm_type == \"simclr\":\n", + " assert feature_source in (\"latent\", \"trunk\", \"concatenated\"), (\n", + " f\"Unsupported feature_source {feature_source} for a SimCLR \"\n", + " \"encoder (there is no separate recognition factor).\")\n", + " is_categorical = model.rpm_type == \"categorical\"\n", + " is_simclr = model.rpm_type == \"simclr\"\n", + "\n", + " if feature_source == \"latent\":\n", + "\n", + " @jax.jit\n", + " def latent_posterior_mean_fn(params: Params,\n", + " batch_views: jax.Array) -> jax.Array:\n", + " main_output, _ = model.apply({\"params\": params}, batch_views)\n", + " if is_simclr:\n", + " # SimCLREncoder's primary output *is* the latent (projection)\n", + " # feature, already shaped like aux[\"trunk\"] (view axis first).\n", + " return main_output\n", + " elif is_categorical:\n", + " # variational.probs is (N, K) — no view axis; add one so\n", + " # compute_features_and_labels can index [0].\n", + " return jnp.expand_dims(main_output[\"variational\"].probs,\n", + " axis=0)\n", + " else:\n", + " # FIX: the original read `main_output[\"variational\"].mean`,\n", + " # but GaussianNatParams/DiagGaussianNatParams have no `.mean`\n", + " # attribute -- only precision/precision_weighted_mean. The\n", + " # actual mean lives in the separately-computed\n", + " # `variational_mean` (a GaussianMeanParams/\n", + " # DiagGaussianMeanParams, which does have `.mean`).\n", + " #\n", + " # variational_mean.mean is (N, dim_latent) — no view axis;\n", + " # add one so compute_features_and_labels can index [0].\n", + " return jnp.expand_dims(\n", + " main_output[\"variational_mean\"].mean, axis=0)\n", + "\n", + " return latent_posterior_mean_fn\n", + " elif feature_source == \"recognition\":\n", + "\n", + " @jax.jit\n", + " def recognition_posterior_mean_fn(params: Params,\n", + " batch_views: jax.Array) -> jax.Array:\n", + " rpm_outputs, _ = model.apply({\"params\": params}, batch_views)\n", + " if is_categorical:\n", + " # factors[0].probs is (N, K) — add view axis.\n", + " return jnp.expand_dims(rpm_outputs[\"factors\"][0].probs,\n", + " axis=0)\n", + " else:\n", + " # recognition[0] selects factor 0; convert to mean params.\n", + " # The result is (N, dim_latent) — add view axis.\n", + " rec_mean = diag_mean_params(rpm_outputs[\"recognition\"][0]).mean\n", + " return jnp.expand_dims(rec_mean, axis=0)\n", + "\n", + " return recognition_posterior_mean_fn\n", + " elif feature_source == \"trunk\":\n", + "\n", + " @jax.jit\n", + " def trunk_posterior_mean_fn(params: Params,\n", + " batch_views: jax.Array) -> jax.Array:\n", + " _, aux = model.apply({\"params\": params}, batch_views)\n", + " return aux[\"trunk\"]\n", + "\n", + " return trunk_posterior_mean_fn\n", + " elif feature_source == \"concatenated\":\n", + "\n", + " @jax.jit\n", + " def concatenated_posterior_mean_fn(params: Params,\n", + " batch_views: jax.Array) -> jax.Array:\n", + " # FIX: the original did `_, aux = model.apply(...)`, discarding\n", + " # the tuple element that actually holds the latent mean, then\n", + " # tried `aux[\"variational\"]` -- but `aux` never has that key (it\n", + " # only ever holds {\"trunk\", \"projection\"}). Now both elements are\n", + " # kept and the latent is read from the right one.\n", + " main_output, aux = model.apply({\"params\": params}, batch_views)\n", + " trunk = aux[\"trunk\"] # (J, N, trunk_dim)\n", + " if is_simclr:\n", + " latent = main_output # (J, N, projection_dim)\n", + " else:\n", + " # (N, dim_latent) — broadcast to match trunk's view dim.\n", + " latent_2d = main_output[\"variational_mean\"].mean\n", + " latent = jnp.broadcast_to(\n", + " latent_2d[jnp.newaxis],\n", + " (trunk.shape[0],) + latent_2d.shape)\n", + " return jnp.concatenate([latent, trunk], axis=-1)\n", + "\n", + " return concatenated_posterior_mean_fn\n", + " else:\n", + " raise ValueError(\n", + " f\"Unsupported feature_source {feature_source}. \"\n", + " f\"Supported sources: ('latent', 'recognition', 'trunk', 'concatenated')\"\n", + " )\n", + "\n", + "\n", + "def compute_features_and_labels(\n", + " posterior_mean_fn: PosteriorMeanFn,\n", + " params: Params,\n", + " data_loader,\n", + ") -> tuple[jax.Array, jax.Array]:\n", + " \"\"\"Extract posterior-mean features for every image, in mini-batches.\"\"\"\n", + " iter_ds = iter(data_loader)\n", + " features = []\n", + " labels_list = []\n", + " for (_, views, label) in iter_ds:\n", + " feats = posterior_mean_fn(params, views)\n", + " # We assume that we do not have randomness during eval.\n", + " feats = feats[0]\n", + " features.append(feats)\n", + " labels_list.append(label)\n", + " features = jnp.concatenate(features, axis=0)\n", + " labels = jnp.concatenate(labels_list, axis=0)\n", + " return features, labels\n", + "\n", + "\n", + "def run_probe(\n", + " model: nn.Module,\n", + " model_params: Params,\n", + " probe_model: nn.Module,\n", + " train_dataloader,\n", + " test_dataloader,\n", + " init_rng: jax.Array,\n", + " output_dir: Path | None,\n", + " lr: float = 0.01,\n", + " weight_decay: float = 0.0,\n", + " num_epochs: int = 20,\n", + " batch_size: int = 256,\n", + " probe_type: str = \"linear\",\n", + " mlp_hidden_dims: Sequence[int] = (512, 256),\n", + " mlp_dropout: float = 0.1,\n", + " step: int = 0,\n", + " feature_source: FeatureSource = \"latent\",\n", + ") -> float:\n", + " \"\"\"Train a probe classifier on frozen features, return best val acc.\"\"\"\n", + " posterior_mean_fn = create_feature_extractor(model,\n", + " feature_source=feature_source)\n", + " probe_tag = f\"{probe_type}-probe\"\n", + " probe_dir_name = f\"{probe_type}_probe\"\n", + " hidden_dims = tuple(int(d) for d in mlp_hidden_dims)\n", + "\n", + " print(\"Pre-computing features for training set...\")\n", + " train_feats, train_labels = compute_features_and_labels(\n", + " posterior_mean_fn=posterior_mean_fn,\n", + " params=model_params,\n", + " data_loader=train_dataloader,\n", + " )\n", + " train_feats = jax.lax.stop_gradient(train_feats)\n", + "\n", + " print(\"Pre-computing features for validation set...\")\n", + " val_feats, val_labels = compute_features_and_labels(\n", + " posterior_mean_fn=posterior_mean_fn,\n", + " params=model_params,\n", + " data_loader=test_dataloader,\n", + " )\n", + " val_feats = jax.lax.stop_gradient(val_feats)\n", + "\n", + " optimizer = optax.sgd(\n", + " learning_rate=lr,\n", + " momentum=0.9,\n", + " nesterov=True,\n", + " )\n", + " if weight_decay > 0:\n", + " optimizer = optax.chain(optax.add_decayed_weights(weight_decay),\n", + " optimizer)\n", + "\n", + " init_rngs = {\"params\": init_rng, \"dropout\": init_rng}\n", + "\n", + " if feature_source == \"latent\":\n", + " probe_feat_dim = model.latent_dim\n", + " elif feature_source == \"recognition\":\n", + " probe_feat_dim = model.latent_dim\n", + " elif feature_source == \"trunk\":\n", + " probe_feat_dim = model.trunk_dim\n", + " elif feature_source == \"concatenated\":\n", + " probe_feat_dim = model.latent_dim + model.trunk_dim\n", + " else:\n", + " raise ValueError(f\"Unknown feature source: {feature_source}\")\n", + "\n", + " dummy_x = jnp.zeros((1, probe_feat_dim), dtype=jnp.float32)\n", + " init_probe_vars = probe_model.init(init_rngs, dummy_x, train=True)\n", + " probe_params = init_probe_vars[\"params\"]\n", + "\n", + " opt_state = optimizer.init(probe_params)\n", + "\n", + " @jax.jit\n", + " def train_step(\n", + " params,\n", + " opt_state,\n", + " features,\n", + " labels,\n", + " dropout_key,\n", + " ):\n", + "\n", + " def loss_fn(curr_params):\n", + " logits = probe_model.apply(\n", + " {\"params\": curr_params},\n", + " features,\n", + " train=True,\n", + " rngs={\"dropout\": dropout_key},\n", + " )\n", + " loss = jnp.mean(\n", + " optax.softmax_cross_entropy_with_integer_labels(logits, labels))\n", + " acc = jnp.mean((jnp.argmax(logits,\n", + " axis=-1) == labels).astype(jnp.float32))\n", + " return loss, acc\n", + "\n", + " (loss, acc), grads = jax.value_and_grad(loss_fn, has_aux=True)(params)\n", + " updates, new_opt_state = optimizer.update(grads, opt_state, params)\n", + " new_params = optax.apply_updates(params, updates)\n", + " return new_params, new_opt_state, loss, acc\n", + "\n", + " @jax.jit\n", + " def eval_step(params, features, labels):\n", + " logits = probe_model.apply({\"params\": params}, features, train=False)\n", + " loss = jnp.mean(\n", + " optax.softmax_cross_entropy_with_integer_labels(logits, labels))\n", + " acc = jnp.mean((jnp.argmax(logits,\n", + " axis=-1) == labels).astype(jnp.float32))\n", + " return loss, acc\n", + "\n", + " best_val_acc = 0.0\n", + " best_epoch = 0\n", + " rng = jax.random.PRNGKey(1)\n", + " n_train = train_feats.shape[0]\n", + " n_val = val_feats.shape[0]\n", + " for epoch in range(num_epochs):\n", + " train_correct = 0\n", + " train_total = 0\n", + " train_loss_sum = 0.0\n", + "\n", + " for i in range(0, n_train, batch_size):\n", + " batch_x = train_feats[i:i + batch_size]\n", + " batch_y = train_labels[i:i + batch_size]\n", + " rng, step_rng = jax.random.split(rng)\n", + " probe_params, opt_state, loss, acc = train_step(\n", + " probe_params, opt_state, batch_x, batch_y, step_rng)\n", + "\n", + " batch_count = batch_y.shape[0]\n", + " train_loss_sum += float(loss) * batch_count\n", + " train_total += batch_count\n", + " train_correct += int(round(float(acc) * batch_count))\n", + "\n", + " val_correct = 0\n", + " val_total = 0\n", + " val_loss_sum = 0.0\n", + "\n", + " for i in range(0, n_val, batch_size):\n", + " batch_x = val_feats[i:i + batch_size]\n", + " batch_y = val_labels[i:i + batch_size]\n", + "\n", + " loss, acc = eval_step(probe_params, batch_x, batch_y)\n", + "\n", + " batch_count = batch_y.shape[0]\n", + " val_loss_sum += float(loss) * batch_count\n", + " val_total += batch_count\n", + " val_correct += int(round(float(acc) * batch_count))\n", + "\n", + " train_acc = train_correct / max(train_total, 1)\n", + " val_acc = val_correct / max(val_total, 1)\n", + " train_loss = train_loss_sum / max(train_total, 1)\n", + " val_loss = val_loss_sum / max(val_total, 1)\n", + "\n", + " if val_acc > best_val_acc:\n", + " best_val_acc = val_acc\n", + " best_epoch = epoch + 1\n", + " if output_dir is not None:\n", + " probe_ckpt_dir = output_dir / probe_dir_name\n", + " probe_ckpt_dir.mkdir(parents=True, exist_ok=True)\n", + " checkpoint_path = (\n", + " probe_ckpt_dir /\n", + " f\"best_{probe_type}_probe_step_{step}\").resolve()\n", + " payload = {\n", + " \"epoch\": int(best_epoch),\n", + " \"val_acc\": float(best_val_acc),\n", + " \"probe_type\": probe_type,\n", + " \"hidden_dims\": tuple(int(v) for v in hidden_dims),\n", + " \"dropout\": float(mlp_dropout),\n", + " \"params\": serialization.to_state_dict(probe_params),\n", + " }\n", + " checkpointer = ocp.PyTreeCheckpointer()\n", + " checkpointer.save(str(checkpoint_path), payload, force=True)\n", + "\n", + " print(f\"[{probe_tag}] epoch={epoch + 1:03d}/{num_epochs:03d} \"\n", + " f\"train_loss={train_loss:.5f} val_loss={val_loss:.5f} \"\n", + " f\"train_acc={train_acc * 100:.2f}% val_acc={val_acc * 100:.2f}% \"\n", + " f\"best_val_acc={best_val_acc * 100:.2f}%\")\n", + "\n", + " print(\n", + " f\"[{probe_tag}] best_val_acc={best_val_acc * 100:.2f}% at epoch {best_epoch}\"\n", + " )\n", + " return best_val_acc\n", + "\n", + "\n", + "print(\"classification defined.\")\n" + ], + "metadata": { + "id": "cell-017" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-018", + "cell_type": "markdown", + "source": [ + "## Part 8 — Training-step builders and checkpointing (`train_helpers`)\n", + "\n", + "`_get_loss_fn` dispatches on `model.rpm_type` to pick the right free-energy\n", + "function — this is the one piece of glue that makes the *same* generic\n", + "training loop below work for Gaussian RPM, Categorical RPM, and SimCLR.\n" + ], + "metadata": { + "id": "cell-018" + } + }, + { + "id": "cell-019", + "cell_type": "code", + "source": [ + "@struct.dataclass\n", + "class TrainState:\n", + " params: Params\n", + " opt_state: optax.OptState\n", + " rng: jax.Array\n", + " step: int\n", + "\n", + "\n", + "def _get_loss_fn(model):\n", + " \"\"\"Return the (params, model, batch, rng, beta) -> (free_energy, metrics)\n", + " function for ``model``. Despite the name, this covers non-RPM models too\n", + " (e.g. SimCLR): \"free_energy\" here just means \"the quantity being\n", + " maximized\", matching create_train_step's ``loss = -free_energy``.\n", + " \"\"\"\n", + " if model.rpm_type == \"categorical\":\n", + " return free_energy_categorical\n", + " elif model.rpm_type == \"gaussian\":\n", + " return free_energy\n", + " elif model.rpm_type == \"simclr\":\n", + " return nt_xent_loss\n", + " else:\n", + " raise ValueError(f\"Unknown RPM type: {model.rpm_type}\")\n", + "\n", + "\n", + "def create_train_step(\n", + " model,\n", + " optimizer: optax.GradientTransformation,\n", + " beta_schedule: float | optax.Schedule = 1.0,\n", + "):\n", + " \"\"\"Build a jitted single optimization step that maximizes the free energy.\"\"\"\n", + " free_energy_fn = _get_loss_fn(model)\n", + " beta_schedule = (beta_schedule if callable(beta_schedule) else\n", + " optax.constant_schedule(beta_schedule))\n", + "\n", + " @jax.jit\n", + " def train_step(\n", + " state: TrainState,\n", + " batch_views: jax.Array) -> tuple[TrainState, dict[str, Array]]:\n", + " rng, loss_rng = jax.random.split(state.rng)\n", + " beta_t = beta_schedule(state.step)\n", + "\n", + " def loss_fn(params, rng):\n", + " vars_dict = {\"params\": params}\n", + " free_energy_val, aux_terms = free_energy_fn(\n", + " vars_dict,\n", + " model,\n", + " batch_views,\n", + " rng=rng,\n", + " beta=beta_t,\n", + " )\n", + " # We maximize the free energy, so minimize its negative.\n", + " loss = -free_energy_val\n", + " return loss, (free_energy_val, aux_terms)\n", + "\n", + " (loss, (free_energy_val,\n", + " aux_terms)), grads = value_and_grad(loss_fn,\n", + " has_aux=True)(state.params,\n", + " loss_rng)\n", + "\n", + " # Raw global gradient norm (before any clipping). A sudden spike here is\n", + " # the early-warning signal for the divergence that blows a healthy run\n", + " # into the collapsed (free energy -> 0) state.\n", + " grad_norm = optax.global_norm(grads)\n", + "\n", + " updates, new_opt_state = optimizer.update(grads, state.opt_state,\n", + " state.params)\n", + " new_params = optax.apply_updates(state.params, updates)\n", + "\n", + " new_state = state.replace(\n", + " params=new_params,\n", + " opt_state=new_opt_state,\n", + " rng=rng,\n", + " step=state.step + 1,\n", + " )\n", + "\n", + " metrics = {\n", + " \"loss\": loss,\n", + " \"free_energy\": free_energy_val,\n", + " \"grad_norm\": grad_norm,\n", + " \"beta\": beta_t,\n", + " }\n", + " metrics.update(aux_terms)\n", + "\n", + " return new_state, metrics\n", + "\n", + " return train_step\n", + "\n", + "\n", + "def create_eval_step(model, beta_schedule: float | optax.Schedule = 1.0):\n", + " \"\"\"Build a jitted function returning the free energy for a batch.\"\"\"\n", + " free_energy_fn = _get_loss_fn(model)\n", + "\n", + " beta_schedule = (beta_schedule if callable(beta_schedule) else\n", + " optax.constant_schedule(beta_schedule))\n", + "\n", + " @jax.jit\n", + " def eval_step(params: Params, batch_views: jax.Array, step: int,\n", + " rng: Array) -> Array:\n", + " beta_t = beta_schedule(step)\n", + " free_energy_val, _ = free_energy_fn({\"params\": params},\n", + " model,\n", + " batch_views,\n", + " rng=rng,\n", + " beta=beta_t)\n", + " return free_energy_val\n", + "\n", + " return eval_step\n", + "\n", + "\n", + "def _checkpoint_metadata(history: dict, best_val_fe: float) -> dict:\n", + " \"\"\"Build a JSON-serializable metadata payload for Orbax checkpoints.\"\"\"\n", + " return {\n", + " \"history\": history,\n", + " \"best_val_fe\": float(best_val_fe),\n", + " }\n", + "\n", + "\n", + "def save_rpm_params_checkpoint(\n", + " ckpt_manager: ocp.CheckpointManager,\n", + " step: int,\n", + " params: Params,\n", + " history: dict,\n", + " best_val_fe: float,\n", + ") -> None:\n", + " \"\"\"Save model params (not full training state) plus lightweight metadata.\"\"\"\n", + " ckpt_manager.save(\n", + " int(step),\n", + " args=ocp.args.Composite(\n", + " params=ocp.args.StandardSave(params),\n", + " metadata=ocp.args.JsonSave(\n", + " _checkpoint_metadata(history, best_val_fe)),\n", + " ),\n", + " )\n", + "\n", + "\n", + "def save_checkpoint(\n", + " ckpt_manager: ocp.CheckpointManager,\n", + " current_state: TrainState,\n", + " current_history: dict,\n", + " current_best_val_fe: float,\n", + ") -> None:\n", + " save_rpm_params_checkpoint(\n", + " ckpt_manager=ckpt_manager,\n", + " step=int(current_state.step),\n", + " params=current_state.params,\n", + " history=current_history,\n", + " best_val_fe=current_best_val_fe,\n", + " )\n", + "\n", + "\n", + "def plot_training_metrics(\n", + " history,\n", + " *,\n", + " val_every=None,\n", + " probe_every=None,\n", + " title=None,\n", + "):\n", + " \"\"\"Plot train free energy, val free energy, and best probe accuracy.\"\"\"\n", + "\n", + " def _infer_interval(total_steps, count):\n", + " if count <= 0 or total_steps <= 0:\n", + " return None\n", + " interval = int(round(total_steps / count))\n", + " return interval if interval > 0 else None\n", + "\n", + " steps = np.asarray(history.get(\"step\", []), dtype=int)\n", + "\n", + " if \"free_energy\" in history and len(history[\"free_energy\"]) > 0:\n", + " free_energy_hist = np.asarray(history[\"free_energy\"][1:], dtype=float)\n", + " else:\n", + " losses = np.asarray(history.get(\"loss\", [])[1:], dtype=float)\n", + " free_energy_hist = -losses\n", + "\n", + " if steps.size == 0:\n", + " steps = np.arange(1, free_energy_hist.size + 1)\n", + "\n", + " val_free_energy = np.asarray(history.get(\"val_free_energy\", []),\n", + " dtype=float)\n", + " probe_val_acc = np.asarray(history.get(\"probe_val_acc\", []), dtype=float)\n", + "\n", + " val_steps = np.array([], dtype=int)\n", + " if val_free_energy.size > 0:\n", + " interval = val_every\n", + " if interval is None and steps.size > 0:\n", + " interval = _infer_interval(int(steps[-1]),\n", + " int(val_free_energy.size))\n", + " if interval is None:\n", + " interval = 1\n", + " val_steps = np.arange(1, val_free_energy.size + 1) * interval\n", + "\n", + " probe_steps = np.array([], dtype=int)\n", + " if probe_val_acc.size > 0:\n", + " interval = probe_every\n", + " if interval is None and steps.size > 0:\n", + " interval = _infer_interval(int(steps[-1]), int(probe_val_acc.size))\n", + " if interval is None:\n", + " interval = 1\n", + " probe_steps = np.arange(0, probe_val_acc.size) * interval\n", + "\n", + " acc_label = \"best val probe acc\"\n", + " if probe_val_acc.size > 0 and probe_val_acc.max() <= 1.0:\n", + " probe_val_acc = 100.0 * probe_val_acc\n", + " acc_label = \"best val probe acc (%)\"\n", + "\n", + " fig, axes = plt.subplots(2, 1, figsize=(8, 6), sharex=True)\n", + " axes[0].plot(steps[1:], free_energy_hist, label=\"train free energy\")\n", + " if val_free_energy.size > 0:\n", + " axes[0].plot(val_steps, val_free_energy, label=\"val free energy\")\n", + " axes[0].set_ylabel(\"free energy\")\n", + " axes[0].grid(True, alpha=0.3)\n", + " axes[0].legend()\n", + "\n", + " if probe_val_acc.size > 0:\n", + " axes[1].plot(probe_steps, probe_val_acc, label=acc_label)\n", + " else:\n", + " axes[1].text(\n", + " 0.5,\n", + " 0.5,\n", + " \"No probe results logged\",\n", + " ha=\"center\",\n", + " va=\"center\",\n", + " transform=axes[1].transAxes,\n", + " )\n", + " axes[1].set_ylabel(acc_label)\n", + " axes[1].set_xlabel(\"training step\")\n", + " axes[1].grid(True, alpha=0.3)\n", + "\n", + " if title:\n", + " fig.suptitle(title)\n", + "\n", + " fig.tight_layout()\n", + " return fig, axes\n", + "\n", + "\n", + "print(\"train_helpers defined.\")\n" + ], + "metadata": { + "id": "cell-019" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-020", + "cell_type": "markdown", + "source": [ + "## Part 9 — Generic training loop (`train_rpm`)\n", + "\n", + "This is the same function for both experiments: it dispatches its loss via\n", + "`_get_loss_fn(model)` (Part 8), so passing a `GaussianRPM` or a\n", + "`SimCLREncoder` in the `model=` argument is the only thing that changes\n", + "between the two experiments below.\n" + ], + "metadata": { + "id": "cell-020" + } + }, + { + "id": "cell-021", + "cell_type": "code", + "source": [ + "def train_rpm(\n", + " model,\n", + " train_dataloader,\n", + " eval_train_dataloader,\n", + " eval_test_dataloader,\n", + " optimizer_name,\n", + " probe_type,\n", + " output_dir: Path = Path(\"outputs\"),\n", + " learning_rate: float | optax.Schedule = 1e-3,\n", + " weight_decay: float = 0.05,\n", + " grad_clip_norm: float | None = 1.0,\n", + " batch_size: int = 64,\n", + " num_steps: int = 1000,\n", + " log_every: int = 100,\n", + " val_every: int = 500,\n", + " val_steps: int = 10,\n", + " probe_every: int = 500,\n", + " probe_epochs: int = 30,\n", + " probe_lr: float = 0.1,\n", + " seed: int = 42,\n", + " beta: float | optax.Schedule = 1.0,\n", + " probe_feature_source: FeatureSource = \"latent\",\n", + " checkpoint_every: int = 10000,\n", + " save_best: bool = True,\n", + " max_checkpoints_to_keep: int = 3,\n", + ") -> tuple[TrainState, dict]:\n", + " \"\"\"Train ``model`` and return the final state and a metrics history.\"\"\"\n", + " output_dir.mkdir(parents=True, exist_ok=True)\n", + " ckpt_dir = (output_dir / \"checkpoints\").absolute()\n", + " best_ckpt_dir = (output_dir / \"best_checkpoint\").absolute()\n", + " ckpt_dir.mkdir(parents=True, exist_ok=True)\n", + " best_ckpt_dir.mkdir(parents=True, exist_ok=True)\n", + "\n", + " rng = jax.random.PRNGKey(seed)\n", + " rng, init_rng = jax.random.split(rng)\n", + "\n", + " _, init_batch, _ = next(iter(train_dataloader))\n", + " variables = model.init(init_rng, init_batch)\n", + " lr_schedule = (learning_rate if callable(learning_rate) else\n", + " optax.constant_schedule(learning_rate))\n", + " tx = getattr(optax, optimizer_name)(learning_rate=lr_schedule,\n", + " weight_decay=weight_decay)\n", + " if grad_clip_norm is not None:\n", + " tx = optax.chain(optax.clip_by_global_norm(grad_clip_norm), tx)\n", + " opt_state = tx.init(variables[\"params\"])\n", + " state = TrainState(\n", + " params=variables[\"params\"],\n", + " opt_state=opt_state,\n", + " rng=rng,\n", + " step=0,\n", + " )\n", + " beta_schedule = beta if callable(beta) else optax.constant_schedule(beta)\n", + " train_step_fn = create_train_step(model=model,\n", + " optimizer=tx,\n", + " beta_schedule=beta_schedule)\n", + " eval_step_fn = create_eval_step(model=model, beta_schedule=beta_schedule)\n", + "\n", + " history = {\n", + " \"step\": [],\n", + " \"loss\": [],\n", + " \"lr\": [],\n", + " \"grad_norm\": [],\n", + " \"beta\": [],\n", + " \"val_free_energy\": [],\n", + " \"probe_val_acc\": [],\n", + " \"neg_kl_qp\": [],\n", + " \"neg_kl_qf\": [],\n", + " \"neg_log_gamma\": [],\n", + " \"q_z_mean_std\": [],\n", + " \"q_z_var_mean\": [],\n", + " \"entropy\": [],\n", + " \"prior_xent\": [],\n", + " \"factors_xent\": [],\n", + " \"denom_xent\": [],\n", + " }\n", + " # FIX: the original initialized `best_val_fe = float(\"inf\")` and saved a\n", + " # new \"best\" checkpoint whenever `mean_val_fe < best_val_fe` -- but free\n", + " # energy is *maximized* during training (verified empirically: it climbs\n", + " # steadily over training steps), so tracking a minimum silently saved the\n", + " # *worst* validation checkpoint as \"best\". Track the maximum instead.\n", + " best_val_fe = float(\"-inf\")\n", + " main_options = ocp.CheckpointManagerOptions(\n", + " max_to_keep=max_checkpoints_to_keep,\n", + " create=True,\n", + " )\n", + " best_options = ocp.CheckpointManagerOptions(\n", + " max_to_keep=1,\n", + " create=True,\n", + " )\n", + " print(f\"Starting training for {num_steps} steps, \")\n", + " start_time = time.time()\n", + " probe_model = ProbeClassifier(\n", + " probe_type=probe_type,\n", + " num_classes=train_dataloader.num_classes,\n", + " )\n", + " with (\n", + " ocp.CheckpointManager(ckpt_dir,\n", + " options=main_options) as ckpt_manager,\n", + " ocp.CheckpointManager(best_ckpt_dir, options=best_options) as\n", + " best_ckpt_manager,\n", + " ):\n", + " # Random initialization probe.\n", + " accuracy = run_probe(\n", + " model=model,\n", + " model_params=state.params,\n", + " probe_model=probe_model,\n", + " train_dataloader=eval_train_dataloader,\n", + " test_dataloader=eval_test_dataloader,\n", + " output_dir=output_dir,\n", + " num_epochs=probe_epochs,\n", + " lr=probe_lr,\n", + " weight_decay=0.0,\n", + " batch_size=batch_size,\n", + " step=int(state.step),\n", + " feature_source=probe_feature_source,\n", + " init_rng=init_rng,\n", + " )\n", + " print(f\"Probe Validation Accuracy: {accuracy * 100:.2f}%\")\n", + " history[\"probe_val_acc\"].append(float(accuracy))\n", + "\n", + " if checkpoint_every > 0:\n", + " save_checkpoint(ckpt_manager=ckpt_manager,\n", + " current_state=state,\n", + " current_history=history,\n", + " current_best_val_fe=best_val_fe)\n", + "\n", + " iter_train_ds = iter(train_dataloader)\n", + "\n", + " for step in range(num_steps):\n", + " state_rng, _ = jax.random.split(state.rng)\n", + " state = state.replace(rng=state_rng)\n", + "\n", + " # Get the batch, cycling to a new epoch once the current one is\n", + " # exhausted (train_dataloader only iterates for one epoch).\n", + " # Re-iterating draws a fresh shuffle, not a repeat of the same\n", + " # order: the DataLoader is built with shuffle=True and a seeded\n", + " # torch.Generator that it reuses across iterations, and that\n", + " # generator's internal state advances on every draw -- so each\n", + " # new epoch's permutation differs from the last, even though the\n", + " # run as a whole is still reproducible from the original seed.\n", + " try:\n", + " _, batch_views, _ = next(iter_train_ds)\n", + " except StopIteration:\n", + " iter_train_ds = iter(train_dataloader)\n", + " _, batch_views, _ = next(iter_train_ds)\n", + "\n", + " state, metrics = train_step_fn(state, batch_views)\n", + "\n", + " if (step + 1) % log_every == 0 or step == 0:\n", + " elapsed = time.time() - start_time\n", + " current_lr = float(lr_schedule(int(state.step)))\n", + "\n", + " print(f\"Step {step + 1}/{num_steps} | \"\n", + " f\"Loss: {float(metrics['loss']):.4f} | \"\n", + " f\"Free Energy: {float(metrics['free_energy']):.4f} | \"\n", + " f\"grad_norm: {float(metrics['grad_norm']):.4f} | \"\n", + " f\"beta: {float(metrics['beta']):.4f} | \"\n", + " f\"LR: {current_lr:.5f} | \"\n", + " f\"Elapsed: {elapsed:.4f}\")\n", + " history[\"step\"].append(step + 1)\n", + " history[\"lr\"].append(current_lr)\n", + " for key, value in metrics.items():\n", + " history.setdefault(key, []).append(float(value))\n", + "\n", + " if (step + 1) % val_every == 0:\n", + " val_fes = []\n", + " eval_ds = iter(eval_test_dataloader)\n", + " eval_rng, state_rng = jax.random.split(state.rng)\n", + " for _ in range(val_steps):\n", + " eval_rng, _ = jax.random.split(eval_rng)\n", + " _, v_views, _ = next(eval_ds)\n", + " val_fe = eval_step_fn(state.params, v_views, state.step,\n", + " eval_rng)\n", + " val_fes.append(float(val_fe))\n", + "\n", + " mean_val_fe = float(jnp.mean(jnp.array(val_fes)))\n", + " print(f\"Validation Free Energy: {mean_val_fe:.4f}\")\n", + " history[\"val_free_energy\"].append(mean_val_fe)\n", + "\n", + " # FIX: `>` instead of the original `<` -- see best_val_fe note above.\n", + " if save_best and mean_val_fe > best_val_fe:\n", + " best_val_fe = mean_val_fe\n", + " print(\n", + " f\"New best validation free energy: {best_val_fe:.4f}. \"\n", + " \"Saving best checkpoint...\")\n", + " save_checkpoint(\n", + " ckpt_manager=best_ckpt_manager,\n", + " current_state=state,\n", + " current_history=history,\n", + " current_best_val_fe=best_val_fe)\n", + "\n", + " if (step + 1) % probe_every == 0:\n", + " print(\"Running probe evaluation\")\n", + " accuracy = run_probe(\n", + " model=model,\n", + " model_params=state.params,\n", + " probe_model=probe_model,\n", + " train_dataloader=eval_train_dataloader,\n", + " test_dataloader=eval_test_dataloader,\n", + " output_dir=output_dir,\n", + " num_epochs=probe_epochs,\n", + " lr=probe_lr,\n", + " weight_decay=0.0,\n", + " batch_size=batch_size,\n", + " step=int(state.step),\n", + " feature_source=probe_feature_source,\n", + " init_rng=state.rng,\n", + " )\n", + " print(f\"Probe Validation Accuracy: {accuracy * 100:.2f}%\")\n", + " history[\"probe_val_acc\"].append(float(accuracy))\n", + "\n", + " if checkpoint_every > 0 and ((step + 1) % checkpoint_every == 0 or\n", + " (step + 1) == num_steps):\n", + " print(f\"Saving checkpoint at step {step + 1}...\")\n", + " save_checkpoint(ckpt_manager=ckpt_manager,\n", + " current_state=state,\n", + " current_history=history,\n", + " current_best_val_fe=best_val_fe)\n", + "\n", + " print(\"Saving final checkpoint...\")\n", + " save_checkpoint(ckpt_manager=ckpt_manager,\n", + " current_state=state,\n", + " current_history=history,\n", + " current_best_val_fe=best_val_fe)\n", + "\n", + " return state, history\n", + "\n", + "\n", + "print(\"train_rpm defined.\")\n" + ], + "metadata": { + "id": "cell-021" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "ff6945b8", + "cell_type": "markdown", + "source": [ + "## Visualize data views\n", + "\n", + "Load a small batch from CIFAR-10 with both preprocessing strategies and\n", + "display a few sample images alongside their generated views.\n" + ], + "metadata": { + "id": "ff6945b8" + } + }, + { + "id": "bde0e3c8", + "cell_type": "code", + "source": [ + "NUM_SAMPLES_TO_SHOW = 6\n", + "NUM_VIEWS_VIS = 4\n", + "\n", + "# --- Load a small batch with each strategy ---\n", + "_vis_simclr = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"train\", batch_size=NUM_SAMPLES_TO_SHOW,\n", + " seed=0, preprocess_type=\"simclr\", num_views=NUM_VIEWS_VIS, is_training=True,\n", + ")\n", + "_vis_masking = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"train\", batch_size=NUM_SAMPLES_TO_SHOW,\n", + " seed=0, preprocess_type=\"lejepa_masking\", num_views=NUM_VIEWS_VIS,\n", + " is_training=True,\n", + ")\n", + "\n", + "_raw_simclr, _views_simclr, _labels_simclr = next(iter(_vis_simclr))\n", + "_raw_masking, _views_masking, _labels_masking = next(iter(_vis_masking))\n", + "\n", + "def _unnormalize(view: np.ndarray) -> np.ndarray:\n", + " \"\"\"Undo per-channel normalization and clip to [0, 1] for display.\"\"\"\n", + " return np.clip(view * CIFAR10_STD + CIFAR10_MEAN, 0.0, 1.0)\n", + "\n", + "CIFAR10_CLASSES = [\n", + " \"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\",\n", + " \"dog\", \"frog\", \"horse\", \"ship\", \"truck\",\n", + "]\n", + "\n", + "for strategy_name, raw_imgs, views_batch, labels in [\n", + " (\"SimCLR\", _raw_simclr, _views_simclr, _labels_simclr),\n", + " (\"LeJEPA masking\", _raw_masking, _views_masking, _labels_masking),\n", + "]:\n", + " n_samples = raw_imgs.shape[0]\n", + " n_views = views_batch.shape[1]\n", + " fig, axes = plt.subplots(\n", + " n_samples, 1 + n_views,\n", + " figsize=(2.2 * (1 + n_views), 2.2 * n_samples),\n", + " )\n", + " fig.suptitle(f\"Views: {strategy_name}\", fontsize=14, fontweight=\"bold\", y=1.01)\n", + "\n", + " for i in range(n_samples):\n", + " label_str = CIFAR10_CLASSES[labels[i]]\n", + " # Original image\n", + " axes[i, 0].imshow(raw_imgs[i])\n", + " axes[i, 0].set_title(f\"Original\\n({label_str})\", fontsize=9)\n", + " axes[i, 0].axis(\"off\")\n", + "\n", + " # Augmented / masked views\n", + " for v in range(n_views):\n", + " view_img = _unnormalize(views_batch[i, v])\n", + " axes[i, v + 1].imshow(view_img)\n", + " axes[i, v + 1].set_title(f\"View {v + 1}\", fontsize=9)\n", + " axes[i, v + 1].axis(\"off\")\n", + "\n", + " plt.tight_layout()\n", + " plt.show()\n", + "\n", + "del _vis_simclr, _vis_masking\n", + "del _raw_simclr, _views_simclr, _labels_simclr\n", + "del _raw_masking, _views_masking, _labels_masking\n" + ], + "metadata": { + "id": "bde0e3c8" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "raise ValueError('ok')" + ], + "metadata": { + "id": "9hIAsh8BMdxj" + }, + "id": "9hIAsh8BMdxj", + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-022", + "cell_type": "markdown", + "source": [ + "---\n", + "# Experiment 1 — Gaussian RPM on CIFAR-10\n", + "\n", + "A ResNet-18-style backbone, `\"constrained\"` auxiliary construction, diagonal\n", + "precision. Mirrors the original `examples/cifar10.py` script's defaults,\n", + "just with plain Python variables instead of `absl` flags.\n", + "\n", + "**Defaults below are a quick smoke run** (`num_steps=2000`) so you can\n", + "confirm everything works end-to-end before committing to a long run — bump\n", + "`NUM_STEPS` (and the schedule-related steps below it) up for a real result.\n", + "On a Colab T4 GPU this backbone does roughly 5-10 steps/sec; on CPU expect\n", + "well under 1 step/sec, so keep steps low unless you have a GPU runtime.\n" + ], + "metadata": { + "id": "cell-022" + } + }, + { + "id": "cell-023", + "cell_type": "code", + "source": [ + "def set_global_seed(seed: int) -> None:\n", + " \"\"\"Set seeds for reproducibility across common Python ML libraries.\"\"\"\n", + " random.seed(seed)\n", + " os.environ[\"PYTHONHASHSEED\"] = str(seed)\n", + " np.random.seed(seed)\n", + "\n", + "\n", + "def get_resnet18_cfg(projection_features):\n", + " return {\n", + " \"stage_sizes\": (2, 2, 2, 2),\n", + " \"stage_widths\": (64, 128, 256, 512),\n", + " \"stem_width\": 64,\n", + " \"stem_kernel_size\": (3, 3),\n", + " \"stem_stride\": 1,\n", + " \"use_max_pool\": False,\n", + " \"block_type\": \"basic\",\n", + " \"projection_features\": projection_features,\n", + " }\n", + "\n", + "\n", + "# --- Config (mirrors examples/cifar10.py's flag defaults) ---\n", + "SEED = 42\n", + "BATCH_SIZE = 128\n", + "EVAL_BATCH_SIZE = 256\n", + "NUM_VIEWS = 8\n", + "NUM_WORKERS = 0 # Colab: keep at 0: DataLoader workers + notebooks\n", + " # sometimes deadlock. Increase only in a plain script.\n", + "LATENT_DIM = 64\n", + "CACHE = False\n", + "PREPROCESS_TYPE = \"simclr\" # \"simclr\" or \"lejepa_masking\"\n", + "# PREPROCESS_TYPE = \"lejepa_masking\" # \"simclr\" or \"lejepa_masking\"\n", + "\n", + "NUM_STEPS = 100_000 # quick smoke run; original default is 100_000\n", + "PROBE_EVERY = 10000\n", + "VAL_EVERY = 5000\n", + "PROBE_EPOCHS = 10\n", + "PROBE_LR = 0.1\n", + "BETA_INIT = 0.1\n", + "BETA_FINAL = 1.0\n", + "USE_BETA_SCHEDULE = True\n", + "BETA_STEPS = min(20_000, NUM_STEPS)\n", + "USE_LR_SCHEDULE = False\n", + "CONSTANT_LR = 1e-4\n", + "INIT_LR = 1e-4\n", + "PEAK_LR = 0.1\n", + "MIN_LR = 1e-5\n", + "WARMUP_STEPS = 10_000\n", + "WEIGHT_DECAY = 1e-6\n", + "GRAD_CLIP_NORM = 1.0\n", + "PROBE_FEATURE_SOURCE = \"latent\"\n", + "PROJECTION_FEATURES = (2048,)\n", + "OPTIMIZER_NAME = \"adamw\"\n", + "PROBE_TYPE = \"linear\"\n", + "LOG_EVERY = 2000\n", + "VAL_STEPS = 10\n", + "CHECKPOINT_EVERY = 20_000\n", + "AUXILIARY_METHOD = \"constrained\"\n", + "SHARE_RECOGNITION = True\n", + "PRECISION_TYPE = \"diag\"\n", + "FIX_PRIOR = True\n", + "N_SAMPLES = 100\n", + "SAVE_BEST = True\n", + "MAX_CHECKPOINTS_TO_KEEP = 10\n", + "\n", + "set_global_seed(SEED)\n", + "print(sys.executable if hasattr(sys, \"executable\") else \"n/a\")\n", + "print(jax.devices())\n" + ], + "metadata": { + "id": "cell-023" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-024", + "cell_type": "code", + "source": [ + "# --- Data ---\n", + "rpm_train_cifar10 = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"train\", batch_size=BATCH_SIZE, seed=SEED,\n", + " cache=CACHE, preprocess_type=PREPROCESS_TYPE, num_views=NUM_VIEWS,\n", + " num_workers=NUM_WORKERS, is_training=True,\n", + ")\n", + "rpm_eval_train_cifar10 = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"train\", batch_size=EVAL_BATCH_SIZE, seed=SEED,\n", + " cache=CACHE, preprocess_type=PREPROCESS_TYPE, num_views=NUM_VIEWS,\n", + " num_workers=NUM_WORKERS, is_training=False,\n", + ")\n", + "rpm_eval_test_cifar10 = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"test\", batch_size=EVAL_BATCH_SIZE, seed=SEED,\n", + " cache=CACHE, preprocess_type=PREPROCESS_TYPE, num_views=NUM_VIEWS,\n", + " num_workers=NUM_WORKERS, is_training=False,\n", + ")\n", + "\n", + "_train_batch = next(iter(rpm_train_cifar10))\n", + "_train_raw, _train_views, _train_labels = _train_batch\n", + "print(\"Train raw images:\", _train_raw.shape, _train_raw.dtype)\n", + "print(\"Train views:\", _train_views.shape, _train_views.dtype)\n", + "print(\"Train labels:\", _train_labels.shape, _train_labels.dtype)\n", + "print(\"Train dataset size:\", rpm_train_cifar10.dataset_size)\n", + "print(\"Num classes:\", rpm_train_cifar10.num_classes)\n" + ], + "metadata": { + "id": "cell-024" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-025", + "cell_type": "code", + "source": [ + "# --- Model + schedules ---\n", + "if USE_LR_SCHEDULE:\n", + " rpm_lr_schedule = optax.warmup_cosine_decay_schedule(\n", + " init_value=INIT_LR, warmup_steps=WARMUP_STEPS,\n", + " decay_steps=NUM_STEPS, peak_value=PEAK_LR, end_value=MIN_LR,\n", + " )\n", + "else:\n", + " rpm_lr_schedule = optax.constant_schedule(value=CONSTANT_LR)\n", + "\n", + "if USE_BETA_SCHEDULE:\n", + " rpm_beta_schedule = optax.linear_schedule(\n", + " init_value=BETA_INIT, end_value=BETA_FINAL,\n", + " transition_steps=BETA_STEPS, transition_begin=0,\n", + " )\n", + "else:\n", + " rpm_beta_schedule = optax.constant_schedule(value=BETA_INIT)\n", + "\n", + "gaussian_model = GaussianRPM.create(\n", + " auxiliary_method=AUXILIARY_METHOD,\n", + " n_factors=NUM_VIEWS,\n", + " dim_latent=LATENT_DIM,\n", + " encoder_arch=\"resnet\",\n", + " encoder_params=get_resnet18_cfg(PROJECTION_FEATURES),\n", + " share_recognition=SHARE_RECOGNITION,\n", + " precision_type=PRECISION_TYPE,\n", + " fix_prior=FIX_PRIOR,\n", + " n_samples=N_SAMPLES,\n", + ")\n", + "print(\"Gaussian RPM model constructed.\")\n" + ], + "metadata": { + "id": "cell-025" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-026", + "cell_type": "code", + "source": [ + "# --- Train ---\n", + "rpm_output_dir = Path(\"/tmp/outputs/cifar10_gaussian_rpm\") / time.strftime(\"%Y%m%d-%H%M%S\")\n", + "\n", + "rpm_final_state, rpm_history = train_rpm(\n", + " model=gaussian_model,\n", + " train_dataloader=rpm_train_cifar10,\n", + " eval_train_dataloader=rpm_eval_train_cifar10,\n", + " eval_test_dataloader=rpm_eval_test_cifar10,\n", + " optimizer_name=OPTIMIZER_NAME,\n", + " probe_type=PROBE_TYPE,\n", + " output_dir=rpm_output_dir,\n", + " learning_rate=rpm_lr_schedule,\n", + " weight_decay=WEIGHT_DECAY,\n", + " grad_clip_norm=GRAD_CLIP_NORM,\n", + " batch_size=BATCH_SIZE,\n", + " num_steps=NUM_STEPS,\n", + " log_every=LOG_EVERY,\n", + " val_every=VAL_EVERY,\n", + " val_steps=VAL_STEPS,\n", + " probe_every=PROBE_EVERY,\n", + " probe_epochs=PROBE_EPOCHS,\n", + " probe_lr=PROBE_LR,\n", + " seed=SEED,\n", + " beta=rpm_beta_schedule,\n", + " probe_feature_source=PROBE_FEATURE_SOURCE,\n", + " checkpoint_every=CHECKPOINT_EVERY,\n", + " save_best=SAVE_BEST,\n", + " max_checkpoints_to_keep=MAX_CHECKPOINTS_TO_KEEP,\n", + ")\n" + ], + "metadata": { + "id": "cell-026" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "markdown", + "source": [], + "metadata": { + "id": "xTvc7RMw53GC" + }, + "id": "xTvc7RMw53GC" + }, + { + "cell_type": "code", + "source": [ + "rpm_history['probe_val_acc']\n" + ], + "metadata": { + "id": "hTXDwFx253Rz" + }, + "id": "hTXDwFx253Rz", + "execution_count": null, + "outputs": [] + }, + { + "id": "B5uKyeM4Yr4m", + "cell_type": "code", + "source": [ + "print('da')" + ], + "metadata": { + "id": "B5uKyeM4Yr4m" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "NvUI1G6j2fXU", + "cell_type": "markdown", + "source": [], + "metadata": { + "id": "NvUI1G6j2fXU" + } + }, + { + "id": "LSUMuIIX2fsy", + "cell_type": "code", + "source": [ + "print('da')" + ], + "metadata": { + "id": "LSUMuIIX2fsy" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-027", + "cell_type": "code", + "source": [ + "_, _ = plot_training_metrics(\n", + " rpm_history, val_every=VAL_EVERY, probe_every=PROBE_EVERY,\n", + " title=\"Gaussian RPM (CIFAR-10, diag) training metrics\")\n", + "plt.savefig(rpm_output_dir / \"training_metrics.png\")\n", + "plt.show()\n" + ], + "metadata": { + "id": "cell-027" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-028", + "cell_type": "markdown", + "source": [ + "---\n", + "# Experiment 2 — SimCLR on CIFAR-10\n", + "\n", + "Same ResNet-18 backbone and the same `train_rpm` training loop as\n", + "Experiment 1 — only the model (`SimCLREncoder` instead of `GaussianRPM`) and\n", + "its loss (NT-Xent instead of free energy, dispatched automatically via\n", + "`model.rpm_type`) differ. Standard SimCLR uses `num_views=2`.\n" + ], + "metadata": { + "id": "cell-028" + } + }, + { + "id": "cell-029", + "cell_type": "code", + "source": [ + "# --- Config (mirrors examples/cifar10_simclr.py's flag defaults) ---\n", + "SIMCLR_SEED = 42\n", + "SIMCLR_BATCH_SIZE = 128\n", + "SIMCLR_EVAL_BATCH_SIZE = 256\n", + "SIMCLR_NUM_VIEWS = 8\n", + "SIMCLR_NUM_WORKERS = 0\n", + "SIMCLR_CACHE = False\n", + "SIMCLR_PREPROCESS_TYPE = \"simclr\" # \"simclr\" or \"lejepa_masking\"\n", + "# SIMCLR_PREPROCESS_TYPE = \"lejepa_masking\" # \"simclr\" or \"lejepa_masking\"\n", + "\n", + "SIMCLR_NUM_STEPS = 100_000 # quick smoke run; original default is 100_000\n", + "SIMCLR_PROBE_EVERY = 10_000\n", + "SIMCLR_VAL_EVERY = 5000\n", + "SIMCLR_PROBE_EPOCHS = 10\n", + "SIMCLR_PROBE_LR = 0.1\n", + "SIMCLR_USE_LR_SCHEDULE = False\n", + "SIMCLR_CONSTANT_LR = 1e-4\n", + "SIMCLR_INIT_LR = 1e-4\n", + "SIMCLR_PEAK_LR = 0.1\n", + "SIMCLR_MIN_LR = 1e-5\n", + "SIMCLR_WARMUP_STEPS = 10_000\n", + "SIMCLR_WEIGHT_DECAY = 1e-6\n", + "SIMCLR_GRAD_CLIP_NORM = 1.0\n", + "SIMCLR_PROBE_FEATURE_SOURCE = \"trunk\"\n", + "SIMCLR_PROJECTION_FEATURES = (2048,)\n", + "SIMCLR_OPTIMIZER_NAME = \"adamw\"\n", + "SIMCLR_PROBE_TYPE = \"linear\"\n", + "SIMCLR_LOG_EVERY = 2000\n", + "SIMCLR_VAL_STEPS = 10\n", + "SIMCLR_CHECKPOINT_EVERY = 20_000\n", + "SIMCLR_SAVE_BEST = True\n", + "SIMCLR_MAX_CHECKPOINTS_TO_KEEP = 10\n", + "SIMCLR_TEMPERATURE = 0.5\n", + "\n", + "set_global_seed(SIMCLR_SEED)\n", + "print(jax.devices())\n" + ], + "metadata": { + "id": "cell-029" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-030", + "cell_type": "code", + "source": [ + "# --- Data ---\n", + "simclr_train_cifar10 = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"train\", batch_size=SIMCLR_BATCH_SIZE,\n", + " seed=SIMCLR_SEED, cache=SIMCLR_CACHE, preprocess_type=SIMCLR_PREPROCESS_TYPE,\n", + " num_views=SIMCLR_NUM_VIEWS, num_workers=SIMCLR_NUM_WORKERS, is_training=True,\n", + ")\n", + "simclr_eval_train_cifar10 = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"train\", batch_size=SIMCLR_EVAL_BATCH_SIZE,\n", + " seed=SIMCLR_SEED, cache=SIMCLR_CACHE, preprocess_type=SIMCLR_PREPROCESS_TYPE,\n", + " num_views=SIMCLR_NUM_VIEWS, num_workers=SIMCLR_NUM_WORKERS, is_training=False,\n", + ")\n", + "simclr_eval_test_cifar10 = load_dataset(\n", + " dataset_name=\"cifar10\", split=\"test\", batch_size=SIMCLR_EVAL_BATCH_SIZE,\n", + " seed=SIMCLR_SEED, cache=SIMCLR_CACHE, preprocess_type=SIMCLR_PREPROCESS_TYPE,\n", + " num_views=SIMCLR_NUM_VIEWS, num_workers=SIMCLR_NUM_WORKERS, is_training=False,\n", + ")\n", + "\n", + "_train_batch = next(iter(simclr_train_cifar10))\n", + "_train_raw, _train_views, _train_labels = _train_batch\n", + "print(\"Train raw images:\", _train_raw.shape, _train_raw.dtype)\n", + "print(\"Train views:\", _train_views.shape, _train_views.dtype)\n", + "print(\"Train labels:\", _train_labels.shape, _train_labels.dtype)\n", + "print(\"Train dataset size:\", simclr_train_cifar10.dataset_size)\n", + "print(\"Num classes:\", simclr_train_cifar10.num_classes)\n" + ], + "metadata": { + "id": "cell-030" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-031", + "cell_type": "code", + "source": [ + "# --- Model + schedule ---\n", + "if SIMCLR_USE_LR_SCHEDULE:\n", + " simclr_lr_schedule = optax.warmup_cosine_decay_schedule(\n", + " init_value=SIMCLR_INIT_LR, warmup_steps=SIMCLR_WARMUP_STEPS,\n", + " decay_steps=SIMCLR_NUM_STEPS, peak_value=SIMCLR_PEAK_LR,\n", + " end_value=SIMCLR_MIN_LR,\n", + " )\n", + "else:\n", + " simclr_lr_schedule = optax.constant_schedule(value=SIMCLR_CONSTANT_LR)\n", + "\n", + "simclr_model = SimCLREncoder(\n", + " n_views=SIMCLR_NUM_VIEWS,\n", + " encoder_arch=\"resnet\",\n", + " encoder_params=get_resnet18_cfg(SIMCLR_PROJECTION_FEATURES),\n", + " temperature=SIMCLR_TEMPERATURE,\n", + ")\n", + "print(\"SimCLR model constructed.\")\n" + ], + "metadata": { + "id": "cell-031" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-032", + "cell_type": "code", + "source": [ + "# --- Train ---\n", + "simclr_output_dir = Path(\"/tmp/outputs/cifar10_simclr\") / time.strftime(\"%Y%m%d-%H%M%S\")\n", + "\n", + "simclr_final_state, simclr_history = train_rpm(\n", + " model=simclr_model,\n", + " train_dataloader=simclr_train_cifar10,\n", + " eval_train_dataloader=simclr_eval_train_cifar10,\n", + " eval_test_dataloader=simclr_eval_test_cifar10,\n", + " optimizer_name=SIMCLR_OPTIMIZER_NAME,\n", + " probe_type=SIMCLR_PROBE_TYPE,\n", + " output_dir=simclr_output_dir,\n", + " learning_rate=simclr_lr_schedule,\n", + " weight_decay=SIMCLR_WEIGHT_DECAY,\n", + " grad_clip_norm=SIMCLR_GRAD_CLIP_NORM,\n", + " batch_size=SIMCLR_BATCH_SIZE,\n", + " num_steps=SIMCLR_NUM_STEPS,\n", + " log_every=SIMCLR_LOG_EVERY,\n", + " val_every=SIMCLR_VAL_EVERY,\n", + " val_steps=SIMCLR_VAL_STEPS,\n", + " probe_every=SIMCLR_PROBE_EVERY,\n", + " probe_epochs=SIMCLR_PROBE_EPOCHS,\n", + " probe_lr=SIMCLR_PROBE_LR,\n", + " seed=SIMCLR_SEED,\n", + " probe_feature_source=SIMCLR_PROBE_FEATURE_SOURCE,\n", + " checkpoint_every=SIMCLR_CHECKPOINT_EVERY,\n", + " save_best=SIMCLR_SAVE_BEST,\n", + " max_checkpoints_to_keep=SIMCLR_MAX_CHECKPOINTS_TO_KEEP,\n", + ")\n" + ], + "metadata": { + "id": "cell-032" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "cell-033", + "cell_type": "code", + "source": [ + "_, _ = plot_training_metrics(\n", + " simclr_history, val_every=SIMCLR_VAL_EVERY, probe_every=SIMCLR_PROBE_EVERY,\n", + " title=\"SimCLR (CIFAR-10, ResNet-18) training metrics\")\n", + "plt.savefig(simclr_output_dir / \"training_metrics.png\")\n", + "plt.show()\n" + ], + "metadata": { + "id": "cell-033" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "ADu8q2dTYogH", + "cell_type": "code", + "source": [ + "print('da')" + ], + "metadata": { + "id": "ADu8q2dTYogH" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-header", + "cell_type": "markdown", + "source": [ + "---\n", + "# Latent Space Analysis\n", + "\n", + "Comprehensive comparison of the learned RPM and SimCLR representations.\n", + "\n", + "1. Feature extraction (trunk features + labels for both models)\n", + "2. Dimensionality reduction (PCA 2D + explained-variance curve, t-SNE)\n", + "3. k-NN accuracy (nearest-neighbor classification without a learned head)\n", + "4. Alignment & Uniformity (Wang & Isola, 2020)\n", + "5. Inter-class / intra-class structure (centroid distance heatmap + per-class variance)\n", + "6. RPM-specific uncertainty analysis (per-class posterior uncertainty, calibration)\n", + "7. CKA (Centered Kernel Alignment) cross-method comparison\n" + ], + "metadata": { + "id": "analysis-header" + } + }, + { + "id": "analysis-extract", + "cell_type": "code", + "source": [ + "from sklearn.decomposition import PCA\n", + "from sklearn.manifold import TSNE\n", + "from sklearn.neighbors import KNeighborsClassifier\n", + "from sklearn.metrics import adjusted_rand_score, normalized_mutual_info_score\n", + "from sklearn.cluster import KMeans\n", + "from scipy.spatial.distance import pdist, squareform\n", + "\n", + "# =====================================================================\n", + "# 1. Extract features for both models\n", + "# =====================================================================\n", + "ANALYSIS_FEATURE_SOURCE = \"latent\" # \"trunk\", \"latent\", or \"concatenated\"\n", + "\n", + "# --- RPM features ---\n", + "rpm_feat_fn = create_feature_extractor(gaussian_model,\n", + " feature_source=ANALYSIS_FEATURE_SOURCE)\n", + "rpm_train_feats, rpm_train_labels = compute_features_and_labels(\n", + " rpm_feat_fn, rpm_final_state.params, rpm_eval_train_cifar10)\n", + "rpm_test_feats, rpm_test_labels = compute_features_and_labels(\n", + " rpm_feat_fn, rpm_final_state.params, rpm_eval_test_cifar10)\n", + "\n", + "# --- SimCLR features ---\n", + "simclr_feat_fn = create_feature_extractor(simclr_model,\n", + " feature_source=ANALYSIS_FEATURE_SOURCE)\n", + "simclr_train_feats, simclr_train_labels = compute_features_and_labels(\n", + " simclr_feat_fn, simclr_final_state.params, simclr_eval_train_cifar10)\n", + "simclr_test_feats, simclr_test_labels = compute_features_and_labels(\n", + " simclr_feat_fn, simclr_final_state.params, simclr_eval_test_cifar10)\n", + "\n", + "# Convert to numpy for sklearn\n", + "rpm_train_feats_np = np.array(rpm_train_feats)\n", + "rpm_test_feats_np = np.array(rpm_test_feats)\n", + "rpm_train_labels_np = np.array(rpm_train_labels)\n", + "rpm_test_labels_np = np.array(rpm_test_labels)\n", + "\n", + "simclr_train_feats_np = np.array(simclr_train_feats)\n", + "simclr_test_feats_np = np.array(simclr_test_feats)\n", + "simclr_train_labels_np = np.array(simclr_train_labels)\n", + "simclr_test_labels_np = np.array(simclr_test_labels)\n", + "\n", + "print(f\"RPM train features: {rpm_train_feats_np.shape}\")\n", + "print(f\"RPM test features: {rpm_test_feats_np.shape}\")\n", + "print(f\"SimCLR train features: {simclr_train_feats_np.shape}\")\n", + "print(f\"SimCLR test features: {simclr_test_feats_np.shape}\")\n" + ], + "metadata": { + "id": "analysis-extract" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-dimred", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# 2. Dimensionality Reduction: PCA (2D + explained variance) and t-SNE\n", + "# =====================================================================\n", + "N_VIS = 5000 # subsample for t-SNE speed\n", + "\n", + "CIFAR10_CLASSES = [\n", + " \"airplane\", \"automobile\", \"bird\", \"cat\", \"deer\",\n", + " \"dog\", \"frog\", \"horse\", \"ship\", \"truck\",\n", + "]\n", + "\n", + "def plot_2d_scatter(feats_2d, labels, title, class_names, ax):\n", + " for c in range(len(class_names)):\n", + " mask = labels == c\n", + " ax.scatter(feats_2d[mask, 0], feats_2d[mask, 1],\n", + " s=3, alpha=0.4, label=class_names[c])\n", + " ax.set_title(title, fontsize=11)\n", + " ax.set_xticks([]); ax.set_yticks([])\n", + "\n", + "fig, axes = plt.subplots(2, 2, figsize=(14, 12))\n", + "\n", + "# --- PCA 2D ---\n", + "for col, (name, feats, labels) in enumerate([\n", + " (\"RPM\", rpm_test_feats_np, rpm_test_labels_np),\n", + " (\"SimCLR\", simclr_test_feats_np, simclr_test_labels_np),\n", + "]):\n", + " pca2 = PCA(n_components=2)\n", + " z2 = pca2.fit_transform(feats)\n", + " plot_2d_scatter(z2, labels, f\"PCA 2D - {name}\", CIFAR10_CLASSES, axes[0, col])\n", + "\n", + "# --- t-SNE ---\n", + "for col, (name, feats, labels) in enumerate([\n", + " (\"RPM\", rpm_test_feats_np, rpm_test_labels_np),\n", + " (\"SimCLR\", simclr_test_feats_np, simclr_test_labels_np),\n", + "]):\n", + " idx = np.random.choice(len(feats), min(N_VIS, len(feats)), replace=False)\n", + " tsne = TSNE(n_components=2, perplexity=30, random_state=42, init=\"pca\",\n", + " learning_rate=\"auto\")\n", + " z2 = tsne.fit_transform(feats[idx])\n", + " plot_2d_scatter(z2, labels[idx], f\"t-SNE - {name}\", CIFAR10_CLASSES,\n", + " axes[1, col])\n", + "\n", + "handles, leg_labels = axes[0, 0].get_legend_handles_labels()\n", + "fig.legend(handles, leg_labels, loc=\"lower center\", ncol=5, fontsize=9,\n", + " markerscale=3)\n", + "plt.tight_layout(rect=[0, 0.05, 1, 1])\n", + "plt.suptitle(\"Dimensionality Reduction\", fontsize=14, fontweight=\"bold\", y=1.01)\n", + "plt.show()\n", + "\n", + "# --- PCA explained variance curve ---\n", + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "for name, feats in [(\"RPM\", rpm_test_feats_np),\n", + " (\"SimCLR\", simclr_test_feats_np)]:\n", + " pca_full = PCA().fit(feats)\n", + " cumvar = np.cumsum(pca_full.explained_variance_ratio_)\n", + " ax.plot(cumvar, label=name)\n", + " n90 = np.searchsorted(cumvar, 0.90) + 1\n", + " n95 = np.searchsorted(cumvar, 0.95) + 1\n", + " print(f\"{name}: dims for 90% var = {n90}, 95% var = {n95}, \"\n", + " f\"total dims = {feats.shape[1]}\")\n", + "ax.axhline(0.90, ls=\"--\", color=\"gray\", alpha=0.5, label=\"90%\")\n", + "ax.axhline(0.95, ls=\"--\", color=\"gray\", alpha=0.3, label=\"95%\")\n", + "ax.set_xlabel(\"Number of PCA components\")\n", + "ax.set_ylabel(\"Cumulative explained variance\")\n", + "ax.set_title(\"Effective Dimensionality\")\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "plt.show()\n" + ], + "metadata": { + "id": "analysis-dimred" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-knn", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# 3. k-NN Accuracy\n", + "# =====================================================================\n", + "print(\"k-NN accuracy (test set):\")\n", + "print(f\"{'k':>4s} {'RPM':>8s} {'SimCLR':>8s}\")\n", + "print(\"-\" * 26)\n", + "for k in [1, 5, 10, 20]:\n", + " knn_rpm = KNeighborsClassifier(n_neighbors=k, metric=\"cosine\")\n", + " knn_rpm.fit(rpm_train_feats_np, rpm_train_labels_np)\n", + " rpm_knn_acc = knn_rpm.score(rpm_test_feats_np, rpm_test_labels_np)\n", + "\n", + " knn_simclr = KNeighborsClassifier(n_neighbors=k, metric=\"cosine\")\n", + " knn_simclr.fit(simclr_train_feats_np, simclr_train_labels_np)\n", + " simclr_knn_acc = knn_simclr.score(simclr_test_feats_np, simclr_test_labels_np)\n", + "\n", + " print(f\"{k:4d} {rpm_knn_acc*100:7.2f}% {simclr_knn_acc*100:7.2f}%\")\n" + ], + "metadata": { + "id": "analysis-knn" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-align-uniform", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# 4. Alignment & Uniformity (Wang & Isola, 2020)\n", + "# =====================================================================\n", + "\n", + "def compute_alignment(feats, labels, alpha=2):\n", + " \"\"\"Mean pairwise distance between same-class representations (L2, l2-normalized).\"\"\"\n", + " feats_norm = feats / (np.linalg.norm(feats, axis=1, keepdims=True) + 1e-8)\n", + " total = 0.0\n", + " count = 0\n", + " classes = np.unique(labels)\n", + " for c in classes:\n", + " z_c = feats_norm[labels == c]\n", + " n = len(z_c)\n", + " if n < 2:\n", + " continue\n", + " # Pairwise distances within class\n", + " dists = pdist(z_c, metric=\"sqeuclidean\")\n", + " total += np.sum(dists ** (alpha / 2))\n", + " count += len(dists)\n", + " return total / max(count, 1)\n", + "\n", + "\n", + "def compute_uniformity(feats, t=2, max_samples=5000):\n", + " \"\"\"Log of expected pairwise Gaussian kernel (l2-normalized features).\"\"\"\n", + " feats_norm = feats / (np.linalg.norm(feats, axis=1, keepdims=True) + 1e-8)\n", + " if len(feats_norm) > max_samples:\n", + " idx = np.random.choice(len(feats_norm), max_samples, replace=False)\n", + " feats_norm = feats_norm[idx]\n", + " sq_dists = pdist(feats_norm, metric=\"sqeuclidean\")\n", + " return np.log(np.mean(np.exp(-t * sq_dists)))\n", + "\n", + "\n", + "print(\"Alignment & Uniformity (lower is better for both):\")\n", + "print(f\"{'Metric':<15s} {'RPM':>10s} {'SimCLR':>10s}\")\n", + "print(\"-\" * 40)\n", + "for name, feats, labels in [\n", + " (\"RPM\", rpm_test_feats_np, rpm_test_labels_np),\n", + " (\"SimCLR\", simclr_test_feats_np, simclr_test_labels_np),\n", + "]:\n", + " align = compute_alignment(feats, labels)\n", + " uniform = compute_uniformity(feats)\n", + " print(f\" {name:<13s} align={align:.4f} uniform={uniform:.4f}\")\n" + ], + "metadata": { + "id": "analysis-align-uniform" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-interclass", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# 5. Inter-class / Intra-class Structure\n", + "# =====================================================================\n", + "\n", + "def class_centroids_and_variance(feats, labels, class_names):\n", + " n_classes = len(class_names)\n", + " centroids = np.zeros((n_classes, feats.shape[1]))\n", + " intra_var = np.zeros(n_classes)\n", + " for c in range(n_classes):\n", + " z_c = feats[labels == c]\n", + " centroids[c] = z_c.mean(axis=0)\n", + " intra_var[c] = np.mean(np.sum((z_c - centroids[c]) ** 2, axis=1))\n", + " return centroids, intra_var\n", + "\n", + "fig, axes = plt.subplots(1, 3, figsize=(18, 5))\n", + "\n", + "for idx, (name, feats, labels) in enumerate([\n", + " (\"RPM\", rpm_test_feats_np, rpm_test_labels_np),\n", + " (\"SimCLR\", simclr_test_feats_np, simclr_test_labels_np),\n", + "]):\n", + " centroids, intra_var = class_centroids_and_variance(\n", + " feats, labels, CIFAR10_CLASSES)\n", + " # Cosine distance between centroids\n", + " centroid_norm = centroids / (np.linalg.norm(centroids, axis=1, keepdims=True) + 1e-8)\n", + " cos_sim = centroid_norm @ centroid_norm.T\n", + " cos_dist = 1 - cos_sim\n", + "\n", + " im = axes[idx].imshow(cos_dist, cmap=\"viridis\", vmin=0)\n", + " axes[idx].set_xticks(range(10))\n", + " axes[idx].set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\", fontsize=8)\n", + " axes[idx].set_yticks(range(10))\n", + " axes[idx].set_yticklabels(CIFAR10_CLASSES, fontsize=8)\n", + " axes[idx].set_title(f\"Centroid cosine distance - {name}\")\n", + " plt.colorbar(im, ax=axes[idx], shrink=0.8)\n", + "\n", + "# Intra-class variance comparison\n", + "rpm_centroids, rpm_intra = class_centroids_and_variance(\n", + " rpm_test_feats_np, rpm_test_labels_np, CIFAR10_CLASSES)\n", + "simclr_centroids, simclr_intra = class_centroids_and_variance(\n", + " simclr_test_feats_np, simclr_test_labels_np, CIFAR10_CLASSES)\n", + "\n", + "x = np.arange(10)\n", + "w = 0.35\n", + "axes[2].bar(x - w/2, rpm_intra, w, label=\"RPM\", alpha=0.8)\n", + "axes[2].bar(x + w/2, simclr_intra, w, label=\"SimCLR\", alpha=0.8)\n", + "axes[2].set_xticks(x)\n", + "axes[2].set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\", fontsize=8)\n", + "axes[2].set_ylabel(\"Mean intra-class L2 variance\")\n", + "axes[2].set_title(\"Intra-class spread\")\n", + "axes[2].legend()\n", + "\n", + "plt.suptitle(\"Inter-class / Intra-class Structure\", fontsize=14,\n", + " fontweight=\"bold\", y=1.02)\n", + "plt.tight_layout()\n", + "plt.show()\n" + ], + "metadata": { + "id": "analysis-interclass" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-uncertainty", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# 6. RPM-Specific: Uncertainty Analysis\n", + "# =====================================================================\n", + "# RPM outputs a Gaussian posterior; extract precision (inverse covariance).\n", + "\n", + "@jax.jit\n", + "def rpm_posterior_params(params, views):\n", + " rpm_output, _ = gaussian_model.apply({\"params\": params}, views)\n", + " variational = rpm_output[\"variational\"]\n", + " variational_mean = rpm_output[\"variational_mean\"]\n", + " # For diag precision: uncertainty ~ 1/precision_diag\n", + " return variational_mean.mean, variational.precision_diag\n", + "\n", + "# Collect posterior means and precisions\n", + "rpm_means_list, rpm_precs_list, rpm_ulabels_list = [], [], []\n", + "for _, views, labels in rpm_eval_test_cifar10:\n", + " mean, prec = rpm_posterior_params(rpm_final_state.params, views)\n", + " rpm_means_list.append(np.array(mean))\n", + " rpm_precs_list.append(np.array(prec))\n", + " rpm_ulabels_list.append(np.array(labels))\n", + "\n", + "rpm_means_all = np.concatenate(rpm_means_list, axis=0)\n", + "rpm_precs_all = np.concatenate(rpm_precs_list, axis=0)\n", + "rpm_ulabels_all = np.concatenate(rpm_ulabels_list, axis=0)\n", + "\n", + "# Uncertainty = trace(covariance) = sum(1/precision_diag)\n", + "rpm_uncertainty = np.sum(1.0 / (rpm_precs_all + 1e-8), axis=-1)\n", + "\n", + "# --- Per-class uncertainty ---\n", + "fig, axes = plt.subplots(1, 2, figsize=(14, 5))\n", + "\n", + "class_uncertainties = []\n", + "for c in range(10):\n", + " mask = rpm_ulabels_all == c\n", + " class_uncertainties.append(rpm_uncertainty[mask])\n", + "\n", + "axes[0].boxplot(class_uncertainties, labels=CIFAR10_CLASSES, vert=True)\n", + "axes[0].set_xticklabels(CIFAR10_CLASSES, rotation=45, ha=\"right\", fontsize=9)\n", + "axes[0].set_ylabel(\"Uncertainty (trace of covariance)\")\n", + "axes[0].set_title(\"RPM: Per-class uncertainty distribution\")\n", + "\n", + "# --- Uncertainty calibration: kNN-based correctness vs uncertainty ---\n", + "knn_rpm_cal = KNeighborsClassifier(n_neighbors=5, metric=\"cosine\")\n", + "knn_rpm_cal.fit(rpm_train_feats_np, rpm_train_labels_np)\n", + "rpm_knn_preds = knn_rpm_cal.predict(rpm_test_feats_np)\n", + "rpm_is_correct = (rpm_knn_preds == rpm_test_labels_np)\n", + "\n", + "n_bins = 10\n", + "sorted_idx = np.argsort(rpm_uncertainty)\n", + "bin_size = len(sorted_idx) // n_bins\n", + "bin_accs, bin_uncs = [], []\n", + "for b in range(n_bins):\n", + " start = b * bin_size\n", + " end = start + bin_size if b < n_bins - 1 else len(sorted_idx)\n", + " idx = sorted_idx[start:end]\n", + " bin_accs.append(np.mean(rpm_is_correct[idx]))\n", + " bin_uncs.append(np.mean(rpm_uncertainty[idx]))\n", + "\n", + "axes[1].bar(range(n_bins), bin_accs, alpha=0.7, label=\"Accuracy\")\n", + "ax2 = axes[1].twinx()\n", + "ax2.plot(range(n_bins), bin_uncs, \"r-o\", label=\"Mean uncertainty\")\n", + "axes[1].set_xlabel(\"Uncertainty bin (low -> high)\")\n", + "axes[1].set_ylabel(\"k-NN Accuracy\")\n", + "ax2.set_ylabel(\"Mean uncertainty\", color=\"r\")\n", + "axes[1].set_title(\"RPM: Uncertainty calibration\")\n", + "axes[1].legend(loc=\"upper left\")\n", + "ax2.legend(loc=\"upper right\")\n", + "\n", + "plt.suptitle(\"RPM Uncertainty Analysis\", fontsize=14, fontweight=\"bold\", y=1.02)\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "correct_unc = np.mean(rpm_uncertainty[rpm_is_correct])\n", + "incorrect_unc = np.mean(rpm_uncertainty[~rpm_is_correct])\n", + "print(f\"Mean uncertainty (correct): {correct_unc:.4f}\")\n", + "print(f\"Mean uncertainty (incorrect): {incorrect_unc:.4f}\")\n", + "print(f\"Ratio (incorrect/correct): {incorrect_unc/correct_unc:.2f}x\")\n" + ], + "metadata": { + "id": "analysis-uncertainty" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "analysis-cka", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# 7. CKA (Centered Kernel Alignment)\n", + "# =====================================================================\n", + "\n", + "def linear_CKA(X, Y):\n", + " \"\"\"Linear CKA between two feature matrices (n_samples, n_features).\"\"\"\n", + " X = X - X.mean(axis=0)\n", + " Y = Y - Y.mean(axis=0)\n", + " hsic_xy = np.linalg.norm(X.T @ Y, ord=\"fro\") ** 2\n", + " hsic_xx = np.linalg.norm(X.T @ X, ord=\"fro\") ** 2\n", + " hsic_yy = np.linalg.norm(Y.T @ Y, ord=\"fro\") ** 2\n", + " return hsic_xy / (np.sqrt(hsic_xx * hsic_yy) + 1e-10)\n", + "\n", + "# Ensure same number of samples (use test set, same ordering from same dataset)\n", + "n = min(len(rpm_test_feats_np), len(simclr_test_feats_np))\n", + "cka_score = linear_CKA(rpm_test_feats_np[:n], simclr_test_feats_np[:n])\n", + "print(f\"Linear CKA between RPM and SimCLR ({ANALYSIS_FEATURE_SOURCE} features): {cka_score:.4f}\")\n", + "print(f\" (1.0 = identical geometry, 0.0 = completely different)\")\n", + "\n", + "# Also compute CKA within each method between train and test for sanity\n", + "n_tt = min(len(rpm_train_feats_np), len(rpm_test_feats_np))\n", + "rpm_self_cka = linear_CKA(rpm_train_feats_np[:n_tt], rpm_test_feats_np[:n_tt])\n", + "simclr_self_cka = linear_CKA(simclr_train_feats_np[:n_tt], simclr_test_feats_np[:n_tt])\n", + "print(f\"RPM train-test CKA (self-consistency): {rpm_self_cka:.4f}\")\n", + "print(f\"SimCLR train-test CKA (self-consistency): {simclr_self_cka:.4f}\")\n" + ], + "metadata": { + "id": "analysis-cka" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-header", + "cell_type": "markdown", + "source": [ + "---\n", + "# Experiment 3 -- Shapes3D disentanglement\n", + "\n", + "Train RPM and SimCLR on the **Shapes3D** dataset (480k synthetic 64x64\n", + "images with 6 independent ground-truth factors: floor hue, wall hue,\n", + "object hue, scale, shape, orientation).\n", + "\n", + "The goal is **not** classification -- instead we analyse whether the\n", + "learned latent dimensions capture the underlying generative factors.\n" + ], + "metadata": { + "id": "shapes3d-header" + } + }, + { + "id": "shapes3d-loader", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# Shapes3D data loader\n", + "# =====================================================================\n", + "SHAPES3D_FACTOR_NAMES = [\n", + " \"floor_hue\", \"wall_hue\", \"object_hue\", \"scale\", \"shape\", \"orientation\",\n", + "]\n", + "\n", + "SHAPES3D_MEAN = np.asarray([0.5, 0.5, 0.5], dtype=np.float32)\n", + "SHAPES3D_STD = np.asarray([0.5, 0.5, 0.5], dtype=np.float32)\n", + "\n", + "\n", + "def shapes3d_simclr_augment(image, image_size=64, pass_mask=False):\n", + " image = _random_resized_crop(image, size=image_size)\n", + " image = tf.image.random_flip_left_right(image)\n", + " do_jitter = tf.random.uniform([]) < 0.8\n", + " image = tf.cond(do_jitter, lambda: _color_jitter(image, 0.5), lambda: image)\n", + " image = _random_grayscale(image, p=0.2)\n", + " do_blur = tf.random.uniform([]) < 0.5\n", + " image = tf.cond(do_blur, lambda: _gaussian_blur(image, kernel_size=5),\n", + " lambda: image)\n", + " image = (image - SHAPES3D_MEAN) / SHAPES3D_STD\n", + " if pass_mask:\n", + " image = _append_ones_mask(image)\n", + " return image\n", + "\n", + "\n", + "def shapes3d_masking_augment(image, image_size=64, num_blocks=4,\n", + " scale_range=(0.15, 0.2),\n", + " aspect_ratio_range=(0.75, 1.5),\n", + " pass_mask=False):\n", + " if image_size != 64:\n", + " image = tf.image.resize(image, [image_size, image_size])\n", + " image = (image - SHAPES3D_MEAN) / SHAPES3D_STD\n", + " mask = _sample_multi_block_mask(image_size, image_size,\n", + " num_blocks=num_blocks,\n", + " scale_range=scale_range,\n", + " aspect_ratio_range=aspect_ratio_range)\n", + " mask_3d = tf.cast(mask[:, :, tf.newaxis], tf.float32)\n", + " image = image * (1.0 - mask_3d)\n", + " if pass_mask:\n", + " visibility = 1.0 - mask_3d\n", + " image = tf.concat([image, visibility], axis=-1)\n", + " return image\n", + "\n", + "\n", + "def shapes3d_eval_preprocess(image, image_size=64, pass_mask=False):\n", + " if image_size != 64:\n", + " image = tf.image.resize(image, [image_size, image_size])\n", + " image = (image - SHAPES3D_MEAN) / SHAPES3D_STD\n", + " if pass_mask:\n", + " image = _append_ones_mask(image)\n", + " return image\n", + "\n", + "\n", + "class Shapes3DLoader:\n", + " \"\"\"Shapes3D data loader.\n", + "\n", + " Yields (raw_images, views, label) where label is label_shape (used as\n", + " a dummy for the training loop). Ground-truth factor values are\n", + " available separately via load_shapes3d_factors().\n", + "\n", + " When ``pass_mask=True``, each view has shape [H, W, 4] where the 4th\n", + " channel is a visibility mask (1=visible, 0=masked).\n", + " \"\"\"\n", + "\n", + " def __init__(self, *, split_spec, batch_size, training, num_views=4,\n", + " image_size=64, preprocess_type=\"simclr\", seed=None,\n", + " drop_remainder=None, cache=True, pass_mask=False):\n", + " self._batch_size = batch_size\n", + " self._training = training\n", + " self._num_views = num_views\n", + " self._image_size = image_size\n", + " self._preprocess_type = preprocess_type\n", + " self._pass_mask = pass_mask\n", + " self._num_classes = 4 # shape factor has 4 categories\n", + " if drop_remainder is None:\n", + " drop_remainder = training\n", + " self._drop_remainder = drop_remainder\n", + "\n", + " ds = tfds.load(\"shapes3d\", split=split_spec, shuffle_files=training)\n", + " self._dataset_size = ds.cardinality().numpy()\n", + " if self._dataset_size < 0:\n", + " # cardinality unknown; shapes3d has 480000 total\n", + " self._dataset_size = 480000\n", + "\n", + " if cache:\n", + " ds = ds.cache()\n", + " if training:\n", + " ds = ds.shuffle(buffer_size=min(self._dataset_size, 50_000),\n", + " seed=seed, reshuffle_each_iteration=True)\n", + " ds = ds.map(self._preprocess, num_parallel_calls=tf.data.AUTOTUNE)\n", + " ds = ds.batch(batch_size, drop_remainder=drop_remainder)\n", + " ds = ds.prefetch(tf.data.AUTOTUNE)\n", + " self._ds = ds\n", + "\n", + " @property\n", + " def dataset_size(self):\n", + " return self._dataset_size\n", + "\n", + " @property\n", + " def num_classes(self):\n", + " return self._num_classes\n", + "\n", + " def _preprocess(self, example):\n", + " image = example[\"image\"]\n", + " label = example[\"label_shape\"] # 0-3, used as dummy label\n", + " raw_image = image\n", + " image_f = tf.cast(image, tf.float32) / 255.0\n", + "\n", + " if self._training:\n", + " if self._preprocess_type == \"simclr\":\n", + " views = tf.stack([\n", + " shapes3d_simclr_augment(image_f, self._image_size,\n", + " pass_mask=self._pass_mask)\n", + " for _ in range(self._num_views)\n", + " ], axis=0)\n", + " elif self._preprocess_type == \"lejepa_masking\":\n", + " views = tf.stack([\n", + " shapes3d_masking_augment(image_f, self._image_size,\n", + " pass_mask=self._pass_mask)\n", + " for _ in range(self._num_views)\n", + " ], axis=0)\n", + " else:\n", + " raise ValueError(f\"Unknown preprocess_type: {self._preprocess_type}\")\n", + " else:\n", + " v = shapes3d_eval_preprocess(image_f, self._image_size,\n", + " pass_mask=self._pass_mask)\n", + " views = tf.stack([v] * self._num_views, axis=0)\n", + "\n", + " return raw_image, views, label\n", + "\n", + " def __iter__(self):\n", + " for raw, views, labels in self._ds:\n", + " yield (raw.numpy(), views.numpy(), labels.numpy())\n", + "\n", + " def __len__(self):\n", + " if self._drop_remainder:\n", + " return self._dataset_size // self._batch_size\n", + " return math.ceil(self._dataset_size / self._batch_size)\n", + "\n", + "\n", + "def load_shapes3d_factors(split_spec, max_samples=None):\n", + " \"\"\"Load ground-truth factor values for analysis.\"\"\"\n", + " ds = tfds.load(\"shapes3d\", split=split_spec)\n", + " factors = {n: [] for n in SHAPES3D_FACTOR_NAMES}\n", + " for ex in ds:\n", + " for n in SHAPES3D_FACTOR_NAMES:\n", + " factors[n].append(float(ex[f\"value_{n}\"].numpy()))\n", + " if max_samples and len(factors[\"shape\"]) >= max_samples:\n", + " break\n", + " return {n: np.array(v) for n, v in factors.items()}\n", + "\n", + "\n", + "print(\"Shapes3D loader defined.\")\n" + ], + "metadata": { + "id": "shapes3d-loader" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-config", + "cell_type": "code", + "source": [ + "# --- Shapes3D experiment config ---\n", + "S3D_SEED = 42\n", + "S3D_BATCH_SIZE = 128\n", + "S3D_EVAL_BATCH_SIZE = 256\n", + "S3D_NUM_VIEWS = 8\n", + "S3D_IMAGE_SIZE = 64\n", + "S3D_LATENT_DIM = 128\n", + "S3D_PREPROCESS_TYPE = \"simclr\" # \"simclr\" or \"lejepa_masking\"\n", + "# S3D_PREPROCESS_TYPE = \"lejepa_masking\" # \"simclr\" or \"lejepa_masking\"\n", + "S3D_PASS_MASK = True # if True, append visibility mask as 4th channel\n", + "S3D_NUM_STEPS = 100_000\n", + "S3D_PROBE_EVERY = 20_000\n", + "S3D_VAL_EVERY = 10000\n", + "S3D_LOG_EVERY = 10000\n", + "S3D_VAL_STEPS = 10\n", + "S3D_CHECKPOINT_EVERY = 20_000\n", + "S3D_LR = 1e-4\n", + "S3D_WEIGHT_DECAY = 1e-6\n", + "S3D_GRAD_CLIP = 1.0\n", + "S3D_PROBE_EPOCHS = 10\n", + "S3D_PROBE_LR = 0.1\n", + "S3D_PROJECTION_FEATURES = (2048,)\n", + "S3D_PRECISION_TYPE = \"diag\"\n", + "S3D_TEMPERATURE = 0.5\n", + "\n", + "# Use 90% train, 10% eval\n", + "S3D_TRAIN_SPLIT = \"train[:90%]\"\n", + "S3D_EVAL_SPLIT = \"train[90%:]\"\n", + "\n", + "set_global_seed(S3D_SEED)\n", + "print(\"Shapes3D config set.\")\n", + "\n", + "\n", + "\n", + "rpm_beta_schedule = optax.linear_schedule(\n", + " init_value=0.1, end_value=1.0,\n", + " transition_steps=5000, transition_begin=0,\n", + " )\n" + ], + "metadata": { + "id": "shapes3d-config" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-data", + "cell_type": "code", + "source": [ + "# --- Shapes3D data loaders ---\n", + "s3d_train_loader = Shapes3DLoader(\n", + " split_spec=S3D_TRAIN_SPLIT, batch_size=S3D_BATCH_SIZE,\n", + " training=True, num_views=S3D_NUM_VIEWS, image_size=S3D_IMAGE_SIZE,\n", + " preprocess_type=S3D_PREPROCESS_TYPE, seed=S3D_SEED,\n", + " pass_mask=S3D_PASS_MASK,\n", + ")\n", + "s3d_eval_loader = Shapes3DLoader(\n", + " split_spec=S3D_EVAL_SPLIT, batch_size=S3D_EVAL_BATCH_SIZE,\n", + " training=False, num_views=S3D_NUM_VIEWS, image_size=S3D_IMAGE_SIZE,\n", + " preprocess_type=S3D_PREPROCESS_TYPE, seed=S3D_SEED,\n", + " pass_mask=S3D_PASS_MASK,\n", + ")\n", + "\n", + "_b = next(iter(s3d_train_loader))\n", + "print(\"Shapes3D train batch shapes:\", _b[0].shape, _b[1].shape, _b[2].shape)\n", + "print(f\"Train size: {s3d_train_loader.dataset_size}, \"\n", + " f\"Eval size: {s3d_eval_loader.dataset_size}\")\n" + ], + "metadata": { + "id": "shapes3d-data" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-models", + "cell_type": "code", + "source": [ + "# --- ResNet-18 config adjusted for 64x64 images ---\n", + "def get_resnet18_cfg_64(projection_features):\n", + " return {\n", + " \"stage_sizes\": (2, 2, 2, 2),\n", + " \"stage_widths\": (64, 128, 256, 512),\n", + " \"stem_width\": 64,\n", + " \"stem_kernel_size\": (3, 3),\n", + " \"stem_stride\": 1,\n", + " \"use_max_pool\": True,\n", + " \"max_pool_window\": (3, 3),\n", + " \"max_pool_stride\": (2, 2),\n", + " \"block_type\": \"basic\",\n", + " \"projection_features\": projection_features,\n", + " }\n", + "\n", + "# --- RPM model ---\n", + "s3d_rpm_model = GaussianRPM.create(\n", + " auxiliary_method=\"constrained\",\n", + " n_factors=S3D_NUM_VIEWS,\n", + " dim_latent=S3D_LATENT_DIM,\n", + " encoder_arch=\"resnet\",\n", + " encoder_params=get_resnet18_cfg_64(S3D_PROJECTION_FEATURES),\n", + " share_recognition=True,\n", + " precision_type=S3D_PRECISION_TYPE,\n", + " fix_prior=True,\n", + ")\n", + "print(\"Shapes3D RPM model constructed.\")\n", + "\n", + "# --- SimCLR model ---\n", + "s3d_simclr_model = SimCLREncoder(\n", + " n_views=S3D_NUM_VIEWS,\n", + " encoder_arch=\"resnet\",\n", + " encoder_params=get_resnet18_cfg_64(S3D_PROJECTION_FEATURES),\n", + " temperature=S3D_TEMPERATURE,\n", + ")\n", + "print(\"Shapes3D SimCLR model constructed.\")\n" + ], + "metadata": { + "id": "shapes3d-models" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-train-rpm", + "cell_type": "code", + "source": [ + "# --- Train RPM on Shapes3D ---\n", + "s3d_rpm_output_dir = Path(\"/tmp/outputs/shapes3d_rpm\") / time.strftime(\"%Y%m%d-%H%M%S\")\n", + "\n", + "s3d_rpm_state, s3d_rpm_history = train_rpm(\n", + " model=s3d_rpm_model,\n", + " train_dataloader=s3d_train_loader,\n", + " eval_train_dataloader=s3d_eval_loader,\n", + " eval_test_dataloader=s3d_eval_loader,\n", + " optimizer_name=\"adamw\",\n", + " probe_type=\"linear\",\n", + " output_dir=s3d_rpm_output_dir,\n", + " learning_rate=optax.constant_schedule(S3D_LR),\n", + " weight_decay=S3D_WEIGHT_DECAY,\n", + " grad_clip_norm=S3D_GRAD_CLIP,\n", + " batch_size=S3D_BATCH_SIZE,\n", + " num_steps=S3D_NUM_STEPS,\n", + " log_every=S3D_LOG_EVERY,\n", + " val_every=S3D_VAL_EVERY,\n", + " val_steps=S3D_VAL_STEPS,\n", + " probe_every=S3D_PROBE_EVERY,\n", + " probe_epochs=S3D_PROBE_EPOCHS,\n", + " probe_lr=S3D_PROBE_LR,\n", + " seed=S3D_SEED,\n", + " # beta=optax.constant_schedule(1.0),\n", + " beta=rpm_beta_schedule,\n", + " probe_feature_source=\"trunk\",\n", + " checkpoint_every=S3D_CHECKPOINT_EVERY,\n", + " save_best=True,\n", + " max_checkpoints_to_keep=3,\n", + ")\n", + "print(\"Shapes3D RPM training done.\")\n" + ], + "metadata": { + "id": "shapes3d-train-rpm" + }, + "execution_count": null, + "outputs": [] + }, + { + "cell_type": "code", + "source": [ + "print('da')" + ], + "metadata": { + "id": "IlkJsKidPXeI" + }, + "id": "IlkJsKidPXeI", + "execution_count": null, + "outputs": [] + }, + { + "id": "7glHhDLazC6c", + "cell_type": "code", + "source": [ + "s3d_rpm_history.keys()" + ], + "metadata": { + "id": "7glHhDLazC6c" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "H0YUempdy-dl", + "cell_type": "code", + "source": [ + "plt.plot(s3d_rpm_history['val_free_energy'])" + ], + "metadata": { + "id": "H0YUempdy-dl" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-train-simclr", + "cell_type": "code", + "source": [ + "# --- Train SimCLR on Shapes3D ---\n", + "s3d_simclr_output_dir = Path(\"/tmp/outputs/shapes3d_simclr\") / time.strftime(\"%Y%m%d-%H%M%S\")\n", + "\n", + "s3d_simclr_state, s3d_simclr_history = train_rpm(\n", + " model=s3d_simclr_model,\n", + " train_dataloader=s3d_train_loader,\n", + " eval_train_dataloader=s3d_eval_loader,\n", + " eval_test_dataloader=s3d_eval_loader,\n", + " optimizer_name=\"adamw\",\n", + " probe_type=\"linear\",\n", + " output_dir=s3d_simclr_output_dir,\n", + " learning_rate=optax.constant_schedule(S3D_LR),\n", + " weight_decay=S3D_WEIGHT_DECAY,\n", + " grad_clip_norm=S3D_GRAD_CLIP,\n", + " batch_size=S3D_BATCH_SIZE,\n", + " num_steps=S3D_NUM_STEPS,\n", + " log_every=S3D_LOG_EVERY,\n", + " val_every=S3D_VAL_EVERY,\n", + " val_steps=S3D_VAL_STEPS,\n", + " probe_every=S3D_PROBE_EVERY,\n", + " probe_epochs=S3D_PROBE_EPOCHS,\n", + " probe_lr=S3D_PROBE_LR,\n", + " seed=S3D_SEED,\n", + " probe_feature_source=\"trunk\",\n", + " checkpoint_every=S3D_CHECKPOINT_EVERY,\n", + " save_best=True,\n", + " max_checkpoints_to_keep=3,\n", + ")\n", + "print(\"Shapes3D SimCLR training done.\")\n" + ], + "metadata": { + "id": "shapes3d-train-simclr" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "zACJyEMDy65l", + "cell_type": "code", + "source": [ + "plt.plot(s3d_simclr_history['loss'])" + ], + "metadata": { + "id": "zACJyEMDy65l" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-extract", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# Extract features and ground-truth factors for analysis\n", + "# =====================================================================\n", + "\n", + "# --- RPM latent: variational posterior mean (aggregated across views) ---\n", + "@jax.jit\n", + "def _rpm_latent_fn(params, views):\n", + " out, _ = s3d_rpm_model.apply({\"params\": params}, views)\n", + " return out[\"variational_mean\"].mean # (batch, latent_dim)\n", + "\n", + "s3d_rpm_latent_list, s3d_rpm_labels_list = [], []\n", + "for _, views, labels in s3d_eval_loader:\n", + " feats = _rpm_latent_fn(s3d_rpm_state.params, views)\n", + " s3d_rpm_latent_list.append(np.array(feats))\n", + " s3d_rpm_labels_list.append(np.array(labels))\n", + "\n", + "s3d_rpm_latent_np = np.concatenate(s3d_rpm_latent_list, axis=0)\n", + "s3d_rpm_labels_np = np.concatenate(s3d_rpm_labels_list, axis=0)\n", + "\n", + "# --- RPM trunk: backbone features before projection head ---\n", + "s3d_rpm_trunk_fn = create_feature_extractor(s3d_rpm_model, feature_source=\"trunk\")\n", + "s3d_rpm_trunk_feats, _ = compute_features_and_labels(\n", + " s3d_rpm_trunk_fn, s3d_rpm_state.params, s3d_eval_loader)\n", + "s3d_rpm_trunk_np = np.array(s3d_rpm_trunk_feats)\n", + "\n", + "# --- SimCLR trunk: backbone features before projection head ---\n", + "s3d_simclr_trunk_fn = create_feature_extractor(s3d_simclr_model,\n", + " feature_source=\"trunk\")\n", + "s3d_simclr_feats, s3d_simclr_labels = compute_features_and_labels(\n", + " s3d_simclr_trunk_fn, s3d_simclr_state.params, s3d_eval_loader)\n", + "s3d_simclr_trunk_np = np.array(s3d_simclr_feats)\n", + "s3d_simclr_labels_np = np.array(s3d_simclr_labels)\n", + "\n", + "for arr, label in [(s3d_rpm_latent_np, \"RPM latent\"),\n", + " (s3d_rpm_trunk_np, \"RPM trunk\"),\n", + " (s3d_simclr_trunk_np, \"SimCLR trunk\")]:\n", + " assert arr.ndim == 2, f\"Expected 2D for {label}, got shape {arr.shape}\"\n", + "\n", + "# Load ground-truth factor values (same split/ordering as eval loader)\n", + "s3d_factors = load_shapes3d_factors(S3D_EVAL_SPLIT,\n", + " max_samples=len(s3d_rpm_latent_np))\n", + "\n", + "# Convenience dict for iterating over all three representations\n", + "S3D_REPR = {\n", + " \"RPM latent\": s3d_rpm_latent_np,\n", + " \"RPM trunk\": s3d_rpm_trunk_np,\n", + " \"SimCLR trunk\": s3d_simclr_trunk_np,\n", + "}\n", + "\n", + "print(\"Feature shapes:\")\n", + "for name, arr in S3D_REPR.items():\n", + " print(f\" {name}: {arr.shape}\")\n", + "for fn in SHAPES3D_FACTOR_NAMES:\n", + " print(f\" {fn}: {len(s3d_factors[fn])} values, \"\n", + " f\"range [{s3d_factors[fn].min():.3f}, {s3d_factors[fn].max():.3f}]\")\n" + ], + "metadata": { + "id": "shapes3d-extract" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-r2", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# Disentanglement: R^2 of each latent dim vs each factor\n", + "# =====================================================================\n", + "from sklearn.linear_model import LinearRegression\n", + "\n", + "def compute_r2_matrix(feats, factors, factor_names):\n", + " \"\"\"R^2 between each latent dimension and each ground-truth factor.\"\"\"\n", + " n_dims = feats.shape[1]\n", + " n_factors = len(factor_names)\n", + " r2 = np.zeros((n_factors, n_dims))\n", + " for fi, fn in enumerate(factor_names):\n", + " y = factors[fn]\n", + " for di in range(n_dims):\n", + " x = feats[:, di:di+1]\n", + " reg = LinearRegression().fit(x, y)\n", + " r2[fi, di] = max(0, reg.score(x, y))\n", + " return r2\n", + "\n", + "def compute_factor_r2(feats, factors, factor_names):\n", + " \"\"\"Per-factor R^2 using ALL latent dims (full linear model).\"\"\"\n", + " results = {}\n", + " for fn in factor_names:\n", + " y = factors[fn]\n", + " reg = LinearRegression().fit(feats, y)\n", + " results[fn] = max(0, reg.score(feats, y))\n", + " return results\n", + "\n", + "# Compute R^2 for all three representations\n", + "s3d_r2_mats = {}\n", + "s3d_full_r2 = {}\n", + "for name, feats in S3D_REPR.items():\n", + " print(f\"Computing R^2 for {name} ({feats.shape[1]} dims)...\")\n", + " s3d_r2_mats[name] = compute_r2_matrix(feats, s3d_factors, SHAPES3D_FACTOR_NAMES)\n", + " s3d_full_r2[name] = compute_factor_r2(feats, s3d_factors, SHAPES3D_FACTOR_NAMES)\n", + "\n", + "# --- Heatmap: per-dim R^2 (top 20 dims) ---\n", + "n_repr = len(S3D_REPR)\n", + "fig, axes = plt.subplots(1, n_repr, figsize=(7 * n_repr, 4))\n", + "for ax, (name, r2_mat) in zip(axes, s3d_r2_mats.items()):\n", + " max_r2_per_dim = r2_mat.max(axis=0)\n", + " top_dims = np.argsort(max_r2_per_dim)[-20:]\n", + " im = ax.imshow(r2_mat[:, top_dims], aspect=\"auto\", cmap=\"YlOrRd\", vmin=0, vmax=1)\n", + " ax.set_yticks(range(len(SHAPES3D_FACTOR_NAMES)))\n", + " ax.set_yticklabels(SHAPES3D_FACTOR_NAMES, fontsize=9)\n", + " ax.set_xlabel(\"Latent dimension (top 20 by max R²)\")\n", + " ax.set_title(f\"{name}: per-dim R²\")\n", + " plt.colorbar(im, ax=ax, shrink=0.8)\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "# --- Bar chart: full-model R^2 per factor ---\n", + "fig, ax = plt.subplots(figsize=(12, 4))\n", + "x = np.arange(len(SHAPES3D_FACTOR_NAMES))\n", + "w = 0.25\n", + "repr_names = list(S3D_REPR.keys())\n", + "colors = [\"#1f77b4\", \"#2ca02c\", \"#ff7f0e\"]\n", + "for i, name in enumerate(repr_names):\n", + " vals = [s3d_full_r2[name][f] for f in SHAPES3D_FACTOR_NAMES]\n", + " ax.bar(x + (i - 1) * w, vals, w, label=name, alpha=0.8, color=colors[i])\n", + "ax.set_xticks(x)\n", + "ax.set_xticklabels(SHAPES3D_FACTOR_NAMES, rotation=30, ha=\"right\")\n", + "ax.set_ylabel(\"R² (linear regression)\")\n", + "ax.set_title(\"Factor predictability from full representation\")\n", + "ax.set_ylim(0, 1.05)\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "print(\"\\nFull-model R² per factor:\")\n", + "header = f\"{'Factor':<15s}\" + \"\".join(f\" {n:>14s}\" for n in repr_names)\n", + "print(header)\n", + "print(\"-\" * len(header))\n", + "for fn in SHAPES3D_FACTOR_NAMES:\n", + " row = f\"{fn:<15s}\" + \"\".join(f\" {s3d_full_r2[n][fn]:14.4f}\" for n in repr_names)\n", + " print(row)\n" + ], + "metadata": { + "id": "shapes3d-r2" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-tsne", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# t-SNE colored by each ground-truth factor\n", + "# =====================================================================\n", + "N_TSNE = 5000 # subsample for speed\n", + "\n", + "idx = np.random.choice(len(s3d_rpm_latent_np),\n", + " min(N_TSNE, len(s3d_rpm_latent_np)), replace=False)\n", + "\n", + "# Compute t-SNE for all three representations\n", + "s3d_tsne = {}\n", + "for name, feats in S3D_REPR.items():\n", + " print(f\"Running t-SNE for {name}...\")\n", + " s3d_tsne[name] = TSNE(n_components=2, perplexity=30, random_state=42,\n", + " init=\"pca\", learning_rate=\"auto\"\n", + " ).fit_transform(feats[idx])\n", + "\n", + "n_factors = len(SHAPES3D_FACTOR_NAMES)\n", + "n_repr = len(S3D_REPR)\n", + "fig, axes = plt.subplots(n_repr, n_factors, figsize=(4 * n_factors, 4 * n_repr))\n", + "\n", + "repr_names = list(S3D_REPR.keys())\n", + "for row, name in enumerate(repr_names):\n", + " z2 = s3d_tsne[name]\n", + " for fi, fn in enumerate(SHAPES3D_FACTOR_NAMES):\n", + " fv = s3d_factors[fn][idx]\n", + " ax = axes[row, fi]\n", + " sc = ax.scatter(z2[:, 0], z2[:, 1], c=fv, s=2, alpha=0.5, cmap=\"viridis\")\n", + " ax.set_title(f\"{name}: {fn}\", fontsize=9)\n", + " ax.set_xticks([]); ax.set_yticks([])\n", + " if fi == n_factors - 1:\n", + " plt.colorbar(sc, ax=ax, shrink=0.7)\n", + "\n", + "plt.suptitle(\"t-SNE colored by ground-truth factors\", fontsize=14,\n", + " fontweight=\"bold\", y=1.01)\n", + "plt.tight_layout()\n", + "plt.show()\n" + ], + "metadata": { + "id": "shapes3d-tsne" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "shapes3d-pca-traversal", + "cell_type": "code", + "source": [ + "# =====================================================================\n", + "# PCA explained variance + factor-latent correlation matrix\n", + "# =====================================================================\n", + "\n", + "# --- Explained variance ---\n", + "fig, ax = plt.subplots(figsize=(8, 4))\n", + "for name, feats in S3D_REPR.items():\n", + " pca = PCA().fit(feats)\n", + " cumvar = np.cumsum(pca.explained_variance_ratio_)\n", + " ax.plot(cumvar, label=name)\n", + " n90 = np.searchsorted(cumvar, 0.90) + 1\n", + " n95 = np.searchsorted(cumvar, 0.95) + 1\n", + " print(f\"{name}: dims for 90% var = {n90}, 95% var = {n95}\")\n", + "ax.axhline(0.90, ls=\"--\", color=\"gray\", alpha=0.5)\n", + "ax.axhline(0.95, ls=\"--\", color=\"gray\", alpha=0.3)\n", + "ax.set_xlabel(\"Number of PCA components\")\n", + "ax.set_ylabel(\"Cumulative explained variance\")\n", + "ax.set_title(\"Shapes3D: Effective dimensionality\")\n", + "ax.legend()\n", + "plt.tight_layout()\n", + "plt.show()\n", + "\n", + "# --- Correlation matrix: PCA components vs factors ---\n", + "n_repr = len(S3D_REPR)\n", + "fig, axes = plt.subplots(1, n_repr, figsize=(7 * n_repr, 5))\n", + "for ax, (name, feats) in zip(axes, S3D_REPR.items()):\n", + " pca = PCA(n_components=min(20, feats.shape[1])).fit_transform(feats)\n", + " corr = np.zeros((len(SHAPES3D_FACTOR_NAMES), pca.shape[1]))\n", + " for fi, fn in enumerate(SHAPES3D_FACTOR_NAMES):\n", + " fv = s3d_factors[fn][:len(pca)]\n", + " for pc in range(pca.shape[1]):\n", + " corr[fi, pc] = np.abs(np.corrcoef(fv, pca[:, pc])[0, 1])\n", + " im = ax.imshow(corr, aspect=\"auto\", cmap=\"Blues\", vmin=0, vmax=1)\n", + " ax.set_yticks(range(len(SHAPES3D_FACTOR_NAMES)))\n", + " ax.set_yticklabels(SHAPES3D_FACTOR_NAMES, fontsize=9)\n", + " ax.set_xlabel(\"PCA component\")\n", + " ax.set_title(f\"{name}: |corr| of PCA dims vs factors\")\n", + " plt.colorbar(im, ax=ax, shrink=0.8)\n", + "plt.tight_layout()\n", + "plt.show()\n" + ], + "metadata": { + "id": "shapes3d-pca-traversal" + }, + "execution_count": null, + "outputs": [] + }, + { + "id": "T_kAyWX8p0KY", + "cell_type": "code", + "source": [], + "metadata": { + "id": "T_kAyWX8p0KY" + }, + "execution_count": null, + "outputs": [] + } + ], + "metadata": { + "colab": { + "provenance": [ + { + "file_id": "1jD49he9xC3CwoNFLiZrpRxIPPNjy47TH", + "timestamp": 1784643358045 + } + ], + "last_runtime": { + "build_target": "//gdm/science/msgym/exploratory:msgym_brain_ext_notebook", + "kind": "private" + }, + "private_outputs": true, + "toc_visible": true, + "collapsed_sections": [ + "cell-028" + ] + }, + "kernelspec": { + "name": "python3", + "display_name": "Python 3" + }, + "language_info": { + "name": "python" + }, + "accelerator": "GPU" + }, + "nbformat_minor": 5, + "nbformat": 4 +}