Distillation

Activation capture, closed-form / fitted regression, and metrics for distilling a target block into a single layer.

polyweave.distill

Activation-space distillation: fit a single layer to a module's I/O map.

The future GPT-2 / Pythia / SwiGLU experiment replaces a transformer feed-forward sub-block with one position-wise layer, trained by regressing the block's cached (input, output) activation pairs. This sub-package is the model-agnostic machinery for that:

capture    — forward-hook collection of a submodule's (input, output) pairs
metrics    — relative MSE and R^2 (coefficient of determination)
regression — ``fit_layer``: MSE-regress any nn.Module onto cached pairs,
             tracking the multiplicative-recruitment gate if the layer has one

Nothing here imports transformers; capture works on any nn.Module, so the same harness fits a synthetic target in the tests and a real FFN later.

IOCapture

IOCapture(module: Module, *, max_rows: Optional[int] = None, device: str = 'cpu', flatten_leading: bool = True)

Context manager that records inputs/outputs of module during forward.

Usage::

with IOCapture(model.transformer.h[6].mlp) as cap:
    for batch in loader:
        model(batch)          # forward passes drive the hook
X, Y = cap.pairs()            # [N, in], [N, out]

Leading dims (batch, sequence, ...) are flattened so each token becomes one independent training example, matching the position-wise nature of an FFN.

Parameters:

Name	Type	Description	Default
module	Module	the submodule to tap.	required
max_rows	Optional[int]	stop accumulating once this many flattened rows are collected (None = unbounded). Useful to cap a large activation cache.	None
device	str	where to store captured tensors (default "cpu" to spare GPU memory during a long caching pass).	'cpu'
flatten_leading	bool	flatten all but the last dim into rows (default True).	True

Source code in polyweave/distill/capture.py

def __init__(
    self,
    module: nn.Module,
    *,
    max_rows: Optional[int] = None,
    device: str = "cpu",
    flatten_leading: bool = True,
) -> None:
    self.module = module
    self.max_rows = max_rows
    self.device = device
    self.flatten_leading = flatten_leading
    self._xs: List[torch.Tensor] = []
    self._ys: List[torch.Tensor] = []
    self._rows = 0
    self._handle = None

pairs

pairs() -> Tuple[torch.Tensor, torch.Tensor]

Return the concatenated (X, Y) activation pairs.

Source code in polyweave/distill/capture.py

def pairs(self) -> Tuple[torch.Tensor, torch.Tensor]:
    """Return the concatenated ``(X, Y)`` activation pairs."""
    if not self._xs:
        raise RuntimeError("no activations captured; run a forward pass inside the context")
    return torch.cat(self._xs, dim=0), torch.cat(self._ys, dim=0)

DistillResult dataclass

DistillResult(train_losses: List[Tuple[int, float]] = list(), val_mse: float = float('nan'), val_rel_mse: float = float('nan'), val_r2: float = float('nan'), val_rmse: float = float('nan'), val_cosine: float = float('nan'), recruit_curve: List[Tuple[int, float]] = list(), num_params: int = 0)

Outcome of fitting one layer to activation pairs.

recruit_delta property

recruit_delta: Optional[float]

Final − initial recruitment gate, or None if not tracked.

collect_io

collect_io(module: Module, run_fn: Callable[[], None], *, max_rows: Optional[int] = None, device: str = 'cpu') -> Tuple[torch.Tensor, torch.Tensor]

Convenience wrapper: run run_fn under an :class:IOCapture and return pairs.

run_fn should drive whatever forward passes exercise module (e.g. a loop feeding text batches through a language model).

Source code in polyweave/distill/capture.py

def collect_io(
    module: nn.Module,
    run_fn: Callable[[], None],
    *,
    max_rows: Optional[int] = None,
    device: str = "cpu",
) -> Tuple[torch.Tensor, torch.Tensor]:
    """Convenience wrapper: run ``run_fn`` under an :class:`IOCapture` and return pairs.

    ``run_fn`` should drive whatever forward passes exercise ``module`` (e.g. a
    loop feeding text batches through a language model).
    """
    with IOCapture(module, max_rows=max_rows, device=device) as cap:
        run_fn()
    return cap.pairs()

