DARTS: Decoder-Aware Representation Tuning via Surgery for Model Merging
2026-08-28 • Machine Learning
Machine Learning
AI summaryⓘ
The authors study a problem called representation bias that happens when combining multiple trained language models into one without retraining. Unlike earlier work focused on models that look at data all at once, they analyze decoder models, which process data step-by-step and face unique challenges like errors building up over positions and some steps being more critical than others. To fix this, they introduce DARTS, a method that corrects important positions more carefully and adjusts errors depending on token position without adding many extra parameters. Their tests show that DARTS improves performance on tasks like code generation, math reasoning, and instruction following.
model merginglarge language modelsrepresentation biasdecoder modelsautoregressiveentropy-weighted lossL1 lossposition-dependent correctioncode generationmathematical reasoning
Authors
Aaryan Ajay Sharma, Sai Nishanth Padala, Seganrasan Subramanian
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).