[토픽 모델링] 깁스 샘플링의 병렬화 & GPU 위에서 돌리기

깁스 샘플링(Gibbs Sampling, 기브스 표집)을 이용한 LDA 모델 추정 과정은 예전 포스팅에서 소개한 바가 있습니다. 다시 한번 간략하게 짚어보자면, 깁스 샘플링은 복잡한 확률분포

에서 표본을 추출하기가 어려우므로 그 대신

와 같이 나머지 변수를 모두 고정한 상태에서 각각의 변수의 표본을 추출하고 이를 합쳐서 전체 확률 분포의 표본을 근사하자는 아이디어입니다. 반복 횟수(iteration)이 늘어날 수록 위 근사는 원래 확률 분포를 따라간다고 이미 알려져 있기에 마음 놓고 이렇게 사용하는 것이죠.

LDA에서 단어의 주제 분포를 계산하는 경우, 모든 문헌의 단어 개수가 N개이고, 주제 개수가 K개라고 하면, 각각의 단어가 주제에 배정될 수 있는 모든 조합의 개수는 K^N개가 됩니다. 어마어마하게 큰 수이기 때문에 이 분포를 직접 계산하는 것은 불가하기에 앞서 설명한 깁스 샘플링을 이용해서 단어 한 개씩 주제를 추출하고, 이를 반복하여 전체 문헌의 모든 단어의 주제를 추출하는 것을 완성하는 것이죠. LDA의 깁스 샘플링 과정은 다음과 같습니다. (자세한 것 전의 포스팅을 참조!)

1. 각각의 단어에 임의의 주제를 배정한다.
2. (잘못됐겠지만) 모든 단어에 주제가 배정되어있고, 그 결과 문헌별 주제 분포와 주제별 단어분포가 결정되었음. 이를 깁스샘플링을 통해서 개선해 나가자.
3. 깁스 샘플링 시작. n = 1 ~ N까지 다음을 반복한다.
   1. n번째 단어를 골라 빼낸다. (나머지 단어들이 모두 적절하게 배정되어 있다고 가정하고)
   2. 빼낸 단어를 제외하고, 그 단어가 속한 문헌의 주제분포(P(Topic | Document))를 계산하고, 그 단어가 속한 주제의 단어분포(P(Word | Topic))를 계산한다. P(Word | Topic) * P(Topic | Document)에서 Topic을 하나 뽑아 그 단어에 배정한다.
4. 수렴할때까지 3번을 충분히 반복한다.

보시다시피 깁스 샘플링을 통한 추정은 매우 간단하지만, 반복 노동이 많이 필요한, 지루한 작업입니다. 그리고 각 샘플링 단계는 앞선 샘플링 단계에 영향을 받기 때문에 쉽사리 병렬화하기도 뭣하죠. 따라서 일단 코어 하나로 차례대로 돌리는게 가장 쉬운 구현입니다. 하지만 세상은 바야흐로 대 멀티코어의 시대. 코어 하나만 돌리는 알고리즘은 살아남기가 어렵습니다. 당연히도 이미 앞선 연구자들이 이를 병렬화하는 방법을 다 고안해 냈으니, 그것이 이번 포스팅의 주제가 되겠습니다.

멀티코어 CPU에서 빠르게 돌려보자 - AD-LDA^[각주:1]

Newman의 연구진은 깁스 샘플링을 근사하는 LDA의 병렬 버전을 제안했습니다. Approximate Distributed LDA라는 뜻에서 AD-LDA라고 이름을 붙였답니다. (앞으로 읽어도 뒤로 읽어도 똑같은 ADLDA를 노린거 같아요!) 사실 앞에서 밝혔듯 깁스 샘플링은 이전 상태에 영향을 받기 때문에 차례대로 수행하는게 원칙이지만, 병렬화하는 대신에 이걸 약간 위반 때문에 Approximate Distributed가 된것이죠. 이 알고리즘을 간단하게 살펴보겠습니다.

이하 N은 모든 단어의 수, C는 분할 개수(즉 CPU 코어의 개수!)입니다.

먼저 모든 단어를 C개의 조각으로 나눌 겁니다. 각각의 코어에 한 조각씩을 할당해서 깁스 샘플링을 하도록 시킬거에요. P(Topic | Document)와 P(Word | Topic)을 계산하기 위해서는 문헌마다 주제별 단어 수 정보(N_dk: d문헌에 배정된 k주제의 개수)와 주제마다 어휘별 단어 수 정보(N_kv: k문헌에 배정된 v 어휘의 개수 )가 필요합니다. 단어를 하나 빼내는 과정에서 N_dk와 N_kv는 1씩 감소하고, 새로운 k를 골라서 다시 배정하는 과정에서 N_dk와 N_kv는 1씩 증가합니다. 

