Research map / Architectures and efficiency / Pruning and efficient tuning
Fine-Tuning Attention Modules Only: Enhancing Weight Disentanglement in Task Arithmetic
Read the paperCheck your idea against 108,967 papers
AI-extracted summary. Written by AI from the paper's full text. Numbers and claims can be wrong or missing context, so check the paper before you cite or build on it.
Proposes fine-tuning only the attention modules (W_q, W_k, W_v, W_o) of pre-trained Transformers to improve weight disentanglement and accuracy of task arithmetic, while separating representation and task-specific modules in the analysis.
Problem
Task arithmetic merges task vectors into a unified model, but interference between tasks (lack of weight disentanglement) degrades per-task accuracy; NTK linearization improves disentanglement but hurts individual-model accuracy and doubles training cost.
Approach
Attention modules exhibit kernel behavior under post-hoc linearization, so fine-tuning only attention weights (non-linearly) gives both good individual accuracy and strong weight disentanglement at lower cost; the analysis also shows weight disentanglement emerges mainly from the representation module, while task-specific heads hurt it.
Key results
- Attention-only fine-tuning of CLIP ViT-B/32 reaches 78.37% absolute multi-task accuracy vs 76.26% for NTK linearization and 70.00% for non-linear fine-tuning (87.42% normalized).
- On Flan-T5-base GLUE tasks, the method averages 85.30% vs 77.77% for NTK linearization and 73.25% for non-linear fine-tuning.
- Attention-only fine-tuning trains 63.7M of 222M parameters and takes 10:15 minutes vs 40:59 for NTK linearization.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| Fine-tune attention modules only instead of using NTK linearization | ViT-B-32 task arithmetic · 8-task task addition | absolute accuracy (%): 76.26 -> 78.37 | Helped |
| Fine-tune attention modules only instead of using NTK linearization | ViT-B-16 task arithmetic · 8-task task addition | absolute accuracy (%): 79.01 -> 80.44 | Helped |
| Fine-tune attention modules only instead of using NTK linearization | ViT-L-14 task arithmetic · 8-task task addition | absolute accuracy (%): 85.53 -> 87.91 | Helped |
| Fine-tune attention modules only instead of using NTK linearization | Flan-T5-base · GLUE (7-task average) | accuracy (%): 77.77 -> 85.30 | Helped |
| Add attention biases to attention-only fine-tuning (ablation) | ViT task arithmetic · 4-setting ablation | multi-task absolute accuracy (%): 78.37 -> 77.71 | Hurt |
| Add attention and MLP weight fine-tuning (ablation) | ViT task arithmetic · 4-setting ablation | multi-task absolute accuracy (%): 78.37 -> 76.52 | Hurt |
| Add attention and MLP weight fine-tuning with biases (ablation) | ViT task arithmetic · 4-setting ablation | multi-task absolute accuracy (%): 78.37 -> 77.80 | Hurt |
Surprising findings
- Classification heads degrade weight disentanglement; disentanglement arises mainly from the representation module.
- Fine-tuning MLP modules in addition to attention did not improve results over attention-only fine-tuning in the ablation.
Limitations
- Task vectors require identical architectures and initializations due to element-wise operations.
- Evaluated only on CLIP ViTs and Flan-T5-base with 7-8 task sets.
Experimental setup
Datasets: Cars DTD SUN397 EuroSAT GTSRB MNIST SVHN RESISC45 GLUE (CoLA, MNLI, MRPC, QQP, RTE, SST2, STSB) ImageNet
Benchmarks: Task arithmetic 8-task vision benchmark GLUE (7 tasks)
Metrics: absolute average accuracy of unified model normalized accuracy single-task accuracy disentanglement error (KL divergence) training time trainable parameters
Baselines: Non-linear fine-tuning NTK linearization (Ortiz-Jimenez et al.) TIES-Merging TAFT
Builds on: task arithmetic NTK linearization weight disentanglement post-hoc linearization
Closest papers by meaning
- DiTASK: Multi-Task Fine-Tuning with Diffeomorphic Transformations (CVPR 2025 · similarity 0.86)
- Parameter-Efficient Multi-Task Model Fusion with Partial Linearization (ICLR 2024 · similarity 0.86)
- Leveraging Submodule Linearity Enhances Task Arithmetic Performance in LLMs (ICLR 2025 · similarity 0.85)
- NTKMTL: Mitigating Task Imbalance in Multi-Task Learning from Neural Tangent Kernel Perspective (NeurIPS 2025 · similarity 0.85)
- The Hedgehog & the Porcupine: Expressive Linear Attentions with Softmax Mimicry (ICLR 2024 · similarity 0.84)
- Coeff-Tuning: A Graph Filter Subspace View for Tuning Attention-Based Large Models (CVPR 2025 · similarity 0.84)
Writing a paper in this area? scime draft is a free LaTeX editor with AI included, real-time co-authors and a check of your references and venue rules before you submit. It also imports OpenAI Prism projects.