기술 요약: ETS: 훈련 없는 RL 정렬을 위한 에너지 유도 테스트 시간 확장

문제 제기

강화 학습 (RL) 사후 훈련은 대형 언어 모델 (LLM) 을 인간의 의도와 정렬하고, 특히 추론 및 지시 따르기 능력을 향상시키는 데 핵심적인 패러다임이 되었습니다. 그러나 이 접근법은 상당한 실용적 한계를 겪습니다: 비용이 많이 드는 보상 모델링이 필요하고, 대규모 인간 선호 데이터에 의존하며, 불안정한 학습 역학을 보이며, 하이퍼파라미터에 민감합니다. 더 나아가, 보상 설계의 변경은 재훈련을 필요로 하여 막대한 계산 및 인력 비용을 초래합니다.

이론적으로 최적의 RL 분포는 폐쇄형 해를 허용하지만, 기존 방법들은 이 해를 직접 활용하기보다는 반복적인 경사 기반 최적화를 통해 이를 근사하는 데 의존합니다. 이는 근본적인 질문을 제기합니다: 추가적인 훈련 없이 추론 시에 최적의 RL 분포에서 직접 샘플링할 수 있을까요?

방법론

저자들은 최적의 RL 정책에서 직접 샘플링하도록 설계된 훈련 없는 추론 방법인 **에너지 유도 테스트 시간 확장 (ETS)**을 제안합니다. 이 방법은 자기회귀 모델 (ARM) 과 확산 언어 모델 (DLM) 을 모두 포괄하는 통합된 마스크 언어 모델링 (MLM) 프레임워크를 기반으로 구축됩니다.

1. 이론적 기반: 에너지 재가중 역전환

핵심 이론적 통찰은 MLM 프레임워크에 대한 최적의 역전환 커널이 두 가지 요소로 분해된다는 것입니다:

1. 참조 전환 (Reference Transition): 참조 모델 ($p_{ref}$) 이 제공하는 전환 커널.
2. 에너지 항 (Energy Term): 지수화된 보상의 조건부 기댓값으로 작용하는 에너지 항.

구체적으로, 목표 역전환 $p(s | t, y)$는 다음에 비례합니다:
$$p_{ref}(s | t, y) \cdot \mathbb{E}_{p_{ref}(x_0|y,s)} \left[ \exp\left(\frac{r(y, x_0)}{\lambda}\right) \right]$$
여기서 $r(y, x_0)$는 보상 함수이고 $\lambda$는 정규화 상수입니다. 이 공식은 몬테카를로 방법을 통해 에너지 항을 추정함으로써 훈련 없는 샘플링을 가능하게 합니다.

2. 알고리즘: 에너지 유도 테스트 시간 확장 (ETS)

이 알고리즘은 반복적 유도 과정으로 작동합니다:

• 샘플링: 각 유도 단계에서 알고리즘은 참조 모델에서 $M$개의 후보 연속체를 샘플링합니다.
• 에너지 추정: 각 후보에 대해 에너지 항을 몬테카를로 샘플링을 사용하여 추정합니다. 구체적으로, 각 후보에 대해 $K$개의 독립적인 완결본을 생성하여 기대 보상의 자기 정규화 중요도 샘플링 추정자를 계산합니다.
• 선택: 후보들은 추정된 에너지에 따라 재가중되며, 그중 하나가 선택되어 다음 단계로 진행됩니다.
• 보상 설계: 추론 시 정답 보상이 없도록 작동하기 위해, ETS는 자기 일관성 보상을 사용합니다. 이는 $K$개의 완결본에서 최종 답변을 추출하고, 다수결을 합의 (consensus) 로 식별한 후, 합의와 일치하는 완결본에는 1 의 보상을, 그렇지 않은 경우에는 0 의 보상을 부여합니다.

3. 중요도 샘플링을 통한 가속 (ETS-IS)

에너지 추정을 위한 반복적인 모델 쿼리로 인한 지연 병목 현상을 해결하기 위해, 저자들은 ETS-IS를 도입합니다.

