LLM : 논문리뷰 : Reward-Guided Speculative Decoding for Efficient LLM Reasoning

Reward-Guided Speculative Decoding for Efficient LLM Reasoning 학습 노트

Purpose of the Paper:

• 기존 Speculative Decoding (SD) 방법론의 한계점 지적:
  • Strict unbiasedness 요구로 인해 draft model 과 target model 간 토큰 불일치 시 불필요한 계산 발생 및 효율성 저하. 특히 multi-step reasoning task 에서 draft model 이 target model 과 divergence 가 커질 때 문제 심화.
  • High-quality 토큰 (reward function 이 선호하는 토큰)이라도 target model probability 가 낮으면 reject 되어 computation 낭비 및 speedup 잠재력 감소.
  • 장기 reasoning task (math, coding)에서 overhead 증가 및 SD benefits 제한.
• 새로운 framework Reward-Guided Speculative Decoding (RSD) 제시:
  • 엄격한 unbiasedness 제약 완화, controlled bias 도입하여 high-reward outputs 우선순위 부여.
  • Draft model 과 target model dynamic mixing 통해 computation cost 와 output quality 간 trade-off 최적화.
  • 특히 resource-intensive reasoning 시나리오에서 LLM deployment 의 robust & cost-effective approach 목표.

Key Contributions:

• Novel Framework: Reward-Guided Speculative Decoding (RSD)
  • Draft model 과 target model 의 출력을 reward function 에 따라 dynamically mixing 하는 adaptive decoding framework 제시.
  • 기존 SD 의 strict unbiasedness 제약에서 벗어나 reward signal 기반 adaptive selection 도입.
  • Unbiasedness 유지 vs. Efficiency 향상 간 균형점을 찾아 flexible mixture distribution (PRSD) 구축.
  • Rejection sampling 활용하여 draft outputs selectively refine, scalability 확보.
• Theoretical Justification:
  • Threshold-based mixture strategy 가 resource utilization 과 performance 간 optimal balance 달성함을 이론적으로 증명.
  • Weighting function w(.) 및 acceptance criterion Aw 제시, 다양한 variant 및 optimal weighting strategy 탐구.
  • Proposition 2.3: Constrained sampling budget 하에서 reward maximization 을 위한 optimal weighting function 은 binary step function 임을 증명.
• Empirical Validation:
  • Challenging reasoning benchmarks (GSM8K, MATH500, Olympiad Bench 등) 에서 RSD effectiveness 입증.
  • Target model only decoding 대비 최대 4.4x FLOPs 감소 달성.
  • Parallel decoding method (SD) 대비 평균 +3.5 accuracy 향상 달성.
  • 다양한 모델 (Qwen2.5, Llama-3) 및 PRM (Skywork-01-Open-PRM) 조합 실험 통해 RSD robustness 및 generalizability 확인.

Experimental Highlights:

• Datasets: GSM8K, MATH500, MMLU STEM, OlympiadBench, GPQA, GaoKao-2023-En 등 diverse reasoning tasks datasets 활용.
• Baselines:
  • Target model only decoding
  • Draft model with/without PRM (Majority Voting, Best-of-N)
  • Speculative Decoding (SD)
  • Search-based methods (Beam Search, Process Best-of-N)
• Metrics: Accuracy, FLOPs reduction
• Key Results:
  • Accuracy 향상: RSD 가 SD 대비 reasoning accuracy 평균 +3.5 향상. 특히 GPQA benchmark 에서 single target model (32.8) 대비 RSD (1.5B/7B/1.5B) 는 38.4 로 significant improvement 달성.
  • Efficiency 증대: Target model only decoding 대비 최대 4.4x FLOPs 감소. RSD (7B/72B/7B) 는 Target (72B) 대비 notable accuracy (88.0 vs 85.6 on MATH500) 유지하며 FLOPs 대폭 감소.
  • Search-based methods outperform: RSD 가 Beam Search, Process Best-of-N 등 search-based methods 대비 성능 우위. 특히 complex reasoning steps 에서 RSD 의 stepwise guidance 효과 입증.
  • Figure 4: MATH500 dataset 에서 RSD (1.5B/7B/7B) 가 SD (1.5B/7B) 및 Target (7B) 대비 accuracy improvement 와 FLOPs reduction 동시에 달성 시각화.
  • Figure 5: Threshold δ 값 변화에 따른 RSD (1.5B/7B/7B) 성능 변화 분석. δ = 0.7 에서 accuracy peak 달성 및 draft model alone 으로 해결 가능한 question proportion 변화 관찰.

Limitations and Future Work:

