DARTS: Decoder-Aware Representation Tuning via Surgery for Model Merging
Organizations: ServiceNow · University of Twente
Abstract
Model merging combines multiple task-specific fine-tuned LLMs into a single multi-task model without additional training. However, merged models are known to suffer from representation bias: systematic drift between the merged model's hidden states and those of each individual source model. Prior work (Yang et al., 2024a) study and mitigate this bias for encoder-based vision models using a lightweight correction module trained with L1 loss. However, such bias is not studied for decoder models due to their autoregressive nature. We analyze the problem of representation bias in decoder models, and show two challenges absent in encoders: (1) the causal attention mask causes bias to accumulate across token positions, requiring position-dependent correction; and (2) not all token positions are equally important, i.e., high-entropy (decision-critical) positions matter far more than low-entropy ones. To address these challenges, we propose Decoder-Aware Representation Tuning via Surgery (DARTS). DARTS employs a novel entropy-weighted L1 loss to upweight correction at high-entropy positions where errors most affect generation quality, and a per-position additive bias that captures position-dependent error without overparameterization. We perform extensive evaluation on three domains: code generation (HumanEval), mathematical reasoning (GSM8K), and instruction following (AlpacaEval) on Llama-2-7B models, and show DARTS achieves significant improvement over the standard surgery approach while adding negligible parameters ( of total parameters).
Figures & tables
| Merging Method | Surgery Method | HumanEval | GSM8K | AlpacaEval | Avg. |
|---|---|---|---|---|---|
| Weight Averaging | None (baseline) | 27.4 | 31.2 | 37.3 | 32.0 |
| Yang et al. (2024a) | 26.2 | 28.7 | 39.5 | 31.5 | |
| DARTS (ours) | 26.8 | 28.9 | 41.3 | 32.3 | |
| Task Arithmetic | None (baseline) | 23.8 | 41.2 | 49.5 | 38.2 |
| Yang et al. (2024a) | 23.2 | 42.5 | 47.0 | 37.6 | |
| DARTS (ours) | 25.0 | 43.4 | 58.6 | 42.3 |
| Model | HE | GSM | AE |
|---|---|---|---|
| Llama-2-7B-Base | 12.8 | 4.5 | 1.5 |
| Llama-2-7B-Instruct | 15.2 | 24.5 | 58.2 |
| WizardMath-7B-v1.0 | 12.8 | 39.7 | 7.0 |
| Llama-2-7B-Code | 28.1 | 5.2 | 0.2 |
| Merged Model | Domain | Task Score | pp | ||
|---|---|---|---|---|---|
| (%) | No Surg. | DARTS | Change | ||
| TA ( ) | Code | 8.2 | 23.8 | 25.0 | +1.2 |
| Math | 11.3 | 41.2 | 43.4 | +2.2 | |
| Instruct | 42.9 | 49.5 | 58.6 | +9.1 | |
| TA ( ) | Code | 14.2 | 21.9 | 26.8 | +4.9 |
| Math | 14.3 | 46.2 | 42.9 | 3.3 | |
| Seed | DARTS (Avg.) | Standard (Avg.) |
|---|---|---|
| default | 42.3 | 37.6 |
| 7 | 42.2 | 41.7 |
| 29 | 42.3 | 40.4 |
| 71 | 43.3 | 41.8 |
| 164 | 42.2 | 41.2 |
| 190 | 43.0 | 41.1 |
| Loss Function | HE | GSM | AE | Avg. |
|---|---|---|---|---|
| (original) | 23.2 | 42.5 | 47.0 | 37.6 |
| KL on logits | 20.1 | 43.4 | 54.5 | 39.3 |
| Cross-entropy distill. | 18.3 | 33.1 | 40.0 | 30.5 |
| (ours) | 25.6 | 41.7 | 57.0 | 41.4 |
| Module | HE | GSM | AE |
|---|---|---|---|
| Standard | 25.6 | 41.7 | 57.0 |
| + Position corr. (ours) | 25.0 | 43.4 | 58.6 |
| 0.6 | +1.7 | +1.6 |
Appendix figures & tables9 assets
Supplementary material from the paper’s appendix.
Appendix
| Benchmark | Total Samples | Calibration Split | Evaluation Split |
|---|---|---|---|
| HumanEval (Code) | 164 | 50 | 114 |
| GSM8K (Math) | 1,319 | 50 | 1,269 |
| AlpacaEval (Instruct) | 805 | 50 | 755 |
| Phase | Yang et al. (2024a) | DARTS |
|---|---|---|
| Load merged model | 3.37 s | 3.46 s |
| Load target model | 3.43 s | 3.37 s |
| Precompute (2 forward) | 1.81 s | 1.87 s |
| Training loop (500 iter) | 2.40 s | 3.15 s |
| Total wall-clock | 11.01 s | 11.85 s |
| Peak GPU memory | 29.91 GB | 31.49 GB |
| Calibration | HE | GSM | AE |
|---|---|---|---|
| 10 | 24.4 | 42.5 | 54.0 |
| 20 | 24.4 | 44.8 | 58.4 |
| 30 | 26.2 | 43.8 | 56.8 |
| 50 | 25.0 | 43.4 | 58.6 |
| 100 | 26.2 | 42.8 | 55.4 |
| Full Test | 23.17 | 43.52 | 59.45 |
| Loss Function | Surgery Module | HumanEval | GSM8K | AlpacaEval | Avg. |
|---|---|---|---|---|---|
| Standard | 23.2 | 42.5 | 47.0 | 37.6 | |
| + Pos. corr. | 21.3 | 43.1 | 48.8 | 37.7 | |
| Entropy-Weighted | Standard | 25.6 | 41.7 | 57.0 | 41.4 |
| Entropy-Weighted | + Pos. corr. | 25.0 | 43.4 | 58.6 | 42.3 |
| Variant | HE | GSM | AE | Extra Params |
|---|---|---|---|---|
| Standard | 25.6 | 41.7 | 57.0 | 0 |
| Add | 21.9 | 44.2 | 57.9 | 2.1M |
| Parallel | 12.2 | 45.0 | 55.5 | 2.2M |
| Pos. Corr. (Ours) | 25.0 | 43.4 | 58.6 | 2.1M |
| Steps | HE | GSM | AE |
|---|---|---|---|
| 50 | 25.6 | 42.3 | 57.9 |
| 100 | 25.0 | 42.9 | 56.5 |
| 200 | 25.0 | 43.4 | 57.9 |
| 400 | 24.4 | 43.0 | 57.0 |
| 500 (default) | 25.0 | 43.4 | 58.6 |
| 1000 | 25.6 | 44.3 | 58.4 |
| Rank | HE | GSM | AE | Params/domain |
|---|---|---|---|---|
| 16 | 25.0 | 43.4 | 58.6 | 131K + 2.1M |
| 32 | 25.0 | 44.3 | 56.4 | 262K + 2.1M |
| 64 | 25.0 | 42.8 | 56.0 | 524K + 2.1M |
| 128 | 23.8 | 42.6 | 58.1 | 1.04M + 2.1M |
| Seed | HE | GSM | AE |
|---|---|---|---|
| default | 25.0 | 43.4 | 58.6 |
| 7 | 24.4 | 45.0 | 57.3 |
| 29 | 25.0 | 43.6 | 58.3 |
| 71 | 25.0 | 45.0 | 59.9 |
| 164 | 25.0 | 44.4 | 57.1 |
| 190 | 25.0 | 45.0 | 59.0 |
| Task | Mean Diff | -stat | -test -val | -stat | Wilcoxon -val |
|---|---|---|---|---|---|
| HumanEval | +0.30 | 0.690 | 0.521 | 7.500 | 0.688 |
| GSM8K | +1.56 | 2.849 | 0.036 | 0.000 | 0.031 |
| AlpacaEval | +3.95 | 2.459 | 0.057 | 0.000 | 0.031 |
Language Model Merging
Language model merging combines capabilities from multiple models or adapters by composing their weights, updates, or outputs.
See all 110 papers