Learning Adaptive LLM Decoding

이 논문은 고정된 샘플링 하이퍼파라미터 대신 강화학습을 통해 추론 시 계산 자원에 따라 동적으로 샘플링 전략을 선택하는 경량 디코딩 어댑터를 제안하여, 수학 및 코딩 벤치마크에서 고정된 예산 대비 정확도를 크게 향상시킨다는 점을 설명합니다.

핵심

🎒 비유: 똑똑한 학생과 '고정된' 시험지

1. 기존의 방식: "모든 시험에 똑같은 연필을 쓰는 학생"
지금까지 우리는 거대한 AI 모델 (LLM) 이 문제를 풀 때, 반드시 같은 규칙을 사용하게 했습니다.

• 예를 들어, "답을 고를 때 무작위성은 0.7 로 하라"거나 "가장 확률이 높은 5 개만 고르라"는 식의 **고정된 설정 (하이퍼파라미터)**을 모델 전체에 적용했습니다.
• 문제점: 이는 마치 매우 어려운 수학 문제를 풀 때나, 아주 쉬운 국어 문제를 풀 때나 반드시 같은 연필을 들고 같은 속도로 쓰는 것과 같습니다.
  • 어려운 문제에서는 조금 더 신중하게 (확률적으로) 여러 가지 가능성을 생각해봐야 하는데, 너무 단단하게 고정되어 있어 실수를 합니다.
  • 쉬운 문제에서는 너무 많은 시간을 써서 불필요하게 고민합니다.

2. 이 논문의 제안: "상황을 보고 연필을 바꿔 쓰는 스마트한 조교"
이 논문은 모델 자체를 바꿀 필요 없이, **모델 옆에 아주 작고 가벼운 '조교 (Adapters)'**를 붙이는 아이디어를 제안합니다. 이 조교는 AI 가 답을 하나씩 써나가는 순간순간, 상황에 따라 가장 좋은 '답을 고르는 방식'을 실시간으로 결정해 줍니다.

• 어려운 부분 (불확실성이 높은 순간): "이 부분은 헷갈리네? 그럼 잠시 멈추고 여러 가지 가능성을 열어두고 (랜덤하게) 생각해보자!"라고 지시합니다.
• 쉬운 부분 (확실한 순간): "이건 너무 쉬워. 그냥 가장 확실한 답으로 바로 넘어가자!"라고 지시합니다.

---

🧠 두 가지 레벨의 '스마트 조교'

이 논문은 이 조교를 두 가지 방식으로 훈련시켰습니다.

1. 문제 전체를 보는 조교 (Sequence-Level)

• 역할: 문제를 받자마자 "이 문제는 어떤 방식으로 풀어야 할까?"를 한 번에 결정합니다.
• 비유: 시험지를 받자마자 "이건 논술형 문제니까 신중하게, 저건 객관식이라 빠르게 풀자"라고 시험 전략을 세우는 것입니다.
• 효과: 병렬로 여러 번 풀이를 시도할 때 (예: 8 번 시도), 어떤 시도는 신중하게, 어떤 시도는 빠르게 돌리는 등 자원 배분을 잘해서 전체 정확도를 높입니다.

2. 단어 하나하나를 보는 조교 (Token-Level)

• 역할: AI 가 문장을 한 글자씩 써나갈 때, 매 글자마다 "지금 이 글자는 확실히 써야 할까, 아니면 여러 가지 가능성을 열어둘까?"를 결정합니다.
• 비유: 글을 쓸 때, "주인공 이름"을 쓸 때는 확실히 쓰지만, "다음에 무슨 일이 일어날지"를 상상할 때는 여러 가지 가능성을 열어두고 고민하는 식입니다.
• 효과: 같은 문장 안에서도 어려운 부분에서는 머리를 굴리고, 쉬운 부분에서는 빠르게 넘어가는 정교한 제어가 가능해져서, 같은 시간 (컴퓨팅 비용) 안에 훨씬 더 정확한 답을 냅니다.

---

🏆 어떻게 훈련시켰나요? (게임의 규칙)

이 조교들은 사람처럼 가르치지 않고, 게임처럼 훈련시켰습니다.

1. 게임: 수학 문제나 코딩 문제를 풀게 합니다.
2. 보상: 정답을 맞추면 점수를 주고, 틀리면 0 점입니다. (정답이 명확한 문제만 사용)
3. 학습: "어떤 상황에서 어떤 방식 (랜덤하게 풀기 vs 확실히 풀기) 을 선택했을 때 점수가 잘 나왔는지"를 반복해서 학습시킵니다.
4. 결과: 모델 자체는 그대로 둔 채, 이 작은 조교만 학습시켜서 정확도와 비용 사이의 균형을 완벽하게 잡게 만들었습니다.

---

💡 핵심 성과: "똑똑한 자원 관리"