• General-purpose PRM 사용: 실험에서 open-source general PRM (Skywork-01-Open-PRM) 활용. Draft model 과 closely aligned 된 specialized PRM (fine-tuning 또는 training) 개발 시 성능 further enhance 가능성 제시. (Acceptance rate 증가 기대)
• RSD 와 SD 결합: Rejected step 에서 SD (draft+target) 활용하여 step regeneration 하는 RSD+SD combined approach 통해 efficiency 추가 최적화 가능성 언급.
• Weighting function 및 rejection schemes 추가 탐색: Binary step function 외에 다른 weighting functions (Table 1) 및 advanced merging methods (MergeKit 활용 linear merging 실험 진행) 등 다양한 rejection schemes exploration 필요성 제시.
• Specialized PRM training/fine-tuning: Targeted datasets 및 chain-of-thought annotations 활용한 specialized PRM 개발 future work 로 제시. Mathematics, coding, scientific discovery 등 domain-specific applications 에서 reward-guided speculative decoding benefits 극대화 기대.
• 모델 merge 효과 추가 연구: Merging larger models 시 performance improvement 관찰 (+1.4 vs. +0.8). Yadav et al. (2024) 연구 결과와 consistent. Model merging scale 에 대한 deeper investigation 필요.

Overall, RSD framework 는 LLM reasoning inference efficiency 와 accuracy 를 동시에 향상시키는 효과적인 접근 방식이며, 향후 specialized PRM 개발 및 다양한 optimization techniques 탐구를 통해 further 발전 가능성이 높음.

 

 

 

마치 "정답 토큰" (정확한 매칭) 만을 고집하는 대신, "정답에 가까운 토큰" (분포 유사성) 이면서 "훌륭한 답변" (high-reward) 이라면 "정답으로 인정" 해주는 유연한 채점 방식과 유사하다고 볼 수 있습니다.

 

 

 

 

 

 

Abstract

본 논문은 large language models (LLMs)의 추론 효율성을 향상시키는 것을 목표로 하는 새로운 프레임워크인 Reward-Guided Speculative Decoding (RSD)을 소개한다. RSD는 기존의 strict unbiasedness를 강제하는 speculative decoding 방법과는 대조적으로, lightweight draft model과 더 강력한 target model을 시너지적으로 결합하여 high-reward outputs을 우선시하도록 제어된 편향을 통합한다. RSD는 중간 디코딩 단계를 평가하고 계산 비용과 output quality 간의 trade-off를 최적화하기 위해 process reward model을 사용하여 target model을 호출할지 여부를 동적으로 결정한다. 우리는 threshold-based mixture 전략이 리소스 활용률과 성능 사이의 최적의 균형을 달성한다는 것을 이론적으로 보여준다. Olympiad 수준의 task를 포함한 어려운 reasoning 벤치마크에 대한 광범위한 평가 결과, RSD는 target model만으로 디코딩하는 것에 비해 상당한 효율성 향상(최대 4.4배 적은 FLOPs)을 제공하는 동시에 평균적으로 parallel decoding 방법에 비해 훨씬 더 나은 정확도(+3.5까지)를 달성하는 것으로 나타났다. 이러한 결과는 RSD가 resource-intensive 시나리오에서 LLM을 배포하기 위한 강력하고 비용 효율적인 접근 방식임을 강조한다.

 

 

 

Introduction

Scaling laws는 machine learning community에서 large language models 개발의 기본 원칙으로 널리 인정받고 있습니다 (Hestness et al., 2017; Kaplan et al., 2020; Hoffmann et al., 2022). 그들은 model size와 dataset scale을 모두 증가시키면 loss 감소와 결과적으로 일반화 능력이 향상된다는 점을 강조합니다. data와 model size가 엄청난 수준으로 확장되면 성능은 전례 없는 수준에 도달할 수 있습니다. Large models는 다양한 task에서 놀라운 능력을 보여주며 강력한 일반화 및 advanced reasoning skills을 보여줍니다 (Brown et al., 2020; Hurst et al., 2024; Anthropic, 2024; Team et al., 2024).

이러한 발전은 높은 computational 및 경제적 비용을 초래합니다. training은 resource-intensive하지만, scale에서의 inference는 훨씬 더 비용이 많이 들며 수십억 건의 query를 처리하기 위해 광대한 computational infrastructure와 에너지를 필요로 합니다 (Patterson et al., 2021). inference 비용의 기하급수적인 증가는 large model 배포의 핵심 과제가 되며, 에너지 사용을 줄이고 scalability를 보장하기 위한 효율적인 기술의 필요성을 강조합니다 (Frantar et al., 2022; Lin et al., 2024; Xu et al., 2023; Sun et al., 2023; Zhang et al., 2023b; Li et al., 2024a; Xu et al., 2024; Liao & Monz, 2024; Liao et al., 2024a).