cosine_similarity

cosine_similarity(y_true: Tensor, y_pred: Tensor) -> float

Mean per-row cosine similarity between predicted and target vectors.

Each row (token) contributes the cosine of the angle between its predicted and true activation vectors; we average over rows. Complementary to R²: it ignores per-token magnitude and reads only directional agreement, so a layer can track the shape of the output manifold even where it mis-scales. 1.0 is perfect alignment, 0.0 orthogonal.

Source code in polyweave/distill/metrics.py

@torch.no_grad()
def cosine_similarity(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
    """Mean per-row cosine similarity between predicted and target vectors.

    Each row (token) contributes the cosine of the angle between its predicted
    and true activation vectors; we average over rows. Complementary to R²: it
    ignores per-token magnitude and reads only *directional* agreement, so a
    layer can track the shape of the output manifold even where it mis-scales.
    ``1.0`` is perfect alignment, ``0.0`` orthogonal.
    """
    return torch.nn.functional.cosine_similarity(
        y_pred, y_true, dim=-1, eps=torch.finfo(y_true.dtype).eps
    ).mean().item()

r2_score

r2_score(y_true: Tensor, y_pred: Tensor) -> float

Coefficient of determination R² = 1 - SS_res / SS_tot.

Computed over all elements against the global target mean. 1.0 is a perfect fit; 0.0 matches the constant-mean predictor; negative is worse than predicting the mean.

Source code in polyweave/distill/metrics.py

@torch.no_grad()
def r2_score(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
    """Coefficient of determination ``R² = 1 - SS_res / SS_tot``.

    Computed over all elements against the global target mean. ``1.0`` is a
    perfect fit; ``0.0`` matches the constant-mean predictor; negative is worse
    than predicting the mean.
    """
    ss_res = torch.sum((y_true - y_pred) ** 2)
    ss_tot = torch.sum((y_true - y_true.mean()) ** 2).clamp_min(
        torch.finfo(y_true.dtype).eps
    )
    return (1.0 - ss_res / ss_tot).item()

relative_mse

relative_mse(y_true: Tensor, y_pred: Tensor) -> float

Squared error normalised by target energy: ||y - ŷ||² / ||y||².

Scale-free and 1.0 for the trivial all-zeros predictor (when the target has zero mean), so it reads as a fraction of the signal left unexplained.

Source code in polyweave/distill/metrics.py

@torch.no_grad()
def relative_mse(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
    """Squared error normalised by target energy: ``||y - ŷ||² / ||y||²``.

    Scale-free and ``1.0`` for the trivial all-zeros predictor (when the target
    has zero mean), so it reads as a fraction of the signal left unexplained.
    """
    num = torch.sum((y_true - y_pred) ** 2)
    den = torch.sum(y_true ** 2).clamp_min(torch.finfo(y_true.dtype).eps)
    return (num / den).item()

rmse

rmse(y_true: Tensor, y_pred: Tensor) -> float

Root mean squared error in the raw activation units.

Scale-dependent (unlike relative_mse / r2_score), so it is only comparable across candidates fitted to the same target block — which is exactly how the distillation experiment uses it.

Source code in polyweave/distill/metrics.py

@torch.no_grad()
def rmse(y_true: torch.Tensor, y_pred: torch.Tensor) -> float:
    """Root mean squared error in the raw activation units.

    Scale-dependent (unlike ``relative_mse`` / ``r2_score``), so it is only
    comparable across candidates fitted to the *same* target block — which is
    exactly how the distillation experiment uses it.
    """
    return torch.sqrt(torch.mean((y_true - y_pred) ** 2)).item()

fit_closed_form_linear

fit_closed_form_linear(layer: Module, X: Tensor, Y: Tensor, *, val_frac: float = 0.2, device: str = 'cpu') -> DistillResult

Solve an affine map X -> Y in closed form (least squares) and load it into layer (an nn.Linear), reporting the same held-out metrics as fit_layer.

This is the exact linear baseline: on ill-conditioned transformer activations a first-order optimiser can leave a plain nn.Linear badly underfit (it may need tens of thousands of steps to converge), which silently inflates the apparent nonlinearity of the target. The closed-form solution removes that confound — it is the true linear ceiling against which the multiplicative / depth candidates are judged. Uses the same fixed tail validation split as :func:fit_layer, so the numbers are directly comparable.

Source code in polyweave/distill/regression.py

def fit_closed_form_linear(
    layer: nn.Module,
    X: torch.Tensor,
    Y: torch.Tensor,
    *,
    val_frac: float = 0.2,
    device: str = "cpu",
) -> DistillResult:
    """Solve an affine map ``X -> Y`` in closed form (least squares) and load it into
    ``layer`` (an ``nn.Linear``), reporting the same held-out metrics as ``fit_layer``.

    This is the *exact* linear baseline: on ill-conditioned transformer activations a
    first-order optimiser can leave a plain ``nn.Linear`` badly underfit (it may need
    tens of thousands of steps to converge), which silently inflates the apparent
    nonlinearity of the target. The closed-form solution removes that confound — it is
    the true linear ceiling against which the multiplicative / depth candidates are
    judged. Uses the same fixed tail validation split as :func:`fit_layer`, so the
    numbers are directly comparable.
    """
    if not isinstance(layer, nn.Linear):
        raise TypeError("fit_closed_form_linear expects an nn.Linear layer")
    n = X.shape[0]
    n_val = max(1, int(round(n * val_frac)))
    n_train = max(1, n - n_val)
    Xtr, Ytr = X[:n_train].double(), Y[:n_train].double()
    Xva, Yva = X[n_train:].double(), Y[n_train:].double()

    ones = torch.ones(n_train, 1, dtype=torch.float64)
    W = torch.linalg.lstsq(torch.cat([Xtr, ones], dim=1), Ytr).solution  # [in+1, out]
    with torch.no_grad():
        layer.weight.copy_(W[:-1].T.to(layer.weight.dtype))
        layer.bias.copy_(W[-1].to(layer.bias.dtype))

    onev = torch.ones(Xva.shape[0], 1, dtype=torch.float64)
    pred_va = (torch.cat([Xva, onev], dim=1) @ W).to(Yva.dtype)
    result = DistillResult(num_params=sum(p.numel() for p in layer.parameters()))
    result.val_mse = F.mse_loss(pred_va, Yva).item()
    result.val_rel_mse = relative_mse(Yva, pred_va)
    result.val_r2 = r2_score(Yva, pred_va)
    result.val_rmse = rmse(Yva, pred_va)
    result.val_cosine = cosine_similarity(Yva, pred_va)
    layer.to(device)
    return result

fit_layer

fit_layer(layer: Module, X: Tensor, Y: Tensor, *, steps: int = 2000, lr: float = 0.001, batch_size: int = 256, weight_decay: float = 0.0, val_frac: float = 0.2, eval_every: int = 100, device: str = 'cpu', seed: int = 0, log_fn: Optional[Callable[[str], None]] = None) -> DistillResult

MSE-regress layer onto (X, Y) and report held-out fit + recruitment.

The last val_frac of the rows form a fixed validation split (the inputs are assumed already shuffled at capture time — tokens are independent rows).

Parameters:

Name	Type	Description	Default
layer	Module	module mapping [batch, in] -> [batch, out]; trained in place.	required
X	Tensor	input activations, shape [N, in].	required
Y	Tensor	target activations, shape [N, out].	required
steps	int	optimisation steps (minibatches).	2000
lr	float	AdamW learning rate.	0.001
weight_decay	float	AdamW weight decay.	0.0
batch_size	int	minibatch size sampled (with replacement) from the train split.	256
val_frac	float	fraction of rows held out for validation metrics.	0.2
eval_every	int	record train loss + recruitment gate every this many steps.	100
device	str	compute device.	'cpu'
seed	int	seed for minibatch sampling.	0
log_fn	Optional[Callable[[str], None]]	optional progress sink.	None

Returns:

Name	Type	Description
A	DistillResult	class:DistillResult with the train-loss curve, final validation
	DistillResult	rel_mse / R² / rmse / mean per-row cosine, the
	DistillResult	recruitment-gate curve (if the layer exposes pi_scale_mean /
	DistillResult	quad_scale_mean), and the layer's parameter count.

Source code in polyweave/distill/regression.py

def fit_layer(
    layer: nn.Module,
    X: torch.Tensor,
    Y: torch.Tensor,
    *,
    steps: int = 2000,
    lr: float = 1e-3,
    batch_size: int = 256,
    weight_decay: float = 0.0,
    val_frac: float = 0.2,
    eval_every: int = 100,
    device: str = "cpu",
    seed: int = 0,
    log_fn: Optional[Callable[[str], None]] = None,
) -> DistillResult:
    """MSE-regress ``layer`` onto ``(X, Y)`` and report held-out fit + recruitment.

    The last ``val_frac`` of the rows form a fixed validation split (the inputs
    are assumed already shuffled at capture time — tokens are independent rows).

    Args:
        layer: module mapping ``[batch, in] -> [batch, out]``; trained in place.
        X: input activations, shape ``[N, in]``.
        Y: target activations, shape ``[N, out]``.
        steps: optimisation steps (minibatches).
        lr: AdamW learning rate.
        weight_decay: AdamW weight decay.
        batch_size: minibatch size sampled (with replacement) from the train split.
        val_frac: fraction of rows held out for validation metrics.
        eval_every: record train loss + recruitment gate every this many steps.
        device: compute device.
        seed: seed for minibatch sampling.
        log_fn: optional progress sink.

    Returns:
        A :class:`DistillResult` with the train-loss curve, final validation
        ``rel_mse`` / ``R²`` / ``rmse`` / mean per-row ``cosine``, the
        recruitment-gate curve (if the layer exposes ``pi_scale_mean`` /
        ``quad_scale_mean``), and the layer's parameter count.
    """
    g = torch.Generator(device="cpu").manual_seed(seed)
    layer = layer.to(device)
    X = X.to(device)
    Y = Y.to(device)

    n = X.shape[0]
    n_val = max(1, int(round(n * val_frac)))
    n_train = max(1, n - n_val)
    Xtr, Ytr = X[:n_train], Y[:n_train]
    Xva, Yva = X[n_train:], Y[n_train:]

    recruit = _recruit_fn(layer)
    result = DistillResult(num_params=sum(p.numel() for p in layer.parameters()))

    opt = torch.optim.AdamW(layer.parameters(), lr=lr, weight_decay=weight_decay)

    def record(step: int) -> None:
        layer.eval()
        with torch.no_grad():
            train_mse = F.mse_loss(layer(Xtr), Ytr).item()
        result.train_losses.append((step, train_mse))
        if recruit is not None:
            result.recruit_curve.append((step, recruit()))
        if log_fn is not None:
            msg = f"step {step:5d}  train_mse={train_mse:.5e}"
            if recruit is not None:
                msg += f"  gate={result.recruit_curve[-1][1]:.5f}"
            log_fn(msg)
        layer.train()

    record(0)
    layer.train()
    for step in range(1, steps + 1):
        idx = torch.randint(0, n_train, (min(batch_size, n_train),), generator=g)
        xb, yb = Xtr[idx], Ytr[idx]
        opt.zero_grad()
        loss = F.mse_loss(layer(xb), yb)
        loss.backward()
        opt.step()
        if eval_every and step % eval_every == 0:
            record(step)

    if not result.train_losses or result.train_losses[-1][0] != steps:
        record(steps)

    layer.eval()
    with torch.no_grad():
        pred_va = layer(Xva)
    result.val_mse = F.mse_loss(pred_va, Yva).item()
    result.val_rel_mse = relative_mse(Yva, pred_va)
    result.val_r2 = r2_score(Yva, pred_va)
    result.val_rmse = rmse(Yva, pred_va)
    result.val_cosine = cosine_similarity(Yva, pred_va)
    return result