Whitening

igl.whitening

Metric estimation and target whitening.

Whitening a reconstruction target by a metric square root makes plain least squares optimize the second-order expansion of the geometry the metric encodes — distillation as a metric change. The pieces compose::

from igl.whitening import TargetWhitener, WhitenedMSELoss, fisher_pullback

g = fisher_pullback(w_unembed, states)          # metric of the softmax read-out
whitener = TargetWhitener(g).fit(states)        # centering + G^{1/2} + unit scale
loss = WhitenedMSELoss(whitener)                # plugs into MatryoshkaTrainer

The loss whitens in target(), so the trainer's closed-form inner solve, validation, early stopping, and the dimension curve all operate in the whitened geometry with no trainer changes.

TargetWhitener

Whiten targets by a metric square root: y_w = ((y - mu) @ a) / y_scale.

Fitting computes the mean of y, the symmetric square-root pair of the metric, and a scalar RMS scale so whitened targets are unit-scale for the trainer. inverse_transform undoes all three. With metric=None the whitener degrades to centering plus unit scaling.

Parameters:

Name	Type	Description	Default
metric	Tensor | None	Symmetric PSD metric [C, C] over the target space, or None for the identity metric.	None
clamp	float	Eigenvalue floor forwarded to :func:igl.whitening.psd_sqrt_inv.	1e-06

Attributes:

Name	Type	Description
mu_		Target mean [C] (post-fit).
a_		Metric square root [C, C] (post-fit).
a_inv_		Inverse square root [C, C] (post-fit).
y_scale_		Scalar RMS of the centered, whitened targets (post-fit).

is_fitted property

Whether :meth:fit has run.

fit(y)

Fit the whitening constants on targets [N, C].

Parameters:

Name	Type	Description	Default
y	Tensor	Targets [N, C].	required

Returns:

Type	Description
Self	self, for chaining.

from_state_dict(state) classmethod

Rebuild a fitted whitener from :meth:state_dict output.

inverse_transform(y_w)

Undo :meth:transform.

state_dict()

Serialize the fitted constants as a flat tensor dict.

transform(y)

Whiten targets: center, rotate-and-scale by a_, unit-scale.

WhitenedMSELoss

MSE in a whitened target space: second-order distillation as least squares.

Wraps a fitted :class:TargetWhitener; whitening happens in :meth:target, so it flows through every trainer touchpoint — the closed-form inner solve, the per-batch loss, validation, early stopping, and :func:igl.eval_dimension_curve — with no trainer changes. When the whitener's metric is the Fisher pullback of a softmax read-out, the optimized quantity is the second-order expansion of the downstream KL.

Parameters:

Name	Type	Description	Default
whitener	TargetWhitener	A fitted :class:TargetWhitener.	required

Attributes:

Name	Type	Description
higher_is_better	bool	Always False — the metric is whitened MSE.

Raises:

Type	Description
IGLNotFittedError	If whitener is not fitted.

curve_score(pred, target)

Whitened MSE — already non-saturating, so identical to :meth:metric.

target(y)

Whiten raw targets [N, C] (1-D targets are lifted to [N, 1]).

damped_metric(g, m, *, lam=0.1)

Blend a sampled metric with a trace-scaled damping metric.

g + lam * (trace(g) / trace(m)) * m — the Tikhonov damping of natural-gradient practice. A sampled Fisher estimate leaves the near-null spectrum ill-conditioned; a floor metric m (typically :func:logit_metric) regularises it without changing the leading geometry.

Parameters:

Name	Type	Description	Default
g	Tensor	The metric to damp [C, C].	required
m	Tensor	The damping metric [C, C].	required
lam	float	Damping weight relative to the trace ratio.	0.1

Returns:

Type	Description
Tensor	The damped metric [C, C].

fisher_pullback(w, states, *, n_sub=4096, batch=256, generator=None)

Estimate the Fisher pullback of KL through a softmax read-out.

G = w^T E[diag(p) - p p^T] w with p = softmax(states @ w^T), estimated on a subsample of states. This is the second-order expansion of KL(p(h) || p(h_hat)) in h_hat - h, averaged over the state distribution: softmax saturation annihilates the logit-shift direction exactly and down-weights directions that only move improbable logits.

Parameters:

Name	Type	Description	Default
w	Tensor	Read-out matrix [V, C].	required
states	Tensor	States [N, C] defining the empirical distribution.	required
n_sub	int	Subsample size for the expectation.	4096
batch	int	Batch size for the accumulation.	256
generator	Generator | None	Optional RNG for the subsample, for reproducibility.	None

