Pruning and Distillation: Exact Masks and Correct Token Losses

Three different compression operations

Quantization changes numerical representation; pruning chooses structure or connections to remove; distillation trains a student against teacher behavior. Their gains and costs overlap, but they are not interchangeable. The bonsai/apprentice analogy is useful only if we remember that a removed branch can matter and an apprentice does not inherit all the teacher's knowledge automatically.

Pruning tolerance, student quality and quantization loss depend on the model and workload; model size alone does not establish how a model was trained. The right question is a measured size/quality/latency/training-cost frontier for the intended workload. Prerequisites are tensor shapes, causal token loss, optimization and evaluation splits. Outcomes include exact magnitude and2:4 masks, masked shifted KL/CE, temperature derivatives, a complete synthetic CPU training experiment and a resource-bounded, unexecuted reference procedure for larger models.

Magnitude pruning with an exact mask

Select eligible tensors explicitly. “Every two-dimensional parameter containing weight” can accidentally include embeddings and is not equivalent to “Linear layers”. For M eligible scalar weights and target fraction s, choose exactly floor(s*M) entries by increasing absolute magnitude, with a stable tie rule. A quantile threshold with >= may retain too many tied values; at s=1 it can even keep the maximum. Existing zeros mean actual numerical zeros may exceed masked entries, so report both.

Masking a dense tensor does not change its shape, storage format or dense GEMM operation count. It becomes a storage win only after a suitable compressed representation, whose indices/metadata also cost bytes. A dense kernel generally does not skip arbitrary zero values. Structured removal of channels/heads/layers changes shapes but must consistently update linked projections, residual dimensions, normalization and model config. Sparse kernels can accelerate some unstructured workloads; “only2:4 can ever be accelerated” is false.

During recovery fine-tuning, apply the mask in forward computation or mask gradients and reapply after optimizer steps. Momentum/Adam state can regrow zeros unless the effective forward weight remains masked; decide whether state for permanently removed entries is cleared. PyTorch pruning commonly uses weight_orig, a weight_mask buffer and a pre-forward hook. Removing the reparameterization makes the current masked weight ordinary; it does not create a compact sparse checkpoint by magic.

Optimizer-state example: suppose a masked scalar weight is zero, its current gradient is zero, but its momentum buffer is 2. With momentum 0.9 and learning rate 0.1, the next ordinary momentum step writes 0 - 0.1*(0.9*2) = -0.18. Gradient masking alone therefore does not preserve a physical zero. Multiplying by the fixed mask during forward keeps the effective weight zero; reapplying the mask after every step also restores it. Neither turns the dense parameter into compressed storage.

Iterative magnitude pruning repeatedly removes some remaining weights. Removing20% of survivors twice retains0.8²=0.64, hence36% cumulative sparsity, not40%. Lottery-ticket experiments also concern initialization/rewinding and training conditions; arbitrary one-shot masking is not proof of the lottery-ticket hypothesis. Use independent validation to choose a sparsity level and freeze the final test.

2:4 sparsity, hardware and reconstruction-aware pruning

For2:4 along the supported reduction axis, retain at most two nonzeros in every four-weight group. This gives a regular metadata/layout pattern compatible with some Sparse Tensor Core paths. A tensor full of2:4 zeros is not yet a sparse-kernel invocation. Convert to the supported semi-structured representation, verify matrix dimensions, dtype, device and operation support, then compare against a dense reference with the same pruned values before timing. Hardware peak sparse throughput is not an end-to-end2× speed guarantee; attention, KV traffic, normalization, conversion and launches remain.

SparseGPT uses calibration activations to approximately reconstruct layer outputs while pruning and compensating remaining weights. If a layer maps X to WX, a local quadratic objective is ||WX-W_hat X||_F²; its curvature contains XXᵀ. With one weight constrained to zero, inverse-curvature information estimates how other weights can compensate. Damping stabilizes nearly singular calibration curvature. This is an approximation, not a global joint optimum of masks and all network weights. The paper reports one-shot50% sparsity results for large models under its setups; it does not establish a universal cliff or fixed latency improvement.

A resource procedure: obtain a reviewed local checkpoint and calibration split; freeze validation/test; hook only target Linear layers to collect bounded activation batches; execute the pinned reference implementation's reconstruction/pruning pass; validate mask pattern, finite outputs and held-out loss after each stage; serialize the mask/layout/config; then benchmark a compatible sparse backend. Use the published repository instructions rather than assuming pip install sparsegpt identifies a supported distribution. Wanda instead uses weight magnitude weighted by input activation norms. Movement pruning learns importance scores from first-order weight/gradient movement during fine-tuning, not simply the absolute distance between initial/final weights. Neither is universally better across tasks.

