Seeking Neural Nuggets: Knowledge Transfer in Large Language Models from a Parametric Perspective

TLDR; Knowledge transfer. From a larger model to a smaller model. Uses LoRA and all. Without any teacher’s participation.

Knowledge Extraction

This is the only part where the teacher is being used. Let’s say the teacher’s parameter is ${\Theta }_T=[{\theta }_1,\dots ,{\theta }_{N_T}]$, we can isolate each parameter $i$ by null-ing every other parameters, ${\Theta }_{T_i}=[0,\dots ,{\theta }_i,\dots ,0]$.

We then use the Taylor approximation to find out the effect of each parameter in ${\Theta }_T$.

$$L({\Theta }_T)-L({\Theta }_T-{\Theta }_{T_i})\approx {\Theta }_{T_i}^{\top }{\nabla }_{{\Theta }_T}L(x_j)$$

We call the above approximated change of loss / sensitivity as $S_{i,j}$. Then, rank all the teacher layer by the summed sensitivity (i.e. $S_{T_l}={\sum }_{{\theta }_i\in {\Theta }_{T_l}}{S}_i$), and select the top $L_s$ layers.

Next, we do dimensionality reduction, which is just choosing the contiguous submatrices (according to the student’s layer dimension) with the highest cumulative sensitivity.

Knowledge Injection

The extracted matrix is decomposed via SVD.

$$W_{T_i,\text{extract}}=U_i{\Sigma }_i{V}_i^{\top }$$

Then, we can approximate the low-rank matrices using the Eckart-Young-Mirsky thm.

$$W_{T_i,\text{extract}}=U_i[:,:r]{\Sigma }_i[:r,:r]V_i^{\top }[:r,:]$$

We can initialize the row-rank module $B_i=U_i[:,:r]$ and $A_i=V_i^{\top }[:r,:]$.

Initialization w/ LoRA

Instead of initializing naively, with $W_i^{\ast }=W_i+B_i{A}_i$, we do it a bit differently so as to ensure training starts from original pre-trained weights.

$$W_i^{\ast }=W_i-W_{T_i,\text{extract},r}+B_i{A}_i$$

In essense, we are doing a “warm start” where the training is informed by the teacher’s knowledge — and we are doing this without disrupting the student’s pretrained weights.

Thoughts

Well there is no “interpretability” on what each parameter is doing. This is only very geared towards knowledge distillation and nothing else.