거부 샘플링 — 싼 분포로 비싼 분포를 정확히 흉내내기

실행: python 딥러닝/mathnotes/M11_rejection.py 스펙큘러티브 디코딩 편에서 쓴 수락 판정이 왜 근사가 아니라 정확한가를 정리한다.

1. 문제

목표 분포 $q$ 에서 뽑고 싶은데 $q$ 를 구하는 일이 비싸다. 대신 싼 제안 분포 $p$ 가 있다. $p$ 에서 뽑아 쓰면 안 될까? 당연히 안 된다 — 분포가 $p$ 가 되어 버린다.

그런데 뽑은 것을 조건부로 버리면 $q$ 를 정확히 복원할 수 있다.

토큰마다 $p$ 와 $q$ 중 작은 쪽까지는 제안을 그대로 써도 아무 문제가 없다. 문제는 $q$ 가 $p$ 보다 큰 자리다. 그 모자란 몫을 모은 것이 잔차이고, 거부가 일어났을 때 바로 거기서 다시 뽑는다.

x \sim p, \qquad \text{수락 확률} = \min\!\Bigl(1, \frac{q(x)}{p(x)}\Bigr), \qquad \text{거부하면 } x \sim \frac{[\,q - p\,]_{+}}{\sum_j [\,q_j - p_j\,]_{+}}

기호	뜻
$p$	제안 분포 — 싼 초안 모델이 주는 확률
$q$	목표 분포 — 우리가 실제로 원하는(비싼 모델의) 확률
$[\,z\,]_{+}$	$\max(z, 0)$ . 음수는 0으로 자른다
$\mathrm{TV}(p,q)$	전변동거리 $\tfrac12\sum_j \lvert p_j - q_j\rvert$ — 두 분포의 어긋남

왜 정확한가는 한 줄로 확인된다. 토큰 $j$ 가 최종적으로 나올 확률은 수락으로 나올 몫 $p_j\min(1, q_j/p_j) = \min(p_j, q_j)$ 에 거부 후 잔차에서 나올 몫을 더한 것인데, 거부가 일어날 확률이 $\sum_k [\,q_k-p_k\,]_{+}$ 와 같아 분모가 지워지고 $[\,q_j-p_j\,]_{+}$ 만 남는다. 둘을 더하면 $\min(p_j,q_j) + [\,q_j-p_j\,]_{+} = q_j$ 다.

2. 직접 재 보기

어휘 8개짜리 장난감 분포로 40만 번 뽑았다.

$거부 샘플링$

	TV(표본, q)
보정 없이 $p$ 에서 그냥 뽑음	0.1995
수락-거부 + 잔차 재샘플링	0.00094

0.00094는 40만 표본의 통계적 오차 수준이다. 근사가 아니라 정확하다. 초안 모델이 아무리 엉성해도 최종 분포는 큰 모델의 것과 같고, 엉성함은 오직 속도로만 나타난다.

그 속도가 무엇으로 정해지는지가 세 번째 실험이다. $p$ 를 $q$ 쪽으로 조금씩 당기며 수락률을 쟀다.

TV(p, q)	이론 1 − TV	실측 수락률
0.2000	80.00%	79.91%
0.1333	86.67%	86.63%
0.0667	93.33%	93.38%
0.0000	100.00%	100.00%

전부 0.1%p 안에서 맞는다. 기대 수락률은 정확히 $1 - \mathrm{TV}(p,q)$ 다.

3. 왜 중요한가

스펙큘러티브 디코딩 편의 수락률이 어디서 오는지가 이것이다. 초안 모델을 키우는 일은 곧 $\mathrm{TV}(p,q)$ 를 줄이는 일이고, 줄어든 만큼 그대로 수락률이 된다. “얼마나 빨라지나”를 추측하지 않고 두 분포의 어긋남만 재면 알 수 있다는 뜻이다.
품질 논쟁이 필요 없다. 스펙큘러티브 디코딩이 출력 품질을 해치느냐는 질문은 애초에 성립하지 않는다. 분포가 수학적으로 같다. 고를 것은 속도뿐이다.
같은 구조가 잔차를 남기는 모든 근사에 쓰인다 — 값싼 후보를 먼저 내고 비싼 검증자가 받아들이거나 물리는 방식이면 무엇이든.

4. 한 줄 요약

싼 분포 $p$ 에서 뽑고 $\min(1, q/p)$ 로 수락하되 거부하면 잔차 $[\,q-p\,]_{+}$ 에서 다시 뽑으면, 최종 분포는 정확히 $q$ 가 된다(실측 TV 0.00094). 그리고 이때 기대 수락률은 정확히 $1 - \mathrm{TV}(p,q)$ 다 — 초안이 목표와 겹치는 만큼이 그대로 속도가 된다.

연결

스펙큘러티브 디코딩 · 디코딩과 샘플링 · 엔트로피와 KL 발산