xwm.metrics

Linear and k-NN probes, and collapse diagnostics. See Diagnostics for what each one misses.

Metrics: is this representation any good, and has it collapsed?

Modules:

Name	Description
probes	Linear and nearest-neighbour probes.
representation	Diagnosing a representation without labels.

Functions:

Name	Description
knn_probe	k-nearest-neighbour probe in cosine distance.
ridge_probe	Closed-form ridge regression from embeddings to targets.
collapse_report	All of the above at once, for logging alongside the loss.
effective_rank_ratio	rankme divided by the embedding width -- 1.0 is ideal.
feature_std	Mean per-dimension standard deviation. Near zero means collapse.
mean_cosine_similarity	Average pairwise cosine similarity between samples (excluding self).
rankme	RankMe: the effective rank as the entropy of the singular-value spectrum.
singular_values	Singular values of the centred embedding matrix, descending.

knn_probe

knn_probe(z_train: Array, y_train: Array, z_test: Array, y_test: Array, *, k: int = 5, classification: bool = False, n_classes: int | None = None) -> dict[str, Array]

k-nearest-neighbour probe in cosine distance.

Unlike ridge_probe this reads local structure, so the two together distinguish "linearly decodable" from "merely clustered".

Parameters:

Name	Type	Description	Default
classification	bool	treat y as integer labels and report accuracy; otherwise average the neighbours' values and report r2.	False
n_classes	int | None	required when classification is set.	None

Source code in xwm/metrics/probes.py

def knn_probe(
    z_train: Array,
    y_train: Array,
    z_test: Array,
    y_test: Array,
    *,
    k: int = 5,
    classification: bool = False,
    n_classes: int | None = None,
) -> dict[str, Array]:
    """k-nearest-neighbour probe in cosine distance.

    Unlike :func:`ridge_probe` this reads *local* structure, so the two together
    distinguish "linearly decodable" from "merely clustered".

    Args:
        classification: treat ``y`` as integer labels and report accuracy;
            otherwise average the neighbours' values and report ``r2``.
        n_classes: required when ``classification`` is set.
    """
    a = z_train / (jnp.linalg.norm(z_train, axis=-1, keepdims=True) + 1e-8)
    b = z_test / (jnp.linalg.norm(z_test, axis=-1, keepdims=True) + 1e-8)
    neighbours = jnp.argsort(-(b @ a.T), axis=1)[:, :k]  # (N_test, k)

    if classification:
        if n_classes is None:
            raise ValueError("classification=True requires n_classes")
        labels = jnp.asarray(y_train, jnp.int32)[neighbours]  # (N_test, k)
        votes = jnp.sum(jax.nn.one_hot(labels, n_classes), axis=1)
        pred = jnp.argmax(votes, axis=-1)
        return {"accuracy": jnp.mean(pred == jnp.asarray(y_test, jnp.int32))}

    y_train = jnp.atleast_2d(y_train)
    y_test = jnp.atleast_2d(y_test)
    pred = jnp.mean(y_train[neighbours], axis=1)
    mse = jnp.mean(jnp.square(pred - y_test))
    variance = jnp.mean(jnp.square(y_test - jnp.mean(y_test, axis=0, keepdims=True)))
    return {"mse": mse, "r2": 1.0 - mse / (variance + 1e-12)}

ridge_probe

ridge_probe(z_train: Array, y_train: Array, z_test: Array, y_test: Array, *, alpha: float = 0.001) -> dict[str, Array]

Closed-form ridge regression from embeddings to targets.

Parameters:

Name	Type	Description	Default
z_train	Array	(N, D) embeddings.	required
y_train	Array	(N, K) regression targets.	required
alpha	float	ridge penalty, on the scale of the (standardised) features.	0.001

Returns:

Type	Description
dict[str, Array]	{"r2", "mse", "train_mse"}. r2 is computed against the test
dict[str, Array]	mean, so a probe that has learned nothing scores ~0 and a harmful one
dict[str, Array]	scores below 0.

Source code in xwm/metrics/probes.py

