PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Aaryan Ajay Sharma, Sai Nishanth Padala, Seganrasan Subramanian

arXiv:2608.28547v1cs.LG

TL;DR

DARTS addresses representation bias in decoder-based language models, where causal attention creates position-dependent bias and token positions differ in importance. It combines entropy-weighted L1 loss with per-position additive corrections, improving multi-task performance across code, math, and instruction-following domains while adding 0.1% of model parameters.

  • Problem

    Representation bias in decoder models is understudied, despite merged models exhibiting hidden-state drift and autoregressive generation introducing position-dependent bias and unequal token importance.

  • Method

    DARTS combines entropy-weighted L1 loss, which focuses correction on decision-critical positions, with a per-position additive bias module for position-dependent correction.

  • Results

    DARTS significantly improves multi-task performance across code, math, and instruction-following tasks over the original surgery approach while adding 0.1% of model parameters.

  • Takeaways & Limitations

    Entropy-weighted loss accounts for most performance gains, while position correction provides complementary improvements in merged Llama-2-7B models.

  • Takeaways & Limitations

    The evaluation uses only Llama-2-7B, so transfer to larger models or different architectures remains untested.

Abstract

from arXiv · show

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).

1 Introduction

Model merging consolidates specialized fine-tuned models, but merged models can exhibit representation bias and task degradation. DARTS adapts representation surgery to decoder-specific, position-dependent errors and unequal token importance.

  • Motivation: Model merging combines specialized fine-tuned models into one multi-task model without additional training, reducing the need to deploy separate models.This consolidation is motivated by the growing serving cost of maintaining many task-specific models.
  • Representation Bias: Merged models can underperform their individual fine-tuned counterparts, with degradation varying across tasks, merging methods, and model scales.The paper identifies understanding and reducing this performance gap as an open challenge.
  • Representation Bias: Prior work links representation bias in merged encoder models to task degradation, but the original position-agnostic surgery framework targets vision encoders rather than decoder LLMs.Encoder patch representations are roughly uniform because bidirectional attention contextualizes all positions similarly.
  • Decoder-Specific Challenges: Decoder models accumulate representation bias across token positions because causal attention propagates earlier errors into later representations.Figure 1 reports later-position bias often reaching 2–3× that of early positions across code, math, and instruction domains.
  • Decoder-Specific Challenges: Decoder generation also assigns unequal importance to positions, with high-entropy decision-critical tokens affecting output quality more than confident low-entropy tokens.This motivates weighting corrections according to prediction entropy rather than treating all positions equally.
  • DARTS: DARTS combines entropy-weighted L1 loss with a per-position additive correction module to target important positions and structured positional bias.The method is evaluated on merged Llama-2-7B models spanning code generation, mathematical reasoning, and instruction following.
  • Results: 0.1% of parameters is the reported overhead for DARTS, which improves multi-task performance across diverse tasks.The paper presents this as an extensive experimental result over the original surgery approach.

2 Related Work

Related work frames model merging as an alternative to joint multi-task training and representation surgery as a way to correct merger-induced hidden-state discrepancies. Inference-time steering and routing provide complementary but more deployment-complex alternatives.

  • Model Merging for Multi-Task Learning: Model merging offers a practical alternative to centralized data collection and joint multi-task training by combining models trained on different tasks.Task Arithmetic introduced task vectors, while later methods addressed conflicts and redundancy in weight space.
  • Representation Surgery: Representation surgery corrects systematic hidden-state divergence in merged vision models using a lightweight adapter-like module trained with L1 loss.Prior work connected smaller representation bias with better task performance in merged models.
  • Inference-Time Interventions and Routing: Inference-time steering dynamically modifies generation behavior, while mixture-of-experts and routing select among expert subnetworks during inference.These approaches maintain flexibility through dynamic intervention or selection rather than permanent consolidation.
  • Inference-Time Interventions and Routing: DARTS is positioned for low-overhead permanent consolidation, whereas inference-time methods provide greater flexibility at the cost of deployment complexity.The comparison describes the approaches as complementary rather than mutually exclusive.

3 Causal Representation Bias in Model Merging

