Universal Reasoner: A Single, Composable Plug-and-Play Reasoner for Frozen LLMs

## Abstract

Abstract Large Language Models (LLMs) have demonstrated remarkable general capabilities, but enhancing skills such as reasoning often demands substantial computational resources and may compromise their generalization. While Parameter-Efficient Fine-Tuning (PEFT) methods offer a more resource-conscious alternative, they typically requires retraining for each LLM backbone due to architectural dependencies. To address these challenges, here we propose Universal Reasoner (UniR) - a single, lightweight, composable, and plug-and-play reasoning module that can be used with any frozen LLM to endow it with specialized reasoning capabilities.
Specifically, UniR decomposes the reward into a standalone reasoning module that is trained independently using predefined rewards, effectively translating trajectory-level signals into token-level guidance. Once trained, UniR can be combined with any frozen LLM at inference time by simply adding its output logits to those of the LLM backbone. This additive structure naturally enables modular composition: multiple UniR modules trained for different tasks can be jointly applied by summing their logits, enabling complex reasoning via composition. Experimental results on mathematical reasoning and machine translation tasks show that UniR significantly outperforms existing baseline fine-tuning methods using the Llama3.2 model . Furthermore, UniR demonstrates strong weak-to-strong generalization: reasoning modules trained on smaller models effectively guide much larger LLMs. This makes UniR a cost-efficient, adaptable, and robust solution for enhancing reasoning in LLMs without compromising their core capabilities. Code is open-sourced at https://github.com/hangeol/UniR

## 1 Introduction

Recent Large Language Models (LLMs), such as OpenAI-o1 [28], DeepSeek-R1 [8], Kimi-1.5 [37], BOLT [29], and Open-Reasoner-Zero [19], have demonstrated remarkable capabilities in complex reasoning tasks. These models leverage Chain-of-Thought (CoT) prompting [39] to elicit step-by-step reasoning processes from LLMs, thereby significantly enhancing their ability to tackle intricate problems. A key driver of this progress has been Reinforcement Fine-Tuning (RFT), which aligns model outputs with desired criteria, often human preferences, through reinforcement learning (RL) algorithms such as Proximal Policy Optimization (PPO) [33], Direct Preference Optimization (DPO) [31], Group Relative Policy Optimization (GRPO) [34], etc.

Unfortunately, adapting these powerful LLMs to specialize in specific domains or further enhance their reasoning abilities through RFT presents significant hurdles. The process is often prohibitively expensive due to the substantial computational and memory resources required for fine-tuning large-scale models. While Parameter-Efficient Fine-Tuning (PEFT) techniques like LoRA [18] have been proposed to mitigate these costs, they introduce their own limitations: (1) Model Dependency: LoRA modules strongly depend on specific model architectures, limiting their transferability across models of different sizes (e.g., from 3B to 14B). (2) Poor Linearity Support: Combining multiple domain-specific LoRA adapters through linear operations lacks robust theoretical support and frequently yields unpredictable outcomes.

To address these challenges, we introduce UniR (Universal Reasoner), a novel method that decouples the reasoning capability into a specialized reasoning module, independently trained and designed to enhance capabilities of a backbone LLM. This lightweight reasoning module can be seamlessly integrated into an architecture-agnostic backbone model without requiring explicit information about the model’s internal structure. This approach not only fosters modularity, preserving the core capabilities of the base model, but also enables effective weak-to-strong generalization where a compact reasoning module can enhance significantly larger models. Moreover, this modularity opens avenues for combining distinct reasoning modules across multiple domains.

The core innovation of UniR lies in its novel decoupling of reward model training from full policy updates.
Instead of relying on preference datasets, UniR leverages explicit, predefined rewards that are often available for tasks with verifiable solutions, such as mathematical problem-solving or machine translation. We propose to model these typically trajectory-level rewards, or rewards over a full response, as a sum of log-probabilities derived from our decomposed reasoning module. This formulation effectively translates a global reward signal into a sequence of implicit token-level guidance signals that steer the reasoning process. The reasoning module is then trained to maximize this reward using a policy gradient algorithm, ensuring that only the smaller reasoning module is updated while the backbone policy remains unchanged. During inference, reasoning module provides guidance to a main backbone model by combining their outputs, typically by adding their respective logits before sampling.
Furthermore, because the guidance occurs at the logit level, multiple reasoning modules trained for distinct tasks can be composed through simple logit addition—enabling composite probabilistic reasoning without additional training or architectural modification.

Our contributions are summarized as follows:

- •
Computational Efficient Training. We introduce a method to train a decomposed reasoning module by directly optimizing for predefined rewards, circumventing the need for expensive preference dataset creation. Furthermore, our algorithm backpropagates only through a smaller reasoning model while keeping the large backbone model frozen, reducing computational and memory overhead.
- •
Transferability Across Model Sizes.
Although our approach trains only small models, the learned reasoning module can be effectively transferred to larger models at inference time. For example, a 0.5B reasoning module trained alongside a 3B backbone can be used in conjunction with a 14B model, providing token-level guidance that improves the larger model’s performance without requiring any fine-tuning of the 14B model itself.
- •
Composibility.
Our method allows optimal composition of multiple small reasoning modules specialized for specific domains or tasks. These specialized models can then be selectively deployed based on the domain context, enabling modular and efficient inference pipelines where domain expertise is encapsulated in lightweight, task-specific policies.
- •
Superior Performance on Reasoning and Translation Tasks. Our experiments on challenging mathematics reasoning and translation benchmarks show that UniR substantially outperforms conventional fine-tuning methods, showcasing its efficacy in enhancing LLM capabilities.

## 2 Related Work

Enhancing LLM Reasoning.
Enhancing the reasoning capabilities of LLMs is a key research objective to develop toward AI agent [10], where complex reasoning and understanding are needed. One prominent strategy involves Supervised Fine-Tuning (SFT) on specialized, often small, datasets curated to foster reasoning skills; these datasets might be distilled from larger models (e.g., Bespoke-Stratos-7B [5], OpenThinker-7B [1] using outputs from models like DeepSeek-R1 [8]) or meticulously selected from high-quality examples (e.g., S1 [27], LIMO [44]). Recently, RL-based Group Relative Policy Optimization (GRPO) and related methods [34, 45, 24] have shown promise in enhancing LLM reasoning capabilities using rule-based rewards without requiring a separate value function. Similarly, our work also aims to improve the reasoning ability of LLMs using the specialized reasoning module through evaluation on tasks that require complex reasoning.

Efficient Methods for Domain-Specific Tasks.
To mitigate the substantial computational costs associated with both SFT and RL fine-tuning, Parameter-Efficient Fine-Tuning (PEFT) techniques like LoRA [18] are commonly employed. However, LoRA modules are inherently tied to the architecture of the model they were trained on, limiting their transferability across LLMs of different sizes, including the number of layers or size of hidden dimensions.
In contrast, UniR trains a compact, self-contained reasoning module that interacts with a large, frozen backbone model at the logit level. This not only makes the training process highly efficient but also results in a reasoning module that is more flexible and architecturally unconstrained for transfer and guidance.

Guided LLM generation. Guiding the generation of frozen LLMs at inference or decoding time offers a flexible way to improve outputs without full retraining [21]. Established methods like Best-of-N (BoN) sampling [12] or search algorithms such as Monte Carlo Tree Search (MCTS) [23] can enhance generation quality but often incur substantial computational costs due to multiple forward passes. More recent techniques focus on finer-grained, token-level guidance [40, 43]. Especially, GenARM [40] proposes a token-level reward model, trained on human preference data, to provide token-level scores that guide a frozen LLM at test-time. In addition, several works [26, 36, 16] train a value function by temporal difference learning to steer generation towards desired attributes. UniR, while also employing a guiding component, is distinguished by the nature and training of its guidance component. UniR does not involve learning a separate value function; instead, it directly trains the reasoning module using a policy gradient algorithm. Furthermore, UniR trains on a dedicated training dataset using predefined rewards, rather than real-time adaptation at inference time.

Figure: Figure 1: UniR Framework Overview. Our approach trains a lightweight, transferable reasoning module (${\pi_{r}}$) using predefined rewards to guide a frozen backbone model (${\pi_{b}}$), offering (1) transferability across different backbone models or tasks; and (2) composability by combining multiple specialized reasoning modules through reward optimization.
Refer to caption: extracted/6484925/figs/main_figure4.png

## 3 Problem Definition

A significant trend in enhancing LLM capabilities, particularly in reasoning, involves guiding their behavior using predefined rewards rather than relying solely on human preference data. This approach leverages objective metrics like task accuracy or other rule-based criteria, as demonstrated in works like DeepSeek-R1 [8] and methodologies employing GRPO [34]. Following this direction, our research concentrates on tasks where performance can be directly assessed through such predefined rewards.
We denote the predefined reward as $r(x,y)$ for a given prompt $x$ from a dataset $\mathcal{D}$ and a generated output sequence $y$.
Standard Reinforcement Learning (RL) fine-tuning aims to find an optimal policy, ${\pi_{\theta}}^{*}(y|x)$, that maximizes the expected reward while remaining close to a backbone policy (${\pi_{b}}(y|x)$), typically a pre-trained base LLM. This objective is formulated as:

