xwm.objectives

Losses: latent prediction, SIGReg, VICReg, InfoNCE, and the distributional statistics SIGReg is built from. See Collapse.

Training objectives for predictive world models.

Two independent choices make up a JEPA objective:

1. What to predict, and how to score it -- prediction_loss, always computed in latent space.
2. How to prevent collapse -- sigreg (LeJEPA: constrain the embedding distribution), vicreg (constrain its first two moments), info_nce (contrast against negatives), or an EMA teacher (architectural, see xwm.core.ema).

Mixing and matching those two axes is what the model classes in xwm.image, xwm.video and xwm.action expose.

Modules:

Name	Description
prediction	Latent prediction losses -- the primary signal in every JEPA.
regularizers	Variance/covariance regularizers (VICReg) and the InfoNCE contrastive loss.
sigreg	SIGReg -- Sketched Isotropic Gaussian Regularization (the LeJEPA objective).

Functions:

Name	Description
layer_normalize	Parameter-free LayerNorm over the last axis.
prediction_loss	Distance between predicted and target embeddings, averaged over everything.
covariance_loss	Penalise off-diagonal covariance, decorrelating the dimensions.
info_nce	Symmetric InfoNCE over in-batch negatives.
variance_loss	Hinge each dimension's standard deviation up to gamma.
vicreg	Variance-Invariance-Covariance regularization on two views.
cramer_von_mises	CDF distance from u to N(0, 1).
epps_pulley	Characteristic-function distance from u to N(0, 1).
random_directions	(dim, n_proj) directions drawn uniformly from the unit sphere.

layer_normalize

layer_normalize(x: Array, eps: float = 1e-06) -> Array

Parameter-free LayerNorm over the last axis.

Applied to prediction targets it removes the scale degree of freedom that a teacher/student pair can otherwise exploit to shrink the loss without improving prediction. V-JEPA normalises targets this way.

Source code in xwm/objectives/prediction.py

def layer_normalize(x: Array, eps: float = 1e-6) -> Array:
    """Parameter-free LayerNorm over the last axis.

    Applied to prediction *targets* it removes the scale degree of freedom that
    a teacher/student pair can otherwise exploit to shrink the loss without
    improving prediction. V-JEPA normalises targets this way.
    """
    mu = jnp.mean(x, axis=-1, keepdims=True)
    var = jnp.var(x, axis=-1, keepdims=True)
    return (x - mu) / jnp.sqrt(var + eps)

prediction_loss

prediction_loss(pred: Array, target: Array, *, kind: LossKind = 'l1', beta: float = 1.0, normalize_target: bool = False) -> Array

Distance between predicted and target embeddings, averaged over everything.

Parameters:

Name	Type	Description	Default
pred	Array	predicted embeddings, any shape ending in D.	required
target	Array	same shape as pred. Detach it before calling if it comes from a teacher -- this function does not stop gradients for you.	required
kind	LossKind	"l1" (V-JEPA), "l2", "smooth_l1" (I-JEPA), or "cosine".	'l1'
beta	float	transition point of the smooth-L1 / Huber loss.	1.0
normalize_target	bool	LayerNorm the target (and, for symmetry, the prediction) before comparing.	False

Returns:

Type	Description
Array	A scalar.

Source code in xwm/objectives/prediction.py

