AdaLoRA › Training Procedure

The complete AdaLoRA training procedure integrates all the elements described above into a coherent algorithm. Understanding how these pieces fit together illuminates why AdaLoRA works effectively and how its various components interact during optimization. The procedure can be divided into three phases: initialization, the main training loop with progressive pruning, and the final extraction of the rank distribution.

Each phase serves a specific purpose in the overall strategy. Initialization sets up the mathematical structure and ensures the model starts from its pretrained state. The main training loop simultaneously learns the task and discovers the importance hierarchy through gradient observations. The final rank distribution is the algorithm's learned answer to the question "which weight matrix directions matter most for this task?"

Training begins by initializing the SVD-style decomposition for each adapted weight matrix. The initialization strategy matters materially because it affects both training stability and the fairness of early importance assessments.

For each adapted weight matrix $W_k$, we initialize:

1. Initialize $P$ with random orthonormal columns, typically by applying QR decomposition to a random matrix. The QR factorization of a random matrix produces exactly orthonormal columns, giving us a perfect starting point that satisfies the orthonormality constraint exactly at initialization.
2. Initialize $\Lambda$ as the zero matrix. This is identical to LoRA's initialization strategy for the $B$ matrix. With $\Lambda = 0$, the initial weight update $\Delta W = P \cdot 0 \cdot Q = 0$, so the adapted model starts identical to the pretrained model. This is important for fine-tuning stability: the pretrained model's outputs serve as a well-calibrated starting point, and we want to begin training from there rather than from a perturbed state.
3. Initialize $Q$ with random orthonormal rows, again using QR decomposition of a random matrix (applied to the transpose).
4. Set all importance score accumulators $\bar{s}_i = 0$. Starting from zero ensures that initial importance estimates are determined entirely by the first few gradient observations rather than by any prior assumption.

The zero initialization of $\Lambda$ has a subtle interaction with importance scoring: in the first training step, all singular values are zero, so all sensitivity estimates $s_i = |\lambda_i \cdot \partial \mathcal{L}/\partial \lambda_i|$ are also zero. This means no component has demonstrated importance at initialization, which is correct: at step zero, before any training, we do not know which components will turn out to be important. The importance hierarchy emerges from gradient observations as training progresses.

Each training iteration proceeds through a sequence of well-defined steps that together drive both task learning and rank allocation:

1. Forward pass: Compute the adapted weight as $W = W_0 + P \Lambda Q$ and use it in the model's forward computation to produce predictions.
2. Loss computation: Calculate the total loss $\mathcal{L}_{\text{total}} = \mathcal{L}_{\text{task}} + \gamma \cdot R(P, Q)$, which combines the downstream task loss with the orthogonality regularization penalty.
3. Backward pass: Compute gradients $\frac{\partial \mathcal{L}_{\text{total}}}{\partial P}$, $\frac{\partial \mathcal{L}_{\text{total}}}{\partial \Lambda}$, and $\frac{\partial \mathcal{L}_{\text{total}}}{\partial Q}$ via standard backpropagation. The gradient with respect to $\Lambda$ is particularly important because it feeds directly into the sensitivity computation.
4. Sensitivity update: For each singular value $\lambda_i$, compute the instantaneous sensitivity $s_i^{(t)} = |\lambda_i^{(t)} \cdot \frac{\partial \mathcal{L}}{\partial \lambda_i^{(t)}}|$ and update the exponential moving average $\bar{s}_i^{(t)}$.
5. Parameter update: Apply the optimizer (typically AdamW) to update $P$, $\Lambda$, and $Q$ in the direction that reduces the total loss.
6. Pruning step (if past warmup and at a scheduled pruning iteration): Compute full importance scores $S_i^{(t)}$ for all triplets across all adapted matrices, rank them globally, and apply the binary mask to zero out the singular values of the lowest-ranked triplets to meet the current budget target.

The pruning step typically occurs every 10-50 iterations rather than every step. This interval serves two purposes: it reduces the computational overhead of the global ranking operation (which requires sorting all triplets in all layers), and it gives stability by letting the model to adapt to each pruning decision before the next one occurs. Continuous pruning at every step would create a moving target that could interfere with the optimizer's ability to converge.

One important implementation detail is that the gradients with respect to masked singular values are still computed but should not update the mask-zero parameters. In practice, this is handled by zeroing out gradients for masked parameters after the backward pass. This keeps pruned components remain pruned rather than being inadvertently revived by gradient updates.

After training completes, different weight matrices will have different effective ranks. This shows the algorithm's learned assessment of their importance for the task. The patterns that typically emerge are informative:

• Query and value projections in attention layers tend to receive higher effective ranks. These matrices directly control what the model attends to and how attended information is combined, both of which are important for most downstream NLP tasks that require the model to focus on specific parts of the input.
• Key projections often receive moderate ranks. Keys determine which queries each position responds to, which is important but often requires less adaptation than queries and values since the general patterns learned during pretraining remain largely applicable.
• Feed-forward layers tend to receive lower ranks, especially in deeper transformer layers. The feed-forward networks have been interpreted as storing factual knowledge about the world, but for many fine-tuning tasks (especially classification tasks), this stored knowledge needs relatively little modification.
• The specific pattern varies substantially by task. For tasks that require understanding long-range dependencies (like multi-hop reasoning), higher-layer attention heads may receive more rank. For tasks that require specific factual knowledge (like open-domain QA), feed-forward layers may need more rank than usual.

This distribution emerges entirely from the data and task: AdaLoRA discovers it rather than requiring manual specification. The patterns often reveal interesting insights about which parts of the model the specific fine-tuning task relies on most heavily, so AdaLoRA can serve as both a practical training tool and an analytical instrument for understanding model adaptation.