Hierarchical Softmax 자세히 알아보기

word2vec은 단어의 의미를 잘 표현하는 임베딩을 생성하는 알고리즘이며, CBOW와 Skip-gram이라는 두 가지 학습 알고리즘을 통해 생성됩니다. word2vec은 단어를 표현하는 데 있어 뛰어난 성능으로 주목을 받았지만, 밀도 높은 벡터를 확률 분포로 표현하는 과정에서 호율성을 개선해야 한다는 근본적인 과제가 남아있었습니다. 기존에는 소프트맥스 함수가 역할을 담당했지만 데이터의 규모나 어휘 집합의 크기가 급증하면서 효율적인 대안이 요구되었습니다. 그래서 등장한 것이 바로 hierarchical softmax(HS)입니다. 따라서 word2vec 논문에는 Hierarchical Softmax라는 개념이 소개되어 있습니다. 소프트맥스 함수에 대해서는 알고 있지만 Hierarchical이 무엇을 의미하는지 정확히 알지는 못했습니다. 논문에서도 이진 트리를 사용하는 방식이라고는 간단하게 언급이 되어있지만 구체적으로 어떻게 나타나는지가 궁금해서 공부한 내용을 정리해보았습니다.

 

---

 

Softmax의 개념

Hierarchical softmax가 무엇인지 알아보기에 앞서 Softmax의 개념을 간단히 짚고 넘어가겠습니다. 소프트맥스 함수는 softargmax 또는 normalized exponential function으로도 알려져 있습니다. 이 함수는 $K$개의 실수를 $K$개의 가능한 결과의 확률 분포로 나타냅니다. 이는 로지스틱 함수, 또는 시그모이드 함수를 다차원으로 일반화한 것으로, 다항 로지스틱 회귀에 사용됩니다. 시그모이드 함수는 이진 분류에 사용되며 어떤 샘플이 클래스에 속할 확률을 0 또는 1로 구분합니다. 다차원으로 일반화했다는 것은 이 개념을 확장하여 하나의 확률 값을 출력하는 대신 어떤 샘플이 각 클래스에 속할 확률을 모든 요소의 합이 1인 벡터로 나타낸다는 의미입니다.

 

소프트맥스 함수는 다음과 같이 정의됩니다.

 

즉, $K$개의 클래스에 대한 원본 로짓 값이 주어졌을 때, 각각의 클래스에 속할 확률을 계산하게 됩니다. 소프트맥스 함수는 로짓을 다중 클래스의 확률분포로 변환하기 때문에 결과를 해석하기 쉽게 하고, 지수 함수를 사용하기 때문에 로짓의 차이를 극대화하여 예측이 명확해지는 장점이 있습니다.

 

Hierarchical Decomposition의 등장

그렇다면 hierarchical softmax는 어떤 개념인지 알아보겠습니다. 원래 hierarchical softmax는 Frederic Morin과 Yoshua Bengio의 논문 Hierarchical Probabilistic Neural Network Language Model에서 처음 등장합니다. 이 개념은 신경망에 기반한 언어 모델의 학습 속도가 극심하게 느리다는 단점을 해소하기 위해 hierarchical decomposition이라는 이름으로 제안되었습니다. 이 논문에서 저자는 WordNet의 IS-A taxonomy를 사용하여 단어의 의미론적 개념을 트리와 같은 형태의 그래프로 구조화하였습니다. WordNet 트리는 이진 트리가 아니었기 때문에 데이터에 기반하여 K-means 알고리즘을 통해 이진 클러스터링을 수행하였습니다.

 

 

이 개념을 바탕으로 word2vec의 hierarchical softmax는 어떻게 계산되었는지 알아보겠습니다. 먼저 word2vec를 다루는 두 논문 Distributed Representations of Words and Phrases and their Compositionality와 Efficient Estimation of Word Representations in Vector Space에서는 공통적으로 hierarchical softmax의 연산 효율을 강조합니다. 기존 소프트맥스의 연산 복잡도가 $O(n)$인 반면, hierarchical softmax는 $O(\log n)$입니다. 그 이유를 간단히 알아보겠습니다.

 

예를 들어서 어휘 사전(vocabulary)의 크기가 10000인 데이터를 통해 학습한다고 생각해보겠습니다. 기존의 소프트맥스는 각 단어에 대해서 10000번의 연산을 수행해야 합니다. 즉, 단어의 개수에 따라 연산의 수가 선형적으로 증가합니다.

 

 

반면 10000만 개의 단어를 100개의 그룹으로 나누고, 각 그룹 아래에 100개의 단어를 포함하면, 연산은 100 + 100으로 총 200번이 수행됩니다. 여기서는 로그의 밑이 100인 경우로 생각할 수 있지만, 실제로는 이진 트리를 사용하기 때문에 $\log_2$의 복잡도를 갖게 됩니다.

 

 

word2vec에서의 Hierarchical Softmax

논문에서는 hierarchical softmax를 다음과 같이 설명합니다. Hierarchical softmax는 출력 레이어의 이진 트리 표현을 사용합니다. 이 트리는 $W$개의 단어를 말단 노드(leaves)로 하고, 각 노드에는 자식 노드가 선택될 확률이 표시되어 있습니다. 따라서 모든 단어에 대한 확률을 계산하는 것이 아니라 적절한 단어까지의 경로를 결정하는 이진 트리의 각 노드에서만 확률을 계산합니다. 즉 기존의 소프트맥스와 다르게 모든 단어에 대한 확률 분포를 갱신할 필요 없이, 경로에 있는 노드에서 각 자식 노드로 이어질 확률만을 계산하며 되기 때문에 연산 효율이 크게 좋아집니다.

 

