Per-token NLL weights from the cached last-layer hidden states. Standard per-layer PCA accumulates ``Sigma_X = X.T @ X`` with uniform weights, implicitly assuming every token contributes equally to downstream loss. Rare / surprising tokens actually carry far more LM-loss gradient than e
(
model,
final_hiddens: list[torch.Tensor], # list of (1, S, d_model) -- output of last backbone layer
input_ids_list: list[torch.Tensor], # list of (1, S) or (S,) -- token ids
)
| 394 | |
| 395 | @torch.no_grad() |
| 396 | def compute_token_weights_nll( |
| 397 | model, |
| 398 | final_hiddens: list[torch.Tensor], # list of (1, S, d_model) -- output of last backbone layer |
| 399 | input_ids_list: list[torch.Tensor], # list of (1, S) or (S,) -- token ids |
| 400 | ) -> list[torch.Tensor]: |
| 401 | """Per-token NLL weights from the cached last-layer hidden states. |
| 402 | |
| 403 | Standard per-layer PCA accumulates ``Sigma_X = X.T @ X`` with uniform weights, |
| 404 | implicitly assuming every token contributes equally to downstream loss. Rare |
| 405 | / surprising tokens actually carry far more LM-loss gradient than easy ones. |
| 406 | The SparseGPT / OBS prescription is ``Sigma_X = X.T @ diag(w) @ X``, with |
| 407 | ``w_t = NLL_teacher(token_{t+1} | tokens_{<=t})`` approximating the diagonal |
| 408 | of the empirical Hessian of the next-token loss at the final hidden state. |
| 409 | |
| 410 | Implementation: apply ``model.model.norm`` + ``model.lm_head`` + cross-entropy |
| 411 | to the cached final hidden states; weight token ``t`` by the NLL of token |
| 412 | ``t+1`` (last position gets weight 0). Per-sample mean-1 normalisation |
| 413 | preserves the cov scale of the downstream PCA. Works under HF |
| 414 | ``device_map="auto"``: norm / head stay on whichever GPU(s) HF placed them. |
| 415 | """ |
| 416 | norm = model.model.norm |
| 417 | head = model.lm_head |
| 418 | norm_param = next(norm.parameters()) |
| 419 | head_param = next(head.parameters()) |
| 420 | norm_device = norm_param.device |
| 421 | norm_dtype = norm_param.dtype |
| 422 | head_device = head_param.device |
| 423 | |
| 424 | weights: list[torch.Tensor] = [] |
| 425 | for h, ids in zip(final_hiddens, input_ids_list): |
| 426 | h = h.to(device=norm_device, dtype=norm_dtype) |
| 427 | if ids.dim() == 1: |
| 428 | ids = ids.unsqueeze(0) |
| 429 | ids_h = ids.to(head_device) |
| 430 | h_n = norm(h) |
| 431 | if norm_device != head_device: |
| 432 | h_n = h_n.to(head_device) |
| 433 | logits = head(h_n).float() # (1, S, V) |
| 434 | shift_logits = logits[:, :-1, :] |
| 435 | shift_targets = ids_h[:, 1:] |
| 436 | nll = F.cross_entropy( |
| 437 | shift_logits.reshape(-1, shift_logits.shape[-1]), |
| 438 | shift_targets.reshape(-1), |
| 439 | reduction="none", |
| 440 | ).view(1, -1) # (1, S-1) |
| 441 | S = ids_h.shape[-1] |
| 442 | w = torch.zeros(S, dtype=torch.float32, device=head_device) |
| 443 | w[:-1] = nll[0].float() |
| 444 | mean_w = w[:-1].mean().clamp_min(1e-12) |
| 445 | w = w / mean_w |
| 446 | weights.append(w.detach().cpu()) |
| 447 | return weights |
| 448 | |
| 449 | |
| 450 | def diagnose_weights(weights: list[torch.Tensor]) -> dict: |