Decoder-model representation bias is position-dependent because causal attention accumulates earlier errors across later tokens. DARTS addresses this with per-position correction and entropy-weighted training that prioritizes decision-critical positions.

  • Causal Representation Bias: Decoder representation bias increases across token positions because causal attention propagates earlier merging errors forward.Later tokens aggregate their own error with accumulated errors from preceding positions.
  • Causal Representation Bias: Tokens at positions 100–256 often exhibit 2–3× the representation bias of early tokens across code, math, and instruction-following domains.
  • Causal Representation Bias: A position-agnostic correction must compromise between under-correcting later tokens and over-correcting earlier tokens.
  • Position Correction Module: DARTS adds a direct per-position bias to capture systematic offsets while its content correction handles input-dependent errors.The two components decouple position-consistent bias from input-specific bias.
  • Position Correction Module: The direct position correction uses independent free parameters for each position, avoiding network parameterization and requiring no additional hyperparameters.For Llama-2-7B, the modules add approximately 2.2M parameters per domain and 0.1% of the model’s total parameters overall.
  • Entropy-Weighted L1 Loss: Entropy-weighted L1 loss emphasizes positions where the target model has high prediction entropy rather than treating all positions equally.Unit-mean entropy normalization keeps the overall loss scale comparable to standard L1.
  • Entropy-Weighted L1 Loss: High-entropy positions are decision-critical because hidden-state errors can change the next-token choice and cascade through later generation.Low-entropy positions are more robust because moderate errors may leave the argmax unchanged.

4 Experimental Setup

The evaluation studies merged Llama-2-7B models built from instruction, math, and code fine-tunes across three complementary benchmarks. It compares individual models, unspecialized merging, standard surgery, and DARTS while also measuring representation-bias reduction.

  • Models and Merging: The study merges Llama-2-7B-Instruct, WizardMath-7B-v1.0, and Llama-2-7B-Code into unified models.
  • Models and Merging: The experiments sweep over 124 configurations spanning Task Arithmetic, Weight Averaging, TIES-Merging, and DARE.
  • Benchmarks: HumanEval reports Pass@1, GSM8K reports Accuracy, and AlpacaEval reports Win Rate against GPT-3.5-Turbo judged by GPT-4.1.
  • Metrics: L1 Reduction measures the fraction of hidden-state representation bias removed by surgery, with higher values indicating greater bias removal.
  • Experimental Comparisons: The comparison includes individual models, merged models without surgery, Standard + L1 surgery, and DARTS.

5 Results

Table 1 compares Standard + L1 with DARTS across merged-model configurations and the three evaluation benchmarks.

  • Main Results: Table 1 compares the original Standard + L1 surgery approach against DARTS.
  • Main Results: The comparison covers multiple merged-model configurations.
  • Main Results: The comparison spans all three benchmarks used in the evaluation.

6 Analysis

The analysis evaluates DARTS against standard surgery across merged-model configurations and three benchmarks, including multi-seed robustness. DARTS significantly outperforms standard surgery, with gains robust to calibration-data variation.

  • Statistical robustness: 2.0 percentage points is the modest absolute improvement in average score, with DARTS outperforming Standard surgery in every one of 6 seeds.The Mann-Whitney U test gives U = 36.0 and p = 0.002.
  • Statistical robustness: DARTS’s multi-seed advantage is robust to calibration-data variation.The calibration-data shuffle order was controlled across six random seeds.
  • Main results: DARTS significantly outperforms Standard surgery across the evaluated merged-model configurations and HumanEval, GSM8K, and AlpacaEval benchmarks.Table 1 compares Standard + L1 with DARTS across all three benchmarks.
  • Evaluation setup: Individual-model baselines use HumanEval pass@1, GSM8K exact match, and AlpacaEval win rate percentage.These metrics define the benchmark comparisons reported for the individual fine-tuned models.

7 Ablation Studies