$$ $\max_{{\pi_{\theta}}}\mathbb{E}_{x\sim\mathcal{D},y\sim{\pi_{\theta}}(y|x)}[r( x,y)]-\beta{\mathbb{D}_{KL}}[{\pi_{\theta}}(y|x)||{\pi_{b}}(y|x)],$ (1) $$

where ${\mathbb{D}_{KL}}[\cdot||\cdot]$ represents the Kullback-Leibler divergence, and $\beta>0$ controls the strength of the regularization towards the backbone policy ${\pi_{b}}$. The optimal policy ${\pi_{\theta}}^{*}(y|x)$ for this objective can be expressed in a closed form as [48, 31]:

$$ ${\pi_{\theta}}^{*}(y|x)=\frac{1}{Z(x)}{\pi_{b}}(y|x)\exp\left(\frac{1}{\beta}r (x,y)\right),$ (2) $$

where $Z(x)=\sum_{y}{\pi_{b}}(y|x)\exp(\frac{1}{\beta}r(x,y))$ is a partition function. A key challenge in directly applying Eq. ([2](https://arxiv.org/html/2505.19075v2#S3.E2)) for autoregressive generation is that the trajectory-level reward $r(x,y)$ can be obtained only after the full sequence $y$ is complete. However, it does not provide token-specific feedback during the generation process. For fine-grained guidance of an LLM’s reasoning steps, per-token signals are more desirable. Theoretically, the optimal token-level decision rule for the objective in Eq. ([1](https://arxiv.org/html/2505.19075v2#S3.E1)) can be expressed using the optimal soft Q-function, $Q^{*}(s_{t},a_{t})$ where $s_{t}=(x,y_{<t})$ is the current state and $a_{t}=y_{t}$ is the next token. This $Q^{*}(s_{t},a_{t})$ represents the maximum expected future return following the optimal policy ${\pi_{\theta}}^{*}$. The optimal token-level policy is then given by [15, 48]:

$$ $\log{\pi_{\theta}}^{*}(y_{t}|x,y_{<t})=\log{\pi_{b}}(y_{t}|x,y_{<t})+\frac{1}{ \beta}Q^{*}(x,y_{<t},y_{t})-\log Z^{\prime}(x,y_{<t}),$ (3) $$

where $Z^{\prime}(x,y_{<t})$ is the per-step partition function [38, 15]. However, guiding with $Q^{*}$ is impractical since it requires knowledge of the optimal policy ${\pi_{\theta}}^{*}$ itself. This challenge motivates the need to bridge observable trajectory-level rule-based rewards with a practical formulation that enables effective token-level guidance.

## 4 UniR: Universal Reasoner

To circumvent the direct estimation of $Q^{*}$, our core idea is to learn a separate, typically smaller, reasoning module, denoted as ${\pi_{r}}(y|x)$ with parameters $\phi$. This policy is specifically trained to capture the desired reasoning or task-specific behaviors that lead to high rewards. We then use this learned reasoning module to guide a larger LLM.

### 4.1 Modeling Trajectory Rewards via the Reasoning Module

To connect trajectory-level reward $r(x,y)$ to our reasoning module ${\pi_{r}}$ and enable its use in a token-by-token guidance mechanism, we posit that the reward can be represented by the sum of log-probabilities of the reasoning module ${\pi_{r}}$:

$$ $\frac{1}{\beta}r(x,y)=\sum_{t=1}^{|y|}\log{\pi_{r}}(y_{t}|x,y_{<t};\phi).$ (4) $$

Here, $\beta$ is the same temperature parameter from Eq. ([1](https://arxiv.org/html/2505.19075v2#S3.E1)). This formulation implies that sequences $y$ assigned a higher sum of token-level log-probabilities by ${\pi_{r}}$ should correspond to sequences achieving higher rewards $r(x,y)$. The feasibility of representing a reward function via the log-likelihood of a policy is supported by theoretical work in [40],
which claims that all reward equivalence classes can be represented using the parameterization $\log{\pi_{r}}(y\mid x)$ for some probability distribution ${\pi_{r}}(y\mid x)$.
Since $\log{\pi_{r}}(y|x)=\sum_{t=1}^{|y|}\log{\pi_{r}}(y_{t}|x,y_{<t})$ by the chain rule, this suggests that a sufficiently expressive reasoning module ${\pi_{r}}$ can indeed model any trajectory-level reward function $r(x,y)$.
By substituting the definition from Eq. ([4](https://arxiv.org/html/2505.19075v2#S4.E4)) into the optimal policy form of Eq. ([2](https://arxiv.org/html/2505.19075v2#S3.E2)), we derive the structure of our guided policy for autoregressive generation of the next token $y_{t}$. This translates to combining their next-token predictions:

$$ $\log{\pi_{\theta}}(y_{t}|x,y_{<t})=\log{\pi_{b}}(y_{t}|x,y_{<t})+\log{\pi_{r}} (y_{t}|x,y_{<t})-\log Z^{\prime}(x,y_{<t}),$ (5) $$

where $Z^{\prime}(x,y_{<t})$ is the per-step partition function. Accordingly, the logits from the frozen backbone ${\pi_{b}}$ and the learned module ${\pi_{r}}$ can be separable to produce the logits for the guided policy ${\pi_{\theta}}$.
This additive and separable guidance mechanism extends to scenarios involving multiple reward criteria. Instead of training a single monolithic model to jointly optimize for all objectives, UniR allows for the training of several distinct reasoning modules $\{{\pi_{r}}^{1},{\pi_{r}}^{2},\cdots,{\pi_{r}}^{N}\}$, each specialized for a different reward function $r_{i}(x,y)$. At inference, these specialized modules can be combined by taking a weighted sum of their logits and adding them to the backbone model’s logits:

$$ $\log{\pi_{\theta}}(y_{t}|x,y_{1:t-1})\propto\log{\pi_{b}}(y_{t}|x,y_{1:t-1})+ \sum_{i=1}^{N}\alpha_{i}\log{\pi_{r}}^{i}(y_{t}|x,y_{1:t-1}),$ (6) $$

where $\alpha_{i}\geq 0$ controls the influence of each reasoning policy ${\pi_{r}}^{i}$. This approach offers significant flexibility, as the weights $\alpha_{i}$ can be adjusted at inference time to explore various trade-offs between objectives without any retraining. Such a strategy is analogous to solving a weighted KL-regularized optimization problem, aligning with principles in multi-objective policy optimization [2, 9].

### 4.2 Training π r subscript 𝜋 𝑟 {\pi_{r}} italic_π start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT with GRPO

The parameters $\phi$ of the reasoning module ${\pi_{r}}$ are trained to maximize the expected predefined reward $r(x,y)$, not the sum of its own log-probabilities defined in Eq. ([4](https://arxiv.org/html/2505.19075v2#S4.E4)). We employ Group Relative Policy Optimization (GRPO) [34] for this purpose. For each input $x\sim\mathcal{D}$, we generate a group of $G$ candidate responses $\{y^{1},y^{2},\cdots,y^{G}\}$ by sampling from ${\pi_{\theta}}(y^{i}|x;\phi_{old})$, where $\phi_{old}$ denotes the parameters of the reasoning module from the previous iteration. For each generated sequence $y^{i}$, we compute its external reward $r(x,y^{i})$ (denoted $r_{i}$ for brevity) and then calculate the normalized advantage $A_{i}$:

$$ $A_{i}=\left({r_{i}-\text{mean}(\{r_{1},r_{2},\cdots,r_{G}\})}\right)/{\text{ std}(\{r_{1},r_{2},\cdots,r_{G}\})}.$ (7) $$

In GRPO, the objective function for updating the parameters $\phi$ of the reasoning module is then:

$$ $\displaystyle\mathcal{L}_{\text{GRPO}}(\phi)=-\mathbb{E}_{x\sim\mathcal{D},\;y ^{i}\sim\pi_{\theta}(y^{i}|x;\phi_{old})}\Bigg{[}\frac{1}{G}\sum_{i=1}^{G} \Bigg{\{}\min\Bigg{(}\frac{\pi_{\theta}(y^{i}|x)}{\pi_{\theta}(y^{i}|x;\phi_{ old})}A_{i},\;$ (8) $\displaystyle\text{clip}\left(\frac{\pi_{\theta}(y^{i}|x)}{\pi_{\theta}(y^{i}| x;\phi_{old})},\;1-\epsilon,\;1+\epsilon\right)A_{i}\Bigg{)}-\beta_{KL}{ \mathbb{D}_{KL}}[{\pi_{\theta}}(y^{i}|x;\phi)||{\pi_{\theta}}(y^{i}|x;\phi_{0} )]\Bigg{\}}\Bigg{]}.$ $$

Here, $\epsilon$ is the clipping hyperparameter from PPO-clip, $\phi_{0}$ is initial value of $\phi$, and $\beta_{KL}$ controls the KL-regularization term.
The learning direction is determined by the ratio $\frac{{\pi_{\theta}}(y^{i}|x,\phi)}{{\pi_{\theta}}(y^{i}|x;\phi_{old})}$, which simplifies to $\frac{{\pi_{r}}(y^{i}|x,\phi)}{{\pi_{r}}(y^{i}|x;\phi_{old})}$ as the terms from the ${\pi_{b}}$ cancel out, thereby exclusively influencing ${\pi_{r}}$.
Furthermore, the gradient $\nabla_{\phi}\mathcal{L}_{\text{GRPO}}$ only affects the parameters of ${\pi_{r}}$ since ${\pi_{b}}$ is frozen.

KL Divergence Considerations. The original GRPO formulation [34] includes a KL divergence term to regularize the optimized policy ${\pi_{\theta}}$ towards the initial policy. However, in our UniR framework, the reasoning module ${\pi_{r}}$ is often initialized from a general pre-trained small model, which may not possess strong task-specific reasoning capabilities. Regularizing ${\pi_{\theta}}$ towards ${\pi_{b}}$ which is heavily influenced by the initial, potentially ${\pi_{r}}_{\text{old}}$ can impede the learning of high-reward reasoning paths. Therefore, we omit such an explicit KL regularization by $\beta_{KL}=0$, allowing ${\pi_{r}}$ more freedom to explore and learn effective reasoning strategies driven by the predefined rewards. This choice aligns with findings in related works such as DAPO [45], which also found benefits in adjusting or omitting certain regularization components.

### 4.3 Interpreting π r subscript 𝜋 𝑟 {\pi_{r}} italic_π start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT with Q ∗ superscript 𝑄 Q^{*} italic_Q start_POSTSUPERSCRIPT ∗ end_POSTSUPERSCRIPT

We now provide theoretical grounding to understand what the learned reasoning module ${\pi_{r}}$ captures at each token-level decision point. This provides a theoretical justification for interpreting $\log{\pi_{r}}$ as a meaningful guidance signal that drives the policy towards optimality for the trajectory-level objective.

###### Theorem 1 .

(Informal)
If the reasoning module ${\pi_{r}}$ satisfies Eq. ([4](https://arxiv.org/html/2505.19075v2#S4.E4)), such that its summed per-token log-probabilities accurately approximate the scaled reward for trajectories $y$ that are consistent with the optimal policy ${\pi_{\theta}}^{*}$, then

$$ $\log{\pi_{r}}(y_{t}|x,y_{<t})=\frac{1}{\beta}Q^{*}(x,y_{<t},y_{t}).$ (9) $$

###### Proof.

We assume guided policy ${\pi_{\theta}}$ has converged to the optimal policy ${\pi_{\theta}}^{*}$ for all states $s_{t}=(x,y_{<t})$ and actions $a_{t}=y_{t}$ by optimizing the GRPO objective to maximize the reward $r(x,y)$ since the KL-regularized objective typically admits a unique optimum. Given that the trajectory-level policies are identical, their per-token conditional distributions must also match for any $s_{t}$. Comparing Eq. ([5](https://arxiv.org/html/2505.19075v2#S4.E5)) with the standard form of the optimal policy expressed in terms of the soft Q-function (Eq. ([3](https://arxiv.org/html/2505.19075v2#S3.E3))) shows that $\log{\pi_{r}}(y_{t}\mid x,y_{<t})$ corresponds to the scaled optimal soft Q-value.
∎

**Table 1: Zero-shot pass@1 performance on mathematic benchmarks. The out-of-distribution results are obtained using models trained on the MATH training dataset. Models marked with * are taken from the report [42].**
|  | Method | Trained<br>Model | In-distribution | Out-of-distribution | Avg. |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | GSM8K | MATH-500 | AIME24 | Minerva | OlympiadBench |  |  |  |
| Other 7B | Mathstral-7B-v0.1* | 7B | 84.9 | 56.6 | - | 16.2 | 21.5 | - |
| Internlm2-math-plus-7B* | 7B | 84.0 | 54.4 | - | 17.3 | 18.8 | - |  |
| NuminaMath-7B-CoT* | 7B | 75.4 | 55.2 | - | 19.1 | 19.9 | - |  |
| DeepSeekMath-7B-RL* | 7B | 88.2 | 52.4 | - | 20.6 | 19.0 | - |  |
| Llama3.2-3B | Baseline | - | 66.1 | 38.0 | 0.0 | 14.8 | 11.4 | 26.1 |
| Baseline + 1B | - | 64.7 | 35.0 | 10.0 | 14.7 | 11.4 | 27.2 |  |
| GRPO LoRA | 3B | 67.9 | 41.8 | 13.3 | 16.9 | 14.2 | 30.8 |  |
| GRPO Full | 3B | 77.3 | 42.4 | 3.3 | 16.2 | 13.5 | 30.6 |  |
| Ours | 1B | 78.3 | 49.2 | 16.7 | 18.4 | 17.3 | 36.0 |  |
| Qwen2.5-3B | Baseline | - | 75.5 | 46.8 | 6.7 | 23.5 | 25.5 | 35.6 |
| Baseline + 0.5B | - | 72.1 | 41.8 | 6.7 | 16.2 | 22.1 | 31.8 |  |
| GRPO LoRA | 3B | 78.2 | 64.2 | 13.3 | 28.3 | 30.1 | 42.8 |  |
| GRPO Full | 3B | 82.8 | 66.0 | 16.7 | 28.3 | 30.1 | 44.8 |  |
| Ours | 0.5B | 82.2 | 64.8 | 10.0 | 26.8 | 27.0 | 42.2 |  |

###### Theorem 1 .

###### Proof.

## 5 Experiments

### 5.1 Experimental Settings

We evaluate UniR on two tasks: mathematical problem-solving and machine translation. Mathematical problem-solving requires explicit, step-by-step reasoning, while machine translation involves implicit reasoning based on domain-specific knowledge and semantic understanding. These tasks provide complementary benchmarks for assessing our decomposed reasoning module.

Models and Baselines. Our experiments utilize models from Qwen2.5 [41] and Llama3.2 [13]. We train the reasoning module ${\pi_{r}}$ alongside a 3B-scale backbone ${\pi_{b}}$: (1) ${\pi_{b}}$ as Qwen2.5-3B with ${\pi_{r}}$ as Qwen2.5-0.5B, and (2) ${\pi_{b}}$ as Llama3.2-3B with ${\pi_{r}}$ as Llama3.2-1B. For the transferability experiment, we use Qwen2.5-14B and Llama3.1-8B as backbone models. All models are Instruct variants unless stated otherwise. As for baseline training, we apply GRPO to fine-tune the entire backbone model (e.g., the 3B-scale ${\pi_{b}}$ model) directly on each task dataset using the same predefined rewards. This baseline is implemented through either full model tuning or LoRA-based fine-tuning.

Dataset and Reward Formulation.
For mathematical reasoning tasks, we use the Math-12k [22, 24] and GSM8k [7] datasets. The reward $r(x,y)$ is determined by the rule-based correctness of the final answer,
following the methodology in Dr.GRPO [24]. For machine translation, we use the IWSLT 2017 dataset [6] for English-to-German and German-to-English translation, employing BLEU [30] and CometKiwi (^1^11[https://huggingface.co/Unbabel/wmt22-CometKiwi-da](https://huggingface.co/Unbabel/wmt22-CometKiwi-da)) [32] as reward during training. Following MT-R1-Zero [11], we combine lexical and semantic metrics for balanced training.

Evaluation Metrics. We evaluate mathematical reasoning using pass@1 accuracy on five benchmarks: GSM8k [7], MATH-500 [22], AIME24 (^2^22[https://huggingface.co/datasets/AI-MO/aimo-validation-aime](https://huggingface.co/datasets/AI-MO/aimo-validation-aime)), Minerva [20], and OlympiadBench [17]. GSM8k uses its own training split; the others use Math-12K. For the translation task, we report BLEU, CometKiwi, and XComet-XL (^3^33[https://huggingface.co/Unbabel/XCOMET-XL](https://huggingface.co/Unbabel/XCOMET-XL)) [14] scores. Further details on training hyperparameters, optimization settings, evaluation details and prompt templates are provided in Appendix [A](https://arxiv.org/html/2505.19075v2#A1).

### 5.2 Enhancing Reasoning Ability

As shown in Table [2](https://arxiv.org/html/2505.19075v2#S4.T2) , our framework demonstrates substantial improvements over GRPO with LoRA across both in-distribution and out-of-distribution benchmarks. UniR achieves an average score of 36.0, surpassing GRPO Full and GRPO LoRA with the Llama3.2 model. When using the Qwen2.5, UniR does not surpass that of GRPO method. However, it is noteworthy that UniR’s performance gain over the naive logit summation approach (Baseline + 0.5B) was more substantial ($+10.4$) than the improvement GRPO achieved over the original baseline ($+9.2$).

Furthermore, in machine translation, as shown in Table [2](https://arxiv.org/html/2505.19075v2#S4.T2), UniR consistently outperforms GRPO with LoRA in both translation directions (EN$\rightarrow$DE and DE$\rightarrow$EN), while achieving performance comparable to full GRPO. This indicates that UniR effectively generalizes beyond mathematical reasoning, extending its efficacy to implicit reasoning tasks such as translation.

To further verify the importance of reward-driven alignment for the reasoning module, we compare it with a naive ensemble approach, directly combining the logits of the frozen backbone model with those of an intact reasoning module without any reward optimization (e.g., Baseline + 1B). This strategy results in degraded performance compared to even baselines. These findings highlight that merely combining outputs from different models does not lead to enhanced reasoning ability. Instead, our reward-driven policy alignment substantially improves reasoning performance, demonstrating the effectiveness of the proposed UniR framework.

### 5.3 Additional Experiments

Figure: Figure 2: Effectiveness of Reasoning Policy Transfer. Results demonstrate that a trained reasoning module can improve performance when integrated with larger backbone models across diverse mathematical reasoning tasks.
Refer to caption: extracted/6484925/figs/transfer_results_math2.png

Transferability to Other Frozen Backbone LLMs.
Figure [2](https://arxiv.org/html/2505.19075v2#S5.F2) illustrates the transferability of our reasoning module ${\pi_{r}}$, which was originally trained on Math-12k and GSM8k datasets with a 3B-scale backbone. We evaluate its effectiveness when deployed to guide larger frozen backbone models with 8B and 14B parameters. In most benchmarks, the transferred ${\pi_{r}}$ improves average reasoning performance (35.5 and 52.8) compared to the respective backbone baselines (33.7 and 51.1), even without additional fine-tuning. This suggests that ${\pi_{r}}$ is not narrowly specialized to its original backbone but instead functions as a generalizable reasoning augmentation module. These findings highlight the potential of UniR to serve as a scalable and resource-efficient framework, enabling lightweight modules to enhance the reasoning capabilities of stronger models at inference time.

Figure: Figure 3: Reasoning Performance of the ${\pi_{r}}$ Module. (Left) The backbone model ${\pi_{b}}$ and (Middle) the standalone reasoning module ${\pi_{r}}$ produce incorrect, repetitive and logically flawed reasoning. (Right) When combined, it generates coherent reasoning and arrive at the correct solution, showing the effectiveness of the modular guidance.

Combining Specialized Reasoning Modules.
To investigate the advantages of integrating multiple specialized reasoning capabilities enabled by UniR, we design a composite task: translating German math problems from the MGSM dataset [35] into English. We utilize two distinct reasoning modules: $\pi^{\text{Math}}_{r}$, a module specialized for mathematical problem-solving on the GSM8k dataset, and $\pi^{\text{Trans}}_{r}$, a module trained for German-to-English translation. These modules are combined with a frozen base LLM (${\pi_{b}}$) through a weighted summation of their logits:

$$ $\log{\pi_{b}}+\alpha\log\pi^{\text{Math}}_{r}+(1-\alpha)\log\pi^{\text{Trans}} _{r},$ (10) $$

Figure: Figure 4: Performance on a German-to-English Math problem-solving task. The numbers in the figure indicate the value of $\alpha$.
Refer to caption: extracted/6484925/figs/math+trans.png

where $\alpha\in[0,1]$ is a coefficient that balances the influence between both modules. We employ GPT-4.1-nano to evaluate translation quality and the accuracy of the generated output. Detailed configurations for this experiment are provided in Appendix [A](https://arxiv.org/html/2505.19075v2#A1).

As depicted in Figure [4](https://arxiv.org/html/2505.19075v2#S5.F4),
we observe that increasing the weight $\alpha$ generally leads to improved accuracy in solving the mathematical problems, while a higher weight for the translation module tends to improve translation quality. This highlights a crucial capability of UniR: the ability to combine and balance multiple specialized reasoning modules. This modularity allows for precise control over the policy model’s behavior, leading to improved performance on multi-faceted tasks. We illustrate the example of the responses in Figure [14](https://arxiv.org/html/2505.19075v2#A2.F14) in Appendix [B](https://arxiv.org/html/2505.19075v2#A2).

Analysis of ${\pi_{r}}$ Reasoning Performance.
The reasoning capabilities of the ${\pi_{r}}$ module and backbone models in isolation reveal both their individual limitations and the significant synergy achieved through their combination. As shown in Figure [3](https://arxiv.org/html/2505.19075v2#S5.F3)
for mathematical reasoning tasks, the standalone ${\pi_{r}}$ module fails to generate coherent and relevant reasoning. Likewise, the large backbone model alone struggles to reach correct conclusions on complex tasks. However, when the reasoning module is paired with a stronger backbone, it provides valuable guidance to the backbone’s chain-of-thought process, resulting in significantly improved reasoning performance. Additional illustrative examples are presented in Appendix [B](https://arxiv.org/html/2505.19075v2#A2).

### 5.4 Computational Efficiency in Training

Figure: Figure 5: VRAM usage versus batch size under an 80GB constraint. Our method scales to batch size 128, while full fine-tuning and LoRA are limited, demonstrating memory efficiency for large batch.
Refer to caption: extracted/6484925/figs/vram_usage_by_method.png

Due to its structural decomposability, our framework substantially reduces VRAM usage by avoiding gradient updates to the large backbone model. While LoRA also requires minimal memory for gradients and optimizer states by restricting updates to a small subset of parameters, it still incurs substantial memory usage for intermediate activations because it has to save the full computation graph for back-propagation [46]. In contrast, our approach updates only the lightweight reasoning module, drastically minimizing the memory required for these activations. This leads to significantly lower VRAM usage, particularly in scenarios with large batch sizes or long sequence lengths where activation memory typically dominates.

We conduct an experiment examining peak memory usage based on the number of tokens processed in one step. The input length was fixed to 250 tokens (plus up to 5 output tokens), and batch size was progressively increased under an 80GB VRAM limit. As shown in Figure [5](https://arxiv.org/html/2505.19075v2#S5.F5), our method supports batch sizes up to 128, whereas LoRA and full fine-tuning are limited to 64 and 32, respectively. Although sequence length was fixed in this specific experiment, activation memory usage generally grows with both batch size and sequence length. As a result, our method provides greater flexibility in scaling the number of tokens processed per step under memory constraints. Furthermore, UniR remains fully compatible with a wide range of efficiency techniques, including quantization, caching, and adapter-based methods like LoRA applied to the reasoning module itself.

### 5.5 Reward Stabilization and Training Dynamics.

A notable empirical finding during training is that UniR not only achieves higher rewards compared to GRPO but also exhibits a smaller standard deviation in the rewards, as shown in Figure [6](https://arxiv.org/html/2505.19075v2#S5.F6). Since the advantage calculation in GRPO normalizes the unbiased advantage $r_{i}-\text{mean}(\{r_{1},r_{2},\cdots,r_{G}\})$ by their standard deviation (Eq. ([7](https://arxiv.org/html/2505.19075v2#S4.E7))), a lower variance amplifies the magnitude of the effective advantage $A_{i}$. This means even small deviations from the mean reward translate into strong learning signals.

This dynamic relates to the “Question-level difficulty bias” identified by [24], where inputs with highly consistent reward induce strong learning directions and thus lead to unstable policy updates. Indeed, GRPO often suffers from training instabilities when encountering policy collapse that can result in overfitting on “too easy” or “too hard” samples.
UniR, on the other hand, ensures structural stability by maintaining a fixed backbone model ${\pi_{b}}$. This design choice allows UniR to produce consistent responses, which typically exhibit low standard deviations in their rewards. Due to the inherent robustness of architecture, resistance of UniR to severe policy collapse allow it to mitigate amplified signals mode effectively. Instead of leading to instability, these strong learning signals can be constructively used by UniR, as opposed to the unstable state of GRPO.

## 6 Conclusion

In this paper, we present Universal Reasoner, a novel plug-and-play reasoning module that directly optimizes predefined rewards by focusing solely on a compact reasoning module while keeping the larger backbone model frozen, leading to efficient training comparable to other parameter-efficient methods. Our approach demonstrates remarkable transferability; the learned policies can effectively guide models of various sizes at inference without requiring further fine-tuning of the larger models. Furthermore, UniR facilitates the convenient combination of modules specialized for different tasks. This suggests the potential of our framework to freely integrate multiple modules, each trained on domain-specific datasets, to achieve synergistic improvements across a diverse range of tasks.

Limitations. The performance of our framework is inherently bounded by the capabilities of the selected frozen backbone model ${\pi_{b}}$, as ${\pi_{r}}$ primarily guides its existing representations. In addition, our approach employs an additional reasoning module introducing computational overhead at inference time. The extent to which a compact reasoning module ${\pi_{r}}$ can capture and generalize complex, multi-faceted reasoning across diverse scenarios remains an area for ongoing investigation, and its effectiveness can depend on the careful design of ${\pi_{r}}$’s architecture and training regimen.

## References

- [1]
Open thoughts.
[https://open-thoughts.ai](https://open-thoughts.ai), 2025.
Accessed: 2025-05-05.
- [2]
Abbas Abdolmaleki, Sandy Huang, Leonard Hasenclever, Michael Neunert, Francis Song, Martina Zambelli, Murilo Martins, Nicolas Heess, Raia Hadsell, and Martin Riedmiller.
A distributional view on multi-objective policy optimization.
In International conference on machine learning, pages 11–22. PMLR, 2020.
- [3]
Daman Arora and Andrea Zanette.
Training language models to reason efficiently.
arXiv preprint arXiv:2502.04463, 2025.
- [4]
Shuai Bai, Keqin Chen, Xuejing Liu, Jialin Wang, Wenbin Ge, Sibo Song, Kai Dang, Peng Wang, Shijie Wang, Jun Tang, et al.
Qwen2. 5-vl technical report.
arXiv preprint arXiv:2502.13923, 2025.
- [5]
Bespoke Labs.
Bespoke-stratos: The unreasonable effectiveness of reasoning distillation.
[https://www.bespokelabs.ai/blog/bespoke-stratos-the-unreasonable-effectiveness-of-reasoning-distillation](https://www.bespokelabs.ai/blog/bespoke-stratos-the-unreasonable-effectiveness-of-reasoning-distillation), 2025.
Accessed: 2025-05-5.
- [6]
Mauro Cettolo, Marcello Federico, Luisa Bentivogli, Jan Niehues, Sebastian Stüker, Katsuitho Sudoh, Koichiro Yoshino, and Christian Federmann.
Overview of the iwslt 2017 evaluation campaign.
In Proceedings of the 14th International Workshop on Spoken Language Translation, pages 2–14, 2017.
- [7]
Karl Cobbe, Vineet Kosaraju, Mohammad Bavarian, Mark Chen, Heewoo Jun, Lukasz Kaiser, Matthias Plappert, Jerry Tworek, Jacob Hilton, Reiichiro Nakano, et al.
Training verifiers to solve math word problems.
arXiv preprint arXiv:2110.14168, 2021.
- [8]
DeepSeek-AI.
Deepseek-r1: Incentivizing reasoning capability in llms via reinforcement learning, 2025.
- [9]
Jasper Dekoninck, Marc Fischer, Luca Beurer-Kellner, and Martin Vechev.
Controlled text generation via language model arithmetic.
arXiv preprint arXiv:2311.14479, 2023.
- [10]
Zane Durante, Qiuyuan Huang, Naoki Wake, Ran Gong, Jae Sung Park, Bidipta Sarkar, Rohan Taori, Yusuke Noda, Demetri Terzopoulos, Yejin Choi, et al.
Agent ai: Surveying the horizons of multimodal interaction.
arXiv preprint arXiv:2401.03568, 2024.
- [11]
Zhaopeng Feng, Shaosheng Cao, Jiahan Ren, Jiayuan Su, Ruizhe Chen, Yan Zhang, Zhe Xu, Yao Hu, Jian Wu, and Zuozhu Liu.
Mt-r1-zero: Advancing llm-based machine translation via r1-zero-like reinforcement learning.
arXiv preprint arXiv:2504.10160, 2025.
- [12]
Leo Gao, John Schulman, and Jacob Hilton.
Scaling laws for reward model overoptimization.
In International Conference on Machine Learning, pages 10835–10866. PMLR, 2023.
- [13]
Aaron Grattafiori, Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Alex Vaughan, et al.
The llama 3 herd of models.
arXiv preprint arXiv:2407.21783, 2024.
- [14]
Nuno M Guerreiro, Ricardo Rei, Daan van Stigt, Luisa Coheur, Pierre Colombo, and André FT Martins.
xcomet: Transparent machine translation evaluation through fine-grained error detection.
Transactions of the Association for Computational Linguistics, 12:979–995, 2024.
- [15]
Tuomas Haarnoja, Haoran Tang, Pieter Abbeel, and Sergey Levine.
Reinforcement learning with deep energy-based policies.
In International conference on machine learning, pages 1352–1361. PMLR, 2017.
- [16]
Seungwook Han, Idan Shenfeld, Akash Srivastava, Yoon Kim, and Pulkit Agrawal.
Value augmented sampling for language model alignment and personalization.
arXiv preprint arXiv:2405.06639, 2024.
- [17]
Chaoqun He, Renjie Luo, Yuzhuo Bai, Shengding Hu, Zhen Leng Thai, Junhao Shen, Jinyi Hu, Xu Han, Yujie Huang, Yuxiang Zhang, et al.
Olympiadbench: A challenging benchmark for promoting agi with olympiad-level bilingual multimodal scientific problems.
arXiv preprint arXiv:2402.14008, 2024.
- [18]
Edward J Hu, Yelong Shen, Phillip Wallis, Zeyuan Allen-Zhu, Yuanzhi Li, Shean Wang, Lu Wang, Weizhu Chen, et al.
Lora: Low-rank adaptation of large language models.
ICLR, 1(2):3, 2022.
- [19]
Jingcheng Hu, Yinmin Zhang, Qi Han, Daxin Jiang, Xiangyu Zhang, and Heung-Yeung Shum.
Open-reasoner-zero: An open source approach to scaling up reinforcement learning on the base model, 2025.
- [20]
Aitor Lewkowycz, Anders Andreassen, David Dohan, Ethan Dyer, Henryk Michalewski, Vinay Ramasesh, Ambrose Slone, Cem Anil, Imanol Schlag, Theo Gutman-Solo, et al.
Solving quantitative reasoning problems with language models.
Advances in Neural Information Processing Systems, 35:3843–3857, 2022.
- [21]
Yuhui Li, Fangyun Wei, Jinjing Zhao, Chao Zhang, and Hongyang Zhang.
Rain: Your language models can align themselves without finetuning.
arXiv preprint arXiv:2309.07124, 2023.
- [22]
Hunter Lightman, Vineet Kosaraju, Yuri Burda, Harrison Edwards, Bowen Baker, Teddy Lee, Jan Leike, John Schulman, Ilya Sutskever, and Karl Cobbe.
Let’s verify step by step.
In The Twelfth International Conference on Learning Representations, 2023.
- [23]
Jiacheng Liu, Andrew Cohen, Ramakanth Pasunuru, Yejin Choi, Hannaneh Hajishirzi, and Asli Celikyilmaz.
Making ppo even better: Value-guided monte-carlo tree search decoding.
Openreview https://openreview.net/forum?id=QaODpeRaOK, 2023.
- [24]
Zichen Liu, Changyu Chen, Wenjun Li, Penghui Qi, Tianyu Pang, Chao Du, Wee Sun Lee, and Min Lin.
Understanding r1-zero-like training: A critical perspective.
arXiv preprint arXiv:2503.20783, 2025.
- [25]
Pan Lu, Ran Gong, Shibiao Jiang, Liang Qiu, Siyuan Huang, Xiaodan Liang, and Song-Chun Zhu.
Inter-gps: Interpretable geometry problem solving with formal language and symbolic reasoning.
arXiv preprint arXiv:2105.04165, 2021.
- [26]
Sidharth Mudgal, Jong Lee, Harish Ganapathy, YaGuang Li, Tao Wang, Yanping Huang, Zhifeng Chen, Heng-Tze Cheng, Michael Collins, Trevor Strohman, et al.
Controlled decoding from language models.
arXiv preprint arXiv:2310.17022, 2023.
- [27]
Niklas Muennighoff, Zitong Yang, Weijia Shi, Xiang Lisa Li, Li Fei-Fei, Hannaneh Hajishirzi, Luke Zettlemoyer, Percy Liang, Emmanuel Candès, and Tatsunori Hashimoto.
s1: Simple test-time scaling, 2025.
- [28]
OpenAI.
Learning to reason with llms.
[https://openai.com/index/learning-to-reason-with-llms/](https://openai.com/index/learning-to-reason-with-llms/), September 2024.
Accessed: 2025-05-05.
- [29]
Bo Pang, Hanze Dong, Jiacheng Xu, Silvio Savarese, Yingbo Zhou, and Caiming Xiong.
Bolt: Bootstrap long chain-of-thought in language models without distillation.
arXiv preprint arXiv:2502.03860, 2025.
- [30]
Kishore Papineni, Salim Roukos, Todd Ward, and Wei-Jing Zhu.
Bleu: a method for automatic evaluation of machine translation.
In Proceedings of the 40th annual meeting of the Association for Computational Linguistics, pages 311–318, 2002.
- [31]
Rafael Rafailov, Archit Sharma, Eric Mitchell, Christopher D Manning, Stefano Ermon, and Chelsea Finn.
Direct preference optimization: Your language model is secretly a reward model.
In Thirty-seventh Conference on Neural Information Processing Systems, 2023.
- [32]
Ricardo Rei, Marcos Treviso, Nuno M Guerreiro, Chrysoula Zerva, Ana C Farinha, Christine Maroti, José GC De Souza, Taisiya Glushkova, Duarte M Alves, Alon Lavie, et al.
Cometkiwi: Ist-unbabel 2022 submission for the quality estimation shared task.
arXiv preprint arXiv:2209.06243, 2022.
- [33]
John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov.
Proximal policy optimization algorithms.
arXiv preprint arXiv:1707.06347, 2017.
- [34]
Zhihong Shao, Peiyi Wang, Qihao Zhu, Runxin Xu, Junxiao Song, Xiao Bi, Haowei Zhang, Mingchuan Zhang, YK Li, Y Wu, et al.
Deepseekmath: Pushing the limits of mathematical reasoning in open language models.
arXiv preprint arXiv:2402.03300, 2024.
- [35]
Freda Shi, Mirac Suzgun, Markus Freitag, Xuezhi Wang, Suraj Srivats, Soroush Vosoughi, Hyung Won Chung, Yi Tay, Sebastian Ruder, Denny Zhou, Dipanjan Das, and Jason Wei.
Language models are multilingual chain-of-thought reasoners, 2022.
- [36]
Charlie Snell, Ilya Kostrikov, Yi Su, Mengjiao Yang, and Sergey Levine.
Offline rl for natural language generation with implicit language q learning.
arXiv preprint arXiv:2206.11871, 2022.
- [37]
Kimi Team, Angang Du, Bofei Gao, Bowei Xing, Changjiu Jiang, Cheng Chen, Cheng Li, Chenjun Xiao, Chenzhuang Du, Chonghua Liao, et al.
Kimi k1.5: Scaling reinforcement learning with llms.
arXiv preprint arXiv:2501.12599, 2025.
- [38]
Masatoshi Uehara, Yulai Zhao, Chenyu Wang, Xiner Li, Aviv Regev, Sergey Levine, and Tommaso Biancalani.
Inference-time alignment in diffusion models with reward-guided generation: Tutorial and review.
arXiv preprint arXiv:2501.09685, 2025.
- [39]
Jason Wei, Xuezhi Wang, Dale Schuurmans, Maarten Bosma, Fei Xia, Ed Chi, Quoc V Le, Denny Zhou, et al.
Chain-of-thought prompting elicits reasoning in large language models.
Advances in neural information processing systems, 35:24824–24837, 2022.
- [40]
Yuancheng Xu, Udari Madhushani Sehwag, Alec Koppel, Sicheng Zhu, Bang An, Furong Huang, and Sumitra Ganesh.
Genarm: Reward guided generation with autoregressive reward model for test-time alignment.
arXiv preprint arXiv:2410.08193, 2024.
- [41]
An Yang, Baosong Yang, Beichen Zhang, Binyuan Hui, Bo Zheng, Bowen Yu, Chengyuan Li, Dayiheng Liu, Fei Huang, Haoran Wei, et al.
Qwen2. 5 technical report.
arXiv preprint arXiv:2412.15115, 2024.
- [42]
An Yang, Beichen Zhang, Binyuan Hui, Bofei Gao, Bowen Yu, Chengpeng Li, Dayiheng Liu, Jianhong Tu, Jingren Zhou, Junyang Lin, et al.
Qwen2. 5-math technical report: Toward mathematical expert model via self-improvement.
arXiv preprint arXiv:2409.12122, 2024.
- [43]
Shentao Yang, Shujian Zhang, Congying Xia, Yihao Feng, Caiming Xiong, and Mingyuan Zhou.
Preference-grounded token-level guidance for language model fine-tuning.
Advances in Neural Information Processing Systems, 36:24466–24496, 2023.
- [44]
Yixin Ye, Zhen Huang, Yang Xiao, Ethan Chern, Shijie Xia, and Pengfei Liu.
Limo: Less is more for reasoning, 2025.
- [45]
Qiying Yu, Zheng Zhang, Ruofei Zhu, Yufeng Yuan, Xiaochen Zuo, Yu Yue, Tiantian Fan, Gaohong Liu, Lingjun Liu, Xin Liu, et al.
Dapo: An open-source llm reinforcement learning system at scale.
arXiv preprint arXiv:2503.14476, 2025.
- [46]
Longteng Zhang, Lin Zhang, Shaohuai Shi, Xiaowen Chu, and Bo Li.
Lora-fa: Memory-efficient low-rank adaptation for large language models fine-tuning.
arXiv preprint arXiv:2308.03303, 2023.
- [47]
Renrui Zhang, Dongzhi Jiang, Yichi Zhang, Haokun Lin, Ziyu Guo, Pengshuo Qiu, Aojun Zhou, Pan Lu, Kai-Wei Chang, Yu Qiao, et al.
Mathverse: Does your multi-modal llm truly see the diagrams in visual math problems?
In European Conference on Computer Vision, pages 169–186. Springer, 2024.
- [48]
Brian D Ziebart, Andrew L Maas, J Andrew Bagnell, Anind K Dey, et al.
Maximum entropy inverse reinforcement learning.
In Aaai, volume 8, pages 1433–1438. Chicago, IL, USA, 2008.

## Appendix A Experimental Details

### A.1 Prompt template

For all models in the LLaMA family, we modified the default chat template by removing the knowledge cutoff and the generation time, as these fields vary depending on the time of execution and thus hinder reproducibility, and are also irrelevant to the task. For models in the Qwen family, however, we retained the default chat template as provided.

#### A.1.1 Math Reasoning Task.

The prompt template for mathematical problems follows the style of DeepSeek-R1 [8, 24]. Our system prompts were designed to encourage the model to articulate its reasoning process before providing a final answer, illustrated in Figure [7](https://arxiv.org/html/2505.19075v2#A1.F7), [8](https://arxiv.org/html/2505.19075v2#A1.F8). Specifically, for the Math-12k dataset using Qwen model in Figure [8](https://arxiv.org/html/2505.19075v2#A1.F8), we instructed the model to output its answer within a boxed LaTeX environment. This format was adopted to
ensure that various valid answer formats were correctly recognized and evaluated.

#### A.1.2 Translation Task.

For translation, inspired by the findings of MT-R1-Zero [11], which show that training with and without explicit reasoning yields comparable performance, we use direct translation prompts without explicit reasoning instructions-detailed in Figure [9](https://arxiv.org/html/2505.19075v2#A1.F9).

#### A.1.3 Compositional Task.

We instructed the model to solve math problems written in German and to perform reasoning and solving in English. The reasoning process was placed within <think> tags, and the final answer was enclosed within <answer> tags, as illustrated in Figure [10](https://arxiv.org/html/2505.19075v2#A1.F10). Using the prompt shown in Figure [11](https://arxiv.org/html/2505.19075v2#A1.F11), we evaluated the responses with GPT-4.1-nano-2025-04-14 on two dimensions: translation quality and math accuracy. Translation quality measures how well the response was translated into English, while math accuracy assesses accuracy using ground truth for a given problem.

While rule-based evaluation is often used for math accuracy due to the structured format of answers, we chose to use an LLM-based evaluator. This is because the math module is explicitly trained to follow a specific output format, whereas the translation module is not. Therefore, using a rule-based approach would unfairly favor the math side. To ensure a fair comparison between translation and problem-solving accuracy, we opt for GPT-4.1-nano as an impartial judge capable of understanding semantic correctness beyond strict formatting.

Figure: Figure 7: For GSM8K, the prompt specifies a reasoning-then-answer format, where the model is instructed to articulate its thought process within <think> tags, followed by the final answer enclosed in <answer> tags. question denotes the question of the benchmark.

Figure: Figure 8: In the Math-12K prompt format used with the LLaMA and Qwen models, answers follow a reasoning-then-answer structure using <think> and <answer> tags. With LLaMA, we observed frequent failures to generate valid \boxed{} expressions, so we chose to omit this syntax. In contrast, when using Qwen, we explicitly enclose the final answer within \boxed{}, prompting the model to express its answer in a boxed format.

Figure: Figure 9: Prompt used for Translation English-to-German and German-to-English. source indicates the source sentences to be translated.

Figure: Figure 10: System prompt used for English to German Math task: specifies the reasoning-then-answer format with <think> and <answer> tags. question indicates the math question in German.

Figure: Figure 11: Prompt used for scoring English-to-German Math task: Translation quality and Math Accuracy. response and gt_answer represent the model output and the ground truth of each math question, respectively.

### A.2 Hyperparameter Setup

The detailed hyperparameter settings are provided in Table [3](https://arxiv.org/html/2505.19075v2#A1.T3). For the fair comparison, we used the same hyperparameters across all methods, including baseline model, GRPO, GRPO with LoRA, and simple logit-sum. Specifically, we set the KL divergence coefficient $\beta=0.04$ for GRPO and GRPO with LoRA, following the default value in the Hugging Face TRL library (^4^44[https://github.com/huggingface/trl](https://github.com/huggingface/trl)). We observed that the original GRPO framework tends to suffer from unstable generation without KL divergence. Therefore, we adopted $\beta=0.04$ for GRPO-based baselines to ensure stable training.

Among the reward configurations, rule_based_accuracy assigns a reward of 1 if the numeric value between the <answer> and </answer> tags matches the ground truth and 0 otherwise.
The boxed_reward evaluates the value enclosed within the \boxed{} format, assigning a reward of 1 if it matches the correct answer and 0 otherwise. Similarly
the tag_based_reward evaluates the value enclosed within the answer tag.
For the implementation of boxed_reward and tag_based_reward we adopted the official code from the DR.GRPO repository [24](^5^55[https://github.com/sail-sg/understand-r1-zero](https://github.com/sail-sg/understand-r1-zero)). For BLEU, we used the default implementation of BLEU scoring provided by SacreBLEU(^6^66[https://github.com/mjpost/sacrebleu](https://github.com/mjpost/sacrebleu)).

### A.3 Evaluation Detail

For evaluation, BLEU and accuracy were computed using the same reward logic applied during training for each dataset. Translation baselines included outputs from Google Translator, implemented via the GoogleTranslator interface of the deep-translator library(^7^77[https://github.com/nidhaloff/deep-translator](https://github.com/nidhaloff/deep-translator)). GPT-4.1-nano-2025-04-14
was also used for comparison, with decoding performed at a sampling temperature of 0.7. All other model generations were conducted deterministically using do_sample=False.

**Table 3: Hyperparameter Setups.**
| Parameter | GSM8K | Math-12K-Qwen | Math-12K-Llama | Translation |
| --- | --- | --- | --- | --- |
| Training Configuration |  |  |  |  |
| gradient_checkpointing | true | true | true | true |
| gradient_checkpointing_kwargs | use_reentrant: false | use_reentrant: false | use_reentrant: false | use_reentrant: false |
| learning_rate | 1.0e-06 | 1.0e-06 | 1.0e-06 | 1.0e-06 |
| lr_scheduler_type | cosine with min_lr | cosine with min_lr | cosine with min_lr | cosine with min_lr |
| lr_scheduler_kwargs | min_lr_rate: 0.1 | min_lr_rate: 0.1 | min_lr_rate: 0.1 | min_lr_rate: 0.1 |
| warmup_ratio | 0.1 | 0.1 | 0.1 | 0.1 |
| max_steps | 1000 | 1000 | 1000 | 10000 |
| effective_batch_size | 64 | 32 | 32 | 32 |
| max_prompt_length | 512 | 512 | 512 | 512 |
| max_completion_length | 1024 | 2048 | 2048 | 1024 |
| num_generations | 8 | 8 | 8 | 8 |
| temperature | 0.7 | 0.7 | 0.7 | 0.7 |
| Evaluation Settings |  |  |  |  |
| max_prompt_length | 512 | 512 | 512 | 512 |
| max_completion_length | 1024 | 2048 | 2048 | 250 |
| do_sample | False | False | False | False |
| Reward Configuration |  |  |  |  |
| reward_funcs | rule_based_accuracy | boxed_reward | tag_based_reward | BLEU, cometkiwi |
| reward_weights | 1.0 | 1.0 | 1.0 | 1.0, 1.0 |

## Appendix B Qualitative Examples

### B.1 Example of Math Task

Based on Table [2](https://arxiv.org/html/2505.19075v2#S4.T2), fine-tuning with GRPO demonstrates improved performance over the base model; however, it also exhibits a tendency to generate longer outputs, which is alongside the results observed in [45, 37, 3]. In contrast, our proposed method not only surpasses the performance of both the base model and the GRPO-tuned model but also achieves this with a lower number of tokens to reach the final answer. Figure [12](https://arxiv.org/html/2505.19075v2#A2.F12) contains detailed examples that clearly show that our model is able to reach the correct answer in fewer steps.

Figure: Figure 12: Chain-of-thought comparison on a GSM8k example: while the Base and GRPO models both arrive at the correct answer with lengthy explanations, our approach produces the same result with a much more concise reasoning process.

Figure: Figure 13: Example of bidirectional translation (German (DE) and English (EN)). Our approach produces more fluent and grammatically correct outputs in both directions, preserving natural phrasing and accurate compound structures.

### B.2 Example of Translation Task

As shown in Figure [13](https://arxiv.org/html/2505.19075v2#A2.F13), our method generates more fluent and faithful translations in both directions. Compared to the base model and GRPO, our DE→EN translation correctly uses the natural phrase "What about…", aligning with the tone of the original. On the EN→DE side, only our method preserves the correct compound adjective "künstlich intelligenten", which is grammatically accurate and semantically precise. In contrast, baseline and GRPO variants either use awkward phrasing ("artificial intelligent system") or deviate from standard German compound structures. This demonstrates the effectiveness of our method in producing higher-quality, context-aware translations.

### B.3 Example of Combining Specialized Modules

Figure [14](https://arxiv.org/html/2505.19075v2#A2.F14) presents an example of the response to a German math-solving in English task. We compare the results from the UniR framework using the math-specific reasoning module ($\pi^{Math}_{r}$), the DE$\rightarrow$EN translation-specific reasoning module ($\pi^{Trans}_{r}$), and their combination. This example demonstrates that using only the math module fails to generate a response in English, while using only the translation module results in an incorrect solution. In contrast, incorporating both modules achieves both high translation quality and mathematical accuracy.

### B.4 Example of π r subscript 𝜋 𝑟 {\pi_{r}} italic_π start_POSTSUBSCRIPT italic_r end_POSTSUBSCRIPT Reasoning Performance

Figure [15](https://arxiv.org/html/2505.19075v2#A2.F15) illustrates a specific instance where the backbone model and the reasoning module fail to adequately comprehend the problem’s context. In contrast, our combined approach demonstrates a correct and coherent reasoning process for the same problem, highlighting the synergistic benefit of integrating the specialized reasoning module with the backbone model for improved contextual understanding and accurate inference.

Figure: Figure 14: We illustrate outputs for a German math problem solved in English. (Top) Guidance with only the math module solves correctly but responds in German; (Middle) while using only a translation module ensures English output but fails at math. (Bottom) Their combination achieves both target language and correct solution.

Figure: Figure 15: The Backbone Model and Reasoning Module both demonstrate flawed reasoning processes—such as making incorrect assumptions or applying invalid constraints—ultimately leading to confusion or incorrect deductions. In contrast, ours performs accurate symbolic reasoning and produces the correct answer with a clear, concise explanation, demonstrating the effectiveness of our approach in structured reasoning tasks.

## Appendix C Additional Experiments

### C.1 Extension of UniR to Vision Language Models (VLMs)

Dataset and Models: To investigate whether our UniR framework can enhance reasoning in multimodal contexts, we extended it by employing a Vision Language Model (VLM) as the backbone. Specifically, we utilized Qwen2.5-VL-3B-Instruct  [4] as the frozen backbone model ${\pi_{b}}$. This VLM was guided by a 0.5B Qwen-based reasoning module ${\pi_{r}}$, which was previously trained on the text-based GSM8k dataset. For evaluation, we used the MathVerse [47] and Geometry3k [25] benchmarks, which feature mathematical problems accompanied by visual diagrams, thus requiring both visual understanding and mathematical reasoning.

Experimental Setup: To assess the effectiveness of the reasoning module on processing combined visual and textual information, particularly for the MathVerse dataset, we focused on its "Vision Intensive" questions that include images. Although the reasoning module does not process visual tokens, we provided the same multimodal input to the overall UniR system. Evaluation protocols and system prompts followed those used in the GSM8k experiments.

**Table 4: Performance comparison on visually-related mathematical benchmarks. UniR enhances the reasoning capabilities of a frozen VLM backbone using a text-only math reasoning module.**
| Method | MathVerse | Geometry3k |
| --- | --- | --- |
| Qwen2.5-VL-3B-Instruct | 8.25 | 23.96 |
| Ours (+ ${\pi_{r}}$) | 11.29 | 25.12 |

Results: As shown in Table [4](https://arxiv.org/html/2505.19075v2#A3.T4), UniR significantly improves performance over the baseline VLM on both benchmarks. These enhancements are particularly noteworthy because the reasoning module ${\pi_{r}}$ was trained exclusively on text-based mathematical problems and had no direct access to visual information during its training. The results suggest that UniR allows the reasoning module to effectively guide the VLM, leading to better overall problem-solving capabilities in these multimodal mathematical tasks. This outcome highlights the ability of our framework to flexibly adapt to multimodal LLM backbones, allowing the framework to seamlessly extend its reach into multimodal reasoning domains. An example illustrating the qualitative difference in responses is depicted in Figure [16](https://arxiv.org/html/2505.19075v2#A3.F16).

Figure: Figure 16: Illustrative examples of responses from the baseline VLM and our UniR-extended framework on tasks from the (Left) Geometry3k and (Right) MathVerse benchmarks.
Refer to caption: x1.png

### C.2 Impact of Reasoning Model Size

Figure: Figure 17: Transferability of the 0.5B ${\pi_{r}}$ and 1.5B ${\pi_{r}}$ reasoning modules when combined with a 14B frozen backbone model. The 1.5B ${\pi_{r}}$ module demonstrated superior performance.
Refer to caption: extracted/6484925/figs/transfer_results_math3_1.5b.png

**Table 5: Zero-shot Pass@1 performance on mathematical benchmarks comparing reasoning modules of different sizes (0.5B vs. 1.5B) trained with a 3B backbone.**
|  | Method | Trained<br>Model | In-distribution | Out-of-distribution | Avg. |  |  |  |
| --- | --- | --- | --- | --- | --- | --- | --- | --- |
|  | GSM8K | MATH-500 | AIME24 | Minerva | OlympiadBench |  |  |  |
| Qwen2.5-3B | Baseline | - | 75.5 | 46.8 | 6.7 | 23.5 | 25.5 | 35.6 |
| Baseline + 0.5B | - | 72.1 | 41.8 | 6.7 | 16.2 | 22.1 | 31.8 |  |
| Baseline + 1.5B | - | 76.0 | 50.2 | 6.7 | 23.1 | 24.7 | 36.1 |  |
| GRPO LoRA | 3B | 78.2 | 64.2 | 13.3 | 28.3 | 30.1 | 42.8 |  |
| GRPO Full | 3B | 82.8 | 66.0 | 16.7 | 28.3 | 30.1 | 44.8 |  |
| Ours(3b+0.5b) | 0.5B | 82.2 | 64.8 | 10.0 | 26.8 | 27.0 | 42.2 |  |
| Ours(3b+1.5b) | 1.5B | 84.8 | 68.6 | 10.0 | 28.7 | 27.6 | 43.9 |  |

We conducted an ablation study to investigate the impact of reasoning module size, training a Qwen2.5-1.5B-Instruct reasoning module ${\pi_{r}}$ with a 3B backbone. As shown in Table [5](https://arxiv.org/html/2505.19075v2#A3.T5), this larger module achieved improved performance compared to the 0.5B ${\pi_{r}}$. Notably, it even outperformed a model fine-tuned using LoRA, suggesting that a larger reasoning model can learn more comprehensive information for reasoning tasks.
Figure [17](https://arxiv.org/html/2505.19075v2#A3.F17) further illustrates the transferability of the 1.5B reasoning module ${\pi_{r}}$ to a 14B frozen backbone model ${\pi_{b}}$. The transferred 1.5B ${\pi_{r}}$ led to an improved average performance (53.6) compared to the 0.5B ${\pi_{r}}$ (52.8).
Nevertheless, it is important to emphasize that our main contribution lies in demonstrating the ability to achieve strong results with a comparatively small language model during training, prioritizing efficiency.

### C.3 Analysis of Reasoning Module during Training

Figure: Figure 18: Performance comparison on the GSM8K dataset between standalone reasoning modules (${\pi_{r}}$, green) and our UniR framework (red), evaluated against their respective backbone models (orange). (Left) ${\pi_{r}}$ module using LLaMA-3.2-1B with a LLaMA-3.2-3B backbone. (Right) ${\pi_{r}}$ module using Qwen2.5-0.5B with a Qwen2.5-3B backbone.
Refer to caption: extracted/6484925/figs/gsm8k_combined_plot.png

We evaluated the performance of trained ${\pi_{r}}$ reasoning modules, as depicted in Figure [18](https://arxiv.org/html/2505.19075v2#A3.F18). This figure tracks the GSM8k dataset accuracy throughout the training process of reasoning modules designed for two different backbone architectures. The results indicate that the reasoning modules, when used in isolation, exhibit suboptimal performance. However, when these same reasoning modules are combined with their respective backbone models within our proposed framework, the resulting performance surpasses that of the backbone model alone. This highlights the efficacy of our learned reasoning modules in providing valuable guidance when integrated with a larger language model, leading to enhanced reasoning capabilities in the combined system.

## Peer Reviews

confidence: 5
contribution: 2
initial_rating: 4
presentation: 2
questions: Please see the weakness section above for relevant questions. In addition, 

- Can you show failure cases where composed modules hurt performance? What are the symptoms of negative interference?
- Code release plan -- will the code for this paper be released?
- Regarding composability, how do you envision setting the combination weights $\alpha_i$ in a real-world scenario without an expensive sweep for each new task combination? Have you considered methods for dynamically learning or setting these weights?
soundness: 3
strengths: Originality
- Proposes a clean, modular training objective mapping trajectory reward to token-level log-probabilities and uses additive logit fusion at inference which is a conceptually neat design.
- Inference-time composition via logit addition is derived from a KL-regularized multi-objective RL objective (Eq. 6), giving the method theoretical grounding absent in prior adapter ensembles.

Empirical Quality
- Consistent gains across 5 math benchmarks (up to +7.4 avg. over GRPO-full) and 2 translation directions (+1.8 BLEU) with ≤ 1.5 B trainable parameters.
- Strong weak-to-strong transfer: a 1.5 B module trained with a 3 B backbone improves a frozen 14 B Qwen-2.5 by +2.5 avg. points (Fig. 2).
- Composability demo: combining math + translation modules on German math→English task beats either specialist alone (Fig. 3).

Clarity & Reproducibility
- Complete hyper-parameter tables (App. A.1), reward definitions, and prompts (App. D) provided.
- Seeds and std-dev reported (10 runs, Tables 1–2).

Significance
- Enables cheap, plug-and-play specialization of proprietary or edge-deployed models without gradients through the backbone.
summary: The paper introduces UniR, a lightweight reasoning module trained with verifiable rewards that can be plugged into frozen LLMs at inference by simply adding its logits to those of the backbone. The module is (i) architecture-agnostic (only tokenizer alignment required), (ii) composable (multiple UniRs can be summed for multi-task reasoning), and (iii) shown to transfer from small to much larger backbones without further tuning. Experiments on mathematical reasoning, machine translation, vision-language math, and medical risk prediction report gains over full-model GRPO fine-tuning while using far fewer trainable parameters and less memory.
weaknesses: - Tokenizer-alignment prerequisite limits practicality. The paper claims “architecture-agnostic” but requires shared tokenizer (sec. 1, line 6). Is there a way to extend the method with tokenizer-mismatch adaptation strategy (e.g., embedding-level mapping layer trained with UniR) or can we quantify performance drop under vocabulary misalignment.

- The central identity (Eq. (4)), representing a trajectory reward as the sum of token log-probabilities of some policy, is assumed and only informally justified; Theorem 1 is informal and depends on strong assumptions (convergence of $\pi_\theta$ to $\pi_{\theta^*}$, uniqueness of optimum, and that π_r can exactly represent the reward) (Sec. 4.1 & 4.3). Can you provide a formal version of the theorem in the appendix with the details of the assumptions and proofs. The proof relies on the assumption that the guided policy has converged to the optimum. It would be beneficial to discuss the conditions under which this assumption holds. 

- The paper compares to GRPO variants and LoRA but does not compare to other recent inference-time guidance / token-level reward approaches (beyond brief GenARM mention).

- Missing Related Work: The paper fails to cite or compare against a relevant line of work on plug-and-play modules for frozen transformers. Specifically, TART ("TART: A Plug-and-Play Transformer for Zero-Shot Task Adaptation", Bhatia et al., ICLR 2022) introduces a nearly identical architectural paradigm: a small, separately trained module whose outputs are combined with a frozen backbone model to achieve task adaptation. While TART's training objective and application (zero-shot classification) differ, the core concept of a "plug-and-play reasoner for frozen LLMs" is not entirely new.

---

confidence: 3
contribution: 4
initial_rating: 6
presentation: 3
questions: Are there sources that can corroborate that the baseline numbers in Table 1 and Table 2 are correct (i.e. what we should expect for these models)?

Can you go into more detail about the architecture of the universal reasoner? The paper is vague on this point.
soundness: 3
strengths: This model seems very useful both as a research artifact (it seems to tell us that there isn't all that much to reasoning behavior; not many weights are actually needed to represent it "from scratch") as well as to users of AI tech. Enabling a "plug-and-play" ecosystem of foundation models and reasoning models could have a high impact on the community.
summary: The paper proposes a single "universal" model that converts practically any LLM into a reasoning-enabled LLM by adding logits at inference time. This enables reasoning behavior without the cost of fine-tuning and without the reduction in quality that reasoning-fine-tuning often produces.
weaknesses: Some of the baseline numbers are suspicious. For example, this paper reports a "backbone only" performance on MATH of 28.7 for llama3.2 3B, but huggingface here (https://huggingface.co/meta-llama/Llama-3.2-3B-Instruct) reports a 48.0 accuracy and that's consistent with what ArtificialAnalysis reports (https://artificialanalysis.ai/evaluations/math-500).  I'm also having trouble finding the other baseline numbers in prior work or online (GSM8K seems to mostly be 8-shot in evaluations of llama3 as opposed to the 0-shot evaluated in the paper).

---

confidence: 3
contribution: 2
initial_rating: 6
presentation: 3
questions: Regarding the first contribution (Line 75–76): circumventing the need for expensive preference dataset creation is not a major contribution of this work. The tasks studied here do not require such preference dataset construction.

In Figure 3, what is the baseline performance—that is, only the base model without the reasoning module? Showing the baseline performance would help readers better understand this figure.
soundness: 3
strengths: This paper has a clear motivation: reduce computational costs in post-training and improve the generalization of learned reasoning abilities. 

The content is high quality and the writing is clear and easy to understand.

The experiments are extensive, covering various benchmarks, tasks, model sizes, and series. The proposed UniR shows better performance compared with baseline methods and improves performance on larger backbones that were not used to train UniR.

When combined, different UniR modules show adjustable performance improvements on math and translation tasks.
summary: This paper studies LLM reasoning tasks and proposes UniR, a standalone reasoning module that reduces training costs while achieving generalization across models. UniR is composable and trained via GRPO. It provides token-level logit guidance to the backbone LLM. UniR demonstrates better performance than GRPO LoRA and GRPO Full with lower training costs. Additionally, UniR trained on a smaller backbone can improve task performance when combined with a larger backbone.
weaknesses: One weakness is the lack of detailed discussion about time and resource costs. The paper should highlight how UniR differs from baseline methods to demonstrate its computational advantages. 

Another point: the paper needs more analysis of the design choices for the standalone reasoning module. The authors use a smaller LLM as the reasoning module, which is a good initial step. However, have they experimented with using parts of a model instead of a full LLM? Could the reasoning module be further reduced in size?

Additionally, more analysis comparing outputs from baseline GRPO versus UniR would better illustrate UniR's influence. For example, what helpful adjustments does token-level logit modification provide compared to optimizing the model using only the final reward signal?

---

confidence: 3
contribution: 3
initial_rating: 6
presentation: 4
questions: 1. Does the reasoning module require any SFT warm start? Prior studies have shown that directly starting with GRPO, especially on small models, often leads to suboptimal performance.
soundness: 3
strengths: 1. A decoupled reasoning module that is model- and architecture-agnostic (assuming tokenizers can be aligned) with the base model, demonstrating that reasoning can be isolated and trained independently.

2. The reasoning module can be seamlessly integrated by adding its output logits to those of the base model, and it appears to enhance performance even in much larger base models.

3. A linear composition of reasoning modules is possible and effective, generalizing well even to unseen test sets (e.g., mortality prediction).

4. The paper is well presented and demonstrates strong understanding and integration of relevant theoretical concepts.
summary: This paper proposes a method to decouple reasoning from the base model by introducing a standalone, pluggable reasoning module that can be integrated with any model size or architecture by adding its output logits to those of the base model. The paper shows that it is possible to compose even two reasoning modules to combine task-specific abilities, achieving improved performance over both the base model and the GRPO-trained model (in math and translation tasks).
weaknesses: 1. How are the coefficients determined when combining reasoning modules?

2. Is it possible to compose more than two modules, and if so, how does performance scale with additional compositions?


Decision: Reject