効率的なファインチューニングのためのLoRAターゲットモジュール選択の最適化

Optimizing LoRA target module selection for efficient fine tuning

Ablation study clarifies trade-offs between accuracy and efficiency when using low-rank adaptation (LoRA) to fine-tune AI models.

Machine learning

Rushil Anirudh Anjie Fang Bhoomit Vasani March 19, 10:39 AM March 19, 10:39 AM Fine-tuning a large language model (LLM) on a specific task requires updates to billions of parameters across trillions of tokens, with the attendant costs in GPU resources and time.

Low-rank adaptation (LoRA) is a more efficient alternative that freezes the original model weights but introduces lightweight matrices into specific model sublayers, or modules. These matrices (commonly referred to as adapters) modify the modules weights, enabling not only efficient fine tuning but also on-demand model serving, which dramatically lowers inference costs; base-model sharing across GPUs, which cuts memory requirements; lower download overhead; and parallel inference across multiple adapters.

The question is where to insert these adapters across the model. Empirically, targeting more and larger modules tends to boost performance, because it allows more flexibility in customization; but it also increases training and inference costs. Using a smaller, well-chosen subset preserves most gains with significantly better efficiency.

Using Amazons Nova 2.0 Lite multimodal reasoning LLM as our base model, we set ourselves the goal of identifying a subset of standardized target-module configurations that works effectively across the vast majority of customer use cases. Through an ablation study, we identified a module known as o_proj, as the single module where adding an adapter achieves the best trade-off between efficiency and accuracy (o_proj is a linear transformation that mixes representations across attention heads into a single, cohesive form for the rest of the model to understand).

The Transformer architecture

Transformer models the models responsible for all of AIs remarkable recent gains consist largely of blocks that are repeated multiple times. Each block in turn has two main components: an attention mechanism, which determines the relevance of previously seen tokens to the token currently being processed, and a feed-forward network, a conventional neural network that does additional processing on the outputs of the attention mechanism.

The attention mechanism involves three different matrices, which take their names from database design: the query matrix represents how relevant the current token is to the other tokens in the input sequence; the key matrix represents how relevant other tokens are to one another; and the value matrix represents the raw content of those other tokens. Multiplying the three matrices together creates, essentially, a recipe for the Transformer's next output.

To reduce computational complexity, these multiplications take place in a space with reduced dimensions. The matrices themselves and the results of their multiplication then have to be projected back up to the original dimensions of the input.

LoRA approximates weight updates using a product of two smaller matrices, drastically reducing the number of trainable parameters. The technique is typically applied to attention projection layers and feed-forward network layers. These modules are ideal candidates because they constitute the bulk of Transformer parameters, directly govern representation learning, and exhibit natural alignment with low-rank approximations. Empirical evidence shows weight changes in these layers often lie within a low-dimensional subspace during fine tuning.

image LoRA for a generic layer-weight matrix (*W*). The weights are modified by the product of two smaller matrices (*A *and* B*), whose lower dimensions drastically reduce the number of trainable parameters. Target module selection

Selecting the right target modules directly affects accuracy, latency, and computational efficiency. The optimal choice of target modules is primarily a function of (a) the base model being fine-tuned (i.e., its architecture, pre- and post-training data distributions, etc.) and (b) customization domain/modality.

When fine-tuning Nova 2.0 Lite, we balanced two competing objectives:

Maximizing accuracy across diverse tasks and modalities and

Minimizing latency to preserve LoRA's efficiency benefits.

We investigated the application of LoRA to four different modules in each Transformer block: the query, key, and value projection layers ( qkv); the o_proj layer; and two different fully connected layers in the feed-forward network, gate_up_proj and gate_down_proj (referred to as fc1 and fc2). Below are the trade-offs for these modules, both singly and in combination, based on results published in literature and empirical studies.

CombinationExpected accuracyExpected latencyUse case

*qkv* onlyGood (baseline)Lowest- Resource-constrained environments

• Tasks where attention mechanisms are critical (e.g., classification, lightweight generation)
• Prioritizes speed over maximum accuracy

*o_proj* onlyModerateLowest- Ultralow-latency scenarios

• Tasks where refining attention outputs is sufficient (e.g., simple sentiment analysis). Plays an important role in reasoning
• Less effective than qkv, but very efficient

*qkv* + *o_proj*HighLow to moderate (+510%)- Attention-focused tasks (e.g., machine translation, summarization)

• Balances refinement of both attention context ( o_proj) and query/key/value projections ( qkv)
• Best accuracy-to-latency ratio for most NLP tasks

*qkv* + *fc1* / *fc2*Very high (close to full fine tuning)Moderate (+1015%)- Complex generation tasks (e.g., translation, long-form summarization)

• When feed-forward layers ( fc1/ fc2) significantly influence output quality as they store and retrieve factual knowledge
• Prioritizes accuracy over speed

*o_proj* + *fc1* / *fc2*Good to highModerate (+510%)- Tasks requiring adaptation of both attention output ( o_proj) and feed-forward layers (e.g., text classification, sentiment analysis)

• Suitable when qkv adaptation is unnecessary

*qkv* + *o_proj* +* fc1* / *fc2*Highest (near-full fine tuning)High (+1520%)- Maximum accuracy for critical tasks (e.g., research benchmarks, high-stakes generation)

• When all components of the Transformer block need adaptation
• Avoid for production if latency matters

All modules( *qkv*, *o_proj*, *fc1*, *fc2*)MaximumHighest (+2025%)- Prototyping/research with no latency constraints

• Rarely justified in practice; marginal gains over qkv + o_proj + fc1/ fc2

Trade-offs of accuracy and latency across target modules, based on literature review and empirical evidence.

Experimental methodology

We conducted a comprehensive ablation study, training multiple supervised-fine-tuning (SFT) LoRA variants on seven datasets spanning both text and visual data, across reasoning (i.e., the training datasets themselves include reasoning content) and non-reasoning tasks. The datasets covered diverse challenges from simple question answering to long-context summarization and structured JSON extraction.

Dataset<td colspan="1" rowspan="1" width="44" valign="top" style="width:32.7pt;border:solid windowtext 1.0pt; border-left:none;mso-border-left-alt:solid windowte