Distillation loss and its derivative

For teacher logits a and student logits z on the same vocabulary and target position, set q=softmax(a/T) and p=softmax(z/T). A common token loss is

Teacher probabilities are detached; the teacher is in evaluation mode. The soft-term gradient with respect to student logit z_j is T(p_j-q_j) after the T² multiplier. Without it the derivative is (p_j-q_j)/T. At high temperature, centered logits give p_j-q_j≈((z_j-z_mean)-(a_j-a_mean))/(V*T), so the raw gradient is approximately proportional to 1/T². This is the regime motivating the multiplier; it does not keep gradients exactly identical for every T.

T=1 does not make a soft teacher distribution one-hot. As T grows, probabilities approach uniform, but the T²-rescaled logit-matching limit can retain a nonzero gradient. Pick temperature and alpha on validation; “higher T always requires higher alpha” is not a theorem. Hard-label cross-entropy also updates every logit through p-onehot, and one token identity carries up to log2(V) bits, not exactly one bit. A fixed teacher's soft targets do not move merely because the student improves, and teacher errors/noise can still transfer.

For a causal LM, logits at position t predict the token at t+1. Slice student_logits[:,:-1], teacher_logits[:,:-1] and labels[:,1:], then apply the same valid-target mask to CE and KL. Summing KL across (batch,time,vocab) with reduction='batchmean' divides only by batch size and grows with sequence length; use per-token vocabulary sum and divide by the number of valid target tokens. Masking only CE while distilling on padded/prompt positions changes the objective. Reject an all-ignored batch rather than producing a misleading zero/NaN.

Exact CPU pruning and distillation experiment

This trains a tiny student on generated token sequences and a fixed randomly initialized teacher. It demonstrates mechanics, not capability transfer from a pretrained LM. It also proves tie-stable exact pruning,2:4 counts, shifted masking and all-ignored rejection. No weights or data are downloaded.

import math
import torch
import torch.nn as nn
import torch.nn.functional as F
torch.manual_seed(68)
torch.set_num_threads(1)
 
def magnitude_mask(w,sparsity):
    if not math.isfinite(sparsity) or not 0<=sparsity<=1: raise ValueError('sparsity')
    if not torch.isfinite(w).all(): raise ValueError('finite weights required')
    count=math.floor(w.numel()*sparsity)
    order=torch.argsort(w.detach().abs().flatten(),stable=True)
    mask=torch.ones(w.numel(),dtype=torch.bool,device=w.device)
    mask[order[:count]]=False
    return mask.reshape_as(w)
 
def mask24(w):
    if w.ndim!=2 or w.shape[1]%4: raise ValueError('groups of four')
    if not torch.isfinite(w).all(): raise ValueError('finite weights required')
    groups=w.detach().reshape(w.shape[0],-1,4)
    order=torch.argsort(groups.abs(),dim=-1,descending=True,stable=True)
    keep=torch.zeros_like(groups,dtype=torch.bool)
    keep.scatter_(-1,order[...,:2],True)
    return keep.reshape_as(w)
 
def distill_loss(student,teacher,labels,T=2.,alpha=0.7):
    if not math.isfinite(T) or T<=0 or not math.isfinite(alpha) or not 0<=alpha<=1:
        raise ValueError('finite T/alpha required')
    if student.ndim!=3 or student.shape!=teacher.shape or labels.shape!=student.shape[:2]:
        raise ValueError('aligned logits/targets required')
    z=student[:,:-1]; a=teacher.detach()[:,:-1]; y=labels[:,1:]
    valid=y.ne(-100)
    if not valid.any(): raise ValueError('no supervised target')
    if not torch.isfinite(z[valid]).all() or not torch.isfinite(a[valid]).all():
        raise ValueError('finite supervised logits required')
    ce=F.cross_entropy(z[valid],y[valid])
    q=F.softmax(a[valid]/T,-1)
    kl=F.kl_div(F.log_softmax(z[valid]/T,-1),q,reduction='none').sum(-1).mean()*T*T
    return (1-alpha)*ce+alpha*kl
 
