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

TL;DR

UniR tackles the challenge of enhancing reasoning in frozen LLMs by introducing Universal Reasoner (UniR), a lightweight, plug-and-play reasoning module trained on predefined rewards and applied at the logit level to a frozen backbone. The key idea is to decouple reasoning from the backbone and represent trajectory rewards as a sum of token-level signals via ${\pi_r}$, enabling modular composition by additive logits when multiple modules are combined. Training employs Group Relative Policy Optimization (GRPO) with a small, domain-specific module, while the backbone remains fixed, yielding transferability across model sizes and improved efficiency. Empirical results on mathematical problem solving and machine translation show UniR outperforms baseline fine-tuning methods and demonstrates strong weak-to-strong generalization, as well as favorable training efficiency and modularity for combining domain-specific reasoning policies. This framework provides a practical, scalable path to augment reasoning in diverse tasks without full fine-tuning of large LLM backbones.

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

Figures (18)

• 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.
• 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.
• 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.
• Figure 4: Performance on a German-to-English Math problem-solving task. The numbers in the figure indicate the value of $\alpha$.
• 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.

Theorems & Definitions (2)

• Theorem 1
• proof