Source-linked AI summary
TaRA: Training-Aware Low-Rank Adaptation Initialization
Taehyeon Kim, Eunhyeok Park
TL;DR
LoRA’s low-rank bottleneck can misalign optimization with full-rank fine-tuning, motivating initialization methods that better preserve training dynamics. TaRA uses covariance-aware SVD to approximate full-rank local gradients and consistently outperforms existing methods across benchmarks, though it requires task-specific calibration.
Problem
LoRA’s low-rank bottleneck can alter the optimization trajectory relative to full-rank training, while prior initializations do not explicitly align with full-fine-tuning learning dynamics.
Method
TaRA jointly uses activation covariance, gradient covariance, and pretrained weights to construct a covariance-aware low-rank initialization that preserves local full-rank gradient behavior.
Results
TaRA consistently outperforms existing initialization methods and achieves state-of-the-art performance across a wide range of benchmarks.
Takeaways & Limitations
TaRA provides a training-aware LoRA initialization that prioritizes directions influential during optimization under a low-rank constraint.
Takeaways & Limitations
TaRA requires task-specific calibration, adding data access and one-time computation and memory costs before fine-tuning.
Abstract
from arXiv · showhide
Low-Rank Adaptation (LoRA) has become a de facto standard for parameter-efficient fine-tuning (PEFT), yet its performance is highly sensitive to initialization due to the information bottleneck imposed by low-rank decomposition. Existing approaches attempt to construct high-quality LoRA initializations by exploiting principal components of pretrained weights, activations, or gradients. However, these methods do not directly account for the training dynamics of the full-rank model. In this paper, we propose Training-aware Low-Rank Adaptation Initialization (TaRA), a method that initializes LoRA such that the gradients induced by the low-rank factors closely approximate the gradient of the corresponding full-rank weight matrix. Derived from a mathematical formulation, TaRA improves gradient fidelity at the start of training while introducing negligible computational overhead. Across diverse and challenging fine-tuning tasks, TaRA consistently outperforms prior state-of-the-art methods, establishing a simple, robust, and scalable solution for effective LoRA initialization.
1 Introduction
LoRA reduces fine-tuning costs but its low-rank bottleneck can alter optimization relative to full-rank training. TaRA addresses this by initializing adapters around training-relevant gradient directions and reports stronger gradient alignment and downstream performance.
- LoRA is widely adopted because it reduces fine-tuning resources and adds no inference-time overhead after updates merge into base weights.
- The low-rank bottleneck alters the optimization trajectory relative to full-rank training, potentially causing suboptimal convergence even with sufficient capacity.
- Prior initialization methods exploit pretrained weights, joint weight-data statistics, or gradients to improve LoRA fine-tuning.
- TaRA preserves directions important to local full-weight gradient behavior by jointly using activation covariance, gradient covariance, and pretrained weights from one forward-backward pass.
- TaRA achieves higher gradient alignment than prior methods at the same rank and consistently improves performance across mathematical, code-generation, and commonsense reasoning tasks.
2 Prior Work on LoRA Initialization
Prior LoRA initializations select low-rank subspaces from weights, activations, or gradients, but they do not explicitly align with full-fine-tuning learning dynamics. TaRA is motivated by this remaining performance gap.
- Original LoRA randomly initializes A and sets B to zero, providing no meaningful task-related initialization information.This design can slow early optimization and lead to suboptimal convergence.
- Improved methods initialize LoRA while preserving the pretrained function through a frozen residual weight Wres = W0 − BA.With suitable A and B, these initializations can accelerate training and typically converge to better solutions than random initialization.
- PiSSA applies SVD to pretrained weights and uses the top-r singular components for data-agnostic initialization.
- CorDA incorporates input activation covariance with pretrained weights, while PiSSA and CorDA align LoRA subspaces with principal components of weight or input-dependent statistics.Both methods are reported to improve convergence and accuracy.
- LoRA-GA and LoRA-One use gradient information, with LoRA-One applying SVD to the one-step full gradient and retaining its top-r components.These methods are motivated by alignment between LoRA updates and dominant singular directions of the full fine-tuning gradient.
- Prior methods emphasize representational capacity rather than explicit alignment with full-fine-tuning dynamics, leaving a gap that rank increases or hyperparameter tuning do not effectively close.
3 Proposed Idea: TaRA
TaRA formulates initialization as a rank-constrained approximation of full-rank local gradient behavior, using curvature-aware statistics to preserve training-relevant directions. It implements this objective through covariance-weighted SVD and layer-wise LoRA construction.
- 3 Proposed Idea: TaRA: TaRA seeks a low-rank parameterization whose induced local gradient behavior approximates the corresponding full-rank weight under a strict rank constraint.
- 3.1 Motivation: The objective minimizes the deviation between gradients at a rank-constrained parameterization and the full-rank initialization parameter θ0.The task loss is denoted by L.
- 3.1 Motivation: TaRA focuses on the one-step gradient at initialization as an approximation to the dominant early optimization signal.The paper reports that this approximation guides effective initialization through later convergence results.
- 3.2 Training-Relevant Decomposition: A second-order loss view shows that curvature modulates gradient changes, making directions associated with large Hessian or Fisher values especially influential.The method therefore targets curvature-weighted gradient deviation and dominant curvature directions.
- 3.2 Training-Relevant Decomposition: Using the Eckart–Young theorem, TaRA obtains the optimal rank-r approximation by retaining the top-r singular components.This corresponds to preserving dominant curvature directions of the pretrained parameter.
- Implementation: The layer-wise procedure collects activations and gradients from a small calibration dataset, computes ΣX and ΣG, and decomposes ΣGW0ΣX before mapping components back with ΣX^-1.
- Implementation: TaRA initializes A and B from the truncated decomposition, scales them with inverse covariance matrices, and updates the frozen residual to preserve the pretrained function.Diagonal damping stabilizes potentially rank-deficient covariance inversions, using c = 10^-2.
- Implementation: The procedure is applied independently to each linear layer to produce a training-friendly LoRA initialization with training-relevant directions.
4 Experiments
TaRA is evaluated across diverse natural-language generation and understanding tasks, model scales, and LoRA ranks. It consistently improves training behavior, gradient alignment, and downstream performance while requiring modest initialization overhead.
- Experimental Setup: TaRA is evaluated against full fine-tuning and multiple LoRA initialization and PEFT baselines across NLG and NLU tasks.The experiments vary models, datasets, benchmarks, and rank budgets to assess generality and robustness.
- Natural Language Generation Task: TaRA achieves the highest average NLG performance at ranks 128, 64, and 32, with the most stable behavior across ranks.At rank 32, it achieves the best performance across all reported NLG tasks.
- Natural Language Generation Task: TaRA converges more smoothly, reaches lower training loss, and attains rapid, sustained accuracy gains compared with prior initialization methods.The advantage appears at both r = 128 and r = 32 and persists through intermediate and final training stages.
- Natural Language Understanding Task: TaRA achieves the best average commonsense-reasoning performance on DeepSeek-R1-Distill-Qwen-1.5B, LLaMA-2-7B, and LLaMA-3.1-8B.On Qwen-3-8B, it slightly trails LoRA-One on average but remains competitive and outperforms LoRA, PiSSA, and CorDA.
- Natural Language Understanding Task: Across 32 commonsense-reasoning cases, TaRA ranks first in 15 and second in 11, versus 8 and 5 top-ranked cases for LoRA-One and CorDA.This pattern supports TaRA’s reported robustness across models and tasks.
- Gradient Alignment and Initialization Overhead Analysis: TaRA has substantially higher full-rank gradient alignment across ranks and after multiple training steps, while initialization requires only 4–5% of total fine-tuning time.It also remains effective with low-precision covariance collection and small calibration sets.
5 Related Work
Related work reduces PEFT costs through adapters, prompts, and low-rank updates, while LoRA extensions target capacity allocation, structure, pruning, and quantization. These efforts are largely complementary to initialization choices.
- PEFT methods span adapter-based, prompt-based, and low-rank adaptation paradigms for reducing the cost of adapting LLMs.
- LoRA freezes pretrained weights and learns a low-rank update through matrices A and B, preserving efficiency and allowing updates to merge into base weights.
- LoRA extensions improve efficiency and performance through adaptive rank allocation, redesigned update structures, pruning, and quantization.
- These extensions primarily address rank-related capacity and resource constraints while preserving LoRA’s core efficiency benefits.
- Because these advances are largely orthogonal to initialization, they can complement TaRA and potentially provide further performance gains when combined.
6 Conclusion
TaRA prioritizes training-relevant directions through covariance-aware low-rank initialization. It closely matches full-finetuning one-step gradients and consistently outperforms existing methods across benchmarks.
- TaRA uses covariance-aware SVD of activation and gradient statistics to identify training-relevant directions for LoRA initialization.
- TaRA’s resulting low-rank approximation produces an induced one-step gradient that closely matches the full-finetuning gradient.
- TaRA consistently outperforms existing initialization methods and achieves state-of-the-art performance across a wide range of benchmarks.
7 Limitations
TaRA has two main scope and cost boundaries: its gains can be less stable under distribution shift, and task-specific calibration adds data-access, computation, and memory costs.
- TaRA’s gains can be less stable under distribution shift because calibration and training are distribution-specific.HumanEval and MBPP are out-of-distribution relative to CodeFeedback, with non-monotonic performance across ranks and methods.
- TaRA requires task-specific calibration to collect activation and gradient covariance statistics before fine-tuning.This creates additional data-access, one-time computation, and memory costs, although the analyses report them as manageable in practice.
C Loss-Landscape Analysis
In a toy loss-landscape experiment, TaRA follows a more direct optimization trajectory toward a local minimum than competing initializations. It also converges to a lower objective value than PiSSA and CorDA.
- TaRA is initialized in a direction that heads more directly toward the local minimum than Full Fine-Tuning, PiSSA, and CorDA.Full Fine-Tuning and PiSSA follow noticeably curved trajectories, while CorDA has a less circuitous early trajectory.
- TaRA ultimately converges to a lower objective value than PiSSA and CorDA in the toy setting.
- The experiment uses a three-linear-layer network pretrained on odd-class MNIST samples and fine-tuned on even-class samples with LoRA rank and scaling r = α = 16.
D Diagonal Damping Analysis
Diagonal damping stabilizes TaRA’s covariance-based initialization: below 10^-2, training collapses, while damping at or above 10^-2 prevents this failure.
- 0.0 task accuracy occurs when the damping coefficient c is smaller than 10^-2, including c = 0.The paper identifies this regime as a failure to stabilize training and complete training breakdown.
- c ≥10^-2 stabilizes performance in the CoLA-task analysis.
E Mitigating Accuracy Fluctuations on Out-of-Distribution Tasks
TaRA’s training-distribution-aware initialization can fluctuate on out-of-distribution code tasks, but Ledoit–Wolf shrinkage makes performance more consistent across ranks and may improve robustness under distribution shift.
- HumanEval and MBPP accuracies show no consistent trend as LoRA rank increases in the out-of-distribution setting.The training and evaluation dataset distributions differ for these code-generation benchmarks.
- TaRA is explicitly tailored to the training distribution because it estimates activation and gradient statistics from training data.
- Ledoit–Wolf shrinkage regularizes training-derived covariance statistics to address TaRA’s bias toward the calibration distribution.The shrinkage coefficient λ controls how dependent the estimated covariance remains on calibration data.
- λ approaching 1 makes the covariance data-agnostic and yields an initialization coinciding with PiSSA under the paper’s formulation.
- LW shrinkage produces substantially more consistent and nearly monotonic performance across ranks on OOD code-generation tasks.The reported LW results use λ = 0.3 and average three seeds.
F Why One-Step Gradient Alignment Persists Beyond Initialization
TaRA’s one-step gradient alignment can persist because locally stable curvature keeps later gradients within related Hessian-determined directions. Experiments also examine calibration size, LoRA α, and alternative Fisher-based formulations.
- Why alignment persists: TaRA retains substantially higher gradient alignment with Full FT than baselines after multiple optimization steps.This extends beyond the one-step gradient-matching objective used to derive TaRA.
- Why alignment persists: Under locally stable curvature, later gradients arise by repeatedly applying a polynomial in the local Hessian to the initial gradient.Training gradients therefore tend to concentrate in a small, relatively stable subspace associated with dominant Hessian directions.
- Why alignment persists: Capturing training-relevant curvature directions at initialization can preserve useful directions beyond the first optimization step.The paper presents this as an intuitive explanation rather than a formal guarantee for the full nonlinear trajectory.
- Robustness analyses: TaRA remains stable across small calibration sets, showing only minor performance changes at rank 128.This suggests meaningful initialization remains possible when available calibration data is limited.
- Robustness analyses: TaRA performs best when LoRA α equals the LoRA rank at rank 128.
- Formulation comparison: The square-root Fisher formulation from compression yields suboptimal LoRA initialization performance on MATH at rank 128.This comparison supports analyzing TaRA from a training perspective rather than directly transferring the compression formulation.