[딥러닝] 생성형 AI (LLM) 에서 Loss는 어떻게 계산될까

LLM을 그래도 1년 남짓 공부하면서 종종 들었던 의문이 하나 있다.

 

LLM(트랜스포머 아키텍처) 특정 단어 뒤에 어떤 단어가 오는지 확률로 계산한다. 

 

대강 이런 식으로. 그래서 LLM은 가장 확률이 높은 단어를 반환하거나, top_k, top_p 샘플링과 temperature 같은 파라미터를 통해 더욱 다양한 답변을 내놓거나 그렇지 않도록 설정할 수 있다.

 

하지만 LLM을 훈련, 특히 QA 데이터셋 (질문 - 답변)으로 LLM을 훈련시키고 있는 상황을 생각해보자.

 

질문 : 나는 오늘 기분이 정말 별로였어. 내가 무엇을 하면 좋을지 추천해줘.
답변 : 기분을 좋게 만드는 방법엔 음악 듣기, 게임하기, 운동하기 등이 있습니다.

 

위와 같은 형식으로 훈련을 시킨다고 치면, LLM은 가장 자연스러운 답변을 생성하도록 훈련된다.

내가 궁금한건 이 과정에서 Loss가 어떻게 계산되는가? 였다.

예를 들면, '기분을' 이라는 단어 다음에 '좋게' 가 정답인 단어로 정해져 있는데, 모델은 '기분을' 이라는 단어 다음에 나올 수 있는 일종의 확률표를 생성하게 되는데, '좋게'라는 단어와 수많은 단어의 확률표로 어떻게 Loss가 계산되지? 하는 것이었다. 얼핏 생각하면 답은 1개라서 정확히 정해져 있는 건데, 모델이 구한 답은 아주 많은 단어에 대한 확률이기 때문에 그 길이가 다른 것 같았기 때문이다.

 

그래서 혹시 추측했던 게 내가 알고 있는 개념에 한해서 원-핫 인코딩 형식처럼 정답 단어만 1로 놓고 나머진 0으로 놓나? 하는 것이었다. 위의 예시에서 '좋게' -> [0, 0, 1, 0, 0, 0, ... 0, 0] 이라는 벡터로 바꾸면 모델이 구한 답은 [0.01, 0.03, 0.08, 0.02, ... 0.04, 0.06] 이런 식일 텐데, 이게 cross-entropy 같은 걸로 Loss를 구할 수도 있겠는데? 하는 것이었다.

 

몇 가지 자료를 찾아본 결과, 정답은 YES! 였다.

 

https://medium.com/@jeraldteokj/visualising-loss-calculation-in-large-language-models-1af410a9d73d

How do Large Language Models learn?

 

위 포스팅은 내가 궁금했던 점을 정확히 설명해주는, LLM이 어떻게 배우는가? 에 대한 점을 자세히 설명해주고 있다.

물론 내가 궁금했던 점은 Pretrain 과정은 아니고 QA 데이터셋처럼 파인튜닝 과정이었지만 과정 자체는 동일하다. 특이한 것은 나는 단어 하나하나를 학습하는 줄 알았는데, 위 포스팅의 설명을 보면 벡터로 한 번에 입력받고 생성도 한 번에 진행한다. 아무래도 이게 계산과정에서 훨씬 빠르고 자연스럽기 때문에, 어떻게 보면 당연한 수순.

 

최근에 교수님과 면담을 하면서 Auto-Regressive Learning에 대한 질문을 받았다가 위처럼 대답했더니 틀렸다는 소리를 들었다. 내가 처음에 알고 있던 것처럼 단어 하나하나를 학습하는 것이 맞다. 아래 그림에서 <s>, I, love를 보고 dogs를 예측해야하는 상황에서, 그 자리 (position2)에 올 단어 확률표를 예측하는 것이 Auto-Regressive Learning이다. 내가 그림만 보고 그 과정을 오해했다.

 

지금은 Decoder only 모델 (ex : GPT) 예시다. Input Text를 받으면 토큰화 -> 임베딩 -> Decoder Layer -> Linear Layer -> 을 거쳐서 Logit이 계산되어 나오게 된다. 이 Logit은 Softmax를 거쳐서 확률이 되기 전 상태이고, 여기에 Softmax를 거치면 그 확률표가 나오는 것이다. (그림에선 logit이 10개 단어에 대해 계산되지만 이 수치는 단어 사전의 크기에 따라 얼마든지 달라진다.)

 

이 확률표에서 True Label, 즉 정답에 해당하는 인덱스에 해당하는 확률을 쏙쏙 빼온다. 왜 쏙쏙 빼오는지 Cross Entropy 식을 보면 알 수 있다.

 

p(x)는 실제 값인데, 실제값은 단 1개이므로 이를 단어 사전의 길이 (확률표의 길이)와 맞춰주기 위해 one-hot으로 바꿔준다. 즉 실제값만 1로 두고 나머진 다 0으로 둔다는 의미. q(x)는 추측값인데, 여기에 로그를 씌워서 one-hot 벡터와 곱해주면 1과 매칭되는 값만 남고 나머지 0과 곱해지는 q(x)의 값들은 0이 되어 사라지게 된다. 따라서 정답 인덱스의 확률에 log를 씌워주고, 1보다 작은 수에 log를 씌우면 음수니까 -를 곱해줘서 양수로 만들어주면 작업은 끝나는 것이다. 그리고 로그함수의 개형을 생각해보면 0에 가까운 값일수록 로그를 씌우면 아주아주 작은 음수가 되기 때문에, 크면 오차가 크다는 loss의 개념에 적용하기 위해서라도 -를 곱해주는 것이 맞다.

 

이렇게 다 구한 다음 총합을 한 다음에 전체 Input 길이, Sequence length로 나눠주면 Train 과정에서의 Loss가 계산된다.

 

자세한 코드도 있는데 이는 포스팅을 참조.

 

정리 : LLM에서 Loss를 계산할 땐 정답 라벨을 one-hot으로 바꿔줘서 계산이 가능하게 한다.