Table of Contents
Fetching ...

RAIN-Merging: A Gradient-Free Method to Enhance Instruction Following in Large Reasoning Models with Preserved Thinking Format

Zhehao Huang, Yuhang Liu, Baijiong Lin, Yixin Lou, Zhengbao He, Hanling Tian, Tao Li, Xiaolin Huang

TL;DR

This work introduces 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.

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.

RAIN-Merging: A Gradient-Free Method to Enhance Instruction Following in Large Reasoning Models with Preserved Thinking Format

TL;DR

This work introduces 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.

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.
Paper Structure (40 sections, 3 theorems, 32 equations, 14 figures, 16 tables, 1 algorithm)

This paper contains 40 sections, 3 theorems, 32 equations, 14 figures, 16 tables, 1 algorithm.

Key Result

Proposition 1

Let the logits of sample $x$ at thinking special tokens $t \in \Omega_{\text{think}}(x)$ be $z_\theta(x,t)$, and let $\pi_\theta(\cdot \mid x, y_{<t})=\operatorname{softmax}(z_\theta(x,t))$. By a second-order approximation of the softmax–KL divergence with a uniform upper bound, for any perturbation Assuming the model’s intermediate representations are Lipschitz continuous and bounded, there exist

Figures (14)

  • Figure 1: An overview of RAIN-Merging. In the case, the LRM arrives at the correct solution but ignores the required format and specific code. To preserve the reasoning structure, we perform training-free merging by combining a task vector projected onto the null space of the thinking format with instruction-attention guided coefficients. The merged model remains correct while satisfying the specified constraints. See Sec.\ref{['sec:method']} for details.
  • Figure 2: Principal subspace cosine similarity between LRM and ITM task vectors for each layer and submodule. The similarities are consistently low ($<0.1$).
  • Figure 3: Two stages of our RAIN-Merging pipeline. (a) For each submodule, the ITM task vector is projected onto the null space preventing shifts in thinking format. (b) Given the instruction calibration set, we compute the instruction-attention score from attention outputs to obtain merging coefficients.
  • Figure 4: GPU memory usage comparison between different methods under the same configuration as Tab.\ref{['tab:baseline_compare']}.
  • Figure 5: $\mathcal{L}_{\text{think}}$ in Eq. (\ref{['eq:think-constraint']}) (left) on the reasoning calibration validation set, and the proportion of generations missing the closing </think> token (right) on IFEval under the same configuration as Tab.\ref{['tab:baseline_compare']}.
  • ...and 9 more figures

Theorems & Definitions (5)

  • Proposition 1
  • Proposition 1: Insufficiency of parameter-space orthogonality for output invariance
  • proof : Proof of Prop. \ref{['prop:orth-not-inv']}
  • proof : Proof of Prop. \ref{['prop:reasoning-aware-null-space-projection']}
  • Lemma 1: Hessian bound for $\operatorname{lse}$