RAIN-Merging: A Gradient-Free Method to Enhance Instruction Following in Large Reasoning Models with Preserved Thinking Format
RAIN-Merging is a gradient-free method that enhances instruction-following in large reasoning models while preserving their thinking format, with over 8% improvement in benchmarks.
Key Findings
Methodology
The proposed RAIN-Merging leverages analysis of task vectors in parameter space, revealing near orthogonality between LRM and ITM in key modules. It involves two stages: first, projecting the ITM task vector into the null space of features at thinking tokens to maintain reasoning structure; second, estimating instruction attention via small calibration sets to derive module-specific scaling factors that amplify instruction relevance. The entire process is gradient-free, based on feature projection and attention statistics, ensuring the preservation of structured reasoning outputs while improving instruction adherence.
Key Results
- Across four instruction-following and nine reasoning benchmarks, RAIN-Merging significantly outperforms traditional merging and fine-tuning, with instruction adherence gains of 5-10%. For example, in Qwen2.5-7B, instruction-following accuracy increased by 8.2%, while reasoning performance remained stable. The method demonstrated consistent stability across model sizes (1.5B to 32B) and architectures (Qwen, Llama), especially in agentic scenarios, confirming its broad applicability.
- Parameter space analysis confirmed the near orthogonality of task vectors, reducing interference. Null-space projection preserved reasoning format, preventing output format violations. Attention-guided coefficient adjustment improved instruction response quality by emphasizing instruction span responses and suppressing leakage, validated through ablation studies. The combined approach effectively balances reasoning and instruction following without gradient updates.
- Experimental results show that the method achieves substantial improvements in instruction adherence while maintaining reasoning capabilities, with performance gains averaging over 7% across models. It also enhances robustness in interactive agent scenarios, demonstrating scalability and versatility. The approach’s low computational cost and theoretical guarantees make it practical for real-world deployment.
Significance
This research addresses the long-standing challenge of balancing reasoning ability and instruction compliance in large models. By introducing a low-cost, gradient-free fusion technique grounded in parameter space analysis, it overcomes output format mismatches and capability interference issues. The method’s ability to improve instruction following without sacrificing reasoning performance has profound implications for deploying trustworthy AI assistants, professional tools, and autonomous agents. It paves the way for more adaptable, reliable, and scalable large models, fostering broader adoption in industry and research. The theoretical insights into task vector orthogonality and null-space projection also contribute to the fundamental understanding of model modularity and capability integration.
Technical Contribution
RAIN-Merging introduces a novel parameter-space approach based on task vector null-space projection, ensuring the preservation of reasoning structures during model merging. It combines this with an attention-guided, gradient-free coefficient adjustment mechanism that amplifies instruction relevance while controlling leakage. The method provides theoretical guarantees on output format stability via second-order approximations of softmax KL divergence, and employs a quadratic surrogate optimization to efficiently determine merging coefficients. This innovative framework enables low-cost, interpretable, and scalable enhancement of instruction-following in large models, setting a new standard for model fusion techniques.
Novelty
This work is the first to utilize task vector null-space projection for gradient-free model merging, specifically addressing the output format mismatch between structured reasoning models and answer-only instruction-tuned models. Unlike prior methods relying on fine-tuning or data-driven calibration, it exploits geometric properties of parameter space to achieve capability integration without gradient updates. The attention-guided coefficient adjustment further introduces a novel, interpretable mechanism to enhance instruction relevance, representing a significant advancement over existing model merging and fine-tuning approaches.
Limitations
- The method depends on the assumption of strong orthogonality between task vectors, which may not hold in all models or tasks, potentially limiting effectiveness in certain scenarios.
- It requires small calibration datasets for reasoning and instruction, which might be insufficient for highly complex or domain-specific tasks, affecting generalization.
- Currently tailored for Transformer architectures, extending to other neural network types remains an open challenge, possibly requiring adaptation.
Future Work
Future research will focus on improving task vector estimation robustness, possibly through unsupervised or self-supervised learning. Extending the approach to multi-modal models integrating vision, speech, and text is a promising direction. Additionally, exploring adaptive calibration techniques to reduce dataset dependency and applying the method to real-world deployment scenarios, such as autonomous agents and professional AI tools, will be key to broadening its impact.
AI Executive Summary
In the rapidly evolving field of artificial intelligence, large reasoning models (LRMs) have demonstrated remarkable capabilities in complex multi-step reasoning tasks, such as mathematical derivations and program synthesis. However, a persistent challenge remains: these models often fail to faithfully follow explicit instructions regarding output formats, constraints, or specific operational requirements. This disconnect hampers their deployment in real-world applications where reliability and adherence to user specifications are critical.
To address this, researchers have explored supervised fine-tuning (SFT) methods, but these are costly and can degrade general capabilities. An alternative approach, model merging, extracts parameter differences—task vectors—between fine-tuned and pre-trained models, combining them to create a unified model. Yet, naive merging risks disrupting the structured reasoning outputs due to format mismatches.
The paper introduces RAIN-Merging, a novel, gradient-free technique that leverages parameter space analysis. By examining the near orthogonality of task vectors, the method projects the instruction-tuned model’s task vector into the null space of the reasoning model’s features at special tokens, preserving the reasoning structure. It then estimates instruction attention via small calibration sets, deriving module-specific scaling factors that amplify instruction relevance.
Extensive experiments across multiple benchmarks show that RAIN-Merging significantly improves instruction adherence—by over 8% in some cases—while maintaining reasoning quality. The approach is robust across model sizes and architectures, and enhances agentic scenario performance. Its low computational cost and theoretical guarantees make it practical for real-world deployment.
Overall, this work offers a scalable, effective solution to enhance large models’ instruction-following abilities without sacrificing their reasoning prowess, marking a substantial step forward in AI capability integration and reliability.
Deep Dive
Abstract
Large reasoning models (LRMs) excel at a long chain of reasoning but often fail to faithfully follow instructions regarding output format, constraints, or specific requirements. We investigate whether this gap can be closed by integrating an instruction-tuned model (ITM) into an LRM. Analyzing their differences in parameter space, namely task vectors, we find that their principal subspaces are nearly orthogonal across key modules, suggesting a lightweight merging with minimal interference. However, we also demonstrate that naive merges are fragile because they overlook the output format mismatch between LRMs (with explicit thinking and response segments) and ITMs (answers-only). We introduce RAIN-Merging (Reasoning-Aware Instruction-attention guided Null-space projection Merging), a gradient-free method that integrates instruction following while preserving thinking format and reasoning performance. First, with a small reasoning calibration set, we project the ITM task vector onto the null space of forward features at thinking special tokens, which preserves the LRM's structured reasoning mechanisms. Second, using a small instruction calibration set, we estimate instruction attention to derive module-specific scaling that amplifies instruction-relevant components and suppresses leakage. Across four instruction-following benchmarks and nine reasoning & general capability benchmarks, RAIN-Merging substantially improves instruction adherence while maintaining reasoning quality. The gains are consistent across model scales and architectures, translating to improved performance in agent settings.