• 제안 모델 (Proposal Model): 에너지 추정을 위해 대규모 참조 모델 ($p_{ref}$) 에서 샘플링하는 대신, 계산 비용이 더 저렴한 경량의 제안 모델 ($p_{small}$) 을 사용합니다.
• 보정: 제안 모델로 인해 발생하는 편향을 수정하기 위해 중요도 샘플링 가중치를 적용하여, 추정자가 편향되지 않고 목표 분포와 이론적으로 일관되도록 보장합니다.
• 구현:
  • ARM 의 경우: 더 작고 규모가 정렬된 모델 (예: Qwen3-8B 를 위한 Qwen3-1.7B) 을 제안 모델로 사용합니다.
  • DLM 의 경우: 블록 단위 KV 캐싱과 병렬 디코딩을 활용하여 추론을 가속화하는 Fast-dLLM 을 사용합니다.

주요 기여

1. 훈련 없는 정렬: 이 논문은 반복적인 경사 업데이트나 미세 조정이 없이 최적의 RL 정책에서 직접 샘플링하는 방법을 제시하여, 기존 RL 사후 훈련의 불안정성과 비용 문제를 해결합니다.
2. 통합 프레임워크: 이 접근법은 통합된 MLM 공식 하에서 ARM 과 DLM 모두에 적용 가능하여, 훈련 없는 확장을 위한 일반적인 해결책을 제공합니다.
3. 이론적 보장: 저자들은 총변동 거리 (Total Variation Distance) 측면에서 샘플링 분포가 목표 최적 분포로 수렴하는 속도에 대한 엄격한 증명을 제공합니다. 오차 한계는 유도 단계 수에 비례하고 샘플 크기의 제곱근에 반비례하여 확장됨이 입증되었습니다.
4. 효율적인 가속: 중요도 샘플링을 경량 제안 모델 및 현대적 가속 프레임워크 (Fast-dLLM 등) 와 통합함으로써, 이 방법은 샘플링 품질을 유지하면서 추론 지연을 크게 줄입니다.

실험 결과

저자들은 ARM(Qwen3-1.7B, Qwen3-8B) 과 DLM(LLaDA-8B-Instruct) 을 사용하여 추론 (MATH500, GSM8K), 코딩 (HumanEval), 과학 (GPQA) 벤치마크에서 ETS 를 평가했습니다.

• 성능: ETS 는 표준 추론 및 기타 테스트 시간 확장 기준 (예: Beam Search, Best-of-N, Power Sampling) 보다 일관되게 생성 품질을 향상시킵니다.
• RL 과의 비교: 주목할 점은 ETS 가 매개변수 업데이트 없이 사후 훈련된 RL 정책 (GRPO 등) 의 성능과 일치하거나 능가한다는 것입니다.
• 효율성: 가속화 방법 (ETS-IS) 은 지연과 정확도 간의 유리한 균형을 달성합니다. 예를 들어, Qwen3-8B 에서 ETS-IS 는 표준 ETS 와 비교할 만한 정확도를 달성하면서도 지연을 크게 줄입니다.
• 확장성: 실험에 따르면, 완결본 수 ($K$) 를 늘리는 것보다 후보 수 ($M$) 를 늘리는 것이 정확도 향상에 더 효과적이며, 중간 유도 단계 ($I$) 는 정확도에 비례하여 개선을 제공합니다.

중요성과 주장

이 논문은 ETS 가 보상 인식 생성을 위한 실용적이고 훈련 없는 경로를 확립한다고 주장합니다. 비용이 많이 드는 사후 훈련 대신 추론 시간 확장을 통해 고품질 정렬을 달성할 수 있음을 입증함으로써, 이 연구는 자원이 제한된 환경에 대한 비용 효율적이고 안정적인 대안을 제공합니다. 저자들은 이 방법이 정렬 기법을 민주화하고 광범위한 재훈련 주기 필요성을 제거함으로써 LLM 배포의 환경 발자국을 줄인다고 주장합니다. 이 프레임워크는 생성 모델에서의 효율적인 추론 및 테스트 시간 확장에 대한 향후 연구의 기초로 제시됩니다.