The ablations show that DARTS’s performance is driven primarily by entropy-weighted L1 loss, while position correction adds complementary improvements. Performance remains stable across a broad learning-rate range.

  • Loss-function ablation: Loss-function choice matters more than surgery-module architecture for decoder-based surgery.Table 5 isolates loss choice by holding the Standard surgery architecture constant.
  • Surgery-module ablation: Adding position correction to EWL1 yields modest but significant improvements over Standard surgery with EWL1.Table 6 isolates the contribution of the position correction module.
  • Learning-rate sensitivity: Less than 2 percentage points separates DARTS’s average scores across learning rates from 10−4 to 5 × 10−3.No single learning rate causes catastrophic degradation on any benchmark.
  • Learning-rate sensitivity: 10−3 achieves the best AlpacaEval score and highest average, while 5 × 10−4 and 5 × 10−3 remain competitive.The sensitivity analysis reports stable performance across two orders of magnitude.

8 Conclusion

The paper extends representation surgery to decoder language models by addressing position-dependent bias accumulation and unequal token-position importance. DARTS combines entropy-weighted L1 loss with position correction, and experiments show significant gains over original surgery.

  • Conclusion: DARTS extends representation surgery from encoder-based vision models to decoder-based language models.The conclusion identifies decoder-specific representation-bias challenges under autoregressive generation.
  • Conclusion: Causal attention causes bias to accumulate across positions, while language generation makes decision-critical positions more important than others.These are the two challenges identified as specific to the autoregressive setting.
  • Conclusion: Entropy-weighted L1 loss accounts for the majority of performance gains, with position correction providing complementary improvements.The conclusion attributes the dominant contribution to loss design and the additional contribution to position correction.
  • Conclusion: DARTS achieves significant improvement over original surgery on Llama-2-7B merged models across code, math, and instruction-following domains.The evaluated domains correspond to HumanEval, GSM8K, and AlpacaEval.

Limitations

DARTS has limitations in sequence-length extrapolation, model and domain scope, and calibration-resource requirements. Its entropy weighting may also behave unevenly across domains.

  • Sequence length generalization: Positions beyond Tmax = 512 receive zero position correction, although content correction still applies.Future work could explore continuous position functions or relative position encodings.
  • Scale: Evaluation is limited to Llama-2-7B, leaving transfer to larger models and architectures such as Mistral untested.
  • Domain asymmetry: Entropy weighting may deprioritize syntactically deterministic but semantically critical code-generation positions, potentially causing domain-dependent gains.
  • Calibration cost: Calibration requires forward passes through both merged and target models, with model loading dominating per-domain wall-clock cost.On one A100, model loading takes approximately 7 seconds per configuration, while surgery computation takes approximately 4–5 seconds.

F Effect of Training Steps

DARTS remains competitive across training-step budgets, with only modest task-specific gains from longer training. Rank 16 offers the best overall balance, and results are robust across calibration seeds.

  • Effect of Training Steps: 50 training iterations already produce competitive results, while 1000 iterations yield a modest GSM8K improvement without degrading other benchmarks.
  • Surgery rank: Rank 16 achieves the best overall balance; ranks 32 and 64 slightly hurt AlpacaEval, while rank 128 degrades HumanEval.The results suggest that higher-rank modules overfit on the small calibration set.
  • Seed variation: Across 6 random seeds, low standard deviations confirm robustness to calibration-sample choice: HE ±0.2, GSM ±0.8, and AE ±1.0.The seeds control calibration data shuffle order.

I Granular Statistical Analysis of Seed Variation

Seed-level statistical tests show that DARTS improves GSM8K and AlpacaEval relative to Standard surgery, while HumanEval remains statistically indistinguishable. A mixed-effects analysis finds a positive overall DARTS effect across tasks and seeds.

  • Task-level tests: DARTS provides statistically significant improvements on GSM8K and AlpacaEval under the Wilcoxon signed-rank test.Parametric t-tests show significant to marginally significant improvements on these tasks.
  • Task-level tests: HumanEval performance remains statistically indistinguishable from the Standard surgery baseline.
  • Mixed-effects analysis: β = 0.019, z = 3.224, p = 0.001 for DARTS in a linear mixed-effects model across tasks and seeds.The model uses Score ∼Model + Task with Seed as a random intercept.
Loading 2608.28547v1…