AFLoRA: Methodology
Module Structure. Inspired by the framework proposed by Kopiczko et al. ( 2024 ) , we design the LoRA module to encompass four components, namely, the down-projection linear layer ( lora_{A} ), the up-projection linear layer ( lora_{B} ), and two feature transform vectors ( s_{d} , and s_{b} ) placed before and after lora_{B} . However, unlike Kopiczko et al. ( 2024 ) , we keep both the projection matrices ( lora_{A} and lora_{B} ) and vectors trainable at the beginning and keep the rank very low .
The module processes a given input X through these components to produce an output Y . The complete operation for a layer l can be described as follows:

{Y}={W^{l}_{0}}{X}+\Lambda^{l}_{b}B^{l}\Lambda^{l}_{d}A^{l}{X} (1)

Here, A^{l} and B^{l} are the trainable LoRA tensors of lora^{l}_{A} and lora^{l}_{B} , respectively. \Lambda_{d} and \Lambda_{b} are the vectors of s_{d} , and s_{b} , respectively. W^{l}_{0} represents the frozen pre-trained weights.
We use Kaiming Uniform initialization for A^{l} and B^{l} , and follow Kopiczko et al. ( 2024 ) to initialize the vectors.

Adaptive Freezing. In pruning literature (Han et al., 2015 ; Molchanov et al., 2019 ; Zhang et al., 2022 ) , sensitivity is gauged to reflect weight variability, necessitating consideration of both the weights’ magnitudes and their gradients. Small weight values suggest minimal impact, while minor gradient values indicate stability. Taking inspiration from this idea, here we introduce the concept of a "freezing score". However, unlike pruning where both magnitude and gradient play a critical role to identify insignificant weight, we leverage only gradient as a proxy to compute the freezing score. This is because, we assume large magnitude weights with negligible change has same priority to be frozen as that for small magnitude weights. This score quantifies the degree to which weights vary throughout the training process. Consequently, when the expected changes to the weights become negligible, we may consider them to be frozen, thereby saving computational resources and energy.

Following equation describes the freezing score evaluation steps for a low-rank tensor A^{l} .

I_{A^{l}}=\left| abla\mathcal{L}(\boldsymbol{\theta)}\right|,\overline{I}_{A^{l}}^{(t)}=\beta_{1}\overline{I}_{A^{l}}^{(t-1)}+(1-\beta_{1})I_{A^{l}}^{(t)} (2)

U_{A^{l}}^{(t)}=\left|I_{A^{l}}^{(t)}-\overline{I}_{A^{l}}^{(t)}\right|,\overline{U}_{A^{l}}^{(t)}=\beta_{2}\overline{U}_{A^{l}}^{(t-1)}+(1-\beta_{2})U_{A^{l}}^{(t)} (3)

s_{A^{l}}^{(t)}=mean(\overline{I}_{A^{l}}^{(t)}\circ\overline{U}_{A^{l}}^{(t)}) (4)

Here, for each projection tensor at iteration t , we compute a smoothed gradient ( \overline{I}_{A^{l}}^{(t)} ) and uncertainly tensor ( \overline{U}_{A^{l}}^{(t)} ), as shown in Eq. 2 and 3, respectively. We then evaluate the freezing score s_{A^{l}}^{(t)} , as the mean of the tensor generated via Hadamard product ( \circ ) between \overline{I}_{A^{l}}^{(t)} and \overline{U}_{A^{l}}^{(t)} .

To apply thresholding on the LoRA freezing scores, we use the cubic schedule as (Zhang et al., 2022 ) . In specific, we keep the projection matrices trainable for the initial t_{i} training steps, and then progessively freeze them by calculating the freezing fraction r(t) as shown in the Eq. 5 .
Finally all the projection matrices freeze beyond T-t_{f} steps. Note, at a step t , for a computed freezing fraction k , we freeze the lowest k\% projection matrices.

r(t)=\left\{\begin{array}[]{ll}0&\quad 0\leq t<t_{i}\ 1-\left(1-\frac{t-t_{i}}{T-t_{i}-t_{f}}\right)^{3}&\quad t_{i}\leq t<T-t_{f}\ 1&\quadotherwise\end{array}\right

where t refers to current #step, T is the total number of fine-tuning steps.
We set t_{i} to the steps corresponding to one epoch and set t_{f} to 70% of the total training steps.