각각에 코어에서 동시에 수행할 경우 N_dk와 N_kv를 여러 스레드가 동시에 접근하게 될 가능성이 있습니다. 이렇게 되면 멀티 스레딩은 망하는 것이기 때문에, 이를 방지하기 위해서는 lock을 걸거나 atomic 연산을 수행해야 합니다. 근데 이들은 매우 비싼 연산이기에 이걸 사용하면 오히려 단일 코어로 수행할때보다 성능이 더 하락할 가능성이 있습니다. 

따라서 AD-LDA에서는 그 대신 각각의 쓰레드가 N_dk, N_kv 사본을 가지고 있는 것으로 이를 회피했습니다. 각자의 코어에서는 자신만의 N_dk, N_kv를 조작하므로 race condition에는 빠지지 않겠지만, 다른 쓰레드에서 샘플링한 결과가 반영되지 않기 때문에 부정확한 샘플이 되겠죠. 그래서 각자 코어에서 1회 깁스 샘플링을 마무리하면, 각자의 N_dk, N_kv의 변동을 합쳐서 전체 N_dk, N_kv를 업데이트하고, 이를 다시 각각의 코어가 나눠가져서, 전체 동기화를 실시합니다.

즉 정리하면 다음과 같이 되겠어요.

1. 초기화: 모든 단어에 임의의 주제를 배정한다.

2. 깁스 샘플링 시작: 전역 N_dk, N_kv를 C개의 코어별 N_dk, N_kv로 복사한다.

   1. 각각의 코어에 대해서 다음을 수행한다(c in 1 ~ C)
   2. c코어에 할당된 단어에 대해 다음을 반복한다
      1. n번째 단어를 골라 빼낸다.
      2. 빼낸 단어를 제외하고, 그 단어가 속한 문헌의 주제분포(P(k | d))를 계산하고, 그 단어가 속한 주제의 단어분포(P(v | k))를 계산한다. P(v | k) * P(k | d)에서 Topic을 하나 뽑아 그 단어에 배정한다.
3. 모든 코어가 위 작업을 마치기를 기다린다.
4. 동기화: 코어별 사본 N_dk, N_kv가 어떻게 바뀌었는지를 바탕으로 전역 N_dk, N_kv를 업데이트한다.
5. 수렴할때까지 2번으로 돌아가 충분히 반복한다.

