RoKA: Robust Knowledge Alignment for Cross-Width Hidden-State Distillation

Distilling hidden states from a wide teacher to a narrower student requires comparing representations with unequal widths. Learned projections address this mismatch but add parameters and geometric bias, whereas relational objectives such as centered kernel alignment (CKA) avoid coordinate-wise matching altogether. We introduce RoKA, a framework for cross-width hidden-state distillation that separates two coupled design choices: the geometry of dimensional alignment and the influence of the resulting residuals on optimization. RoKA combines fixed, learned, and spectral alignment operators with a Student-t-inspired radial penalty that attenuates the gradient contribution of large operator-induced residuals. We further show that even fixed coordinates introduced by dimensional expansion can alter gradients on the retained student coordinates. Across BERT, BART, and GPT-2 teacher–student pairs, we evaluate RoKA against normal knowledge distillation (Normal KD) and CKA over five seeds. In the encoder-only complete-method comparison, student-anchored SVD achieves the highest unweighted nine-metric average, 80.13 against CKA’s 78.59. On XSUM, one padding outperforms CKA on all ROUGE metrics at 1.05× versus CKA’s 2.30× Normal-KD training time in a single-GPU probe. At equal hidden coefficients, robust residual optimization outperforms the corresponding quadratic objective on XSUM; on MNLI, the residual objectives perform similarly, although learned ν exceeds fixed ν = 1 by 0.91 points despite approaching that bound. CKA remains strongest on CNN/DailyMail, and decoder-only differences are small, indicating that no single alignment operator dominates across settings. Overall, RoKA positions explicit alignment geometry and robust residual influence as practical, interacting design axes for cross-width hidden-state distillation.