Returns:

Type	Description
Tensor	The estimated metric [C, C], symmetric PSD up to sampling noise.

logit_metric(w)

Pull the Euclidean logit geometry back to the target space.

Parameters:

Name	Type	Description	Default
w	Tensor	Read-out matrix [V, C] mapping targets to logits.	required

Returns:

Type	Description
Tensor	G = w^T w of shape [C, C]: every logit direction counts
Tensor	equally, kernel directions of w are free.

psd_sqrt_inv(g, *, clamp=1e-06)

Compute the square root and inverse square root of a symmetric PSD matrix.

A single eigendecomposition of the symmetrised input is used for both outputs, with eigenvalues clamped to clamp * lambda_max so the inverse root is the exact inverse of the root on the retained spectrum. Sampled metrics (e.g. :func:igl.whitening.fisher_pullback) leave the near-null spectrum ill-conditioned; the clamp keeps both factors finite.

Parameters:

Name	Type	Description	Default
g	Tensor	Symmetric PSD matrix [C, C].	required
clamp	float	Eigenvalue floor relative to the largest eigenvalue.	1e-06

Returns:

Type	Description
Tensor	A tuple (a, a_inv) with a = g^{1/2} and a_inv = g^{-1/2},
Tensor	both symmetric [C, C].

Raises:

Type	Description
IGLConfigError	If g is not a square 2-D tensor or clamp is not in (0, 1).

tail_metric(fn, states, g, *, n_probes=64, generator=None)

Pull a metric back through a differentiable map by Hutchinson probing.

Estimates E[J^T g J] with J = d fn / d x averaged over states, without ever forming J: for probe vectors u ~ N(0, I), E_u[(J^T g^{1/2} u)(J^T g^{1/2} u)^T] = J^T g J. One backward pass per probe. In the paper use case fn runs the remaining transformer blocks and g is the head's Fisher metric, giving each layer the metric of its own consumer.

Parameters:

Name	Type	Description	Default
fn	Callable[[Tensor], Tensor]	Differentiable map from [N, C_in] to [N, C_out].	required
states	Tensor	Points [N, C_in] at which to average the pullback.	required
g	Tensor	Metric on the output space [C_out, C_out].	required
n_probes	int	Number of Hutchinson probes.	64
generator	Generator | None	Optional RNG for the probes, for reproducibility.	None

Returns:

Type	Description
Tensor	The mean-field pullback metric [C_in, C_in].

igl.whitening.psd_sqrt_inv(g, *, clamp=1e-06)

Compute the square root and inverse square root of a symmetric PSD matrix.

A single eigendecomposition of the symmetrised input is used for both outputs, with eigenvalues clamped to clamp * lambda_max so the inverse root is the exact inverse of the root on the retained spectrum. Sampled metrics (e.g. :func:igl.whitening.fisher_pullback) leave the near-null spectrum ill-conditioned; the clamp keeps both factors finite.

Parameters:

Name	Type	Description	Default
g	Tensor	Symmetric PSD matrix [C, C].	required
clamp	float	Eigenvalue floor relative to the largest eigenvalue.	1e-06

Returns:

Type	Description
Tensor	A tuple (a, a_inv) with a = g^{1/2} and a_inv = g^{-1/2},
Tensor	both symmetric [C, C].

Raises:

Type	Description
IGLConfigError	If g is not a square 2-D tensor or clamp is not in (0, 1).

igl.whitening.logit_metric(w)

Pull the Euclidean logit geometry back to the target space.

Parameters:

Name	Type	Description	Default
w	Tensor	Read-out matrix [V, C] mapping targets to logits.	required

Returns:

Type	Description
Tensor	G = w^T w of shape [C, C]: every logit direction counts
Tensor	equally, kernel directions of w are free.

igl.whitening.fisher_pullback(w, states, *, n_sub=4096, batch=256, generator=None)

Estimate the Fisher pullback of KL through a softmax read-out.

G = w^T E[diag(p) - p p^T] w with p = softmax(states @ w^T), estimated on a subsample of states. This is the second-order expansion of KL(p(h) || p(h_hat)) in h_hat - h, averaged over the state distribution: softmax saturation annihilates the logit-shift direction exactly and down-weights directions that only move improbable logits.

Parameters:

