Sink-content Aitchison distillation

(1) Integration point: add an auxiliary attention-distillation loss to a student Transformer during supervised training or teacher-student compression. Capture teacher and student attention probabilities after softmax for selected layers and heads, with tensors shaped [batch, heads, query_position, key_position]. Treat key position 0 as the sink; for causal attention, compute the loss only over valid unmasked keys. (2) Pseudocode: for each valid row, set s=clamp(p[...,0],eps,1-eps) and t=clamp(q[...,0],eps,1-eps); set c=p[...,1:]/(1-s) and d=q[...,1:]/(1-t), clamp content probabilities to eps, and renormalize them. Compute z_c=log(c)-mean(log(c),last_dim) and z_d=log(d)-mean(log(d),last_dim). Set L_content=mean(sum((z_c-z_d)^2,last_dim)/m). Set L_sink=mean((logit(s)-logit(t))^2), where logit(x)=log(x/(1-x)). Optimize L=L_task+lambda_content L_content+lambda_sink L_sink. (3) The decomposition and log-ratio coordinates come from the paper; epsilon, layer selection, and loss weights are empirical choices. Masked keys must be excluded before content normalization. (4) First experiment: distill a 6-layer, 256-hidden student from BERT-base or a small decoder teacher on GLUE or WikiText-2. Compare ordinary attention KL, content-only Aitchison loss, and the two-channel loss at equal student FLOPs. Record validation accuracy or perplexity, sink-mass error, content Aitchison distance, and gradient stability. Success is lower validation loss at equal compute, especially for teachers with highly concentrated sink attention.