Efficient Softmax Approximation for GPUs

Photo by Nana Dua on Unsplash

ICML 17년도에 게재된 논문으로 지금까지 157회 인용을 받은 논문이다. softmax 함수 자체가 계산 시간이 꽤 필요한 operation인데 이 계산 시간을 짧게 만들기 위해서 조건부 확률과 chain rule을 이용해 softmax를 근사하는 방법이 알려져 있다. 이 논문에서는 GPU 환경에서는 해당 방법을 어떻게 구성하는 것이 좋은 지에 대한 고찰이 포함되어 있다.

클래스 기반 계층 소프트맥스

뉴럴 랭귀지 모델링에서는 다음 단어의 확률을 예측하기 위해서 사전 안에 있는 모든 단어의 점수를 계산한 뒤, 정규화를 해야 한다. 보통 정규화되지 않은 점수에 softmax라는 함수를 적용하여 정규화를 진행한다.

$$ f(z_w) = \frac{exp(z_w)}{\sum_{w' \in V}exp(z_{w'})} $$

이 함수는 $O(|V|)$만큼의 계산이 필요하다. 그런데 사전의 크기가 커지게 되면, 이 과정은 계산상 오래 걸리게 되고 전체 모델의 계산 시간 중 대부분을 차지하게 된다. 이 계산 비용을 줄이기 위한 간단한 방법으로 사전에 있는 각 단어 $w$들을 어떤 고유한 클래스 $C(w)$로 할당하고 단어들의 확률분포를 아래처럼 분해하는 방법이 있다.

$$ p(w_t | h_t) = p_1(C(w_t)|h_t) \times p_2(w_t|C(w_t), h_t) $$

여기서 p1과 p2는 softmax함수를 이용해 구할 수 있다. 만약 각각의 클래스가 사전 크기의 루트 만큼 가지고 있다면, 계산 비용은 $O(\sqrt{|V|})$로 감소된다.

적응 기반 소프트맥스

이 논문에서는 단어의 확률분포를 계산할 때 그 속도를 향상시킬 수 있는 간단한 방법인 adaptive softmax을 제안한다. 이것은 계산 시간을 줄이기 위해 단어의 클래스를 만드는 클래스 기반 계층 소프트맥스에서 영감을 받았다. 우리의 방법은 뉴럴 네트워크를 학습시키는데 주로 사용되는 GPU에 효율적으로 설계되었다. 좀더 명확히 하기 위해, 우리는 먼저 우리의 방법을 제안하게 된 직관을 간단한 케이스에서 보이도록 하겠다. 우리는 일반적인 케이스에서 분석하기 전에 먼저 2개의 분별력 있는 클러스터로 분리하였다.

행렬곱을 위한 계산 시간 모델링

우리의 방법을 설명하기 전에 우리의 모델의 장점을 제대로 이해하기 위해서 계산 시간에 대한 모델링을 먼저 진행하고자 한다. 모델을 계산 시간의 대부분은 은닉 상태를 나타내는 행렬과 단어의 표현형을 나타내는 행렬간의 곱셈이 차지한다. 따라서 이 곱셈을 모델링할 것이다. 고정된 사이즈 $d$인 은닉층에서 우리는 $g(k, B)$를 위의 행렬곱에 걸리는 시간이라고 표기하고자 한다. 여기서 $k$는 사전의 크기, $B$는 배치의 크기를 의미한다.

두 행렬곱을 했을 때, 걸리는 시간. 배치의 크기는 2560이고 2048차원의 은닉 차원을 사용하고 사전의 크기가 k일 때 걸리는 시간을 보고했다. 1000번의 시행을 반복했고 그것의 평균을 보고했다. 2의 15제곱인 32768일 때, 시간이 0.05초 정도 걸린다. 논문이 오래전 것이라 최근 GPU 아키텍처에 대해서는 나와있지 않다.

배치의 크기를 고정하고 $k$를 바꿔가면서 계산 시간을 측정해보면 위의 그림과 같다. 위 그림을 보면 $k$가 작을 때 (대략 50 미만) 에는 상수 시간 정도 걸리지만 그 이상일 경우 선형으로 증가함을 알 수 있다. 이를 수식으로 나타내면 아래와 같다.

$$ g(k) = c_m + max[0, \lambda(k-k_0)] $$

여기서 $c_m$은 최소로 걸리는 시간 $k_0$는 상수 시간이 걸리는 것 중 제일 큰 크기를 의미한다. 이러한 형태는 배치의 크기에 대해서 분석했을 때에도 같은 양상을 보인다. 즉, 하나의 차원이 작을 때 행렬곱은 비효율적이라는 의미가 된다. (여기서 작다는 것은 50 미만을 의미하는 것 같다.) 이러한 관측은 이진 Huffman code와 같이 노드당 적은 수의 자식을 가지는 계층 구조는 꽤나 좋은 구조와는 거리가 멀다는 것을 나타낸다. (왜냐하면 자식을 적게 가지면 GPU를 비효율적으로 사용하게 되기 때문이다.) 비슷한 이유로, 희소한 단어들로만 구성된 클래스 군집은 낮은 확률과 줄어드는 배치 크기를 가지게 되어 결국 비효율적인 행렬곱을 초래한다. (????? 무슨 뜻인지 아시는 분 구합니다..) 따라서, 우리는 행렬곱에 걸리는 계산 시간의 모델을 다음과 같이 사용하기를 제안한다.