w=torch.ones(2,8,dtype=torch.float64)
assert (~magnitude_mask(w,0.5)).sum()==8
assert magnitude_mask(w,0).all() and not magnitude_mask(w,1).any()
m=mask24(w); assert (m.reshape(2,2,4).sum(-1)==2).all()
assert torch.count_nonzero(w*m)==8
# A dense tensor is still the same number of stored elements.
assert (w*m).numel()*w.element_size()==w.numel()*w.element_size()
 
class TinyLM(nn.Module):
    def __init__(self,width):
        super().__init__()
        self.emb=nn.Embedding(9,width,dtype=torch.float64)
        self.out=nn.Linear(width,9,dtype=torch.float64)
    def forward(self,ids):
        h=self.emb(ids).cumsum(1)
        den=torch.arange(1,ids.shape[1]+1,dtype=h.dtype)[None,:,None]
        return self.out(torch.tanh(h/den))
ids=torch.randint(0,9,(8,6))
labels=ids.clone(); labels[:,:2]=-100; labels[0,-1]=-100
teacher=TinyLM(12).eval(); student=TinyLM(5)
for p in teacher.parameters(): p.requires_grad_(False)
with torch.no_grad(): target=teacher(ids)
opt=torch.optim.Adam(student.parameters(),lr=0.03)
initial=distill_loss(student(ids),target,labels).item()
for _ in range(140):
    opt.zero_grad(); loss=distill_loss(student(ids),target,labels)
    loss.backward(); opt.step()
assert loss.item()<initial
z=student(ids).detach().requires_grad_()
loss=distill_loss(z,target,labels); loss.backward()
assert torch.equal(z.grad[:,-1],torch.zeros_like(z.grad[:,-1]))
assert torch.equal(z.grad[:,0],torch.zeros_like(z.grad[:,0]))
assert torch.equal(z.grad[0,-2],torch.zeros_like(z.grad[0,-2]))
# Changing an ignored predicting position cannot change the loss.
changed=z.detach().clone(); changed[:,0]+=100*torch.randn_like(changed[:,0])
assert torch.allclose(distill_loss(changed,target,labels),loss)
assert not torch.allclose(distill_loss(z,target,labels),distill_loss(z,target,labels.roll(1,1)))
try: distill_loss(z,target,torch.full_like(labels,-100))
except ValueError: pass
else: raise AssertionError('all ignored accepted')
# Exact soft-loss derivative, independent of causal slicing.
T=3.; zz=torch.tensor([0.2,-0.4,1.1],dtype=torch.float64,requires_grad=True)
aa=torch.tensor([1.,0.,-1.],dtype=torch.float64)
q=(aa/T).softmax(0)
soft=T*T*(q*(q.log()-(zz/T).log_softmax(0))).sum(); soft.backward()
assert torch.allclose(zz.grad,T*((zz.detach()/T).softmax(0)-q))
print('initial',initial,'final',loss.item(),'pruning, shift/mask and derivative checks passed')

Full local-model training reference and vocabulary contract

Before online logit distillation, compare tokenizer vocabulary ID mappings, special IDs, normalization, chat templates and sequence alignment—not merely model-family names or equal vocabulary sizes. A shared token ID must mean the same event. Feature projection does not repair unrelated output vocabularies automatically. For incompatible tokenizers, sampled-text distillation lets each student tokenize the teacher's completed text under a separately reviewed SFT objective; more sophisticated cross-tokenizer distribution alignment is a different method.

The following continuation reuses the exact distill_loss above. It expects Transformers4.57.1, two local compatible checkpoints, a reviewed tokenizer and a local trusted tensor file with input_ids, attention_mask, labels. Prepare labels with prompt/pad targets set to-100; rows are right-padded. This larger-model reference was not executed; the generated CPU experiment above was.

# Resource reference: local-model continuation; not executed.
import os
os.environ['HF_HUB_OFFLINE']='1'
os.environ['TRANSFORMERS_OFFLINE']='1'
from transformers import AutoModelForCausalLM,AutoTokenizer
from torch.utils.data import TensorDataset,DataLoader
teacher_path='/srv/models/teacher'; student_path='/srv/models/student'
tt=AutoTokenizer.from_pretrained(teacher_path,local_files_only=True)
st=AutoTokenizer.from_pretrained(student_path,local_files_only=True)
if tt.get_vocab()!=st.get_vocab() or tt.all_special_ids!=st.all_special_ids:
    raise ValueError('token event spaces differ')