실험 결과 (수학 문제 MATH 와 코딩 대회 CodeContests 기준), 이 방식을 적용한 AI 는 다음과 같은 성과를 냈습니다.

• 같은 시간, 더 높은 점수: 정해진 시간 (컴퓨팅 비용) 안에 풀었을 때, 기존 고정 방식보다 최대 10% 이상 더 많은 문제를 맞췄습니다.
• 유연한 대응: 문제가 어렵든 쉽든, 혹은 같은 문제 안에서도 헷갈리는 부분과 확실한 부분을 스스로 구분해서 최적의 전략을 썼습니다.

📝 한 줄 요약

"거대한 AI 모델은 그대로 두고, 그 옆에 '상황을 보고 답을 고르는 방식을 실시간으로 바꿔주는 작은 조교'를 붙여주니, 같은 노력으로 훨씬 더 똑똑한 결과를 얻을 수 있었다!"

이 기술은 앞으로 AI 가 더 적은 전기를 쓰면서도 더 복잡한 문제를 해결하는 데 큰 도움이 될 것입니다.

기술 요약

1. 문제 정의 (Problem Statement)

대규모 언어 모델 (LLM) 의 추론 (Inference) 과정에서 **디코딩 (Decoding)**은 모델이 예측 분포에서 출력 토큰을 샘플링하는 과정입니다. 현재 관행에서는 온도 (temperature), top-k, top-p 와 같은 샘플링 하이퍼파라미터가 전체 모델이나 데이터셋에 대해 **고정된 값 (Fixed)**으로 설정됩니다.

그러나 이러한 고정된 접근 방식은 다음과 같은 한계를 가집니다:

• 과도한 단순화: 프롬프트의 난이도, 추론 스타일, 심지어 개별 토큰 단계의 불확실성 (uncertainty) 에 따라 최적의 디코딩 전략은 크게 달라집니다.
• 비효율성: 최근 연구에 따르면 추론 과정 중 불확실성이 높은 '포크링 토큰 (forking tokens)' 소수만이 결과에 결정적인 영향을 미치는데, 고정된 전략은 이러한 상황에 유연하게 대응하지 못합니다.
• 학습 - 테스트 불일치 (Train-Test Mismatch): 강화학습 (RL) 기반의 추론 향상 기법들 (RLVR 등) 이 대부분 고정된 디코딩 전략 하에서 훈련되지만, 실제 배포 시에는 다른 추론 제약 조건 하에서 평가되는 경우가 많아 성능 저하를 초래할 수 있습니다.

따라서, 계산 자원 (Compute Budget) 을 명시적으로 고려하면서 추론 시점에 동적으로 샘플링 전략을 선택하는 적응형 디코딩 정책을 학습하는 것이 핵심 문제입니다.

---

2. 방법론 (Methodology)

저자들은 언어 모델 (LLM) 자체를 파인튜닝하지 않고, **가벼운 디코딩 어댑터 (Decoding Adapters)**를 강화학습 (RL) 으로 훈련하여 추론 시 디코딩 행동을 제어하는 프레임워크를 제안합니다.

2.1. 기본 구조

• 고정된 LLM: 베이스 LLM 의 파라미터는 고정 (Frozen) 되어 있으며, 어댑터만 학습됩니다.
• 가용성 있는 보상: 수학 및 코딩 문제와 같이 정답 여부를 검증할 수 있는 **검증 가능한 종료 보상 (Verifiable Terminal Rewards)**을 사용하여 학습합니다. (학습된 보상 모델이나 선호도 레이블 불필요)
• 두 가지 수준의 적응 (Adaptation Levels):
  1. 시퀀스 수준 (Sequence-Level):
     • 문제 설정: 컨텍스트 밴딧 (Contextual Bandit) 문제로 모델링.
     • 행동: 각 프롬프트에 대해 단 하나의 디코딩 전략 (예: Greedy, Top-k, Top-p, Min-p 등) 을 선택하여 전체 생성 과정에 적용합니다.
     • 입력: 프롬프트 임베딩 + 병렬 샘플링 예산 (Parallel Sampling Budget, $B$).
     • 목표: 주어진 병렬 샘플 수 내에서 가장 높은 정확도를 내는 전략을 선택.
  2. 토큰 수준 (Token-Level):
     • 문제 설정: 부분 관측 마르코프 결정 과정 (POMDP) 으로 모델링.
     • 행동: 생성의 각 토큰 단계마다 샘플링 전략 (주로 온도 파라미터) 을 동적으로 선택합니다.
     • 입력: 현재 토큰의 내부 모델 특징 (Hidden State) + 남은 토큰 예산 ($b_t$).
     • 목표: 불확실성이 높은 단계에서는 탐색 (Exploration) 을, 안정성이 필요한 단계에서는 결정론적 (Deterministic) 인 샘플링을 수행하여 전체 경로의 정확도를 극대화.

2.2. 학습 알고리즘

