DARTS: Decoder-Aware Representation Tuning via Surgery for Model Merging

AI summary

Overview

• Research area: Model merging for large language models, specifically correcting representation bias in merged multi-task LLMs.
• Technical level: Advanced. The paper assumes familiarity with transformer architectures, causal attention, task vectors, and weight-space merging methods such as TIES-Merging and DARE.
• Scope: The paper analyzes representation bias in decoder-only models, proposes two decoder-specific corrections (entropy-weighted L1 loss and a per-position additive bias), and evaluates them on Llama-2-7B models merged across code, math, and instruction-following domains.

What This Paper Is About

Merging several fine-tuned LLMs into one model is cheap and popular, but merged models perform worse than their individual source models. Prior work (Yang et al., 2024a) explained part of this gap for encoder-based vision models through "representation bias" and fixed it with a small correction module trained with L1 loss. This paper asks whether that fix transfers to decoder-only LLMs, shows it does not transfer cleanly because of causal attention and unequal token importance, and proposes a decoder-aware correction method called DARTS.

Key Contributions

1. Diagnosis of decoder-specific representation bias. The authors analyze representation bias in decoder models and empirically demonstrate two challenges absent in encoders: position-dependent bias accumulation under the causal attention mask, and non-uniform position importance arising from the autoregressive generation process.
2. The DARTS method. They propose Decoder-Aware Representation Tuning via Surgery (DARTS), which combines a novel entropy-weighted L1 loss (upweighting correction at decision-critical, high-entropy positions) with a position correction module that applies a per-position additive bias to capture position-dependent error without network overparameterization.
3. Extensive multi-domain evaluation. Experiments on code generation (HumanEval), mathematical reasoning (GSM8K), and instruction following (AlpacaEval) with Llama-2-7B models merged from three specialized fine-tunes show DARTS improves over the standard surgery approach while adding roughly 0.1% of total parameters.
4. Isolation of what drives the gains. Ablations show that loss function choice matters more than surgery module architecture for decoder-based surgery, with the entropy-weighted L1 loss accounting for the majority of improvement and the position correction module providing complementary gains.

Main Findings

• Bias grows with token position. Per-token L1 representation bias between merged and individual Llama-2-7B models shows a clear increasing trend with position across all three domains, with later positions (100–256) exhibiting 2–3× higher bias than early positions, attributed to error accumulation under the causal mask.
• Loss design dominates architecture. On Task Arithmetic (λ = 0.5) with Standard surgery held fixed, entropy-weighted L1 reached 41.4 average (HumanEval 25.6, GSM8K 41.7, AlpacaEval 57.0), versus 37.6 for the original L1, 39.3 for KL on logits, and 30.5 for cross-entropy distillation.
• Position correction adds complementary gains. With EWL1 loss fixed, adding the position correction module changed HumanEval by −0.6, GSM8K by +1.7, and AlpacaEval by +1.6, for a 42.3 versus 41.4 average.
• Main results across four merging methods (Table 1 averages): Weight Averaging 32.0 (no surgery), 31.5 (Yang et al.), 32.3 (DARTS); Task Arithmetic 38.2, 37.6, 42.3; TIES-Merging 12.9, 14.4, 14.4; DARE 34.0, 33.8, 35.0.
• Largest single gain on Task Arithmetic. Under Task Arithmetic, DARTS raised AlpacaEval from 49.5 to 58.6 and GSM8K from 41.2 to 43.4, and HumanEval from 23.8 to 25.0.
• Bias reduction does not perfectly track task gains. On Task Arithmetic with λ = 0.7, L1 reduction was 14.2 for code (+4.9 points), 14.3 for math (−3.3 points), and 51.0 for instruct (+9.9 points).
• Statistically significant over Standard surgery. Across 6 seeds, DARTS averaged 42.6 ± 0.5 versus 40.6 ± 1.6 for Standard surgery, winning in every run; a Shapiro-Wilk test gave DARTS W = 0.80, p = 0.059 and Standard W = 0.77, p = 0.034, so a Mann-Whitney U test was used (U = 36.0, p = 0.002).
• Calibration data needs are small. Performance was relatively stable from N = 10 to N = 100 calibration samples, with as few as 20 samples achieving results competitive with larger budgets; using the full test set (23.17 HumanEval, 43.52 GSM8K, 59.45 AlpacaEval) was an exploratory upper bound only, kept out of main claims.
• Alternative position modules underperform or destabilize. In the variant comparison, "Add" scored 21.9/44.2/57.9 with 2.1M extra parameters, the "Parallel" variant collapsed on HumanEval at 12.2 with 2.2M extra parameters, and the paper's position correction scored 25.0/43.4/58.6 with 2.1M extra parameters.
• Learning rate robustness. Performance varied by less than 2 percentage points across learning rates from 10⁻⁴ to 5 × 10⁻³, with the default 10⁻³ achieving the best AlpacaEval score and highest average.
• Negligible deployment overhead. On a single NVIDIA A100-SXM4-80GB, total calibration wall-clock was 11.85 s for DARTS versus 11.01 s for Yang et al., with peak GPU memory of 31.49 GB versus 29.91 GB and checkpoint size of 4.46 MB versus 265 KB; surgery computation itself took about 4–5 s, and model loading about 7 s dominated.