$$ g(k, B) = max(c + \lambda k_0 B_0, c + \lambda k B) $$

이 모델이 대충 계산 시간을 모델링 한 것 같지만 이것은 실험적인 관측을 잘 설명하도록 해준다.

2개의 클러스터일 때의 예시

자연어에서 단어의 분포는 지프 분포를 따른다고 악명 높게 알려져 있다. 대부분의 확률 질랑은 사전의 아주 작은 부분들로 구성되어 있다. 예를 들어, Penn treebank라는 문서의 87%가 오직 사전 속의 20%의 단어로 표현되어져 있다. 이러한 정보는 계산 시간을 줄이는데 이용될 수 있다.

전체 계산시간을 줄이는 간단한 전략으로 사전을 두개의 군집으로 나누는 것이다. 하나는 빈번하게 등장하는 단어들로 이루어진 사전과 (이를 head 사전이라고 부른다) 많은 수의 희소한 단어들로 이루어진 사전 (이를 tail 사전이라고 부른다)이다. 분류기는 빈번하게 head 사전에 접근할 것이고 이는 곧 계산을 효율적으로 할 것임을 의미한다. 대조적으로, tail 단어들은 적게 나타날 것이고 상응하는 계산 시간은 느려질 것이다. 이것은 클러스터를 불균형한 크기와 확률로 나누라는 것을 제안한다. 이렇게 2개로 나눈 클러스터를 2가지 방식으로 계층화할 수 있다. 하나는 루트 노드에서 2개의 자식을 가진 뒤, 각각의 자식이 head 사전과 tail 사전에 있는 단어들을 자식으로 가지는 방법과 루트 노드가 head 사전에 있는 단어들과 추가적인 1개의 자식을 가진 뒤, 추가한 1개의 자식이 tail 사전에 있는 단어들을 자식으로 가지는 short-list 기법이 있다.

전자의 경우 성능이 심각하게 떨어지는 경향이 실험적으로 관측되었는데 이는 short-list 기법이 head 사전에 있는 단어들의 확률 분포를 좀더 날카롭게 만들 수 있기 때문이라고 한다. 따라서, 속도가 크게 차이나지 않는다면 short-list 기법을 사용하는 것을 더 선호한다.

이제 softmax를 계산하는 시간을 최소화할 수 있게끔 군집화해보자. 사전의 크기는 $k$라고 하고, 우리는 head 사전의 크기가 몇일 때 계산 시간을 최소화할 수 있을 지 알고 싶다. 먼저, head 사전의 크기를 $k_h$라고 나타내자. 이 head 단어들은 총 분포의 $p_h$를 커버할 것이다. tail 사전은 사전의 나머지 부분으로 구성될 것이고 총 분포의 $p_t=1-p_h$를 커버할 것이다. 루트 노드에서 확률을 계산하기 위한 시간은 $g(k_h+1, B)$가 될 것이고 (short-list 구조일 때), tail 사전을 위해서는 추가적으로 $g(k_t, p_tB)$가 걸릴 것이다. (평균적으로 전체 배치의 $p_t$ 만큼만 tail 사전에 포함될 것이다.) 즉, 전반적인 계산 시간은 다음과 같다.

$$ C = g(k_h+1, B) + g(k_t, p_t B) $$

이제 우리는 head 사전의 크기를 바꿔가며 계산 시간이 가장 작아지는 head 사전의 크기를 찾으면 된다. 이를 Figure 2에 그려보았다.

불가리안 유러피안 데이터에서 2개의 클러스터로 나타냈을 때 걸리는 계산 시간 모델링. 이론상으로 계산한 모델과 실험적으로 측정한 시간이 거의 일치한다. 그냥 소프트맥스를 하는 것보다 5배 정도 더 빨라진 것을 확인할 수 있엇다. 빨간 점선은 두개의 클러스터가 동일한 확률 질량을 가지도록 설계했을 때를 의미한다. 빈번한 단어가 100개도 안되는데 전체 문서의 50%를 커버할 수 있음을 알 수 있다.

이를 보며, 같은 확률을 가지도록 2개의 군집을 가르는 것이 head 군집의 좋은 크기가 아니라는 것을 알 수 있다.

각 클러스터는 다른 클러스터에 대해 독립적으로 접근이 가능하므로 굳이 같은 차원의 크기를 가질 필요는 없다. 빈번한 단어들은 제대로 예측되기 위해서 큰 차원이 필요하다. 반대로, 희소한 단어들은 자주 나타나지 않아서 잘 학습이 되지 않을 것이다. 그래서 이 단어들에게 큰 차원의 벡터를 이용하는 것은 낭비일 것이다. 우리는 이 관측을 이용해서 계산 시간을 좀더 단축시켰다. 하지만 우리는 이전 연구와 다르게 사영 행렬을 적용하여 분류기의 입력 크기를 줄였다. 주로, tail 군집의 차원은 head 차원보다 4배 적게 구성했다.

