PYRA: Parallel Yielding Re-Activation for Training-Inference Efficient Task Adaptation

TL;DR

This paper tackles the challenge of training-inference efficient task adaptation for large-scale vision transformers, where traditional PEFT methods reduce training cost but leave inference efficiency unchanged and model compression often incurs heavy re-training. It proposes PYRA, a token-merging framework that uses Parallel Yielding Adaptive Weights and a Re-Activation strategy to modulate tokens and feature channels before merging, enabling effective downstream adaptation with minimal training parameters and substantial inference speedups. Across VTAB-1k benchmarks and multiple backbone variants (ViT, DeiT, MAE), PYRA consistently achieves comparable performance to uncompressed models at low compression and bridges the adverse compression gap at high compression, outperforming ToMe and other baselines. The method demonstrates strong generalization to different architectures and PEFT settings, offering a practical route to deploy foundation models with reduced cost and maintained accuracy.

Abstract

Recently, the scale of transformers has grown rapidly, which introduces considerable challenges in terms of training overhead and inference efficiency in the scope of task adaptation. Existing works, namely Parameter-Efficient Fine-Tuning (PEFT) and model compression, have separately investigated the challenges. However, PEFT cannot guarantee the inference efficiency of the original backbone, especially for large-scale models. Model compression requires significant training costs for structure searching and re-training. Consequently, a simple combination of them cannot guarantee accomplishing both training efficiency and inference efficiency with minimal costs. In this paper, we propose a novel Parallel Yielding Re-Activation (PYRA) method for such a challenge of training-inference efficient task adaptation. PYRA first utilizes parallel yielding adaptive weights to comprehensively perceive the data distribution in downstream tasks. A re-activation strategy for token modulation is then applied for tokens to be merged, leading to calibrated token features. Extensive experiments demonstrate that PYRA outperforms all competing methods under both low compression rate and high compression rate, demonstrating its effectiveness and superiority in maintaining both training efficiency and inference efficiency for large-scale foundation models. Our code is available at https://github.com/THU-MIG/PYRA.

Figures (6)

• Figure 1: (a) Full fine-tuning trains all parameters on downstream tasks and utilizes the trained model for inference, thereby lacking efficiency in both training and inference stages. (b) Model compression employs pruning to enhance inference efficiency, but the pruned model necessitates extensive re-training on large-scale data. (c) PEFT freezes the model backbone and only fine-tunes a small amount of parameters, yet retains the inference complexity. (d) Our training-inference efficient task adaptation incorporates the advantages of all existing pipelines by training inference-efficient models with minimal tunable parameters.
• Figure 2: Comparisons between simply combining ToMe bolya2022token and LoRA hu2021lora and our proposed PYRA. Red boxes represent the performance drop problem in low compression rates. Green boxes represent the adverse compression in high compression rates. See \ref{['sec:4_experiment']} for more results.
• Figure 3: The pipeline of our PYRA. PYRA conducts token modulation before the MHSA module in each transformer block. Inside PYRA, a pair of learnable modulation weight generators are leveraged to generate adaptive modulation weights parallelly. After that, generated weights modulate tokens through re-activation. The generators in PYRA can be trained along with the LoRA module $\theta_p$ in an end-to-end manner.
• Figure 4: Comparisons between PYRA and ToMe bolya2022token under different compression rates for ViT-Large (MAE) and DeiT-Base. Red boxes: PYRA mitigates the performance drops under low compression rate. Green boxes: PYRA elinimates adverse compression under high compression rate.
• Figure 5: t-SNE van2008visualizing visualization of ToMe bolya2022token and PYRA on the Oxford Pet parkhi2012cats dataset.