특히, large LLM의 Sequential token generation은 더 작은 models에 비해 훨씬 더 높은 computational 비용을 발생시킵니다. 이러한 latency 증가는 실시간 또는 high-throughput 애플리케이션에서의 배포를 저해할 수 있습니다. 이 문제를 해결하기 위해 speculative decoding과 같은 parallel decoding 기술이 효과적인 해결책으로 등장했습니다 (Leviathan et al., 2023). Speculative decoding은 더 작고 가벼운 model을 활용하여 candidate outputs을 생성한 다음 더 큰 model에 의해 검증 및 개선되는 방식으로 작동합니다. 이 접근 방식은 large model에 필요한 decoding tokens의 수를 크게 줄여 전체 process 속도를 높입니다. 더 작은 model은 guide 역할을 하여 더 큰 model이 확인하거나 조정할 수 있는 sequence를 제안하므로 품질 저하 없이 더 빠른 inference 시간을 제공합니다. 또한, speculative decoding은 더 작은 model의 예측을 더 큰 model의 기능에 맞춰 조정함으로써 높은 품질의 outputs을 유지하여 효율성을 보장합니다. 이러한 방법을 통해 작은 model의 predictions과 큰 model의 capabilities 간의 불일치를 최소화하고 inference 중 안정성을 향상시킵니다.

parallel decoding의 발전에도 불구하고, speculative decoding은 특히 multi-step generation과 같은 복잡한 reasoning task에 대해서는 여전히 충분히 활용되지 않고 있습니다. 핵심적인 제한 사항은 final token distribution이 large model의 distribution과 일치하도록 보장하지만 다양한 completions 탐색의 유연성을 제한하는 strict unbiasedness 요구 사항입니다 (Holtzman et al., 2020). unbiasedness는 이론적 충실도를 유지하지만, 특히 draft model이 large model과 다를 때 효율성을 저하시키는 경우가 많습니다. High-quality tokens (예: process reward가 선호하는 tokens)는 large model의 확률이 너무 낮으면 여전히 거부될 수 있으며, 이는 computational 낭비로 이어지고 잠재적인 속도 향상 효과를 무효화합니다. 이러한 의존성은 overhead를 부풀리고 특히 수학 및 코딩과 같은 long-trajectory reasoning task에서 speculative decoding의 이점을 제한합니다. final distribution이 large model과 약간 벗어나는 것을 허용하는 controlled bias는 성능을 향상시킬 수 있습니다. draft token이 정확하지만 large model의 distribution과 정확히 일치하지 않는 경우 strict rejection은 역효과를 냅니다. Reward-guided acceptance는 valuable partial solutions을 유지하고 불필요한 queries를 줄이며 large model의 성능을 능가할 수도 있습니다. 따라서 효율성과 robustness의 균형을 맞추고 더 광범위한 실제 적용 가능성을 보장하기 위해서는 보다 adaptive 접근 방식이 필요합니다.

본 연구에서는 computational적으로 가벼운 "draft" evaluations과 더 강력한 "target" model의 reward-driven refinements을 통합하여 효율성과 정확성의 균형을 맞추는 새로운 프레임워크인 Reward-Guided Speculative Decoding (RSD)를 소개합니다. strict unbiasedness를 엄격하게 적용하는 기존의 speculative decoding과 달리, RSD는 reward signals를 활용하여 일치하지 않는 tokens을 즉시 버리는 대신 high-value draft outputs을 adaptive하게 선택합니다. 이 process는 draft model이 candidate steps을 생성하는 것으로 시작하며, 이는 reward function을 사용하여 평가됩니다. 충분히 높은 reward 점수를 가진 steps은 reasoning trajectory를 계속하기 위해 허용되는 반면, 낮은 점수를 가진 steps은 target model을 사용한 speculative corrections를 트리거합니다. Fig. 1에서 볼 수 있듯이, 이 adaptive 메커니즘은 resource 할당을 최적화하면서 draft model과 target model 간의 distribution shifts 문제에 robust합니다. RSD는 large model을 호출할 시점을 동적으로 조정함으로써 불필요한 computations을 크게 줄이는 동시에 기존의 inference 접근 방식의 품질을 유지하거나 능가합니다. 이 접근 방식은 computational 비용과 정확성의 균형을 맞추는 것이 중요한 long-horizon reasoning task에 특히 적합합니다.