def prediction_loss(
    pred: Array,
    target: Array,
    *,
    kind: LossKind = "l1",
    beta: float = 1.0,
    normalize_target: bool = False,
) -> Array:
    """Distance between predicted and target embeddings, averaged over everything.

    Args:
        pred: predicted embeddings, any shape ending in ``D``.
        target: same shape as ``pred``. Detach it before calling if it comes
            from a teacher -- this function does not stop gradients for you.
        kind: ``"l1"`` (V-JEPA), ``"l2"``, ``"smooth_l1"`` (I-JEPA), or
            ``"cosine"``.
        beta: transition point of the smooth-L1 / Huber loss.
        normalize_target: LayerNorm the target (and, for symmetry, the
            prediction) before comparing.

    Returns:
        A scalar.
    """
    if pred.shape != target.shape:
        raise ValueError(f"shape mismatch: pred {pred.shape} vs target {target.shape}")
    if normalize_target:
        pred, target = layer_normalize(pred), layer_normalize(target)

    if kind == "l1":
        return jnp.mean(jnp.abs(pred - target))
    if kind == "l2":
        return jnp.mean(jnp.square(pred - target))
    if kind == "smooth_l1":
        d = jnp.abs(pred - target)
        return jnp.mean(jnp.where(d < beta, 0.5 * d**2 / beta, d - 0.5 * beta))
    if kind == "cosine":
        p = pred / (jnp.linalg.norm(pred, axis=-1, keepdims=True) + 1e-8)
        t = target / (jnp.linalg.norm(target, axis=-1, keepdims=True) + 1e-8)
        return jnp.mean(1.0 - jnp.sum(p * t, axis=-1))
    raise ValueError(f"unknown loss kind {kind!r}")

covariance_loss

covariance_loss(z: Array) -> Array

Penalise off-diagonal covariance, decorrelating the dimensions.

z is (..., D); see variance_loss on leading axes.

Source code in xwm/objectives/regularizers.py

def covariance_loss(z: Array) -> Array:
    """Penalise off-diagonal covariance, decorrelating the dimensions.

    ``z`` is ``(..., D)``; see :func:`variance_loss` on leading axes.
    """
    z = z.reshape(-1, z.shape[-1])
    n, d = z.shape
    z = z - jnp.mean(z, axis=0, keepdims=True)
    cov = (z.T @ z) / max(n - 1, 1)
    off_diagonal = cov - jnp.diag(jnp.diag(cov))
    return jnp.sum(jnp.square(off_diagonal)) / d

info_nce

info_nce(z_a: Array, z_b: Array, temperature: float = 0.1) -> Array

Symmetric InfoNCE over in-batch negatives.

Parameters:

Name	Type	Description	Default
z_a, z_b		(B, D) embeddings; row i of each is a positive pair.	required

Source code in xwm/objectives/regularizers.py

def info_nce(z_a: Array, z_b: Array, temperature: float = 0.1) -> Array:
    """Symmetric InfoNCE over in-batch negatives.

    Args:
        z_a, z_b: ``(B, D)`` embeddings; row ``i`` of each is a positive pair.
    """
    z_a = z_a / (jnp.linalg.norm(z_a, axis=-1, keepdims=True) + 1e-8)
    z_b = z_b / (jnp.linalg.norm(z_b, axis=-1, keepdims=True) + 1e-8)
    logits = (z_a @ z_b.T) / temperature
    labels = jnp.arange(z_a.shape[0])
    log_p_ab = logits - logsumexp(logits, axis=1, keepdims=True)
    log_p_ba = logits.T - logsumexp(logits.T, axis=1, keepdims=True)
    return -0.5 * (
        jnp.mean(log_p_ab[labels, labels]) + jnp.mean(log_p_ba[labels, labels])
    )

variance_loss

variance_loss(z: Array, gamma: float = 1.0, eps: float = 0.0001) -> Array

Hinge each dimension's standard deviation up to gamma.

z is (..., D); leading axes are all treated as samples, matching xwm.objectives.sigreg, so a token sequence (B, N, D) contributes B * N samples rather than erroring.

Source code in xwm/objectives/regularizers.py

def variance_loss(z: Array, gamma: float = 1.0, eps: float = 1e-4) -> Array:
    """Hinge each dimension's standard deviation up to ``gamma``.

    ``z`` is ``(..., D)``; leading axes are all treated as samples, matching
    :func:`xwm.objectives.sigreg`, so a token sequence ``(B, N, D)`` contributes
    ``B * N`` samples rather than erroring.
    """
    z = z.reshape(-1, z.shape[-1])
    std = jnp.sqrt(jnp.var(z, axis=0) + eps)
    return jnp.mean(jax.nn.relu(gamma - std))

vicreg

vicreg(z_a: Array, z_b: Array, *, sim_coeff: float = 25.0, var_coeff: float = 25.0, cov_coeff: float = 1.0, gamma: float = 1.0) -> tuple[Array, dict[str, Array]]