if tt.chat_template!=st.chat_template: raise ValueError('template mismatch')
teacher=AutoModelForCausalLM.from_pretrained(teacher_path,local_files_only=True,
    trust_remote_code=False,use_safetensors=True).to('cuda:0').eval()
student=AutoModelForCausalLM.from_pretrained(student_path,local_files_only=True,
    trust_remote_code=False,use_safetensors=True).to('cuda:1').train()
for p in teacher.parameters(): p.requires_grad_(False)
data=torch.load('/srv/data/distill-tensors.pt',map_location='cpu',weights_only=True)
ids,mask,labels=(data[k] for k in ('input_ids','attention_mask','labels'))
if ids.shape!=mask.shape or ids.shape!=labels.shape: raise ValueError('shapes')
if not ((mask==0)|(mask==1)).all(): raise ValueError('mask')
if ((mask[:,1:]-mask[:,:-1])>0).any(): raise ValueError('right padding required')
if not (labels[mask==0]==-100).all(): raise ValueError('unmasked padding')
loader=DataLoader(TensorDataset(ids,mask,labels),batch_size=1,shuffle=True,
    generator=torch.Generator().manual_seed(68))
opt=torch.optim.AdamW(student.parameters(),lr=2e-5)
for input_ids,attention_mask,labels in loader:
    with torch.no_grad():
        target=teacher(input_ids=input_ids.to('cuda:0'),
            attention_mask=attention_mask.to('cuda:0'),use_cache=False).logits
    logits=student(input_ids=input_ids.to('cuda:1'),
        attention_mask=attention_mask.to('cuda:1'),use_cache=False).logits
    opt.zero_grad(set_to_none=True)
    loss=distill_loss(logits.float(),target.to('cuda:1').float(),labels.to('cuda:1'))
    if not torch.isfinite(loss): raise ValueError('nonfinite loss')
    loss.backward(); torch.nn.utils.clip_grad_norm_(student.parameters(),1.)
    opt.step()
student.save_pretrained('./distilled-student',safe_serialization=True)
st.save_pretrained('./distilled-student')

Two devices are this reference's explicit placement, not a theoretical requirement: sequential teacher/student execution, cached teacher logits or an adequately sized single device are alternatives. Online full-vocabulary logits can dominate transfer/memory; top-k storage changes the objective unless tail mass is handled. Teacher inference, student optimization and checkpoint/data provenance all count in training cost. Evaluate a held-out split after checkpointing, compare to the unchanged student and a matched-budget hard-label SFT control, then test deployment kernels separately.

Feature, relation, response and on-policy distillation

Response/logit distillation matches output distributions. Sampled-completion distillation uses teacher-generated text as SFT targets and does not expose the full distribution. Alpaca used52K instruction-following examples generated with text-davinci-003; do not relabel that as a specific GPT-3.5 chat endpoint. Generated traces can teach behavior but also transfer incorrect reasoning. Filter and verify final answers without treating a fluent rationale as evidence of correctness.

Feature distillation aligns selected teacher/student hidden states, often using a learned projection when widths differ. Define layer correspondence, normalization and a valid-token denominator. Relation distillation matches relationships such as normalized distances/angles or attention relations; define which samples/tokens and avoid comparing incompatible padding. There is no source-backed universal claim that most LLM methods use logits rather than features or that one loss accounts for90% of DistilBERT, TinyLlama and MiniLM. Those projects have different training recipes.

Off-policy targets come from a fixed corpus/teacher trajectory. On-policy methods query the teacher on student-generated prefixes, reducing the mismatch between training contexts and contexts the student visits at inference. The paper arXiv2306.13649 describes generalized/on-policy knowledge distillation. On-policy generation costs compute and can visit poor/unsafe trajectories; bound lengths and evaluate distribution shift rather than assuming universal sample efficiency.

Forward KL, KL(q_teacher||p_student), penalizes missing teacher-supported modes. Reverse KL, KL(p_student||q_teacher), weights the student's own distribution and can emphasize modes it assigns mass to. In both cases the student is optimized and teacher fixed; reverse KL does not mean training the teacher to match the student. MiniLLM develops a specific reverse-KL/on-policy method, not a guaranteed3–5-point MT-Bench improvement for every student.

Attributed models and compression economics