Name	Type	Description	Default
w	Tensor	Read-out matrix [V, C].	required
states	Tensor	States [N, C] defining the empirical distribution.	required
n_sub	int	Subsample size for the expectation.	4096
batch	int	Batch size for the accumulation.	256
generator	Generator | None	Optional RNG for the subsample, for reproducibility.	None

Returns:

Type	Description
Tensor	The estimated metric [C, C], symmetric PSD up to sampling noise.

igl.whitening.damped_metric(g, m, *, lam=0.1)

Blend a sampled metric with a trace-scaled damping metric.

g + lam * (trace(g) / trace(m)) * m — the Tikhonov damping of natural-gradient practice. A sampled Fisher estimate leaves the near-null spectrum ill-conditioned; a floor metric m (typically :func:logit_metric) regularises it without changing the leading geometry.

Parameters:

Name	Type	Description	Default
g	Tensor	The metric to damp [C, C].	required
m	Tensor	The damping metric [C, C].	required
lam	float	Damping weight relative to the trace ratio.	0.1

Returns:

Type	Description
Tensor	The damped metric [C, C].

igl.whitening.tail_metric(fn, states, g, *, n_probes=64, generator=None)

Pull a metric back through a differentiable map by Hutchinson probing.

Estimates E[J^T g J] with J = d fn / d x averaged over states, without ever forming J: for probe vectors u ~ N(0, I), E_u[(J^T g^{1/2} u)(J^T g^{1/2} u)^T] = J^T g J. One backward pass per probe. In the paper use case fn runs the remaining transformer blocks and g is the head's Fisher metric, giving each layer the metric of its own consumer.

Parameters:

Name	Type	Description	Default
fn	Callable[[Tensor], Tensor]	Differentiable map from [N, C_in] to [N, C_out].	required
states	Tensor	Points [N, C_in] at which to average the pullback.	required
g	Tensor	Metric on the output space [C_out, C_out].	required
n_probes	int	Number of Hutchinson probes.	64
generator	Generator | None	Optional RNG for the probes, for reproducibility.	None

Returns:

Type	Description
Tensor	The mean-field pullback metric [C_in, C_in].

igl.whitening.TargetWhitener

Whiten targets by a metric square root: y_w = ((y - mu) @ a) / y_scale.

Fitting computes the mean of y, the symmetric square-root pair of the metric, and a scalar RMS scale so whitened targets are unit-scale for the trainer. inverse_transform undoes all three. With metric=None the whitener degrades to centering plus unit scaling.

Parameters:

Name	Type	Description	Default
metric	Tensor | None	Symmetric PSD metric [C, C] over the target space, or None for the identity metric.	None
clamp	float	Eigenvalue floor forwarded to :func:igl.whitening.psd_sqrt_inv.	1e-06

Attributes:

Name	Type	Description
mu_		Target mean [C] (post-fit).
a_		Metric square root [C, C] (post-fit).
a_inv_		Inverse square root [C, C] (post-fit).
y_scale_		Scalar RMS of the centered, whitened targets (post-fit).

is_fitted property

Whether :meth:fit has run.

fit(y)

Fit the whitening constants on targets [N, C].

Parameters:

Name	Type	Description	Default
y	Tensor	Targets [N, C].	required

Returns:

Type	Description
Self	self, for chaining.

from_state_dict(state) classmethod

Rebuild a fitted whitener from :meth:state_dict output.

inverse_transform(y_w)

Undo :meth:transform.

state_dict()

Serialize the fitted constants as a flat tensor dict.

transform(y)

Whiten targets: center, rotate-and-scale by a_, unit-scale.

igl.whitening.WhitenedMSELoss

MSE in a whitened target space: second-order distillation as least squares.

Wraps a fitted :class:TargetWhitener; whitening happens in :meth:target, so it flows through every trainer touchpoint — the closed-form inner solve, the per-batch loss, validation, early stopping, and :func:igl.eval_dimension_curve — with no trainer changes. When the whitener's metric is the Fisher pullback of a softmax read-out, the optimized quantity is the second-order expansion of the downstream KL.

Parameters:

Name	Type	Description	Default
whitener	TargetWhitener	A fitted :class:TargetWhitener.	required

Attributes:

Name	Type	Description
higher_is_better	bool	Always False — the metric is whitened MSE.

Raises:

Type	Description
IGLNotFittedError	If whitener is not fitted.

curve_score(pred, target)

Whitened MSE — already non-saturating, so identical to :meth:metric.

target(y)

Whiten raw targets [N, C] (1-D targets are lifted to [N, 1]).