Coverage for cafl4ds/measurements.py: 91%
125 statements
« prev ^ index » next coverage.py v7.15.0, created at 2026-07-17 14:39 +0000
1"""Representation-health instruments (Phase 0 — *Instrument*).
3This is the **measurement apparatus** of the project: the "thermometer" the rest of the
4study reads. Every function here is *standalone* — it depends on nothing but embedding
5tensors (``Z``, shape ``[N, d]``) or a passed-in ``encoder`` callable. There is no stream,
6no SSL model, and no training in this module, by design: per the plan's *Strategic
7ordering*, the instruments are built and verified in isolation first, because everything
8downstream depends on them being correct.
10Signal catalogue implemented here (see ``docs/project-plan/metrics.md`` → *Signal catalogue*):
12Geometry / collapse (label-free):
13 * :func:`rankme` / :func:`effective_rank` — exponential entropy of the normalized
14 singular-value spectrum; the core collapse readout (Roy & Vetterli 2007; Garrido
15 et al. 2023).
16 * :func:`feature_variance` — per-dimension variance; flags dimensions collapsing to a
17 constant (VICReg variance term, Bardes et al. 2022).
18 * :func:`offdiag_covariance` — sum of squared off-diagonal covariances; flags
19 informational (redundancy) collapse (VICReg covariance term, Bardes et al. 2022).
20 * :func:`alignment` / :func:`uniformity` — positive-pair closeness and spread on the
21 hypersphere, for the joint-embedding collapse demo (Wang & Isola 2020).
23Dynamics / stability (label-free):
24 * :func:`linear_cka` / :func:`cka_drift` — representation drift of a fixed probe set
25 across checkpoints (CKA, Kornblith et al. 2019). Rotation/isotropic-scale invariant.
26 * :func:`cosine_drift` — rotation-*sensitive* per-sample churn; the direct readout of
27 "how fast the coordinate frame moves" (the moving-reference problem).
29Downstream / forgetting (labels used HERE ONLY — never in any SSL/training path):
30 * :func:`knn_probe` — nearest-neighbour accuracy on frozen features (Wu et al. 2018;
31 Caron et al. 2021).
32 * :func:`linear_probe` — linear-head accuracy on frozen features (standard SSL eval).
33 * :func:`frozen_baseline` — evaluate a never-updated encoder (baseline **B5**);
34 init-matched (frozen-pretrained *or* frozen-random).
36Conventions:
37 * Embedding matrices are ``[N, d]`` (rows are samples). Inputs may be ``torch.Tensor``
38 or ``numpy.ndarray``; everything is computed in ``float`` on CPU.
39 * Scalar metrics return Python ``float``; per-dimension metrics return a 1-D
40 ``torch.Tensor`` of length ``d``.
41 * An ``encoder`` is any callable ``Tensor -> Tensor`` mapping a batch of inputs to a
42 batch of embeddings.
43"""
45from collections.abc import Callable, Sequence
46from typing import TypeAlias, cast
48import numpy as np
49import torch
50from sklearn.linear_model import LogisticRegression
51from sklearn.neighbors import KNeighborsClassifier
53# Anything we accept as an embedding matrix / label vector.
54TensorLike: TypeAlias = torch.Tensor | np.ndarray
56# A callable mapping a batch of inputs to a batch of embeddings.
57Encoder: TypeAlias = Callable[[torch.Tensor], torch.Tensor]
59# A ``(inputs, labels)`` pair for the probe instruments. Labels are used HERE ONLY.
60Dataset: TypeAlias = tuple[TensorLike, TensorLike]
63def _as_tensor(x: TensorLike) -> torch.Tensor:
64 """Return ``x`` as a detached, contiguous float ``torch.Tensor`` on the CPU.
66 Args:
67 x: A ``torch.Tensor`` or ``numpy.ndarray``.
69 Returns:
70 A 2-D-preserving float tensor (dtype ``float32``) detached from any graph.
71 """
72 if isinstance(x, np.ndarray):
73 t = torch.from_numpy(x)
74 elif isinstance(x, torch.Tensor):
75 t = x.detach()
76 else: # pragma: no cover - defensive; callers pass tensors/arrays.
77 t = torch.as_tensor(x)
78 return t.to(dtype=torch.float32).cpu().contiguous()
81def _to_float(x: torch.Tensor) -> float:
82 """Extract a Python ``float`` from a scalar tensor (keeps mypy strict happy)."""
83 return cast(float, x.item())
86def _check_matrix(z: torch.Tensor, name: str = "Z") -> torch.Tensor:
87 """Validate that ``z`` is a 2-D ``[N, d]`` embedding matrix.
89 Args:
90 z: Candidate embedding matrix.
91 name: Name used in error messages.
93 Returns:
94 ``z`` unchanged, once validated.
96 Raises:
97 ValueError: If ``z`` is not 2-dimensional or is empty.
98 """
99 if z.ndim != 2: 99 ↛ 100line 99 didn't jump to line 100 because the condition on line 99 was never true
100 raise ValueError(f"{name} must be 2-D [N, d]; got shape {tuple(z.shape)}.")
101 if z.shape[0] == 0 or z.shape[1] == 0: 101 ↛ 102line 101 didn't jump to line 102 because the condition on line 101 was never true
102 raise ValueError(f"{name} must be non-empty; got shape {tuple(z.shape)}.")
103 return z
106# --------------------------------------------------------------------------------------
107# Geometry / collapse (label-free)
108# --------------------------------------------------------------------------------------
111def _singular_values(z: TensorLike) -> torch.Tensor:
112 """Return the singular values of an embedding matrix in descending order.
114 Args:
115 z: Embedding matrix ``[N, d]``.
117 Returns:
118 A 1-D tensor of ``min(N, d)`` singular values, sorted descending.
119 """
120 zt = _check_matrix(_as_tensor(z))
121 return cast(torch.Tensor, torch.linalg.svdvals(zt))
124def effective_rank(z: TensorLike) -> float:
125 """Effective rank: exponential of the entropy of the normalized singular-value spectrum.
127 Roy & Vetterli (2007). Given singular values ``s`` of ``Z``, form the probability
128 distribution ``p = s / sum(s)`` and return ``exp(H(p))`` where ``H`` is the Shannon
129 entropy in nats. Reads out how many embedding dimensions are *actually in use*: a
130 rank-1 (fully collapsed) representation scores ~1; ``k`` equal directions score ~``k``.
132 Args:
133 z: Embedding matrix ``[N, d]``.
135 Returns:
136 The effective rank, a float in ``[1, min(N, d)]``.
137 """
138 s = _singular_values(z)
139 total = s.sum()
140 if total <= 0: # pragma: no cover - all-zero embeddings; degenerate but well-defined.
141 return 1.0
142 # Drop exact-zero singular values (from rank-deficient matrices) before the log so they
143 # contribute 0 to the entropy (the -p*log(p) -> 0 limit) rather than producing NaNs.
144 p = s[s > 0] / total
145 entropy = -(p * p.log()).sum()
146 return _to_float(torch.exp(entropy))
149def rankme(z: TensorLike, eps: float = 1e-7) -> float:
150 """RankMe (Garrido et al. 2023): effective rank with the paper's numerical smoothing.
152 Identical in spirit to :func:`effective_rank`, using the RankMe formulation
153 ``p_k = s_k / ||s||_1 + eps`` before taking the exponential entropy. The ``eps`` term
154 keeps the log well-defined and matches the reference implementation.
156 Args:
157 z: Embedding matrix ``[N, d]``.
158 eps: Numerical-stability constant added to each normalized singular value.
160 Returns:
161 The RankMe score, a float in ``[1, min(N, d)]``.
162 """
163 s = _singular_values(z)
164 total = s.sum()
165 if total <= 0: # pragma: no cover - all-zero embeddings.
166 return 1.0
167 p = s / total + eps
168 entropy = -(p * p.log()).sum()
169 return _to_float(torch.exp(entropy))
172def feature_variance(z: TensorLike, unbiased: bool = True) -> torch.Tensor:
173 """Per-dimension variance of the embeddings (VICReg variance term, Bardes et al. 2022).
175 Flags dimensions that have collapsed to a (near-)constant value: such a dimension
176 reads ~0. Returned per-dimension rather than aggregated so callers can inspect *which*
177 dimensions are collapsing.
179 Args:
180 z: Embedding matrix ``[N, d]``.
181 unbiased: Whether to use Bessel's correction (``ddof=1``), as VICReg does.
183 Returns:
184 A 1-D tensor of length ``d`` giving the variance of each embedding dimension.
185 """
186 zt = _check_matrix(_as_tensor(z))
187 return zt.var(dim=0, unbiased=unbiased)
190def offdiag_covariance(z: TensorLike) -> float:
191 """VICReg covariance term: mean squared off-diagonal covariance (Bardes et al. 2022).
193 Computes the covariance matrix ``C`` of the (feature-centered) embeddings and returns
194 ``sum_{i != j} C[i, j]^2 / d``. Flags *informational* (redundancy) collapse: when
195 dimensions become linearly redundant the off-diagonal mass grows; a decorrelated
196 representation scores ~0.
198 Args:
199 z: Embedding matrix ``[N, d]`` with ``N >= 2``.
201 Returns:
202 The mean squared off-diagonal covariance, a non-negative float.
204 Raises:
205 ValueError: If ``N < 2`` (covariance is undefined for a single sample).
206 """
207 zt = _check_matrix(_as_tensor(z))
208 n, d = zt.shape
209 if n < 2:
210 raise ValueError(f"offdiag_covariance needs N >= 2 samples; got N={n}.")
211 zc = zt - zt.mean(dim=0, keepdim=True)
212 cov = (zc.T @ zc) / (n - 1)
213 off_diag_sq = cov.pow(2).sum() - cov.diagonal().pow(2).sum()
214 return _to_float(off_diag_sq / d)
217def _l2_normalize(z: torch.Tensor, eps: float = 1e-12) -> torch.Tensor:
218 """Row-wise L2-normalize an embedding matrix onto the unit hypersphere.
220 Args:
221 z: Embedding matrix ``[N, d]``.
222 eps: Lower bound on the norm to avoid division by zero.
224 Returns:
225 ``z`` with each row scaled to unit L2 norm.
226 """
227 return cast(torch.Tensor, z / z.norm(dim=1, keepdim=True).clamp_min(eps))
230def alignment(
231 pos_a: TensorLike | Sequence[TensorLike],
232 pos_b: TensorLike | None = None,
233 *,
234 alpha: float = 2.0,
235) -> float:
236 """Alignment of positive pairs on the hypersphere (Wang & Isola 2020).
238 ``E[ ||f(x) - f(y)||_2^alpha ]`` over positive pairs ``(x, y)``, both L2-normalized to
239 the unit sphere. Lower is better (positives map close together); it *rises* as a
240 joint-embedding model's positives drift apart.
242 Args:
243 pos_a: Either the first matrix of a positive pair ``[N, d]``, or — when ``pos_b`` is
244 omitted — a ``(z_a, z_b)`` sequence of the two matrices.
245 pos_b: The second matrix of the positive pair ``[N, d]``, aligned row-for-row with
246 ``pos_a``. Omit when passing the pair as ``pos_a``.
247 alpha: Exponent on the pairwise distance (Wang & Isola use ``alpha = 2``).
249 Returns:
250 The mean aligned distance, a non-negative float.
252 Raises:
253 ValueError: If the two matrices differ in shape.
254 """
255 if pos_b is None:
256 pair = cast("Sequence[TensorLike]", pos_a) # unpack the (z_a, z_b) sequence.
257 a_mat, b_mat = pair[0], pair[1]
258 else:
259 a_mat, b_mat = cast(TensorLike, pos_a), pos_b
260 za = _check_matrix(_as_tensor(a_mat), "pos_a")
261 zb = _check_matrix(_as_tensor(b_mat), "pos_b")
262 if za.shape != zb.shape: 262 ↛ 263line 262 didn't jump to line 263 because the condition on line 262 was never true
263 raise ValueError(f"positive pairs must share a shape; got {tuple(za.shape)} vs {tuple(zb.shape)}.")
264 za, zb = _l2_normalize(za), _l2_normalize(zb)
265 dist = (za - zb).norm(dim=1)
266 return _to_float(dist.pow(alpha).mean())
269def uniformity(z: TensorLike, t: float = 2.0) -> float:
270 """Uniformity of embeddings on the hypersphere (Wang & Isola 2020).
272 ``log E_{x,y}[ exp(-t * ||x - y||_2^2) ]`` over all distinct pairs, with rows
273 L2-normalized to the unit sphere. *Lower* (more negative) means the embeddings are
274 spread more uniformly; the value climbs toward 0 as embeddings clump together
275 (a collapse signature).
277 Args:
278 z: Embedding matrix ``[N, d]`` with ``N >= 2``.
279 t: Temperature of the Gaussian potential (Wang & Isola use ``t = 2``).
281 Returns:
282 The uniformity loss, a float (typically negative).
284 Raises:
285 ValueError: If ``N < 2`` (no pairs exist).
286 """
287 zt = _check_matrix(_as_tensor(z))
288 if zt.shape[0] < 2: 288 ↛ 289line 288 didn't jump to line 289 because the condition on line 288 was never true
289 raise ValueError(f"uniformity needs N >= 2 samples; got N={zt.shape[0]}.")
290 zt = _l2_normalize(zt)
291 # pdist gives the condensed vector of pairwise L2 distances (upper triangle, i < j).
292 sq_dist = torch.pdist(zt).pow(2)
293 return _to_float(torch.log(torch.exp(-t * sq_dist).mean()))
296# --------------------------------------------------------------------------------------
297# Dynamics / stability (label-free)
298# --------------------------------------------------------------------------------------
301def linear_cka(x: TensorLike, y: TensorLike) -> float:
302 """Linear Centered Kernel Alignment between two representations (Kornblith et al. 2019).
304 Both matrices describe the *same* set of ``N`` samples under two representations
305 (e.g. a fixed probe set at two checkpoints). Returns a similarity in ``[0, 1]`` that is
306 invariant to orthogonal transforms and isotropic scaling of either representation
307 (that invariance is the point of CKA — it compares representational *content*, not the
308 coordinate frame; see :func:`cosine_drift` for a frame-sensitive companion).
310 Args:
311 x: First representation ``[N, d_x]``.
312 y: Second representation ``[N, d_y]`` (same ``N`` as ``x``).
314 Returns:
315 The linear CKA similarity, a float in ``[0, 1]`` (1.0 = identical up to
316 rotation/scaling).
318 Raises:
319 ValueError: If ``x`` and ``y`` have a different number of samples.
320 """
321 xt = _check_matrix(_as_tensor(x), "x")
322 yt = _check_matrix(_as_tensor(y), "y")
323 if xt.shape[0] != yt.shape[0]: 323 ↛ 324line 323 didn't jump to line 324 because the condition on line 323 was never true
324 raise ValueError(f"CKA needs matching sample counts; got {xt.shape[0]} vs {yt.shape[0]}.")
325 xc = xt - xt.mean(dim=0, keepdim=True)
326 yc = yt - yt.mean(dim=0, keepdim=True)
327 # Linear HSIC via the feature-space cross-covariance: ||Yc^T Xc||_F^2.
328 hsic_xy = (yc.T @ xc).pow(2).sum()
329 hsic_xx = (xc.T @ xc).pow(2).sum()
330 hsic_yy = (yc.T @ yc).pow(2).sum()
331 denom = torch.sqrt(hsic_xx * hsic_yy)
332 if denom <= 0: # pragma: no cover - a constant representation; similarity undefined.
333 return 0.0
334 return _to_float(hsic_xy / denom)
337def cka_drift(z_ref_t0: TensorLike, z_ref_t: TensorLike) -> float:
338 """Representation drift of a fixed probe set via linear CKA (``1 - CKA``).
340 Measures how much the *content* of a fixed probe set's representation has changed
341 between two checkpoints. ``0`` means unchanged (identical up to rotation/scaling);
342 larger means more drift. Because linear CKA is rotation-invariant, a pure change of
343 coordinate frame reads ~0 here — use :func:`cosine_drift` to detect that.
345 Args:
346 z_ref_t0: Probe-set embeddings at the reference checkpoint ``[N, d0]``.
347 z_ref_t: Probe-set embeddings at a later checkpoint ``[N, dt]`` (same ``N``).
349 Returns:
350 The drift ``1 - linear_cka``, a float in ``[0, 1]``.
351 """
352 return 1.0 - linear_cka(z_ref_t0, z_ref_t)
355def cosine_drift(z_ref_t0: TensorLike, z_ref_t: TensorLike) -> float:
356 """Per-sample cosine churn of a fixed probe set (``1 - mean cosine similarity``).
358 Unlike :func:`cka_drift`, this is *sensitive* to rotation of the embedding coordinate
359 frame: it compares each sample's embedding to its own embedding at the reference
360 checkpoint. This is the direct readout of "how fast the coordinate frame moves" — the
361 moving-reference problem that destabilizes a fixed novelty reference. Requires the two
362 checkpoints to share an embedding dimensionality (same coordinate axes).
364 Args:
365 z_ref_t0: Probe-set embeddings at the reference checkpoint ``[N, d]``.
366 z_ref_t: Probe-set embeddings at a later checkpoint ``[N, d]`` (same ``N`` and ``d``).
368 Returns:
369 The mean per-sample cosine drift, a float in ``[0, 2]`` (0 = unchanged direction).
371 Raises:
372 ValueError: If the two matrices do not share a shape.
373 """
374 z0 = _check_matrix(_as_tensor(z_ref_t0), "z_ref_t0")
375 zt = _check_matrix(_as_tensor(z_ref_t), "z_ref_t")
376 if z0.shape != zt.shape: 376 ↛ 377line 376 didn't jump to line 377 because the condition on line 376 was never true
377 raise ValueError(f"cosine_drift needs matching shapes; got {tuple(z0.shape)} vs {tuple(zt.shape)}.")
378 cos = torch.nn.functional.cosine_similarity(z0, zt, dim=1)
379 return _to_float(1.0 - cos.mean())
382# --------------------------------------------------------------------------------------
383# Downstream / forgetting (labels used HERE ONLY)
384# --------------------------------------------------------------------------------------
387def _encode(encoder: Encoder, inputs: TensorLike) -> np.ndarray:
388 """Run ``encoder`` over ``inputs`` (no grad) and return a 2-D ``float`` numpy array.
390 Args:
391 encoder: Callable mapping a batch of inputs to a batch of embeddings.
392 inputs: Model inputs accepted by ``encoder``.
394 Returns:
395 The embeddings as a ``[N, d]`` numpy array.
396 """
397 with torch.no_grad():
398 z = encoder(_as_tensor(inputs))
399 zt = _as_tensor(z)
400 return _check_matrix(zt, "encoder output").numpy()
403def _labels(y: TensorLike) -> np.ndarray:
404 """Coerce a label vector to a 1-D numpy array.
406 Args:
407 y: Label vector (tensor or array).
409 Returns:
410 A 1-D numpy array of labels.
411 """
412 arr = y.detach().cpu().numpy() if isinstance(y, torch.Tensor) else np.asarray(y)
413 return arr.reshape(-1)
416def knn_probe(
417 encoder: Encoder,
418 support: Dataset,
419 query: Dataset,
420 k: int = 20,
421 *,
422 metric: str = "cosine",
423 weights: str = "distance",
424) -> float:
425 """k-NN probe accuracy on frozen features (Wu et al. 2018; Caron et al. 2021).
427 Encodes the ``support`` and ``query`` inputs with the (frozen) ``encoder``, fits a
428 nearest-neighbour classifier on the support embeddings/labels, and reports accuracy on
429 the query set. Labels are used HERE ONLY — this is an evaluation probe, never a
430 training signal.
432 Args:
433 encoder: Frozen encoder callable ``inputs -> embeddings``.
434 support: ``(inputs, labels)`` used as the neighbour database.
435 query: ``(inputs, labels)`` to classify and score.
436 k: Number of neighbours. Clamped to the support-set size if larger.
437 metric: Distance metric passed to scikit-learn (default cosine, the SSL k-NN norm).
438 weights: Neighbour weighting (``"distance"`` or ``"uniform"``).
440 Returns:
441 Top-1 query accuracy in ``[0, 1]``.
442 """
443 xs, ys = _encode(encoder, support[0]), _labels(support[1])
444 xq, yq = _encode(encoder, query[0]), _labels(query[1])
445 n_neighbors = max(1, min(k, xs.shape[0]))
446 clf = KNeighborsClassifier(n_neighbors=n_neighbors, metric=metric, weights=weights)
447 clf.fit(xs, ys)
448 return float((clf.predict(xq) == yq).mean())
451def linear_probe(
452 encoder: Encoder,
453 support: Dataset,
454 query: Dataset,
455 *,
456 max_iter: int = 1000,
457 C: float = 1.0,
458 standardize: bool = True,
459) -> float:
460 """Linear-probe accuracy on frozen features (standard SSL linear-evaluation protocol).
462 Encodes ``support`` and ``query`` with the (frozen) ``encoder`` and fits a multinomial
463 logistic-regression head on the support embeddings, then scores the query set. Only a
464 *linear* head is trained; the encoder is never updated. Labels are used HERE ONLY.
466 Args:
467 encoder: Frozen encoder callable ``inputs -> embeddings``.
468 support: ``(inputs, labels)`` used to fit the linear head.
469 query: ``(inputs, labels)`` to score.
470 max_iter: Maximum solver iterations for the logistic-regression head.
471 C: Inverse L2-regularization strength for the linear head.
472 standardize: Whether to zero-mean/unit-variance the features (fit on support) before
473 the linear head — the usual, more stable linear-probe recipe.
475 Returns:
476 Top-1 query accuracy in ``[0, 1]``.
477 """
478 xs, ys = _encode(encoder, support[0]), _labels(support[1])
479 xq, yq = _encode(encoder, query[0]), _labels(query[1])
480 if standardize: 480 ↛ 486line 480 didn't jump to line 486 because the condition on line 480 was always true
481 mean = xs.mean(axis=0, keepdims=True)
482 std = xs.std(axis=0, keepdims=True)
483 std[std == 0] = 1.0
484 xs = (xs - mean) / std
485 xq = (xq - mean) / std
486 clf = LogisticRegression(max_iter=max_iter, C=C)
487 clf.fit(xs, ys)
488 return float((clf.predict(xq) == yq).mean())
491def frozen_baseline(
492 encoder_init: Encoder,
493 support: Dataset,
494 query: Dataset,
495 *,
496 probe: str = "knn",
497 **probe_kwargs: object,
498) -> float:
499 """Baseline **B5**: evaluate a never-updated encoder (init-matched frozen backbone).
501 B5 is the "no adaptation" floor of the study — a backbone that never sees a gradient
502 step. It is *init-matched*: pass a frozen-pretrained encoder in the pretrained regime,
503 or a frozen-random encoder in the from-scratch regime. This is a thin wrapper that runs
504 the chosen probe on ``encoder_init``; the point is the frozen encoder, not new logic.
506 Args:
507 encoder_init: The frozen encoder (pretrained or randomly initialized) to evaluate.
508 support: ``(inputs, labels)`` for the probe's support/training split.
509 query: ``(inputs, labels)`` for the probe's query/eval split.
510 probe: Which downstream probe to run: ``"knn"`` or ``"linear"``.
511 **probe_kwargs: Forwarded to :func:`knn_probe` / :func:`linear_probe`.
513 Returns:
514 The probe's top-1 query accuracy in ``[0, 1]``.
516 Raises:
517 ValueError: If ``probe`` is not ``"knn"`` or ``"linear"``.
518 """
519 if probe == "knn":
520 return knn_probe(encoder_init, support, query, **probe_kwargs) # type: ignore[arg-type]
521 if probe == "linear":
522 return linear_probe(encoder_init, support, query, **probe_kwargs) # type: ignore[arg-type]
523 raise ValueError(f"probe must be 'knn' or 'linear'; got {probe!r}.")