일반적인 경우

이제 군집을 여러 개 만드는 경우를 고려해보자. 아마 전체적인 구성은 아래 그림과 같아질 것이다.

short-list의 구조를 따르고 있으며 head 군집만 root랑 바로 연결되고 나머지 군집은 한번 더 거치는 구조로 구성되어 있다.

마찬가지로, 계산 시간을 모델링해보자. 지금은 각 사전의 크기와 확률의 함수로써 계산 시간을 분석하기 위해 배치 크기와 은닉 차원을 고정하여 계산했다. 총 $J+1$개의 군집이 있다고 하자. 하나는 head 사전이고 $J$개의 tail 사전이 있다. $p_i$를 각각의 사전의 커버하는 확률 질량이라고 하고 $k_i$를 각 사전의 크기라고 생각하면 계산 시간의 기대값은 아래와 같다.

$$ C = C_h + \sum_i C_i $$

$$ C_h = g(k_h + J, B) $$

$$ \forall i, C_i = g(k_i, p_i B) $$

$$ C = g(J+k_h, B) + \sum_i g(k_i ,p_i B) $$

여기서 각각의 군집이 충분히 크다고 가정하면, 원래 계산 모델에서 constant 부분을 무시하여 간단하게 표현할 수 있게 된다.

$$ C = c + \lambda (J+k_h) B + \sum_i (c + \lambda k_i p_i B) = (J + 1)c + \lambda B [J + k_h + \sum_i p_i k_i] $$

이제 이 식을 통해 각 군집에 단어들이 어떻게 들어가는 것이 좋은지 분석해보자. 이를 위해 우선 각 군집의 크기를 고정하고 생각해보자. 어떤 두 개의 군집 $V_i$와 $V_j$에 대해서, $p_j k_j$를 아래와 같이 다시 적을 수 있다.

$$ p_{i+j} = p_i + p_j $$

$$ p_i k_i + p_j k_j = p_i (k_i - k_j) + p_{i+j} k_j $$

여기서 일반성을 잃어버리지 않고 $k_i$는 $k_j$보다 크다고 가정할 수 있다. 이제 이 식을 최소화하는 것이 결국 전체적인 계산 시간을 최소화할 수 있게 되는데 위 식의 우변의 두 번째 항이 가정에 따라서 상수가 된다. 즉, 우변의 첫 번째 항을 최소화해야 하는데 이것은 즉 더 큰 군집 ($k_i$가 $k_j$보다 크기가 크다)에 더 작은 확률 ($p_i$를 줄여야 된다)을 할당해야 한다는 것을 의미한다. 달리 말하면, 빈번한 단어들이 가장 작은 군집에 포함되어야 한다는 것을 요구한다. 이는 어떠한 군집 2개를 가지고 계산을 해보아도 참이다. 따라서, 결론적으로 군집의 개수와 각 군집의 크기를 고정했을 때, 가장 좋은 전략은 각 군집을 크기 순서대로 정렬한 다음, 확률이 높은 단어 (빈번한 단어) 순서대로 크기가 작은 군집에다가 할당하는 것이다.

다음으로 군집의 개수를 고정했을 때 각 군집의 크기를 결정하는 방법에 대해서 알아보자. 이는 사실 다이나믹 프로그래밍을 통해 완벽하게 결정할 수 있다. 위의 관측을 기반으로 빈번한 단어 순서대로 단어를 정렬한 뒤, $m$개의 군집으로 나누는 문제로 풀 수 있다. (short-list라서 조금 헷갈릴 거 같긴 한데 해보면 할 수 있지 않을까?)

남은 파라미터는 군집의 개수이다. 이는 그냥 개수를 변화시켜가면서 계산 시간을 측정해보았다. 해보니 클러스터의 개수를 10개에서 15개 정도로 잡는게 가장 빠르게 계산되었다. 또한 5개 보다 많이 클러스터로 구성해도 5개로 클러스터를 구성하는 것보다 크게 빨라지지 않는다는 것을 관측했다. 따라서 우리는 2개에서 5개 정도의 클러스터 개수를 이용했다. 이는 그저 실험 결과가 2개에서 5개 정도의 클러스터를 사용했을 때 더 좋은 평가 지표를 얻을 수 있었기 때문이고 이것이 계산 시간을 빠르게 하는 것과 모델의 정확도 간의 상관관계에서 좋은 타협점이라는 것을 의미한다. 뒤의 실험에서 보면 우리의 방법이 softmax를 그냥 사용하는 것보다 더 빠르고 충분히 비슷한 평가 지표를 달성할 수 있음을 알 수 있다.

군집의 개수에 따라서 변화하는 계산 시간. 클러스터의 개수가 증가하면 계산 시간이 빨라지지만 5 이상일 때는 딱히 계산 시간이 많이 증가하지 않는 것을 알 수 있다.

실험

작성중