Variance-Invariance-Covariance regularization on two views.

Parameters:

Name	Type	Description	Default
z_a, z_b		(B, D) embeddings of two views of the same inputs.	required

Returns:

Type	Description
tuple[Array, dict[str, Array]]	(loss, parts) where parts holds the three terms for logging.

Source code in xwm/objectives/regularizers.py

def vicreg(
    z_a: Array,
    z_b: Array,
    *,
    sim_coeff: float = 25.0,
    var_coeff: float = 25.0,
    cov_coeff: float = 1.0,
    gamma: float = 1.0,
) -> tuple[Array, dict[str, Array]]:
    """Variance-Invariance-Covariance regularization on two views.

    Args:
        z_a, z_b: ``(B, D)`` embeddings of two views of the same inputs.

    Returns:
        ``(loss, parts)`` where ``parts`` holds the three terms for logging.
    """
    z_a = z_a.reshape(-1, z_a.shape[-1])
    z_b = z_b.reshape(-1, z_b.shape[-1])
    invariance = jnp.mean(jnp.square(z_a - z_b))
    variance = variance_loss(z_a, gamma) + variance_loss(z_b, gamma)
    covariance = covariance_loss(z_a) + covariance_loss(z_b)
    loss = sim_coeff * invariance + var_coeff * variance + cov_coeff * covariance
    return loss, {
        "invariance": invariance,
        "variance": variance,
        "covariance": covariance,
    }

cramer_von_mises

cramer_von_mises(u: Array) -> Array

CDF distance from u to N(0, 1).

Parameters:

Name	Type	Description	Default
u	Array	(n, P) -- P independent 1-D samples of size n.	required

Returns:

Type	Description
Array	(P,) non-negative statistics (the omega^2 form, so the scale is
Array	independent of n).

Source code in xwm/objectives/sigreg.py

def cramer_von_mises(u: Array) -> Array:
    """CDF distance from ``u`` to ``N(0, 1)``.

    Args:
        u: ``(n, P)`` -- ``P`` independent 1-D samples of size ``n``.

    Returns:
        ``(P,)`` non-negative statistics (the ``omega^2`` form, so the scale is
        independent of ``n``).
    """
    n = u.shape[0]
    cdf = jstats.norm.cdf(jnp.sort(u, axis=0))
    ranks = (2.0 * jnp.arange(1, n + 1) - 1.0) / (2.0 * n)
    return jnp.sum((cdf - ranks[:, None]) ** 2, axis=0) / n + 1.0 / (12.0 * n**2)

epps_pulley

epps_pulley(u: Array, *, n_nodes: int = 32, sigma: float = 1.0) -> Array

Characteristic-function distance from u to N(0, 1).

Parameters:

Name	Type	Description	Default
u	Array	(n, P) -- P independent 1-D samples of size n.	required
n_nodes	int	quadrature nodes; 32 is ample for so smooth an integrand.	32
sigma	float	width of the quadrature weight, i.e. which frequencies the test emphasises. Larger values probe finer structure in the tails.	1.0

Returns:

Type	Description
Array	(P,) non-negative statistics, zero iff the empirical characteristic
Array	function matches exp(-t^2 / 2) on the weighted grid.

The quadrature is a jax.lax.scan rather than one batched einsum: the dense form would allocate n * P * n_nodes floats, which at a real batch size (n in the tens of thousands once tokens are counted) reaches hundreds of megabytes for a term that is only a scalar penalty. Scanning keeps the footprint at n * P.

Source code in xwm/objectives/sigreg.py