(+말보다 코드가 편하신 분들을 위해 https://github.com/bab2min/tm-cpp/blob/master/src/PLDAModel.cpp#L61)

이렇게 수행할 경우 초기화와 동기화 과정만 제외하고 나머지 부분은 전부 병렬화 가능합니다. 그리고 한 가지 묘안으로 전체 단어를 C개의 조각으로 분할할때 문헌을 단위로 분할하게 되면, 한 문헌에 접근하는 쓰레드는 항상 1개가 되므로, 전역 N_dk를 바로 조작해도 충돌이 발생하지 않습니다. 이 경우 문헌 크기가 제각각이라 각 쓰레드별로 할당되는 작업의 크기가 조금 상이하여 먼저 끝나고 놀고 있는 쓰레드가 발생할 수 있지만, 동기화 과정이 단순해지므로 오히려 성능 향상에 도움이 됩니다.

이렇게 병렬화해도 원래 깁스 샘플링 과정과 유사할까요? 저자들에 따르면 Perplexity 측면에 거의 동일한 결과를 내었고, 제가 직접 돌려본 결과도 큰 차이가 없었습니다. 어차피 깁스 샘플링도 원 확률 분포를 근사하는 것이고, 이것은 그 근사의 근사이기 때문에, 도찐개찐이라고 볼수 있겠습니다. 오히려 더 나은 결과가 나올수도 있는데, 멀티 쓰레딩 과정에서 각각의 쓰레드가 깁스 샘플링을 수행하는 순서는 누구도 예측할 수 없는 그야말로 랜덤이기 때문에, 개인적인 견해로 랜덤하게 샘플링을 해야하는 MCMC 과정에서는 오히려 도움이 될수도 있는듯 합니다.

깁스샘플링을 활용한 LDA 알고리즘은 이렇게 병렬화를 성공적으로 수행하여 대 멀티코어 시대에 잘 적응하나 싶었는데, 그의 앞에 다시 나타난 것은 대 GPGPU 시대였습니다.

GPGPU에서 더 빠르게 돌려보자 - PCGS^[각주:2]

범용 GPU의 성능이 날이 갈수록 발전하면서 더 이상 GPU는 게임만을 위한 도구가 아니게 되었습니다. 단순하고 반복적인 연산을 어마어마한 물량으로 처리하는 괴물 같은 컴퓨팅 도구가 되었죠. 깁스 샘플링 역시 단순한 작업의 반복이라 GPU에서 돌리기에 안성맞춤이죠!

현대 GPU의 구조는 단순한 코어 여러개가 그룹을 이루어있는 모양으로 생각할 수 있습니다. 이 덕분에 몇 십 몇 백 쓰레드를 동시에 돌리는 위엄을 보여주는데, 문제는 메모리죠. CPU로부터 메모리를 입력받고 결과를 다시 돌려줘야하는데 그 대역폭도 크지 않고 내부의 GPU 메모리도 넉넉하지 않은 편입니다. 

이 때문에 AD-LDA 알고리즘을 GPU 위에서 그대로 돌리기에는 무리가 있습니다. 쓰레드마다 N_kv 사본을 가지고 있어야하는데 이걸 몇 백 쓰레드가 다 가지고 있기에는 메모리가 부족하거든요...!

Yan의 연구진들은 곰곰히 생각해본 끝에 묘안을 발견했습니다. 데이터를 잘 쪼개면 굳이 매 쓰레드가 N_kv 사본을 갖고 있을 필요가 없거든요. 우리의 데이터에 전체 V 종류의 어휘가 있다고 하면 이걸 C개의 조각으로 자릅니다. 그러면 1~C번째 조각이 생기겠죠. 아까 앞서 전체 단어도 C개의 조각으로 잘랐습니다. 헷갈리지 않기 위해 단어 조각은 N_1, N_2, N_3, ... , N_C 라고 표기하고, 어휘 조각은 V_1, V_2, ... , V_C라고 할게요. AD-LDA와 마찬가지로 1번 코어는 N_1에 있는 단어를 대상으로 샘플링을 진행합니다. 2번 코어는 N_2, 3번 코어는 N_3 ...  을 대상으로하구요.

하지만 이렇게 단순히 단어들만 쪼개서는 겹치는 어휘 V가 생기기 때문에 깁스 샘플링 과정도 C개로 나눕니다.

즉 첫번째 과정에서는 1번 코어는 N_1에 있는 단어 중 V_1에 속하는 것들에 대해서만 샘플링을 수행하고, 2번 코어는 N_2에 있는 단어 중 V_2, 3번은 N_3 중 V_3, ... C번 코어는 N_C에 있는 단어 중 V_C에 속하는 것들에 대해서만 샘플링을 하는 식이죠.

그 다음 두 번째 단계에서는 1번 코어는 N_1에 있는 단어 중 V_2에 속하는 것들에 대해, 2번 코어는 N_2에 있는 단어 중 V_3에 속하는 것들에 대해 ... C번 코어는 N_C에 있는 단어 중 V_1에 속하는 것들에 대해 샘플링합니다.

세번째 단계에서는 1번 코어는 V_3, 2번 코어는 V_4, ... C번 코어는 V_2에 속하는 것들에 대해 샘플링하구요. 이를 C번째 단계까지 반복하면 모든 V_n에 대해 샘플링이 이뤄집니다. 그리고 각 단계에서 동시에 접근하는 V_n도 없죠. 그렇기에 이 방법을 쓰면 매 쓰레드가 N_kv 사본을 굳이 가지고 있을 필요가 없습니다.

코어	1	2	3	...	C
1단계	V_1	V_2	V_3	...	V_C
2단계	V_2	V_3	V_4	...	V_1
3단계	V_3	V_4	V_5	...	V_2
...	...	...	...	...	...
C단계	V_C	V_1	V_2	...	V_(C-1)

정리하면 다음과 같이 되겠습니다.

1. 초기화 : 모든 단어에 임의의 주제를 배정한다.

2. 깁스 샘플링 시작. 전역 N_dk를 C개의 코어별 N_dk로 복사한다.

   1. 총 C 단계의 샘플링을 진행한다(g in 0 ~ C-1)
      1. 전역 N_kv 중 V_(c + g mod C)에 속하는 v들에만 대하여, 코어별 N_kv`로 복사한다
      2. 각각의 코어에 대해서 다음을 수행한다(c in 0 ~ C-1)

      3. c코어에 할당된 단어 중, V_(c + g mod C)에 속한 것들에 대해서 다음을 반복한다

         1. n번째 단어를 골라 빼낸다.
         2. 빼낸 단어를 제외하고, 그 단어가 속한 문헌의 주제분포(P(k | d))를 계산하고, 그 단어가 속한 주제의 단어분포(P(v | k))를 계산한다. P(v | k) * P(k | d)에서 Topic을 하나 뽑아 그 단어에 배정한다.
      4. 모든 코어가 위 작업을 마치기를 기다린다.
      5. 동기화: 코어별 사본 N_dk가 어떻게 바뀌었는지를 바탕으로 전역 N_dk를 업데이트하고, 각각의 코어별 N_kv`를 합쳐서 전역 N_kv를 다시 구성한다.
3. 수렴할때까지 2번으로 돌아가 충분히 반복한다.

앞에서와 마찬가지로 문헌별로 쪼개게 되면 N_dk 사본을 코어별로 가지고 있을 필요가 없습니다. (N_kv도 굳이 코어별로 나누고 합치고를 반복할 필요없이 전역 N_kv를 바로 사용해도 되겠지만, GPU의 특성상 global 변수에 접근하는 것보다는 local 변수가 접근하는 것이 빠르다고 하여 굳이 나눠서 local로 나눠갖도록 해놓았습니다.)

발상의 전환으로 race condition 없이 여러 코어가 동시에 샘플링을 진행할 수 있도록 한것이 굉장히 흥미로웠습니다. 그래서 한 번 구현해보고자 설 연휴 간 부랴부랴 OpenCL을 살펴보면서 코딩을 해보았습니다. 좌충우돌 OpenCL 코딩을 하면서 겪었던 삽질은 다음 포스팅으로 넘기고 먼저 결과부터 공개! 

성능 실험!

실험 데이터 셋은 총 13,123개의 문헌으로 구성되어있고, 16,286개의 서로 다른 어휘, 전체 단어 수는 1,274,895개입니다. 깁스 샘플링 1000회를 실시하는데 걸리는 시간을 측정해보았습니다.

방법	걸린 시간(ms)	속도 비교치(%)
일반 깁스 샘플링(CPU 1코어)	258,076	100%
AD-LDA 깁스 샘플링(CPU 4코어)	96,023	268.8%
PCGS (on GPU)	9,417	2740.5%
AD-LDA 깁스 샘플링(16코어)*	52,985	487.1%

테스트에 사용한 CPU는 Intel i5-6600이고, GPU는 NVIDIA GeForce GTX 960입니다. CPU 4코어를 사용할 경우의 속도 향상이 3배가 약간 안 되었던것에 비해, GPU를 팍팍 갈구니 무려 27배 넘게 빨라져버렸네요...!! (* 그리고 궁금해서 16코어를 돌릴 수 있는 학교 서버와도 비교해보았습니다. 서버 CPU는 Intel Xeon E5-2620 v4인데, 16코어 풀로 돌려도 성능 향상은 5배가 약간 안되네요. 압도적으로 우리집 컴퓨터 그래픽 카드가 승리했습니다.)

Yan의 논문에서는 Intel 4개 코어 @ 2.4GHz인 CPU와 Geforce GTX 280을 비교했을때 약 20~30배의 성능향상이 있었다고 하니, i5-6600과 GTX960사이의 비율도 대략 비슷하게 나온것 같습니다. 게다가 OpenCL을 급하게 살펴보면서 짠 것이라 최적화가 덜 되었을수도 있다는 것을 감안하면 잠재적인 성능 향상도 더 기대해볼 수 있겠습니다.

이 알고리즘은 비단 LDA 토픽 모델링 뿐만 아니라 다른 깁스 샘플링을 사용하는 거의 모든 알고리즘에 적용이 가능할 수 있기 때문에 그 잠재력이 굉장히 크다고 할 수 있겠습니다. 요즘 컴퓨터로 게임을 잘 못해서 놀고 있는 GPU에 일을 시킬 생각에 신이 났는데, 신난김에 LDA 말고 좀더 복잡한 토픽 모델링 코드도 OpenCL로 구현이 가능한지 고민을 해봐야겠습니다.

1. Newman, D., Smyth, P., Welling, M., & Asuncion, A. U. (2008). Distributed inference for latent dirichlet allocation. In Advances in neural information processing systems (pp. 1081-1088). [본문으로]
2. Yan, F., Xu, N., & Qi, Y. (2009). Parallel inference for latent dirichlet allocation on graphics processing units. In Advances in Neural Information Processing Systems (pp. 2134-2142). [본문으로]