지난 논문 리뷰에서 hierarchical softmax에 대해 설명한 부분을 가져와서 예시와 함께 단어가 어떻게 선택되는지 한번 알아보겠습니다. 먼저 논문에서 hierarchical softmax를 기술적으로 어떻게 정의했는지를 알아보겠습니다.

 

트리의 루트에서 적절한 경로를 따라 단어 $w$에 도달할 수 있습니다. $n(w,j)$를 루트에서 $w$로 가는 경로의 $j$번째 노드라고 하고, $L(w)$를 이 경로의 길이라고 하겠습니다. 그러면 $n(w,1)=\textrm{root}$이고 $n(w,L(w))=w$입니다. 그리고 모든 내부 노드 $n$에 대하여 $\textrm{ch}(n)$을 $n$의 임의의 고정된 자식 노드라고 하고 $[\![x]\!]$는 $x$가 참일 경우 1, 그 외의 경우 -1이라고 하겠습니다. 그러면 hierarchical softmax는 $p(w_O|w_I)$를 다음과 같이 정의할 수 있습니다. 식에서 $\sigma(x)=1/(1+\exp(-x))$입니다.

 

 

출처:  https://paperswithcode.com/method/hierarchical-softmax

 

위와 같은 트리에서 Context $C$가 주어졌을 때 루트 노드에서 “I’m”까지의 경로를 생각해보겠습니다.$L(\textrm{"I'm"})=4$이며, $p(\textrm{”I’m”}|w_I)$를 계산해야 합니다. 이 확률은 위에서 언급한 수식에 따라 계산할 수 있습니다. 그런데 논문에서는 $ch(n)$이 임의의 고정된 자식 노드(arbitrary fixed child of $n$)라고 표현됩니다. 즉 사전에 left child와 right child 중 하나를 지정하고, 항상 지정한 방향에 대해서만 생각하겠다는 의미입니다. Left child로 고정하였다면, $[![x]!]$는 순서대로 1, 1, -1이 될 것입니다. 그리고 $v^\prime_{n(w,j)}$는 단어가 아닌 각 노드가 갖는 벡터 표현입니다.

 

첫 번째로 $C$의 벡터와 루트 노드의 곱을 계산합니다. 식으로는 ${v^{\prime}{n(\textrm{"I'm"},1)}}^{\top}v{C}$와 같이 표현될텐데, 이 값을 계산하고 시그모이드 함수를 통과하면 0과 1 사이의 값이 출력됩니다. 그 값이 0.57이라고 하겠습니다. Fixed child를 left child로 지정하였으니, 결국 left child를 선택할 확률이 0.57이 될 것이고, 반대로 right child의 확률은 0.43이 될 것입니다. 그리고 $j=3$이 될 때까지 이를 계산하면, $p(\textrm{”I’m”}|w_I)=0.57 \times 0.68 \times 0.72$이 될 것입니다. 참고로, 각 노드에서 자식 노드가 선택될 확률의 합이 1이 된다는 것은 다음 수식을 통해 증명할 수 있습니다.

 

 

Binary Huffman Tree

논문에서는 hierarchical softmax에서 이진 허프만 트리(binary Huffman tree)를 사용하였다고 했습니다. 이 부분이 각 노드에도 벡터 표현이 할당되는 이유입니다. 정보 이론에서 허프만 코드는 무손실 데이터 압축에 사용되는 접두사 코드입니다. 그리고 이런 코드를 찾거나 사용하는 것을 허프만 코딩이라고 합니다. 각 기호는 발생 빈도를 가중치로 하는 코드 테이블을 통해서 인코딩됩니다. 자세한 내용은 관련 문서를 확인해주세요. 결국 허프만 트리는 각 노드에 하위 노드를 표현할 수 있는 정보가 인코딩되어 있을 것이기 때문에, 어휘 집합을 적절하게 구조화할 수 있고, hierarchical softmax의 연산 효율성을 높이자는 목표를 달성하는데 적합합니다.

 

앞서 언급했듯이 기존 소프트맥스는 복잡도가 $O(n)$인데 반해 hierarchical softmax는 $O(\log n)$의 복잡도를 갖습니다. 그런데 논문에서는 hierarchical softmax의 연산 비용이 $L(w_O)$에 비례한다는 표현도 나타납니다. 이는 hierarchical softmax에 사용하는 트리의 모든 노드가 똑같이 두 개의 자식 노드를 갖지 않기 때문에, 연산량이 각 단어에 이르는 경로에 비례한다는 것을 의미합니다. 만약 완전 이진 트리가 된다면 평균적인 연산량은 $O(\log n)$에 비례할 것입니다.

 

---

 

이처럼 hierarchical softmax는 기존 softmax에 비해서 연산 효율을 크게 높일 수 있습니다. 신경망 훈련을 위한 대규모 데이터셋이 등장하고, word2vec 논문에서는 phrase representations을 특별히 추가하여 어휘 사전의 크기도 커진 만큼, 연산 효율성을 더욱 민감하게 다루었습니다. 저자는 이와 같은 기발한 테크닉을 여럿 적용하여 word2vec의 성능과 속도를 동시에 끌어올릴 수 있었습니다.