DistilBERT's abstract reports40% smaller,97% retained language-understanding capability and60% faster under its evaluations. “At40% of original size” is a different claim. Meta's Llama 3.2 model card states that its 1B/3B models use logits from Llama 3.1 8B and 70B as token-level targets during pretraining, with distillation after pruning to recover performance; this does not prove a universal depth-over-width rule or a precise0.5-point cost for removing layer25. Sheared LLaMA studies structured pruning plus continued pretraining; ShortGPT studies layer redundancy under specified evaluations. Neither grants permission to remove arbitrary layers without validation.

DeepSeek-R1's report distinguishes R1-Zero's RL-only exploration from R1's multistage training. Its released distilled Qwen/Llama students are fine-tuned on800K curated samples; the report's Qwen7B AIME2024 pass@1 is55.5 under the stated evaluation settings. The report uses sampling temperature0.6 and top-p0.95, estimates pass@1 from sampled responses, and separately reports AIME consensus over64samples; these metrics must not be conflated. It is not a reproduced result here, a “free” transfer, or evidence that the student inherits all capabilities. A matched-token scratch comparison must also acknowledge the teacher's prior training and teacher-generation cost; equal student tokens do not equal total compute.

For an illustrative compression calculation,70B BF16 weight payload is140GB. A7B INT4 payload is3.5GB before scales/exemptions:40× smaller, not35× by arithmetic necessity. KV cache/activations and quality do not follow this ratio. If a one-time distillation costs C currency units and saves d per served request, simple break-even volume is C/d for positive d, before maintenance and evaluation. A student held back on a domain may need better data/teacher coverage, greater capacity or its own supervised training—not automatically more KL steps.

Choose quantization when supported low-bit kernels and calibration/quality budgets fit; structured pruning when shapes and recovery are viable; distillation when a suitable smaller student and enough amortized volume justify training; MoE when sparse activation suits a trained architecture and hardware. MoE reduces active compute relative to total parameters but does not automatically reduce resident total weight memory. There is no universal “always quantize first” mandate or guaranteed narrow-domain teacher match. Change one axis at a time, then evaluate interactions.

Worked exercises and operational acceptance

1. 50% magnitude mask with ties? Stable-sort eligible magnitudes and mask exactly floor(M/2); report masked count and actual zeros independently. A quantile threshold is not exact under ties.
2. Why no dense-GPU speedup from zeros? Same dense layout and kernel dimensions; compressed storage and suitable kernels are separate actions. Whole-layer/channel pruning can also accelerate hardware outside2:4 paths.
3. Why T²? It compensates approximate high-temperature gradient shrinkage; the exact derivative is T(p-q). It is not exact temperature invariance and T=1 still permits soft targets.
4. What does response masking mean? Shift logits/labels first, then mask the target positions identically for CE and KL; average by valid targets, not just batch size.
5. Sparsity sweep recipe? Restore the same baseline for each fraction, use the same fixed evaluation windows/token denominator, compare one-shot and equal-budget recovery, and report loss/perplexity, allocated/serialized bytes and actual latency separately. No fixed cliff location is predicted.
6. 2:4 timing recipe? Warm up dense and converted sparse matmuls on compatible hardware, synchronize timing, verify numerical agreement with the pruned dense reference, include conversion separately, repeat across actual serving shapes, and then measure full-model quality/latency. Do not assert1.8× before measuring.
7. Reasoning-distillation stretch? Create a rights-cleared, training-only problem set; generate bounded local teacher traces; verify answers, deduplicate/group-split, train the student and hard-label baseline at matched student budget, and evaluate untouched problems.200 examples need not yield a detectable gain.
8. Is the teacher an upper bound? No universal bound: student data/inductive biases can improve some tasks, but teacher errors and missing coverage are risks. Report per-capability gains and losses, not only perplexity.

Primary sources and boundaries

PyTorch2.9 sparse tensors, pruning tutorial, SparseGPT, Movement Pruning, Hinton distillation, DistilBERT, on-policy distillation, MiniLLM, Sheared LLaMA, and DeepSeek-R1. Primary passages and numerical oracles are separate from the exact CPU experiment. No sparse GPU kernel, large-model recovery, external corpus, teacher service or deployment benchmark was executed.

Course navigation

Series: Deep Learning & LLMs From Scratch — 80 sessions · Session 68 / 80 · Module M11 · ~2 hours

Previous: ← DL S067 · Next: DL S069 →

Optional reading

Use these links for additional background. Version-specific behavior is documented in the primary references above.

• Optional resource 1
• Optional resource 2
• Optional resource 3
• Optional resource 4