def ridge_probe(
    z_train: Array,
    y_train: Array,
    z_test: Array,
    y_test: Array,
    *,
    alpha: float = 1e-3,
) -> dict[str, Array]:
    """Closed-form ridge regression from embeddings to targets.

    Args:
        z_train: ``(N, D)`` embeddings.
        y_train: ``(N, K)`` regression targets.
        alpha: ridge penalty, on the scale of the (standardised) features.

    Returns:
        ``{"r2", "mse", "train_mse"}``. ``r2`` is computed against the *test*
        mean, so a probe that has learned nothing scores ~0 and a harmful one
        scores below 0.
    """
    z_train, z_test = jnp.atleast_2d(z_train), jnp.atleast_2d(z_test)
    y_train, y_test = jnp.atleast_2d(y_train), jnp.atleast_2d(y_test)
    mu = jnp.mean(z_train, axis=0, keepdims=True)
    sigma = jnp.std(z_train, axis=0, keepdims=True) + 1e-6
    a = jnp.concatenate([(z_train - mu) / sigma, jnp.ones((z_train.shape[0], 1))], axis=1)
    b = jnp.concatenate([(z_test - mu) / sigma, jnp.ones((z_test.shape[0], 1))], axis=1)
    d = a.shape[1]
    w = jnp.linalg.solve(a.T @ a + alpha * a.shape[0] * jnp.eye(d), a.T @ y_train)
    pred_test, pred_train = b @ w, a @ w
    mse = jnp.mean(jnp.square(pred_test - y_test))
    variance = jnp.mean(jnp.square(y_test - jnp.mean(y_test, axis=0, keepdims=True)))
    return {
        "r2": 1.0 - mse / (variance + 1e-12),
        "mse": mse,
        "train_mse": jnp.mean(jnp.square(pred_train - y_train)),
    }

collapse_report

collapse_report(z: Array) -> dict[str, Array]

All of the above at once, for logging alongside the loss.

Source code in xwm/metrics/representation.py

def collapse_report(z: Array) -> dict[str, Array]:
    """All of the above at once, for logging alongside the loss."""
    return {
        "rankme": rankme(z),
        "rank_ratio": effective_rank_ratio(z),
        "feature_std": feature_std(z),
        "mean_cosine": mean_cosine_similarity(z),
    }

effective_rank_ratio

effective_rank_ratio(z: Array) -> Array

rankme divided by the embedding width -- 1.0 is ideal.

Source code in xwm/metrics/representation.py

def effective_rank_ratio(z: Array) -> Array:
    """:func:`rankme` divided by the embedding width -- ``1.0`` is ideal."""
    return rankme(z) / z.shape[-1]

feature_std

feature_std(z: Array) -> Array

Mean per-dimension standard deviation. Near zero means collapse.

Source code in xwm/metrics/representation.py

def feature_std(z: Array) -> Array:
    """Mean per-dimension standard deviation. Near zero means collapse."""
    return jnp.mean(jnp.std(_flatten(z), axis=0))

mean_cosine_similarity

mean_cosine_similarity(z: Array) -> Array

Average pairwise cosine similarity between samples (excluding self).

Approaching 1.0 means every input maps to nearly the same direction -- collapse, even if the per-dimension variance still looks healthy.

Source code in xwm/metrics/representation.py

def mean_cosine_similarity(z: Array) -> Array:
    """Average pairwise cosine similarity between samples (excluding self).

    Approaching ``1.0`` means every input maps to nearly the same direction --
    collapse, even if the per-dimension variance still looks healthy.
    """
    z = _flatten(z)
    z = z / (jnp.linalg.norm(z, axis=-1, keepdims=True) + 1e-8)
    n = z.shape[0]
    sim = z @ z.T
    off_diagonal = jnp.sum(sim) - jnp.trace(sim)
    return off_diagonal / max(n * (n - 1), 1)

rankme

rankme(z: Array, eps: float = 1e-07) -> Array

RankMe: the effective rank as the entropy of the singular-value spectrum.

exp(H(p)) where p is the normalised spectrum. Equals D when every direction carries equal energy and 1 when all the energy is in one direction, and unlike a hard rank it responds smoothly to partial collapse -- the failure mode that actually shows up in JEPA training.

Note that the spectrum is taken after centring, so this is blind to a constant offset: an encoder emitting c + tiny_noise still scores a high rank. feature_std and mean_cosine_similarity catch that case, which is why collapse_report reports all three.

Source code in xwm/metrics/representation.py

def rankme(z: Array, eps: float = 1e-7) -> Array:
    """RankMe: the effective rank as the entropy of the singular-value spectrum.

    ``exp(H(p))`` where ``p`` is the normalised spectrum. Equals ``D`` when every
    direction carries equal energy and ``1`` when all the energy is in one
    direction, and unlike a hard rank it responds smoothly to *partial*
    collapse -- the failure mode that actually shows up in JEPA training.

    Note that the spectrum is taken after centring, so this is blind to a
    constant offset: an encoder emitting ``c + tiny_noise`` still scores a high
    rank. :func:`feature_std` and :func:`mean_cosine_similarity` catch that
    case, which is why :func:`collapse_report` reports all three.
    """
    s = singular_values(z)
    p = s / (jnp.sum(s) + eps)
    return jnp.exp(-jnp.sum(p * jnp.log(p + eps)))

singular_values

singular_values(z: Array) -> Array

Singular values of the centred embedding matrix, descending.

Source code in xwm/metrics/representation.py

def singular_values(z: Array) -> Array:
    """Singular values of the centred embedding matrix, descending."""
    z = _flatten(z)
    z = z - jnp.mean(z, axis=0, keepdims=True)
    return jnp.linalg.svd(z, compute_uv=False)