Research map / Architectures and efficiency / Pruning and efficient tuning
When is Task Vector Provably Effective for Model Editing? A Generalization Analysis of Nonlinear Transformers
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.
Provides the first generalization analysis of task arithmetic (addition, negation, out-of-domain generalization, low-rank/sparse task vectors) for a one-layer nonlinear Transformer trained by SGD, validated on Colored-MNIST and unlearning with Phi-1.5.
Problem
Task arithmetic (adding/negating task vectors) works empirically for model editing, but its theoretical basis for nonlinear Transformers is unexplained; the paper asks when and why multi-task learning, unlearning, and out-of-domain generalization via task vectors succeed.
Approach
Analyzes SGD feature learning in a one-layer single-head Transformer on binary majority-vote tasks defined by discriminative patterns, deriving conditions on arithmetic coefficients λ as a function of inter-task correlation α (aligned, irrelevant, contradictory) and out-of-domain correlations γ_i, plus justification for low-rank and magnitude-pruned task vectors.
Key results
- On Colored-MNIST, irrelevant tasks allow multi-task learning with lambda>=1 and unlearning with lambda<=0 (estimated correlation alpha=0.164), matching the theory.
- Unlearning HP1 with lambda=-1 on Phi-1.5 drops Rouge-L on HP1 by 37.23% and on the aligned HP2 by 34.71%, versus 15.13% on the less-aligned PP.
- With Phi-3-small (7B) using LoRA task vectors, unlearning HP1 at lambda=-1 reduces HP1 Rouge-L by 55.61% and HP2 by 52.29%.
Measured effects reported in the paper
| What was changed | Where | Measured | Effect |
|---|---|---|---|
| task addition (positive lambda on second task vector) | irrelevant and aligned tasks · Colored-MNIST multi-task | test accuracy | Helped |
| task negation (negative lambda on second task vector) | irrelevant tasks · Colored-MNIST unlearning | T2 accuracy drop while T1 retained | Helped |
| task negation | contradictory tasks · Colored-MNIST unlearning | T2 accuracy | Helped |
| task addition for contradictory tasks | contradictory tasks · Colored-MNIST multi-task | test accuracy | Hurt |
| searched arithmetic coefficients on two irrelevant-ish task vectors for out-of-domain target | out-of-domain target task · Colored-MNIST | target task accuracy: 91.74 vs 88.10 (best single vector) and 45.06 | Helped |
| negative lambda (unlearning HP1) on full-rank task vector | Phi-1.5 HP1 unlearning · Rouge-L on HP1, HP2, PP | Rouge-L: HP1 down 37.23%, HP2 down 34.71%, PP down 15.13% at lambda=-1 | Helped |
| low-rank (rank 32) task vector instead of full-rank | Phi-1.5 HP1 unlearning · Rouge-L on HP1, HP2, PP at lambda=-1 | Rouge-L: HP1 down 35.53% vs 37.23% full-rank; HP2 down 24.11% vs 34.71%; PP down 13.83% vs 15.13% | No clear effect |
| larger model (Phi-3-small 7B) with LoRA rank-32 task vector | HP1 unlearning · Rouge-L on HP2 at lambda=-1 | Rouge-L decrease on aligned task HP2: HP2 decrease 52.29% vs 24.11% for Phi-1.5 low-rank; HP2 unlearning stronger with 7B | Helped |
Surprising findings
- Unlearning a task also substantially degrades the aligned task (HP2), showing collateral forgetting beyond the target.
Limitations
- Theory is restricted to a simplified one-layer single-head Transformer with binary tasks and a single discriminative pattern per task.
- NLP experiments are small-scale, use Rouge-L, and search over few lambda values.
- Relies on estimated task correlations and a practical margin replacing the Hinge-loss margin in the out-of-domain conditions.
Experimental setup
Datasets: Colored-MNIST Harry Potter 1 Harry Potter 2 Pride and Prejudice
Benchmarks: Colored-MNIST parity tasks Harry Potter 1/2 and Pride and Prejudice unlearning (Rouge-L)
Metrics: test accuracy Rouge-L generalization error
Baselines: fine-tuned task vectors Psi*_T1 and Psi*_T2 joint fine-tuning TIES DARE
Builds on: task arithmetic feature learning framework for Transformers neural tangent kernel LoRA
Closest papers by meaning
- Efficient Model Editing with Task-Localized Sparse Fine-tuning (ICLR 2025 · similarity 0.88)
- Training-Free Task Vectors for LLM Behavioral Control (NeurIPS 2026 · similarity 0.87)
- Task Generalization with Autoregressive Compositional Structure: Can Learning from $D$ Tasks Generalize to $D^T$ Tasks? (ICML 2025 · similarity 0.86)
- Parameter-Efficient Multi-Task Model Fusion with Partial Linearization (ICLR 2024 · similarity 0.86)
- AdaMerging: Adaptive Model Merging for Multi-Task Learning (ICLR 2024 · similarity 0.86)
- Task Vector Geometry Underlies Dual Modes of Task Inference in Transformers (NeurIPS 2026 · similarity 0.86)
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.