• 정책 경사 (Policy Gradient, REINFORCE): 검증 가능한 종료 보상 (예: 수학 문제 정답 여부) 을 기반으로 정책을 업데이트합니다.
• 행동 공간 선택:
  • 시퀀스 수준: 다양한 디코딩 설정 (Temperature, Top-k, Top-p, Min-p 조합) 의 후보 풀에서, 검증 데이터에서 '최고의 조합 (Best-of-S)' 성능을 최대화하는 소수의 전략 집합을 그리디 커버리지 (Greedy Coverage) 방식으로 선별합니다.
  • 토큰 수준: 해석 가능성과 효율성을 위해 온도 (Temperature) 파라미터를 조절하는 행동 공간에 집중합니다.
• 안정화 기법: 토큰 수준 학습 시 고변동성 (High Variance) 문제를 해결하기 위해, 보상 신호가 희소한 프롬프트를 필터링하거나, 이미 확률이 집중된 토큰 (최대 확률 > 0.95) 에 대한 학습을 마스킹합니다.

---

3. 주요 기여 (Key Contributions)

1. 추론 제어의 정책 학습 프레임워크: 프롬프트 수준과 토큰 수준 모두에서 명시적인 계산 예산 하에서 적응형 디코딩을 수행하는 통합된 강화학습 프레임워크를 제시했습니다.
2. 경량 어댑터 및 검증 가능 보상: 별도의 보상 모델 학습이나 수동 설계된 휴리스틱 없이, 오직 작업 수준의 정답 신호 (Correctness) 만으로 어댑터를 학습시켜 베이스 모델을 변경하지 않고도 성능을 향상시켰습니다.
3. 예산 인식 (Budget-Aware) 학습: 추론 시 제약 조건 (병렬 샘플 수, 토큰 예산) 을 정책 입력에 포함시킴으로써, 다양한 추론 환경에 강건한 디코딩 행동을 학습했습니다.

---

4. 실험 결과 (Results)

MATH (수학) 및 CodeContests (코딩) 벤치마크에서 실험을 수행했습니다.

• 시퀀스 수준 어댑터:
  • 고정된 최선의 전략 (Best Static Baseline) 대비 Pass@1 정확도에서 2~3% 향상 (CodeContests 의 경우 최대 33% 향상).
  • 병렬 샘플링 예산을 입력으로 포함할 때 성능이 더욱 개선되었으며, 다양한 프롬프트 형식 (CoT 포함/미포함) 에 대해 강건함을 보였습니다.
• 토큰 수준 어댑터:
  • 고정된 토큰 예산 하에서 Pass@1 정확도를 최대 10.2% 까지 향상시켰습니다.
  • 시퀀스 수준 어댑터보다 더 큰 성능 향상을 보였으며, 이는 단일 생성 경로 내에서 미세하게 확률적 요소 (Stochasticity) 를 분배하는 것이 효과적임을 시사합니다.
• 일반화 (Generalization):
  • MATH 데이터셋으로 훈련된 어댑터가 CodeContests 나 AIME 2025 와 같은 다른 도메인에서도 경쟁력 있는 성능을 보여주어, 학습된 전략이 전이 가능한 신호를 포착함을 입증했습니다.
• 분석:
  • 학습된 정책은 불확실성이 높은 토큰 (고엔트로피) 에서는 확률적 샘플링을 유지하고, 낮은 토큰에서는 결정론적 행동을 취하는 경향이 있음을 확인했습니다.
  • 단순한 엔트로피 기반 휴리스틱만으로는 이러한 성능 향상을 재현할 수 없었으며, 학습된 컨텍스트 정보가 필수적입니다.

---

5. 의의 및 결론 (Significance)

이 연구는 LLM 의 성능 향상을 위한 새로운 축을 제시합니다.

• 모델 스케일링 및 파인튜닝의 대안: 모델 자체를 재학습하거나 크기를 늘리는 대신, 추론 시 (Inference-time) 제어를 통해 효율적으로 성능을 끌어올릴 수 있음을 증명했습니다.
• 자원 효율성: 제한된 계산 예산 (Budget) 하에서 최적의 샘플링 전략을 동적으로 선택함으로써, 계산 비용 대비 정확도 (Accuracy-Budget Tradeoff) 를 극대화합니다.
• 적응형 추론의 가능성: 고정된 하이퍼파라미터가 아닌, 학습된 정책이 문제의 난이도와 토큰별 불확실성에 맞춰 유연하게 대응할 수 있음을 보여주어, 향후 LLM 의 추론 능력 향상과 효율적 배포에 중요한 통찰을 제공합니다.

요약하자면, 이 논문은 **"언어 모델을 고정된 채로, 강화학습을 통해 추론 시의 샘플링 전략을 학습함으로써, 제한된 계산 자원 하에서도 고정된 전략보다 훨씬 우수한 추론 성능을 달성할 수 있다"**는 것을 입증한 연구입니다.