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

Motivation

• Known: PEFT 기법(LoRA 등)은 자원 효율적으로 LLM을 특정 도메인에 특화시키지만, 모델 아키텍처에 종속되어 다른 크기의 backbone으로 전이가 어렵고, 여러 어댑터를 선형적으로 결합하는 것도 이론적 근거가 부족해 예측 불가능한 결과를 낳는다. RLVR 기반 방법(GRPO 등)은 verifiable reward로 reasoning 능력을 강화하지만 일반적으로 backbone 전체 또는 그 파라미터를 직접 파인튜닝해야 한다.
• Gap: 기존 방법들은 reasoning 능력 강화를 위해 backbone별로 재학습이 필요하거나(LoRA의 architecture dependency), 조합 시 이론적 지지가 부족한 선형 결합에 의존하며, reward 신호(trajectory-level)를 token-level guidance로 분리해 별도 모듈에 담아 frozen backbone에 순수 추론 시점(inference-time)에서만 결합하는 방법은 부재했다.
• Why: UniR은 backbone LLM의 파라미터를 전혀 건드리지 않고 작은 reasoning module만 학습해 다양한 크기·도메인의 frozen LLM에 재사용할 수 있게 하여, reasoning 강화의 학습 비용과 전이 비용을 동시에 크게 낮출 수 있는 실용적 방향을 제시한다.
• Approach: RLVR의 closed-form 최적 policy(πθ*(y|x) ∝ πb(y|x) exp(r(x,y)/β))에서 착안해, trajectory-level reward를 reasoning module의 log-probability 합으로 재매개변수화하고, 이 reasoning module만 policy gradient로 학습하여 backbone에 logit 합산으로 결합한다.

Achievement

Figure 2. Effectiveness of Reasoning module Transfer. Results demonstrate that a trained reasoning module can improve pe

1. Transferability: 3B backbone과 함께 학습된 1.5B reasoning module이 별도 튜닝 없이 14B backbone에도 적용되어 성능을 향상시킴 (weak-to-strong generalization).
2. Composability: 서로 다른 태스크(예: 수학 reasoning과 번역, 혹은 비전-언어 reasoning과 의료 reasoning)에 대해 학습된 여러 reasoning module을 logit 합산만으로 결합해 복합적 reasoning을 수행할 수 있음을 보임.
3. Superior Performance: 수학 reasoning, 기계 번역, 의료 reasoning, 그리고 vision-language 벤치마크에서 기존 fine-tuning 및 PEFT 기반 방법을 능가하는 성능을 달성.

How

Figure 1. UniR Framework Overview. Our approach trains a lightweight, transferable reasoning module (πr) using predefine

• RLVR 목적함수(Eq.1)의 closed-form 해(Eq.2)를 이용, trajectory-level reward r(x,y)를 reasoning module πr의 로그확률 합으로 근사·분해
• reasoning module πr만 별도로 파라미터화하고 policy gradient(예: GRPO 계열) 알고리즘으로 학습, backbone πb는 완전히 frozen 상태 유지
• 추론 시 backbone과 reasoning module의 output logit을 단순 덧셈하여 결합 후 샘플링
• 서로 다른 태스크로 학습된 다수의 reasoning module을 동일한 backbone에 대해 logit 덧셈으로 동시 적용해 composite reasoning 구현
• tokenizer가 공유되거나 토큰 공간이 정렬된 경우, 크기가 다른 backbone(예: 3B→14B) 혹은 도메인이 다른 모델(vision-language, medical)에도 reasoning module을 그대로 전이 적용

Originality

• reward를 trajectory-level에서 token-level logit guidance로 변환하고 이를 별도의 학습 가능한 reasoning module로 명시적으로 분리(decouple)한 최초의 정식화
• backbone 파라미터나 내부 레이어 접근 없이 순수 logit 덧셈만으로 결합하는 완전히 modular한 inference-time composition 구조 제안
• 기존 RAST(contrastive logit 차분), GenARM(preference 기반 surrogate reward), Reasoning Vectors(parameter 차분) 등과 달리 verifiable reward를 직접 policy gradient로 최적화하는 방식으로 차별화
• weak-to-strong transfer와 multi-domain composability(수학+번역, vision-language+medical)를 단일 additive 프레임워크로 동시에 실증

Limitation & Further Study

• backbone과 reasoning module 간 tokenizer 공유 또는 토큰 공간 정렬이 전제 조건이므로, 이질적인 tokenizer를 가진 모델 조합으로는 확장이 제한적일 수 있음
• logit 덧셈이라는 단순한 결합 방식이 항상 최적의 결합인지에 대한 이론적 보장이 충분히 논의되지 않았을 가능성(β 등 결합 강도 조절의 민감도 분석 필요)
• 여러 reasoning module을 동시에 조합할 때 태스크 간 상충(reward conflict)이나 조합 개수 증가에 따른 성능 저하 가능성에 대한 심층 분석 부족
• 후속 연구로 이질적 tokenizer 간 정렬 기법, 조합 시 자동 가중치 조절, 더 다양한 도메인(코드, 에이전트 태스크 등)으로의 확장 검토가 필요

Evaluation

총평: reward를 token-level logit guidance로 분해해 별도의 경량 모듈로 학습하고 frozen backbone에 additive하게 결합한다는 아이디어는 간단하면서도 강력하며, weak-to-strong transfer와 modular composability를 동시에 실현한 점이 인상적인 연구이다.