Source-linked AI summary
KronA: Parameter Efficient Tuning with Kronecker Adapter
Ali Edalati, Marzieh Tahaei, Ivan Kobyzev, Vahid Partovi Nia, James J. Clark, Mehdi Rezagholizadeh
TL;DR
As PLMs grow, full fine-tuning becomes expensive and low-rank PET methods have limited representation power. The paper replaces low-rank adaptation with Kronecker-based modules, introducing KronA and higher-accuracy parallel variants, and evaluates them on T5 with GLUE. The proposed methods outperform corresponding low-rank baselines, while the residual variant improves results at increased inference cost.
Problem
Growing PLMs make full fine-tuning expensive, while low-rank adaptation imposes a limited representation structure for task-specific updates.
Method
The paper replaces low-rank projections in LoRA with Kronecker product factors to develop KronA, plus parallel FFN variants with an optional learnable residual connection.
Results
On T5 GLUE, KronA and KronAB outperform LoRA and Parallel-Adapter respectively, while KronABres achieves remarkably better results.
Takeaways & Limitations
Kronecker-based adapters provide a PET alternative that can improve accuracy over corresponding low-rank methods, with latency-sensitive and higher-accuracy variants.
Takeaways & Limitations
KronA’s merged form imposes increased inference time and computation.
Abstract
from arXiv · showhide
Fine-tuning a Pre-trained Language Model (PLM) on a specific downstream task has been a well-known paradigm in Natural Language Processing. However, with the ever-growing size of PLMs, training the entire model on several downstream tasks becomes very expensive and resource-hungry. Recently, different Parameter Efficient Tuning (PET) techniques are proposed to improve the efficiency of fine-tuning PLMs. One popular category of PET methods is the low-rank adaptation methods which insert learnable truncated SVD modules into the original model either sequentially or in parallel. However, low-rank decomposition suffers from limited representation power. In this work, we address this problem using the Kronecker product instead of the low-rank representation. We introduce KronA, a Kronecker product-based adapter module for efficient fine-tuning of Transformer-based PLMs. We apply the proposed methods for fine-tuning T5 on the GLUE benchmark to show that incorporating the Kronecker-based modules can outperform state-of-the-art PET methods.
1 Introduction
Growing PLMs make full fine-tuning costly to train, store, and deploy, motivating PET methods that update only small parameter sets. KronA replaces low-rank adaptation with Kronecker-based modules, while KronAB and KronABres trade additional latency for higher accuracy on GLUE.
- Motivation: Large PLMs make full fine-tuning increasingly expensive and require a separate full checkpoint for each downstream application.In-context learning also adds inference overhead because training examples must be processed for each sample.
- Existing PET methods: PET methods freeze most or all pretrained parameters and insert a small trainable parameter set, reducing fine-tuning resources and per-task storage.Soft prompts and adapters can nevertheless increase inference computation or latency.
- Existing PET methods: LoRA preserves inference latency by merging task-specific weights, but its low-rank structure usually reduces accuracy relative to full fine-tuning.Kronecker decomposition is presented as an alternative that does not rely on the low-rank assumption and has outperformed low-rank factorization in compression settings.
- Proposed methods: KronA replaces LoRA’s low-rank projections with Kronecker factors, while KronAB applies the module in parallel to the FFN and KronABres adds a learnable residual connection.KronA targets latency-critical use, whereas the latter variants accept increased inference latency for accuracy gains.
- Evaluation: The methods are evaluated on T5 with GLUE, comparing GLUE score, training time, and inference latency against state-of-the-art PET baselines.The evaluation is intended to assess the performance and efficiency impact of Kronecker-based modules.
2 Related Works and Baselines
Related PET baselines differ in where trainable modules are inserted and whether their updates can be merged into the pretrained model. LoRA avoids inference overhead after merging, whereas adapter-style methods can be slower.
- PET methods: BitFit freezes PLM weights and tunes biases or a bias subset, offering parameter efficiency and speed but usually weaker performance than state-of-the-art methods.It is included as a related PET approach rather than among the baselines selected for this comparison.
- Adapter methods: Adapters freeze the PLM and insert sequential down-projection, nonlinearity, up-projection, and residual modules after FFN or attention blocks.Parallel-Adapter instead places the module in parallel, uses a scaling factor, and removes its internal residual connection.
- Adapter methods: Compacter reconstructs weight matrices from Kronecker products of factors and achieves good GLUE results, but is notably slow during training and inference.Its factors are themselves formed by multiplying two sub-factors.
- Low-rank methods: LoRA trains low-rank modules while freezing PLM weights and merges them into the original matrices during inference, avoiding added inference time.The authors recommend placing LoRA modules in parallel to query and value matrices.
- Baselines: LoRA, Parallel-Adapter, and Compacter are used as comparison baselines because they had outperformed many earlier PET methods.The listed earlier methods include Pfeiffer-Adapter, AdapterDrop, VGLM-Adapter, prompt tuning, and prefix tuning.
3 Methodology
The methodology replaces low-rank adaptation with Kronecker-factor modules and develops KronA for weight matrices and KronAB for parallel FFN use. These designs target stronger representation while balancing training and inference efficiency.
- Kronecker Product: Kronecker products preserve input-matrix rank and can compute module outputs without reconstructing the full product matrix.The latter avoids reconstruction of W to reduce required FLOPS.
- KronA: KronA replaces LoRA’s projections with Kronecker factors and applies the resulting module in parallel to PLM weight matrices.A fixed scale factor s controls the added update, and the factors can later be merged into the original weights.
- KronA: After fine-tuning, KronA’s Kronecker factors are multiplied, scaled, and merged into the original PLM weights, avoiding added inference time.This merging follows the same latency-preserving design goal as LoRA.
- KronAB: KronAB applies a KronA-style module in parallel to the FFN block, while its residual variant adds a scaled, tunable residual connection.The residual scale is initialized to one and tuned during fine-tuning.
- KronAB: Because the FFN nonlinearity prevents merging KronAB into the pretrained model, KronAB increases inference-time computation and latency.The method is therefore distinct from KronA’s mergeable weight-matrix adaptation.
4 Results and Discussion
On T5 with GLUE, KronA variants outperform low-rank and other state-of-the-art baselines on average and most tasks, while trading modest training-time increases for improved accuracy. KronA and related methods avoid added inference latency, whereas KronAB incurs latency from its residual connection and Kronecker computation.
- KronA and KronAB outperform LoRA and PA, their respective low-rank counterparts, on T5 GLUE evaluation.
- All proposed modules outperform other state-of-the-art baselines on average and most GLUE tasks.
- KronAB’s learnable residual connection achieves remarkably better results.
- KronA, LoRA, FT, and BitFit do not increase inference latency, while KronAB is slower than PA but faster than Compacter and Adapter.KronAB’s additional latency reflects Kronecker-product computation and its learnable residual connection.
- KronA variants require slightly more training time than low-rank counterparts but remain significantly faster than full fine-tuning.The training-time increase is described as not remarkable.
5 Conclusion and Future Directions
The paper develops Kronecker-based adapters by replacing low-rank projections with Kronecker products and evaluates them against state-of-the-art baselines using T5 on GLUE.
- Kronecker-based adapters replace the low-rank projections of state-of-the-art parameter-efficient tuning methods with Kronecker products.
- The proposed adapters are evaluated through comparisons of training and inference time and fine-tuning T5 on the GLUE benchmark.
- Table 3 provides details about the Kronecker factors.
B Ablation Study
The ablation study finds that KronA initialization affects performance, with a mixed zero and Kaiming-uniform initialization outperforming two Normal-distribution initializations.
- KronA factor initialization affects module performance.
- Initializing one Kronecker factor with Kaiming-uniform and the other with zero performs significantly better than initializing both factors from the Normal distribution.
B.2 Step by Step Improvement of KronAB
KronAB was refined through sequential design changes involving placement, scaling, residual control, and nonlinearity. The experiments favored parallel insertion with scaling and without the tested nonlinear or sigmoid modifications.
- Initial design: KronAB was initially inserted sequentially after the FFN and attention blocks without scaling, then modified to improve performance.The initial design followed a normal adapter configuration.
- Placement and scaling: Parallel insertion of KronAB improved performance significantly compared with sequential insertion.The comparison is reported in the step-by-step experiments on GLUE.
- Placement and scaling: Adding a scaling factor increased the GLUE score, while placing two modules in each FFN outperformed placing modules in both FFN and attention blocks.
- Nonlinearity: SiLU was the best tested nonlinearity on QNLI, but it decreased KronAB’s GLUE score, so the nonlinearity was removed.The nonlinearity was tested between the two Kronecker multiplications.
- Residual scaling: Adding a Sigmoid to the KronAB residual reduced performance and increased latency, so the Sigmoid function was removed.The tested variant constrained the residual scale between 0 and 1.
C Details of Measuring Training and Inference Time
Inference latency was measured with repeated GPU warm-up and timed runs, while training time was normalized under matched experimental settings.
- Inference latency: Inference latency used a batch size of one and sequence length ten, with 150 GPU warm-up iterations followed by 200 timed iterations.
- Inference latency: Latency was averaged across three repetitions, then normalized for reporting in Table 2.
- Training time: Training-time comparisons used the same epochs, batch size, GPU count, and gradient accumulation step for every technique.
D.1 Datasets
The evaluation uses GLUE tasks spanning language understanding categories, with task-specific metrics and constructed test sets because the original GLUE test set is unavailable.
- Benchmark: The GLUE evaluation covers natural language inference, linguistic acceptability, similarity and paraphrasing, and sentiment classification.The listed tasks include MNLI, RTE, QNLI, CoLA, MRPC, QQP, and SST-2.
- Dataset construction: Because GLUE’s original test set is unpublished, the study generated test sets from evaluation and training data.
- Training-time evaluation: Table 8 reports normalized training time for each technique across the GLUE tasks.
- Metrics: CoLA, MRPC, STS-B, and other GLUE tasks use Matthew correlation coefficient, F1, averaged Pearson/Spearman correlations, and accuracy, respectively.
D.2 Experimental Setup
Experiments compare parameter-efficient methods on T5-base under approximately matched trainable-parameter budgets, using selected Kronecker-factor configurations and tuned task settings.
- Platform and backbone: All experiments used one NVIDIA Tesla V100, PyTorch, Hugging Face Transformers, and T5-base as the backbone model.
- Compared methods: The compared methods were configured with roughly equal numbers of trainable parameters for fairness.
- KronA configuration: For KronA, Kronecker-factor shapes were selected using the best development results among tested options.
- Module placement: The study selected the PLM block for each technique and tuned bias counts to maximize each task’s score.
- Hyperparameters: LoRA and KronA scaling factors were fixed at one to isolate scaling-factor effects in their comparison.
- Hyperparameters: LoRA and PA used ranks one and two, while KronAB and KronABres factor shapes were not tuned because of time and resource limitations.