Affine-to-Axis Distillation

Apply this to teacher-student distillation or to two independently trained models that must later be merged. Choose corresponding hidden layers, or define a depth window I(a) when teacher and student depths differ. On a calibration batch, collect teacher preactivations z^A_j and student preactivations z^B_i. First fit a global affine map E(z) = Az + c by least squares from student to teacher, producing weak-alignment residual omega_l^2 = E||E(z^B_l) - z^A_l||_2^2. Then build a cost matrix C_{ji} = min_{alpha in H_phi} E[(z^B_i - alpha z^A_j)^2]/(Var(z^A_j) + epsilon); for ReLU use H_phi = (0,infinity) and estimate alpha = max(epsilon, E[z^A_j z^B_i]/E[(z^A_j)^2]). Solve a Hungarian assignment, retain pairs with cost below theta^2, and add lambda L_axis to the task loss. Recompute the assignment every 100-500 steps, stop-gradient through the assignment, and optionally initialize the student's next-layer weights by the matched rescaling W^B[:,pi(j)] <- W^B[:,pi(j)]/alpha_j. Estimate m_l(epsilon) by the pruning procedure in the first idea; estimate the theorem's conditioning constant empirically from the smallest singular value of the fitted affine map. The first cheap test should use two 3-layer ReLU MLPs on CIFAR-10, with one student trained from scratch under ordinary logit distillation and the other using logit plus axis loss. Measure validation accuracy, calibration, assignment stability, and post-hoc weight permutation or merging quality. A positive signal is faster student convergence or improved accuracy at equal parameters, plus lower merged-model loss after aligning hidden units.