OLLM: Options-based Large Language Models

OLLM: Options-based Large Language Models 기술 요약

1. 문제 정의 (Problem)

기존의 대규모 언어 모델 (LLM) 은 다음 토큰을 단일 확률 분포로 예측하며, 다양성을 확보하기 위해 온도 (temperature) 설정이나 샘플링 휴리스틱 (top-k/p) 에 의존합니다. 그러나 자연어 생성, 특히 수학 추론과 같은 복잡한 작업에서는 많은 위치에서 여러 개의 그럴듯한 다음 토큰 (다중 모드, multimodality) 이 존재합니다.

• 한계점: 단일 소프트맥스 (softmax) 분포 내에서 경쟁하는 토큰들은 확률 질량을 나누게 되어, 모델이 지배적인 모드에 과적합되거나 중요한 대안적 추론 경로를 놓칠 수 있습니다.
• 관찰: 논문은 수학 데이터셋 분석을 통해 약 58% 의 토큰 위치가 높은 엔트로피 (모호함) 를 가지며, 이는 명시적인 옵션 (options) 표현이 필요함을 시사합니다.

2. 방법론 (Methodology)

저자들은 **Options LLM (OLLM)**을 제안합니다. 이는 표준 LLM 의 단일 다음 토큰 예측을 대신하여, 이산적 잠재 변수 (discrete latent variable) 로 인덱싱된 학습된 '옵션' 집합을 통해 다음 토큰을 예측하는 방법입니다.

• 아키텍처 (Architecture):

  • 플러그인 방식: 기존 사전 학습된 LLM 백본 (Backbone) 을 변경하지 않고, 출력 헤드 (lm_head) 직전에 **인코더 (Encoder)**와 디코더 (Decoder) 두 개의 경량 모듈만 추가합니다.
  • 인코더 ($E_\psi$): 컨텍스트 ($h_t$) 와 목표 토큰 임베딩을 입력받아, 해당 토큰에 적합한 잠재 옵션 ($z \in \{1, ..., K\}$) 의 분포를 예측합니다.
  • 디코더 ($D_\omega$): 선택된 잠재 변수 $z$와 컨텍스트 $h_t$를 입력받아 편향 벡터 ($\Delta h_t$) 를 생성하여, LLM 은닉 상태에 더합니다. 이를 통해 토큰 예측 로짓 (logits) 이 해당 옵션에 맞춰 조정됩니다.
  • 추론 (Inference): 학습 중에는 인코더가 정답 토큰을 보고 옵션을 예측하지만, 추론 시에는 **경량 정책 네트워크 (Policy $\pi_\phi$)**가 컨텍스트만 보고 적절한 잠재 변수 $z$를 선택하여 생성을 제어합니다.

• 학습 목표 (Training Objectives):

  1. 토큰 예측 손실 (Cross-Entropy): 주어진 잠재 옵션 $z$ 하에서 정답 토큰을 예측하도록 학습.
  2. 잠재 공간 정규화 (KL Divergence): 인코더가 특정 옵션으로 붕괴 (collapse) 되는 것을 방지하기 위해 균일 분포 (uniform prior) 에 대한 KL 발산을 적용 (자동 적응 스케일링 사용).
  3. 정책 학습 (Behavioral Cloning): 인코더가 학습한 옵션 매핑을 모방하도록 정책 네트워크를 학습 (SFT 기반).

3. 주요 기여 (Key Contributions)

1. 명시적 옵션 모델링: 토큰 생성의 불확실성을 단일 분포가 아닌, 잠재 공간의 이산적 옵션 집합으로 명시적으로 모델링하여 다양한 추론 경로를 분리합니다.
2. 경량 및 범용성: 기존 LLM 의 백본을 수정하지 않고 소수의 파라미터 (약 1.56%) 만을 학습 가능한 플러그인으로 구현하여, 거의 모든 사전 학습 모델에 적용 가능합니다.
3. 잠재 공간 기반 정책 학습: 전체 어휘 (Vocabulary) 가 아닌 저차원의 잠재 옵션 공간에서 정책을 학습함으로써, 샘플 효율성을 극대화하고 언어 전환 (language switching) 이나 퇴화된 추론 (degenerate reasoning) 과 같은 정렬 (alignment) 문제를 자연스럽게 해결합니다.
4. 추가 손실 함수 불필요: 구조적 설계 (옵션 분리) 를 통해 KL 페널티나 수동으로 설계된 정렬 손실 없이도 모델의 안정성과 다양성을 확보합니다.

4. 실험 결과 (Results)

• 실험 설정: 17 억 파라미터 (1.7B) 의 Qwen2.5 Deepseek R1 기반 모델을 OpenMathReasoning 데이터셋으로 학습하고 OmniMath 데이터셋으로 평가했습니다.
• 성능 비교:
  • 기존 LoRA 기반 모델: 최종 정답 정확도 (Final Answer Correctness) 가 약 **51%**에서 정체되었습니다.
  • OLLM (최적 옵션 선택 시): 잠재 변수를 최적으로 선택했을 때 최종 정답 정확도가 약 **70%**까지 상승했습니다.
  • 토큰 단위 정확도: OLLM 은 약 83% 의 토큰 정확도를 달성하여, 단일 분포 모델의 과적합 문제를 해결하고 다양한 가능성을 유지함을 보였습니다.
• 분석: 단순한 정책 (Policy) 만으로도 저차원 잠재 공간에서 작동하여 최종 정확도를 크게 향상시킬 수 있음을 입증했습니다.

5. 의의 및 결론 (Significance)

• 통제 가능성과 견고성: OLLM 은 수학 추론과 같은 복잡한 작업에서 모델의 제어 가능성 (controllability) 과 견고성 (robustness) 을 크게 향상시킵니다.
• 강화학습 (RL) 을 위한 새로운 방향: 전체 어휘 공간이 아닌 저차원 잠재 옵션 공간에서 강화학습 (RL) 을 수행할 수 있는 기반을 마련했습니다. 이는 기존 RLHF 의 샘플 비효율성과 정렬 실패 (reward hacking) 문제를 해결할 수 있는 유망한 접근법입니다.
• 해석 가능성: 서로 다른 잠재 변수가 서로 다른 논리적 추론 경로에 대응할 수 있어, 모델의 내부 작동 원리를 해석하고 다양한 해결책을 탐색하는 데 유용합니다.

이 논문은 LLM 의 생성 메커니즘을 단순한 확률 예측에서 '옵션 선택'의 관점으로 전환함으로써, 효율성과 성능을 동시에 개선할 수 있는 새로운 패러다임을 제시합니다.