코드로 간단하게 구현해 보는 RNNLM

RNNLM

RNN 기반 언어 모델은 Transformer의 등장 이전에 많이 사용된 언어 모델 형태로 기존 고정 길이 컨텍스트 방식의 문제점을 해결하기 위해 고안된 모델 형태입니다.
RNN에서는 이전 단계에서의 은닉 상태 값인 $h_{t-1}$과 현재 단계의 입력값인 $i_t$를 합쳐서 출력값 $o_t$를 구하는 방식으로 이전 단계들에서 등장한 값들을 버리지 않고 이후의 단어를 예측하는 데 사용합니다.
참고로 이 글은 Recurrent neural network based language model 논문을 바탕으로 작성한 글입니다.

기존 고정 길이 컨텍스트 방식의 문제점과 RNN의 해결 방법

기존의 고정 길이 컨텍스트 방식은 항상 고정된 길이의 컨텍스트로 예측을 진행하였습니다. 예를 들자면
나는 사과를 먹었다. 하지만 나는 사과보다 바나나를 더 좋아한다.라는 문장이 있을 때 3개의 컨텍스트로 다음 단어를 예측한다면
나는 사과를 먹었다.
사과를 먹었다. 하지만
먹었다. 하지만 나는
이런 식으로 항상 3개의 컨텍스트만을 이용해서 다음 단어를 예측합니다. 하지만 이런 방식은 문제가 많을 수밖에 없습니다.

만약 저 문장 뒤에 그러므로 나는 다음번에는 바나나를 살 것이다라는 문장이 온다고 생각을 해보면
그러므로 나는 다음번에는이라는 컨텍스트만 가지고 바나나를 살 것이다라는 문장을 예측해야 하는 것입니다. 어떤 힌트도 없이 제가 바나나를 살 거라는 사실을 예측하기는 힘들겠죠.

따라서 RNN에서는 이전에 나온 값들을 버리지 않고 다음 예측에 사용하는 것으로 이를 해결합니다.
그러므로 나는 다음번에는이라는 문장 뒤를 예측하기에 앞서 제가 사과보다 바나나를 더 좋아한다는 사실을 알면 그 뒤를 예측하는데 도움이 되겠죠. 이것이 RNN의 기본 컨셉입니다.
이제 RNNLM이 어떤 것인지 알아보았으니 직접 구현해 보는 단계로 넘어가겠습니다.

RNNLM의 모델 구조

RNNLM의 모델 구조는 매우 간단합니다. RNNLM은 3개의 layer를 가지고 있는데 각각 input layer인 $x$, hidden layer인 $s$ 그리고 output layer인 $y$입니다. 또한 각 시간 $t$에서의 input 값은 $x(t)$, hidden layer의 state 값은 $s(t)$ 그리고 output 값은 $y(t)$로 나타냅니다.

def __init__(self, vocab_size, embedding_dim, hidden_dim):
	super(MyRNNLM, self).__init__()
	self.embedding = nn.Embedding(vocab_size, embedding_dim)
	self.hidden_dim = hidden_dim
	self.U = nn.Parameter(torch.randn(embedding_dim + hidden_dim, hidden_dim))
	self.V = nn.Parameter(torch.randn(hidden_dim, vocab_size))
	self.f = nn.Sigmoid()
	self.g = nn.Softmax(dim=1)

위 코드는 모델의 요소들을 생성하는 코드로 각각 input layer와 hidden layer 간의 weight($U$), hidden layer와 output layer 간의 weight($V$), embedding layer 그리고 각각의 연산에 필요한 함수들(Sigmoid, Softmax)입니다.

RNN의 연산 과정

RNN의 학습 과정에서는 위에서 설명한 3가지 값이 사용됩니다. 입력값 $x(t)$, hidden layer의 state 값 $s(t)$, 출력값 $y(t)$입니다. 이 중 $y(t)$에 대해서는 설명할 필요가 없지만 $x(t)$와 $s(t)$의 관계에 대해서는 설명해보겠습니다.

RNN에서는 이전 단계에서 등장한 값을 버리지 않고 이후의 단어를 예측하는 데 사용한다고 하였습니다.
(1) $x(t) = w(t) + s(t − 1)$
위의 식은 이를 나타낸 수식으로 어느 한 시점 $t$에서의 $x(t)$는 이전 단계에서의 $s(t)$ 값(=$s(t-1)$)과 현 단계에서의 입력된 단어의 token embedding $w(t)$를 합한 것입니다.

word_embedding = self.embedding(word)
x = torch.cat((word_embedding, hidden), dim=1)

그리고 이 $x(t)$를 이용해서 hidden layer에서 수행하는 연산은
(2) $s_j(t) = f(∑_ix_i(t)u_{ji})$
입니다. 여기서 함수 $f$는 Sigmoid 함수입니다.
위 수식의 각 요소들을 자세히 알아보면 $s_j(t)$는 $s(t)$의 $j$번째 element, $x_i(t)$는 $x(t)$의 $i$번째 element이고 $u_{ji}$는 input $i$의 hidden unit $j$에 대한 weight입니다.
여기서 $∑_ix_i(t)u_{ji}$은 모든 $i$에 대해 $u_ji$와 $x_i(t)$의 곱을 구하고 그 합을 구하는 것입니다. 또한 이 식은 $s_j(t)$에 대한 식이므로 $s(t)$에 대한 식으로 생각할 경우 두 matrix의 각 행과 열의 원소들을 곱한 뒤 그 합을 구하는 것이므로 행렬 곱 연산(torch.mm)과 동일한 연산이 됩니다.

hidden = self.f(torch.mm(x, self.U))

마지막으로 예측값 출력 과정을 수식으로 나타내면
(3) $y_k(t)=g(∑_js_j(t)v_{kj})$
입니다. 여기서 함수 $g$는 Softmax입니다.
위 수식도 각 요소들을 자세히 알아보면 $s_j(t)$는 $s(t)$의 $j$번째 element이고 $v_{kj}$ hidden unit $j$의 output $k$에 대한 weight 입니다.
여기서도 $∑_js_j(t)v_{kj}$은 모든 $j$에 대해 $s_j(t)$와 $v_{kj}$의 곱을 구하고 그 합을 구하는 것입니다. 또한 (2) 연산과 같은 원리로 이 연산은 행렬 곱 연산과 동일한 연산이 됩니다.

output = self.g(torch.mm(hidden, self.V))

이제 각 연산의 수식과 그를 표현한 코드를 모두 보았으니 이를 모두 합치면 아래와 같습니다.

def forward(self, word, hidden):
	word_embedding = self.embedding(word)
	x = torch.cat((word_embedding, hidden), dim=1)
	hidden = self.f(torch.mm(x, self.U))
	output = self.g(torch.mm(hidden, self.V))
	return output, hidden

후기

이렇게 RNNLM을 간단하게나마 코드로 구현해 보는 시간을 가졌습니다. 간단한 코드지만 ChatGPT와 Perplexity, 그 외 다른 블로그들의 도움을 많이 받았는데 아직 공부할 게 굉장히 많다는 생각이 들었습니다. 그동안은 논문을 읽고 모델의 컨셉을 이해하고 수학적인 내용을 해석해 보는 것 정도에서 그쳤었는데 그쳤었는데 이제는 그것을 어떻게 코드로 구현할 수 있을지도 고민해봐야 할 것 같습니다.