Methodology in Plain English

The researchers start from an existing recipe: take several models fine-tuned from the same base (here Llama-2-7B, with Llama-2-7B-Instruct, WizardMath-7B-v1.0, and Llama-2-7B-Code), combine their weights with standard merging methods (Task Arithmetic, Weight Averaging, TIES-Merging, DARE) using mergekit, and then train a small add-on module that nudges the merged model's internal hidden states back toward those of each individual fine-tune.

They first measure how far apart those hidden states are, token by token, and find the gap widens for later tokens. Two changes follow. First, instead of a single uniform penalty, they weight the correction loss at each position by how uncertain the target model is there, so uncertain ("decision-critical") positions get corrected harder; the weights are normalized by the mean entropy so the overall loss scale stays comparable. Second, they add a simple lookup table of one bias vector per token position, which learns the systematic, input-independent offset at each spot without a network. The table and the low-rank content correction are both zero-initialized so training starts from an identity correction.

Each domain gets its own surgery module, trained for 500 iterations with Adam at a learning rate of 10⁻³ on 50 calibration prompts (tokenized to a maximum length of 512) held out from the evaluation sets, with hidden states precomputed in bfloat16. Evaluation uses HumanEval (164 problems, Pass@1), GSM8K (1,319 problems, chain-of-thought zero-shot, accuracy), and AlpacaEval (805 instructions, win rate against GPT-3.5-Turbo judged by GPT-4.1), run through lm-eval-harness with greedy decoding for the first two and sampling (temperature 0.7, top-p 0.9, max 1024 tokens) for AlpacaEval. They also sweep more than 124 merging configurations across the four merging families to pick the best representative baseline for each.

Why This Matters

• Impact on research: The paper extends representation surgery from encoder-based vision models to autoregressive decoders, identifies why the encoder assumptions break, and reports that loss weighting matters more than module architecture in this setting. It also provides a negative result for the same relationship between hidden-state bias reduction and task score that holds in encoders.
• Real-world applications:
  • Consolidating separate code, math, and instruction-tuned checkpoints into one deployable model without retraining.
  • Serving a single model footprint instead of multiple specialized models, since DARTS adds only about 0.1% parameters and applies a static correction with latency identical to the base model.
  • Low-cost offline calibration on a single commodity GPU (about 36 seconds total for all three domains in the reported setup).
  • Drop-in replacement scenarios where a standard single-checkpoint model is required and no specialized inference engine is available.
• Industry relevance: Practitioners in the open-source LLM ecosystem already merge checkpoints routinely; DARTS targets exactly the performance gap those practitioners observe, and the reported calibration costs and checkpoint sizes (265 KB–4.46 MB on disk, under 0.04% of backbone parameters) make the approach practical for teams without large training budgets.

Future Directions

• Sequence length generalization. The position bias table is sized to T_max = 512 and cannot extrapolate; positions beyond that receive zero position correction. The authors suggest exploring continuous position functions or relative position encodings.
• Scaling and architecture transfer. Evaluation was exclusively on Llama-2-7B; whether DARTS transfers to larger models (13B, 70B) or different architectures such as Mistral remains untested.
• Domain asymmetry. The entropy weighting may deprioritize positions critical for code generation, such as syntactically deterministic but semantically critical tokens, which the authors offer as a possible explanation for domain-dependent gains.
• Comparison and combination with inference-time methods. The paper positions DARTS as complementary to activation steering, mixture-of-experts routing, and expert composition, leaving open how the static correction interacts with those flexible but higher-latency approaches.

Target Audience

Researchers and engineers working on model merging, multi-task consolidation of LLMs, and efficient post-training adaptation. It is most useful to readers already comfortable with transformer internals, causal attention, task vectors, and the mergekit ecosystem, since the paper's core argument depends on how hidden states propagate under autoregressive masking.

Authors’ 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 ($0.1\%$ of total parameters).