DARTS: Decoder-Aware Representation Tuning via Surgery for Model Merging
DARTS uses entropy-weighted L1 and position biases to correct autoregressive decoder biases, improving multi-task model merging with only 0.1% parameter overhead.
Key Findings
Methodology
This work analyzes position-dependent representation bias in decoder models, revealing bias accumulation along token positions due to causal attention. It introduces DARTS, combining an entropy-weighted L1 loss that emphasizes correction at high-entropy (decision-critical) tokens, and a per-position additive bias to capture systematic position-dependent errors. Extensive experiments on Llama-2-7B models merged from task-specific fine-tunes across code, math, and instruction domains demonstrate that DARTS significantly reduces bias and improves task performance, with negligible parameter increase (0.1%). The approach involves calibrating a small correction module using limited samples and evaluating on multiple benchmarks, including HumanEval, GSM8K, and AlpacaEval, showing consistent gains over baseline surgery methods.
Key Results
- DARTS achieves over 15% average performance improvement across tasks, with Pass@1 reaching 85% in code, 78% accuracy in math reasoning, and a 72% win rate in instruction following. The L1 bias reduction exceeds 40%, and the total parameter overhead remains at just 0.1%. Ablation studies confirm that entropy-weighted loss and position bias contribute approximately 8-10% each to the overall gains, with their combination yielding the best results.
- Compared to traditional uniform correction, DARTS effectively addresses the position-dependent bias accumulation, especially in later tokens, leading to more accurate and stable generation. The method’s low parameter cost makes it suitable for deployment in resource-constrained environments, and its modular design allows easy integration with existing merging pipelines.
- The experiments demonstrate that focusing correction efforts on high-entropy, decision-critical positions yields the most substantial improvements, validating the importance of position-aware bias correction in autoregressive models.
Significance
This research advances understanding of representation bias in decoder models, highlighting the importance of position-dependent correction. By addressing bias accumulation and position importance, DARTS offers a practical, low-cost solution to improve the robustness and accuracy of merged multi-task models. Its implications extend to deploying more reliable large language models in real-world applications, reducing performance degradation caused by model merging. The approach bridges a critical gap in model fine-tuning and merging, fostering more efficient multi-task learning frameworks that can scale to larger models and diverse tasks, ultimately contributing to more intelligent and adaptable AI systems.
Technical Contribution
The paper introduces a novel bias correction framework tailored for autoregressive decoders, combining an entropy-weighted L1 loss with a position-specific bias module. This dual mechanism effectively captures the position-dependent nature of representation bias, which traditional uniform correction methods overlook. The position bias parameters are learned independently for each token position, preventing overparameterization while maintaining flexibility. The entropy weighting dynamically emphasizes correction at decision-critical tokens, aligning correction capacity with task importance. These innovations enable high-fidelity bias mitigation with minimal parameter overhead, setting a new standard for model merging techniques in autoregressive architectures.
Novelty
This work is the first to systematically analyze and address position-dependent representation bias in decoder models during model merging. Unlike prior methods designed for bidirectional encoders, DARTS explicitly models the recursive bias accumulation and importance heterogeneity inherent in autoregressive models. Its combination of entropy-weighted loss and position bias parameters represents a significant departure from existing uniform correction strategies, providing a scalable, low-parameter solution that effectively mitigates bias and enhances multi-task performance.
Limitations
- The approach has been validated primarily on Llama-2-7B models; its effectiveness on larger models or different architectures remains to be tested, which could reveal scalability issues.
- The correction relies on limited calibration samples, potentially limiting robustness in highly diverse or out-of-distribution tasks.
- While parameter-efficient, the method may still face challenges in extremely complex or specialized tasks where bias patterns are more intricate and require more sophisticated modeling.
Future Work
Future directions include developing adaptive, dynamic bias correction mechanisms that can adjust during inference, exploring integration with reinforcement learning for bias mitigation, and extending the framework to multi-modal models. Additionally, scaling the approach to larger models and more diverse tasks, along with automated hyperparameter tuning for entropy weighting, could further enhance its applicability and robustness.
AI Executive Summary
The rapid growth of large language models (LLMs) has revolutionized AI capabilities across diverse tasks. However, merging multiple task-specific fine-tuned models into a single multi-task model often introduces systematic representation biases, especially in decoder architectures. Unlike encoder models, decoder-based models process tokens autoregressively, causing biases to accumulate along token positions. This positional bias leads to performance degradation, particularly in later tokens, and is compounded by the fact that not all tokens are equally important for output quality.
Addressing this challenge, the paper introduces DARTS, a novel bias correction framework tailored for autoregressive models. DARTS employs an entropy-weighted L1 loss that emphasizes correction at high-entropy, decision-critical tokens, ensuring the model focuses on positions where errors most impact generation quality. Additionally, it introduces a position bias parameter for each token position, capturing the systematic, position-dependent bias without overparameterization. This combination allows for efficient, targeted bias mitigation.
Extensive experiments on Llama-2-7B models merged from three specialized fine-tunes—covering code, math, and instruction tasks—demonstrate that DARTS significantly reduces representation bias, with over 40% bias reduction and performance improvements exceeding 15% across benchmarks. The parameter overhead is minimal (0.1%), making it practical for deployment. Ablation studies confirm the importance of entropy weighting and position bias, validating the core design principles.
This work advances understanding of autoregressive bias dynamics and provides a scalable, low-cost solution for multi-task model merging. Its implications span improved robustness, accuracy, and deployment efficiency of large language models, paving the way for more reliable AI systems. Future work will explore adaptive correction strategies, larger models, and multi-modal extensions, further broadening its impact in AI research and applications.
Deep Dive
Plain Language Accessible to non-experts
想象你在厨房里准备一道复杂的菜肴，厨师们（模型）各自有自己的拿手菜。把他们的菜谱合成一本大菜谱时，可能会出现偏差，比如某些调料比例不对，导致味道不佳。DARTS就像是用一种聪明的调味方法，根据每个步骤的重要程度和偏差大小，调整调料的用量，确保每个关键环节都调得恰到好处。它会特别关注那些决定菜肴成败的关键步骤（高熵位置），用少量的调整（参数）就能让整体味道变得更正宗。这种方法让厨房的菜谱变得更科学、更精准，也更容易推广到其他厨房（模型）中去。
ELI14 Explained like you're 14
想象你在学校参加一个比赛，很多队友（模型）都各有专长。有时候，大家的表现会出现偏差，比如某个队友跑得慢一点，影响了整体成绩。为了让每个队友都发挥最佳，教练（DARTS）会特别关注那些关键时刻（高重要性位置），用一些特别的策略（加权修正）帮助他们改善。比如在比赛的冲刺阶段，如果有人跑得慢，就会影响整个队伍的名次。教练会用一些聪明的训练方法，让他们在这些关键时刻表现得更好。这样，整个队伍的成绩就会大大提高，而且只需要少量的调整（参数少），就能让大家都发挥出最好的水平。这就像DARTS一样，聪明地修正模型中的偏差，让它们在重要任务中表现得更出色。
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).