Contributions. 우리는 특히 reasoning task에 대해 LLM inference를 가속화하는 새로운 접근 방식인 Reward-Guided Speculative Decoding을 제안합니다. 이는 각 step에서 output 품질을 평가하는 reward function에 의해 guide되는 draft model과 target model의 outputs을 동적으로 혼합하는 adaptive decoding 프레임워크를 도입합니다. 이를 통해 reward-based weighting을 통해 효율성과 정확성의 균형을 맞춘 flexible mixture distribution, PRSD를 구성하여 효율적이고 high-quality reasoning을 가능하게 합니다. RSD는 rejection sampling을 사용하여 draft outputs을 선택적으로 개선하여 scalability를 보장합니다. 이론적으로 우리는 computational 제약 조건 하에서 최적의 weighting 전략을 도출하여 성능 저하 없이 효율성을 극대화합니다. GSM8K, MATH500, Olympiad Bench, GPQA, MMLU STEM 및 GaoKao2023-En에 대한 광범위한 실험 결과, RSD는 평균적으로 SD보다 reasoning 정확도를 최대 3.5까지 향상시킬 뿐만 아니라 target model만 사용하는 경우와 비교하여 최대 4.4배 적은 FLOPs로 inference computation을 크게 줄이는 것으로 나타났습니다.

 

Introduction 정리 노트

배경

• Scaling laws에 따라 model 크기와 dataset 크기를 키우면 LLM의 성능이 향상되지만, inference 비용이 매우 커져 large model 배포에 어려움이 있음.
• 특히, sequential token generation은 작은 model에 비해 높은 computational 비용을 발생시켜 latency를 증가시킴.
• Speculative decoding과 같은 parallel decoding 기술이 latency를 줄이는 효과적인 대안으로 제시됨.

문제점

• Speculative decoding은 strict unbiasedness 요구 사항으로 인해 다양한 completions 탐색에 제약이 있음.
• Draft model과 large model 간의 차이가 클 경우 high-quality tokens이 large model의 확률이 낮다는 이유로 거부되어 computational 낭비가 발생할 수 있음.
• 기존 방식은 long-trajectory reasoning task에서 효율성이 떨어짐.

제안하는 방법 (RSD)

• Reward-Guided Speculative Decoding (RSD) 프레임워크를 제안하여 효율성과 정확성 균형을 맞춤.
• Reward signals를 활용하여 high-value draft outputs을 adaptive하게 선택하고, 일치하지 않는 tokens을 즉시 버리는 대신 개선함.
• Draft model에서 생성된 candidate steps을 reward function으로 평가하고, 높은 reward 점수를 가진 steps은 reasoning trajectory를 계속하고 낮은 점수를 가진 steps은 target model을 사용하여 수정함.
• Draft model과 target model 간의 distribution shifts 문제에 robust하며, resource 할당을 최적화함.
• Long-horizon reasoning task에 적합함.

핵심 기여

• LLM inference, 특히 reasoning task를 가속화하는 새로운 접근 방식인 RSD를 제안.
• 각 step에서 output 품질을 평가하는 reward function에 따라 draft model과 target model의 outputs을 동적으로 혼합하는 adaptive decoding 프레임워크를 도입하여 효율적이고 고품질의 reasoning을 가능하게 함.
• Rejection sampling을 사용하여 draft outputs을 선택적으로 개선하여 scalability를 보장.
• Computational 제약 조건 하에서 최적의 weighting 전략을 도출하여 성능 저하 없이 효율성을 극대화함.

쉬운 설명:

LLM은 모델 크기가 클수록 성능이 좋아지지만, 사용하는 데 비용이 많이 듭니다. 특히 텍스트를 한 글자씩 순서대로 생성하는 과정이 오래 걸립니다. Speculative decoding은 이 과정을 빠르게 하기 위해 작은 모델이 먼저 대략적인 텍스트를 생성하고, 큰 모델이 이를 확인하고 수정하는 방식입니다. 하지만 기존의 speculative decoding은 작은 모델과 큰 모델의 예측이 정확히 일치해야만 수정하지 않기 때문에, 작은 모델이 좋은 텍스트를 생성했더라도 큰 모델의 기준에 맞지 않으면 버려지는 비효율적인 부분이 있었습니다.

RSD는 이러한 문제점을 해결하기 위해 reward function이라는 평가 도구를 사용하여 작은 모델이 생성한 텍스트의 가치를 판단합니다. 가치가 높은 텍스트는 큰 모델의 예측과 정확히 일치하지 않더라도 수정하지 않고 그대로 사용하고, 가치가 낮은 텍스트만 큰 모델을 사용하여 수정합니다. 이를 통해 불필요한 계산을 줄여 효율성을 높이고, 정확도 또한 유지할 수 있습니다. 즉, RSD는 speculative decoding의 장점을 살리면서 단점을 보완한 똑똑한 방법이라고 할 수 있습니다.