# Affine-to-Axis Distillation - ID: 583 - Canonical URL: https://synthcore.org/idea/583/affine-to-axis-distillation - API JSON: https://synthcore.org/api/idea/583.json - API Markdown: https://synthcore.org/api/idea/583.md - Verification status: mechanism_failed - Source: [arXiv:2607.08561](https://arxiv.org/abs/2607.08561) - Category: regularization - Solves: accuracy, generalization, sample-efficiency - ML areas: distillation, fine-tuning, mlp, regularization - Math tags: geometry, linear-algebra, approximation-theory, optimization - Ratings: usefulness 6/10; difficulty 5/10; novelty 6/10 ## Idea description Use a teacher and student with comparable hierarchical layers, first matching their representations through a learned affine map and then explicitly encouraging one-to-one coordinate correspondence up to permutation and positive rescaling. This operationalizes the paper's weak-to-strong alignment mechanism for representation distillation, model merging, or stabilizing independently trained replicas. ## Mathematical statement The networks use z^N_l(x) = W^N_l h^N_{l-1}(x) + b^N_l and h^N_l(x) = phi(z^N_l(x)), with N in {A,B} denoting teacher and student and phi typically ReLU or Softplus. Weak alignment is an injective affine correspondence between layer representations; strong alignment means individual axes satisfy z^B_{pi(j)} = alpha_j z^A_j on the task set for a permutation or injection pi and allowed scale alpha_j. The depth-warped axis score in the paper is DWAxisAlign_{a,I(a),phi}(A,B) = (1/d^A_a) max_pi # {j: z^B_{pi(j)} = alpha_j z^A_{a,j} on Omega_0 for some alpha_j in H_phi}. The asymptotic theorem bounds alignment below by the used-axis fraction minus a quadratic weak-alignment error: AxisAlign_l^theta >= m_l(epsilon)/d^A_l - (||W^B_{l+1}|| omega_l + omega_{l+1})^2/(kappa_K(theta)^2 d^A_l). Here omega_l and omega_{l+1} measure affine mismatch, kappa_K(theta) > 0 is the compact-set conditioning constant, and theta is the tolerated axis residual. The practical loss replaces exact equality by normalized squared residuals and computes pi with Hungarian matching. ## Key formulas - $$z^N_{\ell}(x)=W^N_{\ell}h^N_{\ell-1}(x)+b^N_{\ell},\qquad h^N_{\ell}(x)=\varphi(z^N_{\ell}(x)).$$ - $$\operatorname{DWAxisAlign}_{a,I(a),\varphi}(A,B)=\frac{1}{d_a^A}\max_{\pi}\#\left\{j:z^B_{\pi(j)}=\alpha_jz^A_{a,j}\ \text{on }\Omega_0\ \text{for some }\alpha_j\in H_\varphi\right\}.$$ - $$\operatorname{AxisAlign}_{\ell}^{\theta}(A,B)\geq\frac{m_{\ell}(\varepsilon)}{d_{\ell}^{A}}-\frac{\bigl(\|W_{\ell+1}^{B}\|\omega_{\ell}(p)+\omega_{\ell+1}(p)\bigr)^{2}}{\kappa_{\mathcal K}(\theta)^{2}d_{\ell}^{A}}.$$ - $$\mathcal L_{\mathrm{axis}}=\frac{1}{|\Omega_0|}\sum_{x\in\Omega_0}\sum_{j\in J_A}u_j\frac{\left(z^B_{\pi(j)}(x)-\alpha_jz^A_j(x)\right)^2}{\operatorname{Var}_{\Omega_0}(z^A_j)+\epsilon}.$$ ## Implementation notes 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. ## Verification - Status: mechanism_failed - Mechanism evidence: yes - Mechanism confirmed: no - Verdict: Built a reproducible Affine-to-Axis Distillation MVP with positive-scale fitting, Hungarian matching, normalized axis loss, affine residual measurement, and a two-hidden-layer PyTorch distillation test. The synthetic check reproduced the intended signal: affine-related representations recovered all axes, recovery degraded smoothly with noise, and unrelated representations scored much worse. The task experiment showed only a small one-run accuracy/loss improvement while learned axis quality was substantially worse than baseline, so the practical effect is not demonstrated. ### Mechanism check - Verdict: Built a reproducible Affine-to-Axis Distillation MVP with positive-scale fitting, Hungarian matching, normalized axis loss, affine residual measurement, and a two-hidden-layer PyTorch distillation test. The synthetic check reproduced the intended signal: affine-related representations recovered all axes, recovery degraded smoothly with noise, and unrelated representations scored much worse. The task experiment showed only a small one-run accuracy/loss improvement while learned axis quality was substantially worse than baseline, so the practical effect is not demonstrated. - Confidence: 8/10 - Limitations: Only one fixed-seed run was performed on sklearn digits rather than CIFAR-10. Repeated seeds, confidence intervals, calibration, merging quality, assignment stability, depth windows, next-layer rescaling, and full global-affine-map training were not tested. ## Artifacts - [affine_axis_distill.py](https://synthcore.org/code/1397/affine_axis_distill.py) - [report.md](https://synthcore.org/code/1397/report.md) - [results.json](https://synthcore.org/code/1397/results.json) - [Download all files as ZIP](https://synthcore.org/download/1397) ## Disclaimer AI-generated research hypothesis, automatically tested. Not peer-reviewed.