def epps_pulley(u: Array, *, n_nodes: int = 32, sigma: float = 1.0) -> Array:
    """Characteristic-function distance from ``u`` to ``N(0, 1)``.

    Args:
        u: ``(n, P)`` -- ``P`` independent 1-D samples of size ``n``.
        n_nodes: quadrature nodes; 32 is ample for so smooth an integrand.
        sigma: width of the quadrature weight, i.e. which frequencies the test
            emphasises. Larger values probe finer structure in the tails.

    Returns:
        ``(P,)`` non-negative statistics, zero iff the empirical characteristic
        function matches ``exp(-t^2 / 2)`` on the weighted grid.

    The quadrature is a :func:`jax.lax.scan` rather than one batched einsum:
    the dense form would allocate ``n * P * n_nodes`` floats, which at a real
    batch size (``n`` in the tens of thousands once tokens are counted) reaches
    hundreds of megabytes for a term that is only a scalar penalty. Scanning
    keeps the footprint at ``n * P``.
    """
    t, w = _quadrature(n_nodes, sigma)

    def node(_, tw):
        t_k, w_k = tw
        phase = t_k * u  # (n, P)
        re = jnp.mean(jnp.cos(phase), axis=0) - jnp.exp(-0.5 * t_k**2)
        im = jnp.mean(jnp.sin(phase), axis=0)
        return _, w_k * (re**2 + im**2)

    _, per_node = jax.lax.scan(node, None, (t, w))  # (nodes, P)
    return jnp.sum(per_node, axis=0)

random_directions

random_directions(key: PRNGKey, dim: int, n_proj: int) -> Array

(dim, n_proj) directions drawn uniformly from the unit sphere.

Source code in xwm/objectives/sigreg.py

def random_directions(key: PRNGKey, dim: int, n_proj: int) -> Array:
    """``(dim, n_proj)`` directions drawn uniformly from the unit sphere."""
    v = jr.normal(key, (dim, n_proj))
    return v / (jnp.linalg.norm(v, axis=0, keepdims=True) + 1e-8)

sigreg

Note

xwm.objectives.sigreg is both a submodule and the function it exports. The function is documented here under its canonical path; xwm.objectives.sigreg(...) is the way to call it.

Sketched isotropic-Gaussian regularizer for a batch of embeddings.

Parameters:

Name	Type	Description	Default
z	Array	(..., D). Every leading axis is treated as a sample, so token sequences (B, N, D) contribute B * N samples.	required
key	PRNGKey	RNG for the projection directions. Redraw it each step -- fresh directions are what make the sketch cover all of R^D over training rather than only a fixed subspace.	required
n_proj	int	number of random directions.	512
statistic	Statistic	which goodness-of-fit test to use.	'epps_pulley'
n_nodes, sigma		quadrature settings for "epps_pulley".	required
center	bool	subtract the batch mean before testing. Off by default, because driving the mean to zero is part of the job.	False

Returns:

Type	Description
Array	A scalar, minimised when the embeddings look isotropic Gaussian.

Source code in xwm/objectives/sigreg.py

def sigreg(
    z: Array,
    key: PRNGKey,
    *,
    n_proj: int = 512,
    statistic: Statistic = "epps_pulley",
    n_nodes: int = 32,
    sigma: float = 1.0,
    center: bool = False,
) -> Array:
    """Sketched isotropic-Gaussian regularizer for a batch of embeddings.

    Args:
        z: ``(..., D)``. Every leading axis is treated as a sample, so token
            sequences ``(B, N, D)`` contribute ``B * N`` samples.
        key: RNG for the projection directions. Redraw it each step -- fresh
            directions are what make the sketch cover all of ``R^D`` over
            training rather than only a fixed subspace.
        n_proj: number of random directions.
        statistic: which goodness-of-fit test to use.
        n_nodes, sigma: quadrature settings for ``"epps_pulley"``.
        center: subtract the batch mean before testing. Off by default,
            because driving the mean to zero is part of the job.

    Returns:
        A scalar, minimised when the embeddings look isotropic Gaussian.
    """
    z = z.reshape(-1, z.shape[-1])
    if center:
        z = z - jnp.mean(z, axis=0, keepdims=True)
    u = z @ random_directions(key, z.shape[-1], n_proj)  # (n, n_proj)
    if statistic == "epps_pulley":
        stats = epps_pulley(u, n_nodes=n_nodes, sigma=sigma)
    elif statistic == "cramer_von_mises":
        stats = cramer_von_mises(u)
    else:
        raise ValueError(f"unknown statistic {statistic!r}")
    return jnp.mean(stats)