자기 개선 생성을 위한 테스트 시간 정책 학습을 통한 적응적 디코딩

Adaptive Decoding via Test-Time Policy Learning for Self-Improving Generation

디코딩 전략은 대규모 언어 모델(LLM)의 출력 품질을 결정하는 데 중요한 역할을 하지만, 일반적으로 사용되는 휴리스틱 방법인 탐욕(greedy) 방식이나 고정된 온도/top-p 디코딩은 정적이며 종종 작업에 독립적이어서, 스타일이나 구조의 유연성이 요구되는 다양한 분야에서 최적이거나 일관성 없는 생성 품질을 초래합니다. 본 연구에서는 강화 학습 기반의 디코더 샘플러를 제안합니다. 이 샘플러는 디코딩을 순차적 의사 결정 문제로 간주하고, LLM의 가중치를 고정한 상태에서 테스트 시간에 샘플링 파라미터를 조정하는 경량 정책을 학습합니다. Granite-3.3-2B 및 Qwen-2.5-0.5B 모델을 사용하여 BookSum, arXiv, WikiHow 등 요약 데이터셋에 대한 실험을 수행했습니다. 제안하는 정책 샘플러는 탐욕 방식 및 고정 기반 모델을 지속적으로 능가하며, BookSum(Granite 모델 기준)에서 최대 +88%, WikiHow(Qwen 모델 기준)에서 최대 +79%의 성능 향상을 달성했습니다. 보상 요소 분석 결과, 중복만 고려하는 목표는 복합 보상에 비해 성능이 떨어지며, 구조화된 보상 요소(길이, 커버리지, 반복, 완전성)는 안정적이고 지속적인 성능 향상을 가능하게 합니다. 이러한 결과는 강화 학습이 디코딩 과정에서의 테스트 시간 적응을 위한 실용적인 메커니즘임을 보여주며, 대규모 모델을 재학습하지 않고도 도메인에 특화된 및 사용자 제어 가능한 생성을 가능하게 합니다.

Decoding strategies largely determine the quality of Large Language Model (LLM) outputs, yet widely used heuristics such as greedy or fixed temperature/top-p decoding are static and often task-agnostic, leading to suboptimal or inconsistent generation quality across domains that demand stylistic or structural flexibility. We introduce a reinforcement learning-based decoder sampler that treats decoding as sequential decision-making and learns a lightweight policy to adjust sampling parameters at test-time while keeping LLM weights frozen. We evaluated summarization datasets including BookSum, arXiv, and WikiHow using Granite-3.3-2B and Qwen-2.5-0.5B. Our policy sampler consistently outperforms greedy and static baselines, achieving relative gains of up to +88% (BookSum, Granite) and +79% (WikiHow, Qwen). Reward ablations show that overlap-only objectives underperform compared to composite rewards, while structured shaping terms (length, coverage, repetition, completeness) enable stable and sustained improvements. These findings highlight reinforcement learning as a practical mechanism for test-time adaptation in decoding, enabling domain-aware and user-controllable generation without retraining large models.