Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo 논문 리뷰

Probabilistic Inference in Language Models via Twisted Sequential Monte Carlo

---

Problem Setting

인간처럼 말하기

언어모델의 궁극적인 목적은 인간처럼 말하기이다. 그렇기 때문에 대부분의 LLM은 RLHF 기반의 보상 모델링을 통해서 자연스러운 응답을 만들도록(alignment를 잘 하도록) 학습되고는 한다.

그러나 어디까지나 두 가지 목표를 두고 있는데:

• 제어 가능성: 유해한 말을 하지 않으며, 주어진 지시를 정확히 따르는 안전하고 예측 가능한 능력을 갖추는 것 (PPO)
• 다양성 및 창의성: 상황에 맞게 새롭고 흥미로우며 자연스러운 응답을 생성하는 능력 (PPO와 반대될 수 있음)

어떻게 보면 이 두 목표는 현재의 RLHF로 추구하기에는 모순스러운 일이다. 이에 대한 이유는 다음과 같다:

GPT 같은 auto-regressive 모델은 단어를 하나씩 순차적으로 출력하는 인과적 방식으로 작동하지만, 우리는 이 단어가 만들어질 때마다 학습하는 것이 아닌, 이 문장이 완성되고 나서야 좋음/나쁨을 평가할 수 있다. 생성은 인과적(casual)이지만 평가는 비인과적(non-casual)한 것이다.

따라서 이 문제를 해결하기 위해서는, 순차적으로 생성되는 도중에 응답을 평가할 수 있어야 하지만, 다음과 같은 문제들이 있다:

1. 다루기 힘든 주변화(intractable marginalization) :

표준적인 인과적 모델은 $P(s_t | s_{1:t-1})$를 계산하여 토큰을 순차적으로 샘플링한다. 그러나 $T$시점에 생성이 종료된다고 하는 경우, $t$시점의 토큰이 발생할 올바른 확률은

좌변의 경우, 전체적인 목표 $\sigma$의 관점에서 $1:t-1$까지의 시퀀스가 주어졌을 때, 현재 토큰으로 $s_t$를 선택하는 것이 얼마나 좋은가를 나타내는 점수 또는 확률이다.

 

현재 토큰 $s_t$의 점수를 알려면, $s_t$ 뒤에 이어질 수 있는 모든 미래의 문장들을 하나씩 다 만들어보고, 각각의 완성된 문장들이 얼마나 좋은지 점수를 매긴 다음, 그 점수들을 전부 다 더해야 한다. 그것이 우변에 관한 식이다.

 

이것이 다루기 어려운 이유는, 경우의 수 때문이다. 이를 조합적 폭발(Combinational Explosion)이라고 부른다. 당연하게도, 각 단위당 가능한 단어의 종류의 수를 남은 문장의 길이에 대해 제곱을 시켜주어야 하기 때문에, 지수적으로 경우의 수가 지나치게 늘어나는 것을 쉽게 생각할 수 있을 것이다.

 

현실적으로, 이러한 모든 경우의 수에 대해서 전부 다루는 것은 불가능에 가까울 것이다.

 

2. 진보적 해법의 시도: Sequential Monte Carlo (SMC)

SMC는 일단 끝까지 생성하고 나중에 평가하는, Importance Sampling의 문제를 극복하고자 한다. 핵심은 여러 개의 후보 시퀀스(particle)를 동시에 탐색하며, 각 단계마다 유망한 토큰을 골라내는 것이다.

토큰을 생성하는 매 단계마다, 각 후보 시퀀스 $s_{1:t}$가 얼마나 목표에 부합하는지를 평가(분자에 있는 $\pi$)하고, 그것이 이전 단계의 평가에 비해서 얼마나 좋은지 계산한 후 자연스러움의 정도($q$)로 나누어주면 해당 경로의 유망함/유망하지 않음에 대해서 가중치를 갱신할 수 있다.

 

갱신된 가중치는에 따라 가능성이 낮은 후보 시퀀스는 점점 버려지고, 가중치가 높은 시퀀스는 더 많이 채택되도록 한다.

그러나 SMC도 여전히 근본적인 문제를 안고 있는데, 이는 다음과 같다:

• SMC는 선택과 집중을 위한 리샘플링 단계를 거치면서, 몇 단계만 지나면 분포의 다양성이 급격히 줄어드는(skew) 현상이 발생한다. 따라서 다양성을 확보하려는 초기의 목적과 어긋나며, 추가적으로 잘못된 경로를 선택하는 경우 회복이 불가능해진다.
• 1. 문제에서 미래의 제약 조건을 내다보지 못한다는 문제점이 있다. 식의 구조와 SMC에 대해서 가볍게 살펴보면 알 수 있게 될 것이다. 미래를 내다보지 못한다면, 소위 근시안적 평가로 경로가 퇴화되거나, 잘못된 선택의 증폭이 일어나는 문제점이 발생할 수 있다.

이를 해결하기 위해서, 본 논문은 핵심 내용인 twisted SMC를 제안한다.

---

Twisted SMC

우선 일반적인 SMC가 과거의 상태 + 현재까지만 상황을 보는 것이라면, Twisted SMC는 매 평가 시에 미래에 대한 힌트를 섞어준다고 생각하면 된다.

 

Twisted SMC는 Twist Function인 $\psi$를 도입한다. 현재까지의 경로를 보고, 이 경로가 최종 목표인 $\phi$에 도달할 잠재력이 얼마나 되는지를 요약해서 알려준다. 따라서 중간 목표 $\pi$는 다음과 같이 정의된다:

해당 수식의 뜻은 새로운 중간 점수 = (기존 모델의 자연스러움) * (미래 잠재력 점수)이다. 즉 기존의 SMC처럼 단순히 지금까지 얼마나 자연스러운가?가 아닌, 미래에 성공할 가능성이 얼마나 되는지까지 고려한다고 생각하면 된다.

 

이제는 이 Twist Function $\psi$를 구해야 한다.

이 식이 애초에 말도 안되는 것을 확인할 수 있다. 왜냐하면 결국 처음에 문제제기가 되었던, 다루기 힘든 주변화(intractable marginalization) 문제에 다시 돌아왔기 때문이다. 따라서 이상적인  $\psi^{*}$은 구할 수 없고, 다른 방법을 사용해야 한다.

 

Derive twisted function $\psi$

본 논문의 핵심 내용들이 점점 나온다. 직접 구할 수 없는 $\psi^{*}$는 우선 포기하고, 다만 이러한 $\psi$를 적절히 학습함으로써 문제를 해결하면 되는 것이다. 3장에서는 Twisted SMC 프레임워크의 이론을 확장하는 단계라, 어렵기도 하여 넘어가도록 하겠다.

---

Learning the Twist Functions

본 논문은 대조 학습(Contrastive Learning)에 기반한 새로운 접근법인 CTL(Contrastive Twist Learning)을 제안한다.

 

우선, 어떻게 학습할 것인가를 알기 전에 무엇을 목표로 학습할 것인가를 먼저 정의해야 한다.

용어 정리부터 다시 하고 가자:

• $\sigma$ : 도달하고 싶은 진짜 목표 분포
• $\pi$ : 모델($p$와 학습 가능한 $\psi$의 조합)이 만들어내는 분포

즉 학습의 목표는 $\pi$가 만들어내는 분포와 진짜 목표 분포를 최대한 똑같이 만드는 것이다.

이 gradient의 형태는 진짜 목표 $\sigma$에서 뽑은 샘플들의 확률을 높이고, $\pi$는 나타날 확률을 낮추도록 하는 것이다.

 

다만 현실적인 문제로, 긍정 샘플을 어떻게 구할 것이냐는 것이다. 여기서 논문은 다음과 같은 해결책을 제시한다:

• 정확한 샘플링: BDMC 트릭이나 Rejection Sampling같은 방법으로 운 좋게 완벽한 샘플을 구할 수 있다면 그것을 사용한다
• 근사적 샘플링: 완벽한 샘플을 구하기 어렵다면, SMC나 IS를 이용하여 꽤 괜찮은 긍정 샘플을 근사적으로 생성하여 학습에 사용한다.

---

Evaluating Inference in Language Models

결국 이렇게 잘 뽑은 샘플링 된 $q$가 우수함을 알 수 있는 방법으로는, KL divergence가 있다. 이때 $\log Z$는 $\sigma$의 모든 가능한 상태에 대한 점수이다.

문제는 KL divergence를 계산하기 위해서는 로그 정규화 상수 $\log Z$의 값을 알아야 하는데, 애초에 그것을 전부 다 계산하기는 현실적으로 불가능하기 때문에, 적절한 범위 내에서 추론할 필요가있다.

 

이러한 boundary를 제시할 수 있는 방법으로, 논문은 bidirectional SMC bound라는 방법을 제안한다:

하한의 경우 $q_{SMC}$를 통해서 도출한 것을 볼 수 있다. 일반적으로 어떠한 미래에 대한 최적의 힌트 없이 경로를 찾아가는 것은 최적보다는 값이 낮은 것이 당연하다. 상한의 경우 정답 샘플을 하나 알고 있는 상태에서 SMC를 여러 번 실행했을 때 나오는 점수들의 평균인데, 이는 정답을 반드시 하나 알고 있기 때문에 반드시 $\log Z_{\sigma}$보다는 크다.

 

이 $\log Z_{\sigma}$의 범위 덕분에 $q$와 $\sigma$의 분포 차이가 어느 정도인지 가늠할 수 있고, 따라서 샘플링 된 $q$의 우수성 여부를 판단할 수 있는 것이다.

---

Experiments

Twisted SMC가 적은 샘플만으로도 경계를 빠르게 좁히는 것을 확인할 수 있다. 따라서 이 논문의 proposal은 성공!

• $D_{KL}(q||\sigma)$: 제안된 모델 $q$가 목표에 없는 샘플을 얼마나 만들지 않는가를 측정한다. 점수가 낮을수록 안전하다.
• $D_{KL}(\sigma ||q)$: 제안된 모델이 얼마나 좋은 실제 목표를 빼놓지 않고 모두 생성하는가를 측정한다. 점수가 낮을수록 좋은 결과물이 다양하고 창의적이게 생성된다는 뜻이다.

Constrastive Twist Learning이 세 가지 테스크에서 좋은 결과물을 내는 것을 확인할 수 있다.

---

Conclusion

이 논문은 순차적 생성 모델의 평가와 모델의 output 제어 기능 두 가지를 해결할 수 있다는 것에서 의의가 있다. 단순히 output을 사람의 선호에 바탕으로 제한적으로 두는 것이 아닌, 다양성을 보존하면서 안전한 output을 낼 수 있다는 것을 보여준다.

다만 매 단계마다 Twisted function을 실행해야 하므로, 생성 속도 자체는 더 느려질 수 있겠다는 결론이 들어 이에 대한 해결책이 필요해보인다.

 

무엇보다 느낀 점... 너무 모르는게 많다 내가... 부족함을 깨닫고 더 열심히 공부해야겠다.